fix: more robust action matching

2 files changed, 11 insertions, 6 deletions

diff --git a/actions.py b/actions.py
index a6051ed..7d2b61c 100644
--- a/actions.py
+++ b/actions.py
@@ -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']}",
 )
diff --git a/llm.py b/llm.py
index db72ba3..7577f9f 100755
--- a/llm.py
+++ b/llm.py
@@ -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 (