#!/usr/bin/env python3 # desktop-assistant: local text assistant daemon. # Copyright (C) 2026 Danilo M. # # This program is free software; you can redistribute it and/or modify it # under the terms of the GNU General Public License version 2 as published # by the Free Software Foundation. # # This program is distributed in the hope that it will be useful, but # WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General # Public License for more details. # # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. """Questions arrive as JSON lines on a Unix socket, the loaded llama-server model answers with tools, replies stream back on the same socket, and every finished turn is kept in a SQLite history with full-text search. Run: python3 assistant.py (--test runs the self-check) """ import base64 import glob import hashlib import http.client import json import mimetypes import os import re import signal import socket import sqlite3 import subprocess import sys import threading import time import tomllib import traceback import urllib.parse import urllib.request CONFIG_DIR = os.path.join(os.environ.get("XDG_CONFIG_HOME") or os.path.expanduser("~/.config"), "desktop-assistant") DB_PATH = os.path.join(os.environ.get("XDG_DATA_HOME") or os.path.expanduser("~/.local/share"), "desktop-assistant", "history.db") DEFAULTS = { "llm": "http://localhost:8181/v1", "system": ("You are a desktop assistant. Answer in the language the user writes, " "Italian or English. Markdown is rendered: use lists, tables and fenced " "code blocks when they help, and keep answers as short as the question " "allows. Facts about the user's computer, devices, accounts, mail, " "calendar or money come only from a tool: call one, and if none can " "tell, say you do not know. Never guess them."), "tool_rounds": 3, "temperature": None, # None: the llama-server preset's "thinking": True, # False skips the model's reasoning "env": {}, } OPEN = {"env"} # free-form tables: replaced whole, not merged def merge(base, over, where=""): # A typo must fail loudly, not fall back to the default in silence. out = dict(base) for k, v in over.items(): if k not in base: raise SystemExit(f"config: unknown key {where}{k}") if isinstance(base[k], dict) and k not in OPEN: if not isinstance(v, dict): raise SystemExit(f"config: {where}{k} must be a table") v = merge(base[k], v, f"{where}{k}.") out[k] = v return out def load_config(path): try: with open(path, "rb") as f: return merge(DEFAULTS, tomllib.load(f)) except FileNotFoundError: return DEFAULTS def load_tools(d): tools = {} for path in sorted(glob.glob(os.path.join(d, "*.toml"))): with open(path, "rb") as f: t = tomllib.load(f) if not isinstance(t.get("description"), str) or not t.get("command"): raise SystemExit(f"{path}: needs description and command") tools[os.path.basename(path)[:-5]] = t return tools def schema(tools): return [{"type": "function", "function": { "name": name, "description": t["description"], "parameters": {"type": "object", "properties": { p: {"type": s.get("type", "string"), "description": s.get("description", "")} for p, s in t.get("params", {}).items()}}}} for name, t in tools.items()] CFG = load_config(os.path.join(CONFIG_DIR, "config.toml")) TOOLS_DIR = os.path.join(CONFIG_DIR, "tools") TOOLS = load_tools(TOOLS_DIR) TOOL_SCHEMA = schema(TOOLS) or None SOCK = os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/tmp"), "desktop-assistant.sock") def log(*a): print(time.strftime("%H:%M:%S"), *a, flush=True) # History: one SQLite file, written only by this process. turns.messages is # what the model is sent for that turn; user, reply and tools are what the # canvas draws. The FTS index follows turns through triggers, and a # conversation delete cascades into turns and so into the index. SCHEMA = """ CREATE TABLE IF NOT EXISTS conversations(id INTEGER PRIMARY KEY, title TEXT, created INTEGER, updated INTEGER); CREATE TABLE IF NOT EXISTS turns(id INTEGER PRIMARY KEY, conv INTEGER REFERENCES conversations ON DELETE CASCADE, seq INTEGER, user TEXT, reply TEXT, tools TEXT, messages TEXT); CREATE INDEX IF NOT EXISTS turns_conv ON turns(conv, seq); CREATE VIRTUAL TABLE IF NOT EXISTS turns_fts USING fts5(user, reply, content='turns', content_rowid='id', tokenize='unicode61 remove_diacritics 2'); CREATE TRIGGER IF NOT EXISTS turns_ai AFTER INSERT ON turns BEGIN INSERT INTO turns_fts(rowid, user, reply) VALUES (new.id, new.user, new.reply); END; CREATE TRIGGER IF NOT EXISTS turns_ad AFTER DELETE ON turns BEGIN INSERT INTO turns_fts(turns_fts, rowid, user, reply) VALUES ('delete', old.id, old.user, old.reply); END; """ def db_open(path): if path != ":memory:": os.makedirs(os.path.dirname(path), exist_ok=True) os.chmod(os.path.dirname(path), 0o700) # also when something else made it old = os.umask(0o177) # 0600; SQLite gives the WAL files the db's mode try: conn = sqlite3.connect(path, check_same_thread=False) finally: os.umask(old) if path != ":memory:": os.chmod(path, 0o600) # an existing file keeps its old mode otherwise conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys=ON") conn.executescript(SCHEMA) # Added with attachments; a database from before them gains the column. if "files" not in {r[1] for r in conn.execute("PRAGMA table_info(turns)")}: conn.execute("ALTER TABLE turns ADD COLUMN files TEXT") return conn def save_turn(conn, conv, user, reply, tools, messages, files=()): # Returns the conversation id: the row is made at its first finished turn, # so a New that is never used leaves nothing behind. A question that is # only files is titled by their names. now = int(time.time()) title = " ".join(user.split()[:6]) or ", ".join(f["name"] for f in files) with conn: if conv is None: conv = conn.execute("INSERT INTO conversations(title, created, updated) " "VALUES (?, ?, ?)", (title, now, now)).lastrowid else: conn.execute("UPDATE conversations SET updated = ? WHERE id = ?", (now, conv)) seq = conn.execute("SELECT COALESCE(MAX(seq), 0) + 1 FROM turns WHERE conv = ?", (conv,)).fetchone()[0] conn.execute("INSERT INTO turns(conv, seq, user, reply, tools, messages, files) " "VALUES (?, ?, ?, ?, ?, ?, ?)", (conv, seq, user, reply, json.dumps(tools), json.dumps(messages), json.dumps(list(files)))) return conv def list_convs(conn): return [{"id": i, "title": t, "updated": u} for i, t, u in conn.execute( "SELECT id, title, updated FROM conversations ORDER BY updated DESC, id DESC")] def load_conv(conn, conv): # (model history, shown turns), or None for an id that is not there. rows = conn.execute("SELECT user, reply, tools, messages, files FROM turns " "WHERE conv = ? ORDER BY seq", (conv,)).fetchall() if not rows: return None return ([json.loads(m) for _, _, _, m, _ in rows], [{"user": u, "reply": r, "tools": json.loads(t), "files": json.loads(f or "[]")} for u, r, t, _, f in rows]) def search(conn, q): # The query is one quoted phrase with a prefix star on its last word, so # FTS syntax typed by the user is matched as text and cannot widen it. # ponytail: rank per turn, first hit per conversation; 30 results. q = q.replace("\0", "").strip() # NUL ends the FTS string early if not q: return [] seen, out = set(), [] for conv, title, snip in conn.execute( "SELECT t.conv, c.title, snippet(turns_fts, -1, '**', '**', '…', 12) " "FROM turns_fts JOIN turns t ON t.id = turns_fts.rowid " "JOIN conversations c ON c.id = t.conv " "WHERE turns_fts MATCH ? ORDER BY rank LIMIT 300", ('"' + q.replace('"', '""') + '"*',)): if conv not in seen: seen.add(conv) out.append({"id": conv, "title": title, "snippet": snip}) return out[:30] def delete_conv(conn, conv): with conn: conn.execute("DELETE FROM conversations WHERE id = ?", (conv,)) # Attachments. Paths come from the canvas; the daemon runs as the user and # reads them itself. Text, and a PDF's text, go inline and are stored that # way, so a reopened conversation still has them. Images are stored as a # path and hash and read again when sent, so the history does not grow by # every picture, and a file changed since is reported gone rather than # shown as something it no longer is. CODE_SUFFIXES = { ".py", ".c", ".h", ".cpp", ".hpp", ".rs", ".go", ".js", ".ts", ".jsx", ".tsx", ".sh", ".bash", ".zsh", ".lua", ".rb", ".pl", ".php", ".java", ".kt", ".swift", ".sql", ".toml", ".yaml", ".yml", ".json", ".ini", ".cfg", ".conf", ".md", ".rst", ".txt", ".csv", ".xml", ".html", ".css", ".diff", ".patch", ".slackbuild", ".info", ".desktop", ".service", ".qml", } IMAGE_TYPES = {"image/png", "image/jpeg", "image/gif", "image/webp"} MAX_FILES = 8 MAX_IMAGE = 20 << 20 MAX_PAGE_IMAGES = 4 # scanned PDF pages sent as images PAGE_PX = 1600 # longest side of a page render class Refused(Exception): """A file or request the user has to change; logged without a trace.""" def kind_of(path): mime = mimetypes.guess_type(path)[0] if mime in IMAGE_TYPES: return "image" if mime == "application/pdf": return "pdf" if (mime or "").startswith("text/") or os.path.splitext(path)[1].lower() in CODE_SUFFIXES: return "text" if mime is None and looks_like_text(path): # no telling name: a script, a Makefile return "text" return "other" def looks_like_text(path): # The first 4 KB: no NUL and valid UTF-8 (a character cut at the end of # the window is allowed). try: with open(path, "rb") as f: head = f.read(4096) except OSError: return False if b"\0" in head: return False try: head.decode("utf-8") except UnicodeDecodeError as e: return len(head) == 4096 and e.reason == "unexpected end of data" return True def sha256(path): h = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(1 << 20), b""): h.update(chunk) return h.hexdigest() def pdf_pages(path): # ponytail: the whole text is read, then cut to the share; a page cap # (-l) if huge PDFs ever matter. try: out = subprocess.run(["pdftotext", "-layout", path, "-"], capture_output=True, text=True, timeout=60) except subprocess.TimeoutExpired: raise Refused(f"cannot read {os.path.basename(path)}: timed out") from None except OSError as e: # poppler missing raise Refused(f"cannot read {os.path.basename(path)}: {e.strerror}") from None if out.returncode: last = out.stderr.strip().splitlines()[-1:] or ["pdftotext failed"] raise Refused(f"cannot read {os.path.basename(path)}: {last[0]}") pages = out.stdout.split("\f") return pages[:-1] if out.stdout.endswith("\f") else pages # each page ends in \f def attach(paths, budget): # (text blocks, image references, chips) for the files; Refused says # what is wrong with the first bad one. Text files and PDFs share the # budget of characters equally. if len(paths) > MAX_FILES: raise Refused(f"at most {MAX_FILES} files") loaded = [] for p in paths: # A newline in a name would close the file block early. name = " ".join((os.path.basename(p) or p).splitlines()) if not (os.path.isabs(p) and os.path.isfile(p) and os.access(p, os.R_OK)): raise Refused(f"cannot read {name}") kind = kind_of(p) if kind == "other": raise Refused(f"{name}: chat supports text, images and PDF") if kind == "image" and os.path.getsize(p) > MAX_IMAGE: raise Refused(f"{name}: images over 20 MB are refused") loaded.append((p, name, kind)) share = budget // max(1, sum(k != "image" for _, _, k in loaded)) blocks, refs, chips = [], [], [] for p, name, kind in loaded: chip = {"name": name, "kind": kind, "truncated": False} chips.append(chip) if kind == "image": refs.append({"type": "image_ref", "name": name, "path": p, "sha256": sha256(p)}) continue if kind == "pdf": pages = pdf_pages(p) chip["pages"] = len(pages) body = "\n\n".join(f"[page {i}]\n{t.rstrip()}" for i, t in enumerate(pages, 1)) # A page with no text is a scan: the model reads it as an image. scans = [i for i, t in enumerate(pages, 1) if not t.strip()] blank, unseen = scans[:MAX_PAGE_IMAGES], scans[MAX_PAGE_IMAGES:] digest = sha256(p) if blank else "" refs += [{"type": "image_ref", "name": f"{name} page {n}", "path": p, "sha256": digest, "page": n} for n in blank] else: with open(p, "rb") as f: # no more than the share could use body = f.read(share * 4 + 4).decode("utf-8", "replace") if len(body) > share: body, chip["truncated"] = body[:share], True if kind == "pdf" and unseen: # the model must not think it saw them all body += f"\n[scanned pages not shown: {', '.join(map(str, unseen))}]" chip["truncated"] = True note = " (truncated)" if chip["truncated"] else "" blocks.append(f"--- file: {name}{note} ---\n{body}\n--- end file ---") return blocks, refs, chips def resolve(part, images=True): # An image reference back into an image the model can see, or a note # that it is gone (moved, deleted, changed since it was sent) or that the # model now loaded cannot see it, which would otherwise fail the request. if part.get("type") != "image_ref": return part if not images: return {"type": "text", "text": f"[image {part['name']}: this model cannot see images]"} gone = {"type": "text", "text": f"[image {part['name']} is no longer available]"} try: if "page" in part: if sha256(part["path"]) != part["sha256"]: return gone # 150 dpi, lowered for a large page so its longest side stays # within PAGE_PX: unbounded, a 3000 pt page rendered 6250x6250 at # 171 MB. Never raised: -scale-to also upscales, and a 600x300 pt # scan sent as 2048x1024 with a second image took llama-server's # model process down. n = str(part["page"]) info = subprocess.run(["pdfinfo", "-f", n, "-l", n, part["path"]], capture_output=True, text=True, timeout=30).stdout size = re.search(r"size:\s+([\d.]+) x ([\d.]+)", info) longest = max(map(float, size.groups())) if size else 842.0 # A4 if unreported dpi = min(150.0, PAGE_PX * 72 / longest) out = subprocess.run(["pdftoppm", "-f", n, "-l", n, "-r", f"{dpi:.2f}", "-png", part["path"]], capture_output=True, timeout=60) if out.returncode or not out.stdout or len(out.stdout) > MAX_IMAGE: return gone data, mime = out.stdout, "image/png" else: with open(part["path"], "rb") as f: # once: the bytes hashed are the bytes sent data = f.read() if hashlib.sha256(data).hexdigest() != part["sha256"]: return gone mime = mimetypes.guess_type(part["path"])[0] except (OSError, subprocess.TimeoutExpired): return gone return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{base64.b64encode(data).decode()}"}} def wire(messages, images=True): # What is sent: the stored messages with their image references resolved. return [{**m, "content": [resolve(p, images) for p in m["content"]]} if isinstance(m.get("content"), list) else m for m in messages] # The canvas's feed: JSON lines on a Unix socket, both ways. Every client # gets every event. A client that stalls for 2 s is dropped rather than # holding up the others. clients, clients_lock = [], threading.Lock() now_state = {"state": "idle"} # repeated in hello for a late client ctx = {"model": None, "max": 0, "used": 0} # One lock for the conversation, the turn counter and the database. Emits # that depend on the turn happen under it, so a stopped reply cannot send a # token after its done. lock = threading.RLock() cur = {"id": None, "history": [], "turns": [], "pending": None} inflight = {"conn": None} # the HTTP connection the current reply streams over writing = {"on": False} # a write in flight, which interrupt() also ends turn = 0 # bumped by every ask, stop, new and open: older replies are stale DB = None def emit(event): if event["type"] == "state": now_state.clear() now_state.update({k: v for k, v in event.items() if k != "type"}) line = (json.dumps(event) + "\n").encode() # ponytail: sends run under the caller's lock, so one stalled client holds # everyone up for its 2 s timeout, once, before it is dropped; per-client # send queues if that is ever felt. with clients_lock: for c in clients[:]: try: c.sendall(line) except OSError: # gone, or stalled past its timeout clients.remove(c) c.close() def hello(): return {"type": "hello", **now_state, "model": ctx["model"], "ctx_used": ctx["used"], "ctx_max": ctx["max"], "id": cur["id"], "turns": cur["turns"], "pending": cur["pending"]} def drop(c): with clients_lock: if c in clients: clients.remove(c) c.close() def reader(c): # A timed-out recv is only an idle client; the socket stays usable. buf = b"" while True: try: data = c.recv(65536) except TimeoutError: continue except OSError: break if not data: break buf += data if len(buf) > 1 << 20: # a megabyte with no newline is not a command log("client line over 1 MB, dropped") emit({"type": "error", "message": "line too long"}) buf = b"" continue *lines, buf = buf.split(b"\n") for line in lines: if line.strip(): handle(line) drop(c) def serve(path=SOCK): probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) try: probe.connect(path) except OSError: pass # missing, or left over from a crash else: # Unlinking a live socket would orphan the running daemon's clients. probe.close() raise SystemExit(f"another assistant is serving {path}") probe.close() try: os.unlink(path) except FileNotFoundError: pass srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) old = os.umask(0o177) # 0600: only this user talks to it try: srv.bind(path) finally: os.umask(old) srv.listen() def accept(): while True: try: c, _ = srv.accept() except OSError as e: # EMFILE and friends: log, back off, keep serving log(f"accept failed: {e}") time.sleep(1) continue c.settimeout(2) # Under the lock: no event can fall between hello and the client # joining the list. with lock: try: c.sendall((json.dumps(hello()) + "\n").encode()) except OSError: c.close() continue with clients_lock: clients.append(c) threading.Thread(target=reader, args=(c,), daemon=True).start() threading.Thread(target=accept, daemon=True).start() def interrupt(): # Caller holds lock. A reply in flight ends now: done marks it stopped, # its partial text is not saved, its thread notices the turn and returns. # Its connection is shut too: llama-server stops only when the client # hangs up, and a reply waiting for its first chunk through a long # prompt pass would otherwise keep the GPU busy until that chunk. global turn turn += 1 if writing["on"]: # a write has no pending turn to close writing["on"] = False # Whatever ended it (stop, an ask, a new conversation): the canvas # waits for an outcome, so it gets one. emit({"type": "error", "message": "stopped", "turn": True}) emit({"type": "state", "state": "idle"}) conn = inflight["conn"] if conn is not None and conn.sock is not None: try: conn.sock.shutdown(socket.SHUT_RDWR) except OSError: pass if cur["pending"] is not None: cur["pending"] = None emit({"type": "done", "id": cur["id"], "stopped": True}) emit({"type": "state", "state": "idle"}) def start_new(): interrupt() cur.update(id=None, history=[], turns=[]) def handle(raw): try: cmd = json.loads(raw) kind = cmd["type"] except (ValueError, TypeError, KeyError): log(f"bad command: {raw[:200]!r}") emit({"type": "error", "message": "bad command"}) return # Ids are integers: SQLite would match "1" to row 1 while Python's # comparison with the current id would not, and a missing id would # equal an unsaved conversation's None. bad = ("bad id" if kind in ("open", "delete") and type(cmd.get("id")) is not int else "bad text" if kind == "ask" and not isinstance(cmd.get("text"), str) else "bad files" if kind == "ask" and not ( isinstance(cmd.get("files", []), list) and all(isinstance(f, str) for f in cmd.get("files", []))) else "bad query" if kind == "search" and not isinstance(cmd.get("q"), str) else "bad write" if kind == "write" and not ( isinstance(cmd.get("files"), list) and cmd["files"] and all(isinstance(f, str) for f in cmd["files"]) and isinstance(cmd.get("prompt"), str) and isinstance(cmd.get("dir"), str) and os.path.isabs(cmd["dir"]) and isinstance(cmd.get("name"), str) and cmd["name"].strip(". ") and "/" not in cmd["name"]) else None) if bad: # A refused write ends the canvas's wait for it, so it carries turn. emit({"type": "error", "message": bad, **({"turn": True} if kind == "write" else {})}) return try: with lock: command(kind, cmd) except Exception: # A failed command must not kill the reader: the client would stay # connected, receiving events, with nothing reading what it sends. log(f"{kind} failed:\n" + traceback.format_exc().rstrip()) emit({"type": "error", "message": f"{kind} failed"}) def command(kind, cmd): # Caller holds lock; arguments are already checked. if kind == "write": interrupt() threading.Thread(target=write, daemon=True, args=(turn, cmd["files"], cmd["prompt"], cmd["dir"], cmd["name"])).start() elif kind == "ask": text, files = cmd["text"].strip(), cmd.get("files", []) if text or files: interrupt() threading.Thread(target=reply, args=(turn, text, files), daemon=True).start() elif kind == "stop": interrupt() elif kind == "new": start_new() emit({"type": "conversation", "id": None, "turns": []}) elif kind == "list": emit({"type": "conversations", "items": list_convs(DB)}) elif kind == "open": got = load_conv(DB, cmd["id"]) if got is None: emit({"type": "error", "message": "no such conversation"}) return interrupt() cur.update(id=cmd["id"], history=got[0], turns=got[1]) emit({"type": "conversation", "id": cur["id"], "turns": cur["turns"]}) elif kind == "search": emit({"type": "results", "items": search(DB, cmd["q"])}) elif kind == "delete": delete_conv(DB, cmd["id"]) if cmd["id"] == cur["id"]: start_new() emit({"type": "conversation", "id": None, "turns": []}) emit({"type": "conversations", "items": list_convs(DB)}) else: log(f"unknown command: {kind!r}") emit({"type": "error", "message": f"unknown command {kind}"}) def context_size(model): # The router can swap models, so ask again whenever the id changes. if ctx["model"] != model: url = f"{CFG['llm'].removesuffix('/v1')}/props?model={urllib.parse.quote(model)}" try: with urllib.request.urlopen(url, timeout=5) as r: ctx["max"] = json.load(r)["default_generation_settings"]["n_ctx"] except (OSError, KeyError, ValueError): ctx["max"] = 0 ctx["model"] = model return ctx["max"] def run_tool(tools, name, arguments, d=TOOLS_DIR): # Whatever goes wrong becomes the tool's answer, so the model can say it # could not find out; a tool never takes the assistant down. t = tools.get(name) if t is None: return f"error: no tool named {name}" try: args = json.loads(arguments or "{}") except json.JSONDecodeError: args = None if not isinstance(args, dict): return "error: arguments are not valid JSON" vals = {p: "" if args.get(p) is None else str(args[p]) for p in t.get("params", {})} # No shell: each {name} fills in inside one argv element. argv = [re.sub(r"\{(\w+)\}", lambda m: vals.get(m.group(1), ""), a) for a in t["command"]] if "/" not in argv[0] and os.path.exists(os.path.join(d, argv[0])): argv[0] = os.path.join(d, argv[0]) env = {**os.environ, **{k: str(v) for k, v in CFG["env"].items()}} try: # Own process group: a timeout kills the script's children too (curl), # which would otherwise keep the pipes open and block us. p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env, start_new_session=True) except OSError as e: return f"error: cannot run {argv[0]}: {e.strerror}" try: out, err = p.communicate(timeout=t.get("timeout", 10)) except subprocess.TimeoutExpired: os.killpg(p.pid, signal.SIGKILL) p.communicate() return "error: timed out" if p.returncode: last = err.strip().splitlines()[-1:] return ("error: " + (last[0] if last else f"exit status {p.returncode}"))[:2000] return out.strip()[:2000] def collect_calls(calls, fragments): # A streamed tool call arrives in pieces keyed by index: id and name in # the first, the JSON arguments spread over the rest. for f in fragments: c = calls.setdefault(f["index"], {"id": "", "name": "", "arguments": ""}) fn = f.get("function") or {} c["id"] = f.get("id") or c["id"] c["name"] = fn.get("name") or c["name"] c["arguments"] += fn.get("arguments") or "" def model_inputs(): # (id, inputs) of the model in use. The router keeps one model resident; # asking for another would swap it out under other clients, so use # whatever is not unloaded. inputs is None when the server does not # report them. with urllib.request.urlopen(CFG["llm"] + "/models", timeout=10) as r: models = json.load(r)["data"] m = next((m for m in models if m["status"]["value"] != "unloaded"), models[0]) return m["id"], (m.get("architecture") or {}).get("input_modalities") def llm_model(): return model_inputs()[0] def ask(messages, calls, tools=None, mine=None, think=None, meter=True): # Yields ("think", text) and ("token", text) as they stream; tool-call # fragments go into calls. mine is the reply's turn: a stop shuts the # connection from the other thread (interrupt), which ends the reading. model = llm_model() context_size(model) body = {"model": model, "messages": messages, "stream": True, "stream_options": {"include_usage": True}, "chat_template_kwargs": {"enable_thinking": CFG["thinking"] if think is None else think}} if tools: body["tools"] = tools if CFG["temperature"] is not None: body["temperature"] = CFG["temperature"] # http.client rather than urllib: it exposes the socket a stop must shut. u = urllib.parse.urlsplit(CFG["llm"] + "/chat/completions") conn = (http.client.HTTPSConnection if u.scheme == "https" else http.client.HTTPConnection)( u.hostname, u.port, timeout=300) with lock: if mine is not None and turn != mine: return # stopped before the request went out inflight["conn"] = conn try: yield from stream(conn, u.path, body, calls, meter) except OSError: if mine is not None and turn != mine: return # the stop hung up under the read raise finally: with lock: if inflight["conn"] is conn: inflight["conn"] = None conn.close() def stream(conn, path, body, calls, meter=True): conn.request("POST", path, json.dumps(body).encode(), {"Content-Type": "application/json"}) r = conn.getresponse() if r.status >= 400: # The body says why (a full context, a bad request); the status alone does not. raise RuntimeError(f"llama-server {r.status}: {r.read().decode(errors='replace')[:300]}") with r: streamed, other = False, b"" for line in r: if not line.startswith(b"data: "): # The router answers a crashed model with a plain line, e.g. # "proxy error: Failed to read connection", not a stream. other = other or line.strip() continue streamed = True if line.strip() == b"data: [DONE]": continue chunk = json.loads(line[6:]) if "error" in chunk: # a failure mid-stream arrives as a chunk err = chunk["error"] raise RuntimeError(err.get("message", str(err)) if isinstance(err, dict) else str(err)) if chunk.get("usage") and meter: # How full the model's context is: history, tool results and # tool definitions all count. ctx["used"] = chunk["usage"]["prompt_tokens"] emit({"type": "ctx", "used": ctx["used"], "max": ctx["max"]}) choices = chunk["choices"] if not choices: continue delta = choices[0]["delta"] collect_calls(calls, delta.get("tool_calls") or []) if delta.get("reasoning_content"): yield "think", delta["reasoning_content"] if delta.get("content"): yield "token", delta["content"] if not streamed: raise RuntimeError(other.decode(errors="replace")[:300] or "llama-server sent nothing") def say(mine, kind, value, done=False): # Sends one piece of the reply, unless the reply went stale; False tells # the reply thread to give up. with lock: if turn != mine: return False p = cur["pending"] if kind == "think": p["think"] += value emit({"type": "think", "delta": value}) elif kind == "token": p["reply"] += value emit({"type": "token", "delta": value}) else: if not done: p["tools"].append(value) emit({"type": "tool", "name": value, "done": done}) emit({"type": "state", "state": "thinking"} if done else {"type": "state", "state": "tool", "tool": value}) return True def free_name(d, name): # The first of name, "stem (1).ext", "stem (2).ext" not taken in d. stem, ext = os.path.splitext(name) path, n = os.path.join(d, name), 1 while os.path.lexists(path): path, n = os.path.join(d, f"{stem} ({n}){ext}"), n + 1 return path def unfence(text): # Models often wrap a whole document in one code fence; the file wants # what is inside it. lines = text.strip().split("\n") if (len(lines) >= 2 and lines[0].startswith("```") and lines[-1].strip() == "```" and not any(l.startswith("```") for l in lines[1:-1])): # one fence, not two blocks return "\n".join(lines[1:-1]) + "\n" return text def write(mine, files, prompt, d, name): # A one-off request whose reply goes to a file, for the canvas's To .md: # no history in, nothing kept in the conversation, no thinking, no tools. tmp = None try: # No history goes with a write, so half the whole context is its share. blocks, refs, _ = attach(files, ctx["max"] // 2 * 3 if ctx["max"] else 24000) model, inputs = model_inputs() images = inputs is None or "image" in inputs if refs and not images: raise Refused(f"{model} cannot see images") with lock: if turn != mine: return writing["on"] = True emit({"type": "state", "state": "thinking"}) body = "\n\n".join(blocks + [prompt]) msgs = [{"role": "user", "content": [{"type": "text", "text": body}] + refs if refs else body}] text, shown = "", 0.0 # meter=False: the chat's context meter is not about this request. for kind, piece in ask(wire(msgs, images), {}, None, mine, think=False, meter=False): if turn != mine: return if kind == "token": text += piece if time.time() - shown >= 1: shown = time.time() emit({"type": "progress", "chars": len(text)}) if turn != mine: # stopped: the hang-up ended the stream return if not text.strip(): raise Refused("the model gave an empty reply; nothing written") os.makedirs(d, exist_ok=True) tmp = os.path.join(d, f".write-{os.getpid()}-{mine}.tmp") with open(tmp, "w", encoding="utf-8") as f: f.write(unfence(text)) with lock: if turn != mine: return path = free_name(d, name) os.rename(tmp, path) tmp = None writing["on"] = False emit({"type": "written", "path": path}) emit({"type": "state", "state": "idle"}) log(f"wrote {path}") except Exception as e: log(f"refused: {e}" if isinstance(e, Refused) else "write failed:\n" + traceback.format_exc().rstrip()) with lock: if turn == mine: writing["on"] = False emit({"type": "error", "message": str(e) or type(e).__name__, "turn": True}) emit({"type": "state", "state": "idle"}) finally: if tmp is not None and os.path.exists(tmp): os.remove(tmp) def reply(mine, text, files=()): try: answer(mine, text, files) except Exception as e: # A refused file is the user's to fix, not a bug: one line, no trace. log(f"refused: {e}" if isinstance(e, Refused) else "reply failed:\n" + traceback.format_exc().rstrip()) with lock: if turn == mine: cur["pending"] = None # turn: this error ended the reply. Command errors carry no # turn and must not end a reply another client is watching. emit({"type": "error", "message": str(e) or type(e).__name__, "turn": True}) emit({"type": "state", "state": "idle"}) def answer(mine, text, files=()): blocks, refs, chips = [], [], [] if files: # Half of what the context has free, at about 3 characters a token. free = (ctx["max"] - ctx["used"]) // 2 * 3 if ctx["max"] else 24000 blocks, refs, chips = attach(files, max(free, 2000)) # A server that does not report inputs is left to decide for itself. model, inputs = model_inputs() images = inputs is None or "image" in inputs if refs and not images: raise Refused(f"{model} cannot see images") with lock: if turn != mine: return history = list(cur["history"]) cur["pending"] = {"user": text, "reply": "", "think": "", "tools": [], "files": chips} emit({"type": "user", "text": text, "files": chips}) emit({"type": "state", "state": "thinking"}) log(f"you: {text}" + "".join(f" [{c['name']}]" for c in chips)) body = "\n\n".join(blocks + [text] if text else blocks) question = {"role": "user", "content": [{"type": "text", "text": body}] + refs if refs else body} # The model has no clock: without the date it makes one up for tools. # Also stamped on the question itself, which history does not keep: with # the time only in the system prompt, "what time is it?" asked again an # hour later copied the old answer from history. Stamping every message # in history fixed that but lost tool arguments (account filter 0/6). now = time.strftime("%A %d %B %Y, %H:%M") # ponytail: last 10 turns, summarise if long sessions need memory msgs = ([{"role": "system", "content": f"{CFG['system']} It is {now}."}] + [m for past in history[-10:] for m in past] + [{"role": "user", "content": [{"type": "text", "text": f"[{now}] {body}"}] + refs if refs else f"[{now}] {body}"}]) first = len(msgs) # this turn's tool messages start here # Resolved once per answer: images re-read and pages re-rendered on every # tool round would cost seconds between rounds, and a file changed # mid-answer would flip from image to gone. sent = wire(msgs, images) said, used = "", [] # A stop shuts this reply's connection (interrupt), so llama-server's slot # frees even mid prompt pass. ponytail: a running tool still runs out. for rnd in range(CFG["tool_rounds"] + 1): calls, content = {}, "" # The last round offers no tools, which forces an answer. for kind, piece in ask(sent, calls, TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None, mine): if not say(mine, kind, piece): log("interrupted") return if kind == "token": content += piece said += content if not calls: break call = {"role": "assistant", "content": content, "tool_calls": [ {"id": c["id"], "type": "function", "function": {"name": c["name"], "arguments": c["arguments"]}} for c in calls.values()]} msgs.append(call) sent.append(call) # ponytail: a stop does not kill a running tool, its result is dropped for c in calls.values(): if not say(mine, "tool", c["name"]): log("interrupted") return used.append(c["name"]) result = run_tool(TOOLS, c["name"], c["arguments"]) log(f"tool {c['name']}({c['arguments']}): {result[:120]!r}") msgs.append({"role": "tool", "tool_call_id": c["id"], "content": result}) sent.append(msgs[-1]) if not say(mine, "tool", c["name"], done=True): log("interrupted") return # History keeps the tool calls: with only what was said, the model saw # facts answered without tools and invented the next ones (no tool call # 9/10 on a replayed session, 0/10 with the calls kept). Results are # replaced by a placeholder, for context: same 10/10, and a follow-up on # an old result calls the tool again (10/10). A stopped exchange is # dropped whole, so user and assistant still alternate. past = [{**m, "content": "(shown to the user at the time)"} if m["role"] == "tool" else m for m in msgs[first:]] entry = [question] + past + [{"role": "assistant", "content": said}] with lock: if turn != mine: log("interrupted") return cur["id"] = save_turn(DB, cur["id"], text, said, used, entry, chips) cur["history"].append(entry) cur["turns"].append({"user": text, "reply": said, "tools": used, "files": chips}) cur["pending"] = None emit({"type": "done", "id": cur["id"]}) emit({"type": "state", "state": "idle"}) log("assistant:", said) def main(): global DB # SIGTERM runs the finally below, so the socket is not left behind. signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) served = False try: DB = db_open(DB_PATH) try: context_size(llm_model()) # so the first hello has it except Exception: pass # llama-server down: asked again per request serve() served = True log(f"ready on {SOCK}") while True: signal.pause() except (KeyboardInterrupt, SystemExit) as e: log(f"exit: {type(e).__name__} {e}".rstrip()) if not served: raise SystemExit(1) # refused or failed before serving: not a clean stop except BaseException: log("crashed:\n" + traceback.format_exc().rstrip()) raise finally: if served and os.path.exists(SOCK): # never another daemon's socket os.remove(SOCK) def test(): import http.server import tempfile c = merge(DEFAULTS, {"tool_rounds": 1, "env": {"K": "v"}}) assert c["tool_rounds"] == 1 and c["env"] == {"K": "v"} assert DEFAULTS["tool_rounds"] == 3 # defaults untouched for bad in ({"card": "X"}, {"whisper": {"port": 1}}): # voice-era keys are gone try: merge(DEFAULTS, bad) assert False, bad except SystemExit as e: assert "config:" in str(e) with tempfile.TemporaryDirectory() as d: with open(os.path.join(d, "echo.sh"), "w") as f: f.write('#!/bin/sh\nprintf "%s|" "$@"\n') os.chmod(os.path.join(d, "echo.sh"), 0o755) with open(os.path.join(d, "now.toml"), "w") as f: f.write('description = "d"\ncommand = ["date"]\n' '[params.x]\ntype = "string"\ndescription = "px"\n') assert load_tools(d) == {"now": {"description": "d", "command": ["date"], "params": {"x": {"type": "string", "description": "px"}}}} assert schema(load_tools(d)) == [{"type": "function", "function": { "name": "now", "description": "d", "parameters": {"type": "object", "properties": { "x": {"type": "string", "description": "px"}}}}}] tools = {"echo": {"command": ["echo.sh", "{a}", "x{b}y"], "params": {"a": {}, "b": {}}}, "slow": {"command": ["sh", "-c", "sleep 5; true"], "timeout": 0.3}, "fail": {"command": ["sh", "-c", "echo oops >&2; exit 3"]}, "gone": {"command": ["no-such-binary-here"]}} # A value is one argv element, never a shell command or a second argument. assert run_tool(tools, "echo", '{"a": "x; rm -rf ~", "b": 2}', d) == "x; rm -rf ~|x2y|" assert run_tool(tools, "echo", "{}", d) == "|xy|" t0 = time.time() assert run_tool(tools, "slow", "{}", d) == "error: timed out" assert time.time() - t0 < 2 # the grandchild sleep was killed too assert run_tool(tools, "fail", "{}", d) == "error: oops" assert run_tool(tools, "gone", "{}", d).startswith("error: cannot run") assert run_tool(tools, "nope", "{}", d) == "error: no tool named nope" assert run_tool(tools, "echo", "{bad", d) == "error: arguments are not valid JSON" assert run_tool(tools, "echo", "[1]", d) == "error: arguments are not valid JSON" calls = {} collect_calls(calls, [{"index": 0, "id": "c1", "type": "function", "function": {"name": "weather", "arguments": "{"}}]) for piece in ['"place', '":', '"', "Rome", '"', "}"]: collect_calls(calls, [{"index": 0, "function": {"arguments": piece}}]) collect_calls(calls, [{"index": 1, "id": "c2", "function": {"name": "now", "arguments": "{}"}}]) assert calls == {0: {"id": "c1", "name": "weather", "arguments": '{"place":"Rome"}'}, 1: {"id": "c2", "name": "now", "arguments": "{}"}} # Storage. db = db_open(":memory:") a = save_turn(db, None, "che tempo fa in città oggi pomeriggio?", "Piove.", ["weather"], [{"role": "user", "content": "q"}]) assert save_turn(db, a, "e domani?", "Sole.", [], [{"role": "user", "content": "q2"}]) == a b = save_turn(db, None, "how do I list files", "Use `ls`.", [], []) assert [x["id"] for x in list_convs(db)] == [b, a] assert list_convs(db)[1]["title"] == "che tempo fa in città oggi" # first six words hist, shown = load_conv(db, a) assert hist == [[{"role": "user", "content": "q"}], [{"role": "user", "content": "q2"}]] assert shown == [{"user": "che tempo fa in città oggi pomeriggio?", "reply": "Piove.", "tools": ["weather"], "files": []}, {"user": "e domani?", "reply": "Sole.", "tools": [], "files": []}] assert load_conv(db, 999) is None hits = search(db, "citta") # no accent typed, accented word found assert [h["id"] for h in hits] == [a] and "**città**" in hits[0]["snippet"] assert [h["id"] for h in search(db, "cit")] == [a] # prefix, as you type assert [h["id"] for h in search(db, "piove")] == [a] # replies are indexed for q in ('"', "*", 'a"b', "NOT", "x OR y", "ls)", "", "a\0b"): # FTS syntax is text search(db, q) assert search(db, "piove OR ls") == [] # unquoted, this would match both delete_conv(db, a) assert load_conv(db, a) is None and search(db, "citta") == [] # Raises if the index still holds the deleted turns: the cascade must # have fired the delete trigger. db.execute("INSERT INTO turns_fts(turns_fts, rank) VALUES ('integrity-check', 1)") with tempfile.TemporaryDirectory() as d: path = os.path.join(d, "sub", "h.db") os.makedirs(os.path.dirname(path), mode=0o755) open(path, "w").close() os.chmod(path, 0o644) # made earlier by something else db_open(path).close() assert os.stat(os.path.dirname(path)).st_mode & 0o777 == 0o700 assert os.stat(path).st_mode & 0o777 == 0o600 # Attachments. import struct import zlib def tiny_pdf(path, pages): # One Helvetica line per page; "" is a page with no text, a scan. A # page may be (text, width, height) in points. objs = ["<< /Type /Catalog /Pages 2 0 R >>", "<< /Type /Pages /Kids [%s] /Count %d >>" % (" ".join(f"{4 + 2 * i} 0 R" for i in range(len(pages))), len(pages)), "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"] for i, t in enumerate(pages): t, w, h = t if isinstance(t, tuple) else (t, 612, 792) stream = f"BT /F1 24 Tf 72 700 Td ({t}) Tj ET" if t else "" objs += [f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {w} {h}] /Resources " f"<< /Font << /F1 3 0 R >> >> /Contents {5 + 2 * i} 0 R >>", f"<< /Length {len(stream)} >>\nstream\n{stream}\nendstream"] out, offsets = b"%PDF-1.4\n", [] for k, o in enumerate(objs, 1): offsets.append(len(out)) out += f"{k} 0 obj\n{o}\nendobj\n".encode() xref = len(out) out += (f"xref\n0 {len(objs) + 1}\n0000000000 65535 f \n".encode() + b"".join(f"{o:010d} 00000 n \n".encode() for o in offsets) + f"trailer\n<< /Size {len(objs) + 1} /Root 1 0 R >>\n" f"startxref\n{xref}\n%%EOF\n".encode()) with open(path, "wb") as f: f.write(out) def tiny_png(path, rgb=b"\xff\x00\x00"): def chunk(t, d): return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d)) with open(path, "wb") as f: f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)) + chunk(b"IDAT", zlib.compress(b"\x00" + rgb)) + chunk(b"IEND", b"")) files_dir = tempfile.TemporaryDirectory() fd = files_dir.name code, png, pdf, other = (os.path.join(fd, n) for n in ("a.py", "red.png", "doc.pdf", "x.bin")) with open(code, "w") as f: f.write("print('hi')\n" * 50) tiny_png(png) tiny_pdf(pdf, ["hello pdf", ("", 3000, 3000)]) with open(other, "wb") as f: f.write(b"\x00\x01") assert [kind_of(p) for p in (code, png, pdf, other, "/x/README.md", "/x/b.SlackBuild")] == \ ["text", "image", "pdf", "other", "text", "text"] # No suffix: the content decides. A script or plain UTF-8 is text; a NUL # byte, or bytes that are not UTF-8, are not. for name, data, want in (("run", b"#!/bin/bash\necho hi\n", "text"), ("notes", "caff\u00e8 alle 9\n".encode(), "text"), ("blob", b"\x7fELF\x00\x01", "other"), ("latin", "caff\u00e8".encode("latin-1"), "other")): with open(os.path.join(fd, name), "wb") as f: f.write(data) assert kind_of(os.path.join(fd, name)) == want, name assert kind_of(os.path.join(fd, "missing")) == "other" with open(os.path.join(fd, "long"), "wb") as f: # a character cut by the window f.write(b"a" * 4095 + "\u00e8 tail".encode()) assert kind_of(os.path.join(fd, "long")) == "text" blocks, refs, chips = attach([code, pdf, png], 100000) assert blocks[0].startswith("--- file: a.py ---\nprint('hi')") and "(truncated)" not in blocks[0] assert "[page 1]\nhello pdf" in blocks[1] and "[page 2]" in blocks[1] assert chips == [{"name": "a.py", "kind": "text", "truncated": False}, {"name": "doc.pdf", "kind": "pdf", "truncated": False, "pages": 2}, {"name": "red.png", "kind": "image", "truncated": False}] # The blank page is a scan: it goes as an image, rendered when sent. assert [r["name"] for r in refs] == ["doc.pdf page 2", "red.png"] and refs[0]["page"] == 2 sent = [resolve(r) for r in refs] assert all(x["type"] == "image_url" for x in sent) assert sent[1]["image_url"]["url"].startswith("data:image/png;base64,") page = base64.b64decode(sent[0]["image_url"]["url"].split(",", 1)[1]) # Page 2, the square one, bounded to PAGE_PX: page 1 would not be square. w, h = struct.unpack(">II", page[16:24]) assert page[:4] == b"\x89PNG" and w == h and PAGE_PX - 2 <= w <= PAGE_PX, (w, h) scans = os.path.join(fd, "scans.pdf") tiny_pdf(scans, ["cover"] + [("", 300, 150)] * 6) b, r, c = attach([scans], 100000) assert len(r) == 4 and c[0]["truncated"] and "[scanned pages not shown: 6, 7]" in b[0] small = base64.b64decode(resolve(r[0])["image_url"]["url"].split(",", 1)[1]) # A small page stays at 150 dpi, not raised to the bound: 300 pt is 625 px. assert struct.unpack(">II", small[16:24]) == (625, 313), struct.unpack(">II", small[16:24]) with open(os.path.join(fd, "exact.txt"), "w") as f: f.write("x" * 24) with open(os.path.join(fd, "over.txt"), "w") as f: f.write("x" * 25) for n, cut in (("exact.txt", False), ("over.txt", True)): b, _, c = attach([os.path.join(fd, n)], 24) assert c[0]["truncated"] is cut and b[0].split("\n")[1] == "x" * 24, n with open(os.path.join(fd, "a\nb.txt"), "w") as f: f.write("x") b, _, c = attach([os.path.join(fd, "a\nb.txt")], 100) assert c[0]["name"] == "a b.txt" and b[0].startswith("--- file: a b.txt ---\n") small = attach([code], 24)[0][0] # the share is the cut: two lines assert "(truncated)" in small and small.count("print") == 2 for bad, why in ((["a.py"], "cannot read a.py"), ([fd + "/nope.txt"], "cannot read nope.txt"), ([fd], "cannot read"), ([other], "x.bin: chat supports text, images and PDF"), ([code] * 9, "at most 8 files")): try: attach(bad, 1000) assert False, bad except Refused as e: assert why in str(e), (bad, e) # Stored as a reference: survives a save and reload, gone once changed. a = save_turn(db, None, "", "A red pixel.", [], [{"role": "user", "content": refs[1:]}], chips[2:]) hist, shown = load_conv(db, a) assert shown[0]["files"] == chips[2:] and list_convs(db)[0]["title"] == "red.png" assert wire(hist[0])[0]["content"][0]["type"] == "image_url" assert wire(hist[0], images=False)[0]["content"][0] == { # a text-only model now "type": "text", "text": "[image red.png: this model cannot see images]"} tiny_png(png, b"\x00\xff\x00") assert wire(hist[0])[0]["content"][0] == { "type": "text", "text": "[image red.png is no longer available]"} os.remove(png) assert resolve(refs[1])["type"] == "text" # deleted reads the same # An old database gains the files column. with tempfile.TemporaryDirectory() as d: old = sqlite3.connect(os.path.join(d, "h.db")) old.executescript(SCHEMA) old.close() assert "files" in {r[1] for r in db_open(os.path.join(d, "h.db")) .execute("PRAGMA table_info(turns)")} # Protocol, against a fake llama-server. global DB, TOOLS, TOOL_SCHEMA script = {"replies": [], "delay": 0, "inputs": ["text", "image"]} class Fake(http.server.BaseHTTPRequestHandler): def log_message(self, *a): pass def do_GET(self): body = ({"data": [{"id": "m", "status": {"value": "loaded"}, "architecture": {"input_modalities": script["inputs"]}}]} if self.path.endswith("/models") else {"default_generation_settings": {"n_ctx": 4096}}) self.send_response(200) self.end_headers() self.wfile.write(json.dumps(body).encode()) def do_POST(self): script["body"] = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) if script.get("hold"): # A long prompt pass: nothing is sent, so only the client # hanging up ends it. Records how soon that happened. import select t0 = time.time() while time.time() - t0 < 5: if select.select([self.connection], [], [], 0.05)[0]: if not self.connection.recv(1, socket.MSG_PEEK): script["hung_up"] = time.time() return self.send_response(200) self.end_headers() try: reply = script["replies"].pop(0) for delta in reply: self.wfile.write(delta.encode() if isinstance(delta, str) else b"data: " + json.dumps({"choices": [{"delta": delta}]}) .encode() + b"\n\n") self.wfile.flush() time.sleep(script["delay"]) # A plain line stands alone, as the router sends it. if not any(isinstance(d, str) and not d.startswith("data:") for d in reply): self.wfile.write(b'data: {"choices": [], "usage": {"prompt_tokens": 42}}\n\n' b"data: [DONE]\n\n") except BrokenPipeError: # a stopped reply hangs up mid-stream pass srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Fake) threading.Thread(target=srv.serve_forever, daemon=True).start() CFG["llm"] = f"http://127.0.0.1:{srv.server_address[1]}/v1" DB = db_open(":memory:") TOOLS = {"now": {"description": "d", "command": ["sh", "-c", "echo noon"]}} TOOL_SCHEMA = schema(TOOLS) with tempfile.TemporaryDirectory() as d: path = os.path.join(d, "s.sock") serve(path) assert os.stat(path).st_mode & 0o777 == 0o600 cl = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) cl.connect(path) cl.settimeout(5) f = cl.makefile("rb") def send(**cmd): cl.sendall((json.dumps(cmd) + "\n").encode()) def until(kind): seen = [] while True: e = json.loads(f.readline()) seen.append(e) if e["type"] == kind: return seen h = json.loads(f.readline()) assert h["type"] == "hello" and h["state"] == "idle" and h["turns"] == [] assert h["pending"] is None and "ctx_max" in h # Thinking, a tool round, then the answer. script["replies"] = [ [{"reasoning_content": "hmm"}, {"tool_calls": [{"index": 0, "id": "c1", "function": {"name": "now", "arguments": "{}"}}]}], [{"content": "It is "}, {"content": "noon."}]] send(type="ask", text="what time is it") got = [(e["type"], e.get("delta") or e.get("name") or e.get("state") or e.get("text")) for e in until("done") if e["type"] != "ctx"] assert got == [("user", "what time is it"), ("state", "thinking"), ("think", "hmm"), ("tool", "now"), ("state", "tool"), ("tool", "now"), ("state", "thinking"), ("token", "It is "), ("token", "noon."), ("done", None)], got assert until("state")[-1]["state"] == "idle" assert script["body"]["chat_template_kwargs"]["enable_thinking"] is True conv = cur["id"] assert load_conv(DB, conv)[1] == [{"user": "what time is it", "reply": "It is noon.", "tools": ["now"], "files": []}] # The tool result went to the model but history keeps a placeholder. assert load_conv(DB, conv)[0][0][2] == {"role": "tool", "tool_call_id": "c1", "content": "(shown to the user at the time)"} # A plain line instead of a stream (the router, when the model died) # is an error carrying that line, not an empty reply. script["replies"] = [["proxy error: Failed to read connection\n"]] send(type="ask", text="anyone?") e = until("error")[-1] assert e["message"] == "proxy error: Failed to read connection" and e["turn"] is True # An error chunk mid-stream reaches the client with its message. script["replies"] = [['data: {"error": {"message": "context full"}}\n\n']] CFG["thinking"] = False send(type="ask", text="too much") e = until("error")[-1] assert e["message"] == "context full" and e["turn"] is True assert script["body"]["chat_template_kwargs"]["enable_thinking"] is False CFG["thinking"] = True # Stop while the server is still reading the prompt: the daemon hangs # up at once rather than at the first chunk, which frees the GPU. script.update(hold=True, hung_up=None) script["replies"] = [[]] send(type="ask", text="a long prompt") until("state") # thinking: the request is out time.sleep(0.3) stopped = time.time() send(type="stop") assert until("done")[-1].get("stopped") is True for _ in range(40): if script["hung_up"]: break time.sleep(0.05) assert script["hung_up"] and script["hung_up"] - stopped < 1, script["hung_up"] script["hold"] = False # Stop mid-stream: no token after done, nothing saved. script["replies"] = [[{"content": f"w{i} "} for i in range(50)]] script["delay"] = 0.02 send(type="ask", text="long one") until("token") c3 = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) c3.connect(path) # mid-reply: hello carries the partial h3 = json.loads(c3.makefile("rb").readline()) assert h3["pending"]["user"] == "long one" and h3["pending"]["reply"].startswith("w0") c3.close() send(type="stop") seen = until("done") assert seen[-1].get("stopped") is True time.sleep(0.3) # the reply thread notices and quits send(type="list") late = until("conversations") assert not [e for e in late if e["type"] == "token"], late assert len(load_conv(DB, conv)[1]) == 1 script["delay"] = 0 # History browsing. send(type="new") assert until("conversation")[-1] == {"type": "conversation", "id": None, "turns": []} send(type="list") assert [x["id"] for x in until("conversations")[-1]["items"]] == [conv] send(type="search", q="tim") assert until("results")[-1]["items"][0]["id"] == conv send(type="open", id=conv) assert until("conversation")[-1]["turns"][0]["reply"] == "It is noon." send(type="open", id=999) assert until("error")[-1]["message"] == "no such conversation" cl.sendall(b"not json\n") assert until("error")[-1]["message"] == "bad command" for cmd, why in (({"type": "open", "id": [1]}, "bad id"), ({"type": "open", "id": str(conv)}, "bad id"), ({"type": "delete"}, "bad id"), ({"type": "ask", "text": None}, "bad text"), ({"type": "search", "q": 3}, "bad query")): send(**cmd) e = until("error")[-1] assert e["message"] == why and "turn" not in e, cmd # a command error assert cur["id"] == conv # none of them touched the conversation real, DB = DB, None # every query now raises send(type="list") assert until("error")[-1]["message"] == "list failed" DB = real send(type="list") # the reader survived it assert [x["id"] for x in until("conversations")[-1]["items"]] == [conv] try: serve(path) # a second daemon on a live socket assert False except SystemExit as e: assert "another assistant" in str(e) send(type="list") # and the first one still serves until("conversations") send(type="delete", id=conv) seen = until("conversations") assert {"type": "conversation", "id": None, "turns": []} in seen # was current assert seen[-1]["items"] == [] # Files: the text inline, the image resolved, chips on the turn. tiny_png(png) script["replies"] = [[{"content": "Red."}]] send(type="ask", text="what is it?", files=[pdf, png]) user = until("user")[-1] assert user["files"][1] == {"name": "red.png", "kind": "image", "truncated": False} until("done") q = script["body"]["messages"][-1]["content"] assert "hello pdf" in q[0]["text"] and q[0]["text"].endswith("what is it?") assert [x["type"] for x in q] == ["text", "image_url", "image_url"] assert cur["turns"][-1]["files"] == user["files"] # Only files, no text. script["replies"] = [[{"content": "Code."}]] send(type="ask", text="", files=[code]) assert until("user")[-1]["text"] == "" until("done") # A model without vision is refused before anything is sent. script["inputs"] = ["text"] send(type="ask", text="see?", files=[png]) e = until("error")[-1] assert e["message"] == "m cannot see images" and e["turn"] is True script["inputs"] = ["text", "image"] send(type="ask", text="x", files=[other]) assert "chat supports" in until("error")[-1]["message"] send(type="ask", text="x", files="not a list") assert until("error")[-1]["message"] == "bad files" # write: the reply goes to a file, not the conversation. before = len(list_convs(DB)) with open(os.path.join(fd, "doc.md"), "w") as fh: # f is the socket's fh.write("taken") script["replies"] = [['data: {"choices": [{"delta": {"content": "```markdown\\n# Title\\n\\nbody\\n"}}]}\n\n', {"content": "```"}]] send(type="write", files=[pdf], prompt="As Markdown.", dir=fd, name="doc.md") w = until("written")[-1] assert w["path"] == os.path.join(fd, "doc (1).md"), w with open(w["path"]) as fh: assert fh.read() == "# Title\n\nbody\n" # the one fence unwrapped assert script["body"]["chat_template_kwargs"]["enable_thinking"] is False assert "tools" not in script["body"] and script["body"]["messages"][-1]["role"] == "user" assert len(list_convs(DB)) == before # nothing in History assert until("state")[-1]["state"] == "idle" # A stopped write leaves nothing, and the state goes back to idle. script.update(hold=True, hung_up=None) script["replies"] = [[]] send(type="write", files=[pdf], prompt="As Markdown.", dir=fd, name="held.md") until("state") send(type="stop") seen = until("state") assert seen[-1]["state"] == "idle" assert {"type": "error", "message": "stopped", "turn": True} in seen # an outcome time.sleep(0.3) assert not os.path.exists(os.path.join(fd, "held.md")) assert not [n for n in os.listdir(fd) if n.startswith(".write-")] script["hold"] = False for bad in ({"dir": "relative", "name": "x.md"}, {"dir": fd, "name": "a/b.md"}, {"dir": fd, "name": ""}): send(type="write", files=[pdf], prompt="p", **bad) e = until("error")[-1] assert e["message"] == "bad write" and e["turn"] is True, bad assert unfence("```\na\n```\nmid\n```\nb\n```") == "```\na\n```\nmid\n```\nb\n```" script["replies"] = [[{"content": " "}]] send(type="write", files=[pdf], prompt="p", dir=fd, name="empty.md") assert "empty reply" in until("error")[-1]["message"] assert not os.path.exists(os.path.join(fd, "empty.md")) # llama-server down: an error, not a dead thread. CFG["llm"] = "http://127.0.0.1:9/v1" send(type="ask", text="anyone?") assert until("error")[-1]["turn"] is True assert until("state")[-1]["state"] == "idle" and cur["pending"] is None # A second client gets hello with the conversation so far. c2 = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) c2.connect(path) assert json.loads(c2.makefile("rb").readline())["type"] == "hello" c2.close() cl.close() srv.shutdown() print("ok") if __name__ == "__main__": test() if sys.argv[1:] == ["--test"] else main()