diff options
| -rw-r--r-- | assistant.py | 208 |
1 files changed, 163 insertions, 45 deletions
diff --git a/assistant.py b/assistant.py index 630f5f8..c91917f 100644 --- a/assistant.py +++ b/assistant.py @@ -138,7 +138,6 @@ def schema(tools): 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") @@ -448,7 +447,7 @@ ctx = {"model": None, "max": 0, "used": 0} # 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} +cur = {"id": None, "history": [], "turns": [], "pending": None, "model": 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 @@ -475,7 +474,7 @@ def emit(event): 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"]} + "pending": cur["pending"], "chat_model": cur["model"]} def drop(c): @@ -584,7 +583,7 @@ def interrupt(): def start_new(): interrupt() - cur.update(id=None, history=[], turns=[]) + cur.update(id=None, history=[], turns=[], model=None) def handle(raw): @@ -605,6 +604,8 @@ def handle(raw): and all(isinstance(f, str) for f in cmd.get("files", []))) else "bad audio" if kind == "ask" and "audio" in cmd and not ( isinstance(cmd["audio"], str) and not cmd["text"].strip()) + else "bad model" if kind == "ask" and "model" in cmd and not ( + isinstance(cmd["model"], str) and cmd["model"]) 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"] @@ -638,12 +639,15 @@ def command(kind, cmd): text, files, audio = cmd["text"].strip(), cmd.get("files", []), cmd.get("audio") if text or files or audio: interrupt() - threading.Thread(target=reply, args=(turn, text, files, audio), daemon=True).start() + # A pick applies from this ask on and stays with the chat. + cur["model"] = cmd.get("model", cur["model"]) + threading.Thread(target=reply, args=(turn, text, files, audio, cur["model"]), + daemon=True).start() elif kind == "stop": interrupt() elif kind == "new": start_new() - emit({"type": "conversation", "id": None, "turns": []}) + emit({"type": "conversation", "id": None, "turns": [], "model": None}) elif kind == "list": emit({"type": "conversations", "items": list_convs(DB)}) elif kind == "open": @@ -652,16 +656,21 @@ def command(kind, cmd): 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"]}) + model = DB.execute("SELECT model FROM conversations WHERE id = ?", + (cmd["id"],)).fetchone()[0] + cur.update(id=cmd["id"], history=got[0], turns=got[1], model=model) + emit({"type": "conversation", "id": cur["id"], "turns": cur["turns"], "model": model}) 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": "conversation", "id": None, "turns": [], "model": None}) emit({"type": "conversations", "items": list_convs(DB)}) + elif kind == "models": + # Off the lock: the router can take seconds to answer. + threading.Thread(target=send_models, daemon=True).start() else: log(f"unknown command: {kind!r}") emit({"type": "error", "message": f"unknown command {kind}"}) @@ -793,14 +802,32 @@ def cloud_models(name): raise SystemExit(f"cloud {name}: {e}") from None -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. +def send_models(): + # The canvas's model dropdown: the router's models, then every cloud + # model whose provider's key loaded. + try: + with urllib.request.urlopen(CFG["llm"] + "/models", timeout=10) as r: + local = [{"id": m["id"], "status": m["status"]["value"], "cloud": False} + for m in json.load(r)["data"]] + except (OSError, KeyError, TypeError, ValueError): + local = [] # router down: the cloud ones still show + emit({"type": "models", "items": local + [ + {"id": f"{n}:{m}", "provider": n, "cloud": True} + for n, p in PROVIDERS.items() for m in p["models"]]}) + + +def model_inputs(model=None): + # (id, inputs) of the model a chat uses: its pick, else whatever the + # router has resident (asking for another would swap it out under other + # clients). inputs is None when not reported, as for every cloud model. + if model and route(model)[3]: + return model, None 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]) + m = (next((m for m in models if m["id"] == model), None) + or next((m for m in models if m["status"]["value"] != "unloaded"), models[0])) + if model and m["id"] != model: + log(f"no local model {model}; using {m['id']}") return m["id"], (m.get("architecture") or {}).get("input_modalities") @@ -808,21 +835,28 @@ def llm_model(): return model_inputs()[0] -def ask(messages, calls, tools=None, mine=None, think=None, meter=True): +def ask(messages, calls, tools=None, mine=None, think=None, meter=True, model=None): # 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}} + # model is the id model_inputs gave, None for the router's loaded one. + url, key, mid, cloud = route(model or llm_model()) + if cloud: + ctx.update(model=model, max=0) # no /props: the meter shows tokens used only + else: + context_size(mid) + body = {"model": mid, "messages": messages, "stream": True, + "stream_options": {"include_usage": True}} + if not cloud: + # llama.cpp's own field; strict providers refuse unknown ones. + body["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") + u = urllib.parse.urlsplit(url + "/chat/completions") conn = (http.client.HTTPSConnection if u.scheme == "https" else http.client.HTTPConnection)( u.hostname, u.port, timeout=300) with lock: @@ -830,7 +864,7 @@ def ask(messages, calls, tools=None, mine=None, think=None, meter=True): return # stopped before the request went out inflight["conn"] = conn try: - yield from stream(conn, u.path, body, calls, meter) + yield from stream(conn, u.path, body, calls, meter, key, cloud or "llama-server") except OSError: if mine is not None and turn != mine: return # the stop hung up under the read @@ -842,12 +876,16 @@ def ask(messages, calls, tools=None, mine=None, think=None, meter=True): conn.close() -def stream(conn, path, body, calls, meter=True): - conn.request("POST", path, json.dumps(body).encode(), {"Content-Type": "application/json"}) +def stream(conn, path, body, calls, meter=True, key=None, who="llama-server"): + headers = {"Content-Type": "application/json"} + if key: + headers["Authorization"] = f"Bearer {key}" + conn.request("POST", path, json.dumps(body).encode(), headers) 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]}") + # The body says why (a full context, a bad request, a spent quota); + # the status alone does not. + raise RuntimeError(f"{who} {r.status}: {r.read().decode(errors='replace')[:2000]}") with r: streamed, other = False, b"" for line in r: @@ -875,12 +913,14 @@ def stream(conn, path, body, calls, meter=True): continue delta = choices[0]["delta"] collect_calls(calls, delta.get("tool_calls") or []) - if delta.get("reasoning_content"): - yield "think", delta["reasoning_content"] + # llama-server sends reasoning_content, OpenRouter reasoning. + think = delta.get("reasoning_content") or delta.get("reasoning") + if think: + yield "think", think if delta.get("content"): yield "token", delta["content"] if not streamed: - raise RuntimeError(other.decode(errors="replace")[:300] or "llama-server sent nothing") + raise RuntimeError(other.decode(errors="replace")[:300] or f"{who} sent nothing") def say(mine, kind, value, done=False): @@ -1024,13 +1064,14 @@ def transcribe(mine, path): return text -def reply(mine, text, files=(), audio=None): +def reply(mine, text, files=(), audio=None, model=None): try: if audio: + # Always the local model: the audio never leaves, only the words. text = transcribe(mine, audio) if text is None: return - answer(mine, text, files) + answer(mine, text, files, model) 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) @@ -1044,15 +1085,21 @@ def reply(mine, text, files=(), audio=None): emit({"type": "state", "state": "idle"}) -def answer(mine, text, files=()): +def answer(mine, text, files=(), chosen=None): 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() + model, inputs = model_inputs(chosen) images = inputs is None or "image" in inputs + # A cloud chat is offered only the tools marked cloud = true: the others + # would send mail, money or calendar to the provider. One of them called + # anyway (copied from history) answers "no tool named". + offered = ({n: t for n, t in TOOLS.items() if t.get("cloud") is True} + if route(model)[3] else TOOLS) + offer = schema(offered) or None if refs and not images: raise Refused(f"{model} cannot see images") with lock: @@ -1087,8 +1134,8 @@ def answer(mine, text, files=()): 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): + for kind, piece in ask(sent, calls, offer if rnd < CFG["tool_rounds"] else None, + mine, model=model): if not say(mine, kind, piece): log("interrupted") return @@ -1109,7 +1156,7 @@ def answer(mine, text, files=()): log("interrupted") return used.append(c["name"]) - result = run_tool(TOOLS, c["name"], c["arguments"]) + result = run_tool(offered, 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]) @@ -1130,7 +1177,7 @@ def answer(mine, text, files=()): log("interrupted") return fresh = cur["id"] is None - cur["id"] = save_turn(DB, cur["id"], text, said, used, entry, chips) + cur["id"] = save_turn(DB, cur["id"], text, said, used, entry, chips, chosen) conv = cur["id"] cur["history"].append(entry) cur["turns"].append({"user": text, "reply": said, "tools": used, "files": chips}) @@ -1139,17 +1186,18 @@ def answer(mine, text, files=()): emit({"type": "state", "state": "idle"}) log("assistant:", said) if fresh: - retitle(mine, conv, text, said) + retitle(mine, conv, text, said, model) -def retitle(mine, conv, user, said): +def retitle(mine, conv, user, said, model=None): # The first six words stand until the model names the conversation. Asked # after done, under the reply's turn: the next ask, stop, new or open # hangs it up, so it never holds a reply back, and the six words stay. try: msgs = [{"role": "user", "content": f"{TITLE}\n\nUser: {user[:2000]}\n\nAssistant: {said[:2000]}"}] - text = "".join(p for k, p in ask(msgs, {}, None, mine, think=False, meter=False) + text = "".join(p for k, p in ask(msgs, {}, None, mine, think=False, meter=False, + model=model) if k == "token") line = next((l for l in text.splitlines() if l.strip()), "") title = line.strip().strip("\"'*#").strip()[:80] @@ -1193,7 +1241,9 @@ def main(): def test(): + import contextlib import http.server + import io import tempfile c = merge(DEFAULTS, {"tool_rounds": 1, "env": {"K": "v"}}) @@ -1432,7 +1482,7 @@ def test(): .execute("PRAGMA table_info(conversations)")} # Protocol, against a fake llama-server. - global DB, TOOLS, TOOL_SCHEMA + global DB, TOOLS script = {"replies": [], "delay": 0, "inputs": ["text", "image"]} class Fake(http.server.BaseHTTPRequestHandler): @@ -1459,6 +1509,7 @@ def test(): self.wfile.write(b"data: " + json.dumps({"choices": [{"delta": { "content": script.get("title", "")}}]}).encode() + b"\n\ndata: [DONE]\n\n") return + script.get("paths", []).append(self.path) script["body"] = body script.get("bodies", []).append(script["body"]) if script.get("hold"): @@ -1515,7 +1566,6 @@ def test(): assert why in str(e), (name, e) 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) @@ -1539,6 +1589,7 @@ def test(): 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 + assert h["chat_model"] is None # Thinking, a tool round, then the answer. script["replies"] = [ [{"reasoning_content": "hmm"}, @@ -1618,7 +1669,8 @@ def test(): script["delay"] = 0 # History browsing. send(type="new") - assert until("conversation")[-1] == {"type": "conversation", "id": None, "turns": []} + assert until("conversation")[-1] == {"type": "conversation", "id": None, "turns": [], + "model": None} send(type="list") assert [x["id"] for x in until("conversations")[-1]["items"]] == [conv] send(type="search", q="tim") @@ -1653,7 +1705,7 @@ def test(): until("conversations") send(type="delete", id=conv) seen = until("conversations") - assert {"type": "conversation", "id": None, "turns": []} in seen # was current + assert {"type": "conversation", "id": None, "turns": [], "model": None} in seen # was current assert seen[-1]["items"] == [] # Files: the text inline, the image resolved, chips on the turn. tiny_png(png) @@ -1770,6 +1822,72 @@ def test(): 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")) + # Cloud: the chat's model picks the provider, and only cloud tools + # go along. A call to another one, copied from history, gets no + # output. The key reaches the provider and nothing else. + saved = TOOLS + TOOLS = {"now": {"description": "d", "command": ["sh", "-c", "echo noon"], "cloud": True}, + "mail": {"description": "m", "command": ["sh", "-c", "echo secret subject"]}} + logged = io.StringIO() + with contextlib.redirect_stdout(logged): + send(type="new") + until("conversation") + script["bodies"], script["paths"] = [], [] + script["replies"] = [ + [{"reasoning": "hm"}, + {"tool_calls": [{"index": 0, "id": "k1", "function": {"name": "mail", + "arguments": "{}"}}]}], + [{"content": "No mail tool."}]] + send(type="ask", text="any mail?", model="fake:vendor/x:free") + seen = until("done") + first, second = script["bodies"] + assert script["paths"] == ["/cloud/v1/chat/completions"] * 2 + assert script["auth"] == "Bearer sk-test" + assert first["model"] == "vendor/x:free" and "chat_template_kwargs" not in first + assert [t["function"]["name"] for t in first["tools"]] == ["now"] + assert second["messages"][-1] == {"role": "tool", "tool_call_id": "k1", + "content": "error: no tool named mail"} + assert ("think", "hm") in [(e["type"], e.get("delta")) for e in seen] + meters = [e for e in seen if e["type"] == "ctx"] + assert meters[-1]["model"] == "fake:vendor/x:free" and meters[-1]["max"] == 0 + assert "sk-test" not in json.dumps(seen) + conv = cur["id"] + assert DB.execute("SELECT model FROM conversations WHERE id = ?", + (conv,)).fetchone()[0] == "fake:vendor/x:free" + # The next ask with no model stays on the chat's. + script["replies"] = [[{"content": "Still cloud."}]] + send(type="ask", text="again") + until("done") + assert script["paths"][-1] == "/cloud/v1/chat/completions" + # Reopened, the chat reports its model. + send(type="new") + assert until("conversation")[-1]["model"] is None + send(type="open", id=conv) + assert until("conversation")[-1]["model"] == "fake:vendor/x:free" + # A local pick goes to the router with every tool and the + # thinking switch. + script["replies"] = [[{"content": "Local."}]] + send(type="ask", text="local now", model="m") + until("done") + assert script["paths"][-1] == "/v1/chat/completions" + assert script["body"]["model"] == "m" and "chat_template_kwargs" in script["body"] + assert [t["function"]["name"] for t in script["body"]["tools"]] == ["now", "mail"] + # The dropdown's list: the router's models, then the cloud ones + # whose key loaded. + send(type="models") + items = until("models")[-1]["items"] + assert {"id": "m", "status": "loaded", "cloud": False} in items + assert {"id": "fake:vendor/x:free", "provider": "fake", "cloud": True} in items + assert not [i for i in items if i["id"].startswith("broken:")] + send(type="ask", text="x", model="") + assert until("error")[-1]["message"] == "bad model" + send(type="ask", text="x", model="broken:y") + e = until("error")[-1] + assert "no key" in e["message"] and e["turn"] is True + send(type="new") + until("conversation") + TOOLS = saved + assert "sk-test" not in logged.getvalue(), logged.getvalue() # llama-server down: an error, not a dead thread. CFG["llm"] = "http://127.0.0.1:9/v1" send(type="ask", text="anyone?") |