diff options
| -rw-r--r-- | assistant.py | 97 |
1 files changed, 96 insertions, 1 deletions
diff --git a/assistant.py b/assistant.py index 37ad5f2..7f31add 100644 --- a/assistant.py +++ b/assistant.py @@ -19,7 +19,8 @@ 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, --defaults -prints the built-in config as JSON, for a settings UI) +prints the built-in config as JSON, for a settings UI, --cloud-models NAME +prints that [[cloud]] provider's model ids as JSON) """ import base64 @@ -30,6 +31,7 @@ import json import mimetypes import os import re +import shlex import signal import socket import sqlite3 @@ -40,6 +42,7 @@ import time import tomllib import traceback import wave +import urllib.error import urllib.parse import urllib.request @@ -720,6 +723,71 @@ def collect_calls(calls, fragments): c["arguments"] += fn.get("arguments") or "" +# Cloud providers, named in [[cloud]]. A key is key_cmd's stdout, read once +# at startup; a provider whose key_cmd fails is left out, so its models are +# not offered and a chat still on one is refused. +PROVIDERS = {} # name -> {"url", "key", "models"} + + +def provider_key(p): + # Only the exit status is reported: a secret tool's stderr is not worth + # the risk of it holding the secret. + try: + r = subprocess.run(shlex.split(p["key_cmd"]), capture_output=True, text=True, + stdin=subprocess.DEVNULL, timeout=10) + except (OSError, ValueError, subprocess.TimeoutExpired) as e: + raise RuntimeError(f"cloud {p['name']}: key_cmd failed ({type(e).__name__})") from None + if r.returncode or not r.stdout.strip(): + raise RuntimeError(f"cloud {p['name']}: key_cmd failed (exit {r.returncode})") + return r.stdout.strip() + + +def load_providers(cloud): + out = {} + for p in cloud: + try: + out[p["name"]] = {"url": p["url"].rstrip("/"), "key": provider_key(p), + "models": p["models"]} + except RuntimeError as e: + log(e) + return out + + +def route(model): + # (base url, key, the id the server knows, provider or None). A prefix + # is a provider only when one by that name is configured, so a local id + # with a colon stays local. + name, sep, rest = (model or "").partition(":") + if sep and name in PROVIDERS: + p = PROVIDERS[name] + return p["url"], p["key"], rest, name + if sep and any(p["name"] == name for p in CFG["cloud"]): + raise Refused(f"cloud {name}: no key (key_cmd failed at startup)") + return CFG["llm"], None, model, None + + +def cloud_models(name): + # For the settings page (--cloud-models): the provider's model ids, + # fetched here so the key never reaches it. + p = next((p for p in CFG["cloud"] if p["name"] == name), None) + if p is None: + raise SystemExit(f"no cloud provider named {name}") + try: + key = provider_key(p) + except RuntimeError as e: + raise SystemExit(str(e)) from None + req = urllib.request.Request(p["url"].rstrip("/") + "/models", + headers={"Authorization": f"Bearer {key}"}) + try: + with urllib.request.urlopen(req, timeout=20) as r: + return sorted(m["id"] for m in json.load(r)["data"]) + except urllib.error.HTTPError as e: + raise SystemExit(f"cloud {name} {e.code}: " + f"{e.read().decode(errors='replace')[:300]}") from None + except (OSError, KeyError, TypeError, ValueError) as e: + 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 @@ -1097,6 +1165,7 @@ def main(): served = False try: DB = db_open(DB_PATH) + PROVIDERS.update(load_providers(CFG["cloud"])) try: context_size(llm_model()) # so the first hello has it except Exception: @@ -1356,6 +1425,7 @@ def test(): pass def do_GET(self): + script["auth"] = self.headers.get("Authorization") body = ({"data": [{"id": "m", "status": {"value": "loaded"}, "architecture": {"input_modalities": script["inputs"]}}]} if self.path.endswith("/models") @@ -1365,6 +1435,7 @@ def test(): self.wfile.write(json.dumps(body).encode()) def do_POST(self): + script["auth"] = self.headers.get("Authorization") body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) if str(body["messages"][0]["content"]).startswith(TITLE): # Answered apart, so a title never takes a reply meant for an ask. @@ -1405,6 +1476,28 @@ def test(): 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" + # Cloud providers: the key from key_cmd, models fetched with it. The + # fake answers /cloud/v1 like /v1, so the path tells them apart. + base = f"http://127.0.0.1:{srv.server_address[1]}" + CFG["cloud"] = [{"name": "fake", "url": base + "/cloud/v1/", "key_cmd": "echo sk-test", + "models": ["vendor/x:free"]}, + {"name": "broken", "url": base + "/v1", "key_cmd": "false", "models": ["y"]}] + PROVIDERS.update(load_providers(CFG["cloud"])) + assert set(PROVIDERS) == {"fake"} and PROVIDERS["fake"]["key"] == "sk-test" + assert route("fake:vendor/x:free") == (base + "/cloud/v1", "sk-test", "vendor/x:free", "fake") + assert route("Gemma:q4") == (CFG["llm"], None, "Gemma:q4", None) # no such provider + try: + route("broken:y") + assert False + except Refused as e: + assert "no key" in str(e) + assert cloud_models("fake") == ["m"] and script["auth"] == "Bearer sk-test" + for name, why in (("broken", "key_cmd failed"), ("nope", "no cloud provider")): + try: + cloud_models(name) + assert False, name + except SystemExit as e: + assert why in str(e), (name, e) DB = db_open(":memory:") TOOLS = {"now": {"description": "d", "command": ["sh", "-c", "echo noon"]}} TOOL_SCHEMA = schema(TOOLS) @@ -1683,5 +1776,7 @@ if __name__ == "__main__": test() elif sys.argv[1:] == ["--defaults"]: print(json.dumps(DEFAULTS)) + elif len(sys.argv) == 3 and sys.argv[1] == "--cloud-models": + print(json.dumps(cloud_models(sys.argv[2]))) else: main() |