diff options
| -rw-r--r-- | actions.py | 13 | ||||
| -rwxr-xr-x | llm.py | 4 |
2 files changed, 11 insertions, 6 deletions
@@ -1,5 +1,5 @@ from pathlib import Path -from typing import Optional, Callable, Dict +from typing import Optional, Callable, Dict, List from dataclasses import dataclass import subprocess import re @@ -17,6 +17,7 @@ class Action: name: str description: str handler: Callable[[JSON, AgentConfig], JSON] + required_keys: List[str] formatter: Callable[[JSON], str] = lambda args: ( f"{args['action']} with " + ",".join([f"{arg} as {args[arg]}" for arg in args if arg != "action"]) @@ -148,30 +149,34 @@ def _webgrep(args: JSON, config: AgentConfig) -> JSON: return {"status": "success", "matches": list(matches)} -ls = Action("ls", _ls.__doc__ or "", _ls) -read = Action("read", _read.__doc__ or "", _read) +ls = Action("ls", _ls.__doc__ or "", _ls, ["path"]) +read = Action("read", _read.__doc__ or "", _read, ["path"]) write = Action( "write", _write.__doc__ or "", _write, + ["path", "content"], formatter=lambda args: f"write {args['path']}", ) -tldr = Action("tldr", _tldr.__doc__ or "", _tldr) +tldr = Action("tldr", _tldr.__doc__ or "", _tldr, ["program"]) shell = Action( "shell", _shell.__doc__ or "", _shell, + ["command"], formatter=lambda args: f"shell command `{args['command']}`", ) mypy = Action( "mypy", _mypy.__doc__ or "", _mypy, + ["path"], formatter=lambda args: f"type check {args['path']}", ) webgrep = Action( "webgrep", _webgrep.__doc__ or "", _webgrep, + ["url", "keywords"], formatter=lambda args: f"search {args['url']}", ) @@ -87,10 +87,10 @@ def query_agent(query: str, config: AgentConfig, state=()): try: action = json.loads(json_str) action_name = action["action"] - except (json.JSONDecodeError, KeyError): + except (json.JSONDecodeError, KeyError, TypeError): continue - f = first(config.actions, lambda f: f.name == action_name) + f = first(config.actions, lambda f: f.name == action_name and all(k in action for k in f.required_keys)) if f is not None: desc = f.formatter(action) if ( |