diff options
| -rw-r--r-- | assistant.py | 89 |
1 files changed, 89 insertions, 0 deletions
diff --git a/assistant.py b/assistant.py index 5dde49b..11333dc 100644 --- a/assistant.py +++ b/assistant.py @@ -21,6 +21,7 @@ Run: python3 assistant.py (--test runs the self-check) """ import ctypes +import glob import json import os import queue @@ -79,9 +80,32 @@ def load_config(path): 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")) WHISPER = f"http://127.0.0.1:{CFG['whisper']['port']}" KLANG = {l: v["lang"] for l, v in CFG["voice"].items()} +TOOLS_DIR = os.path.join(CONFIG_DIR, "tools") +TOOLS = load_tools(TOOLS_DIR) +TOOL_SCHEMA = schema(TOOLS) or None def log(*a): @@ -148,6 +172,43 @@ def split_sentences(buf): return [p.strip() for p in parts[:-1] if p.strip()], parts[-1] +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}") + return out.strip()[:2000] + + def llm_model(): # The router keeps one model resident; asking for another would swap it # out under other clients, so use whatever is not unloaded. @@ -322,6 +383,34 @@ def test(): assert False, bad except SystemExit as e: assert "config:" in str(e) + import tempfile + 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" print("ok") |