feat: trust system

2 files changed, 54 insertions, 24 deletions

diff --git a/actions.py b/actions.py
index 3dc7557..dfeebc3 100644
--- a/actions.py
+++ b/actions.py
@@ -15,6 +15,10 @@ class Action:
     name: str
     description: str
     handler: Callable[[JSON, AgentConfig], JSON]
+    formatter: Callable[[JSON], str] = lambda args: (
+        f"{args['action']} with "
+        + ",".join([f"{arg} as {args[arg]}" for arg in args if arg != "action"])
+    )
 
     def __call__(self, data: JSON, config: AgentConfig) -> JSON:
         return self.handler(data, config)
@@ -84,30 +88,36 @@ def _tldr(args: JSON, config: AgentConfig) -> JSON:
 
 def _shell(args: JSON, config: AgentConfig) -> JSON:
     '{"action": "shell", "command": CMD} - Propose executing a shell command, but interrupts to ask for permission from the user.'
-    user_input = input(f"Run cmd `{args['command']}`? [y/N]: ")
-    if user_input.lower() == "y":
-        res = subprocess.run(
-            args["command"],
-            shell=True,
-            cwd=config.workdir,
-            capture_output=True,
-            text=True,
-        )
-        if res.returncode == 0:
-            return {"status": "success", "stdout": res.stdout}
-        else:
-            return {
-                "status": "error",
-                "reason": f"Shell command exited with code {res.returncode}",
-                "stdout": res.stdout,
-                "stderr": res.stderr,
-            }
+    res = subprocess.run(
+        args["command"],
+        shell=True,
+        cwd=config.workdir,
+        capture_output=True,
+        text=True,
+    )
+    if res.returncode == 0:
+        return {"status": "success", "stdout": res.stdout}
     else:
-        return {"status": "error", "reason": "Permission denied by user."}
+        return {
+            "status": "error",
+            "reason": f"Shell command exited with code {res.returncode}",
+            "stdout": res.stdout,
+            "stderr": res.stderr,
+        }
 
 
 ls = Action("ls", _ls.__doc__ or "", _ls)
 read = Action("read", _read.__doc__ or "", _read)
-write = Action("write", _write.__doc__ or "", _write)
+write = Action(
+    "write",
+    _write.__doc__ or "",
+    _write,
+    formatter=lambda args: f"write {args['path']}",
+)
 tldr = Action("tldr", _tldr.__doc__ or "", _tldr)
-shell = Action("shell", _shell.__doc__ or "", _shell)
+shell = Action(
+    "shell",
+    _shell.__doc__ or "",
+    _shell,
+    formatter=lambda args: f"shell command `{args['command']}`",
+)
diff --git a/llm.py b/llm.py
index b21e81e..649f927 100755
--- a/llm.py
+++ b/llm.py
@@ -27,6 +27,8 @@ class AgentConfig:
 
     workdir: Path
     actions: List["Action"]
+    trustlevels: Dict[str, int]
+    trustlevel: int
 
 
 client = OpenAI(
@@ -105,9 +107,16 @@ def query_agent(query: str, config: AgentConfig, state=()):
 
         f = first(config.actions, lambda f: f.name == action_name)
         if f is not None:
-            result_json = f(action, config)
+            desc = f.formatter(action)
+            if (
+                config.trustlevels[action_name] > config.trustlevel
+                and not input(f"Perform {desc}? [y/N] ") == "y"
+            ):
+                result_json = {"status": "error", "reason": "Permission denied by user"}
+            else:
+                result_json = f(action, config)
             result_jsons.append(result_json)
-            replacements.append((m.start(), m.end(), action_name))
+            replacements.append((m.start(), m.end(), desc))
 
     replaced_res = ""
     i = 0
@@ -133,7 +142,18 @@ def main():
     )
     args = parser.parse_args()
 
-    config = AgentConfig(workdir=args.workdir, actions=[ls, read, write, tldr, shell])
+    config = AgentConfig(
+        workdir=args.workdir,
+        actions=[ls, read, write, tldr, shell],
+        trustlevels={
+            "ls": 0,
+            "read": 0,
+            "tldr": 0,
+            "write": 1,
+            "shell": 2,
+        },
+        trustlevel=0,
+    )
 
     state = (
         {