refactor: add AgentConfig and more type annotations

2 files changed, 151 insertions, 80 deletions

diff --git a/actions.py b/actions.py
index ffd33f7..3dc7557 100644
--- a/actions.py
+++ b/actions.py
@@ -1,6 +1,24 @@
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Callable, Dict
+from dataclasses import dataclass
 import subprocess
+from llm import AgentConfig, JSON
+
+
+@dataclass
+class Action:
+    """Class for LLM actions.
+
+    Should only be called on JSON dicts where "action" -> `name`.
+    """
+
+    name: str
+    description: str
+    handler: Callable[[JSON, AgentConfig], JSON]
+
+    def __call__(self, data: JSON, config: AgentConfig) -> JSON:
+        return self.handler(data, config)
+
 
 def _path_checked(path: Path, cwd=Path.cwd()) -> Optional[Path]:
     """Resolves `path`, returning `None` if it's not a subdirectory of `cwd`."""
@@ -8,6 +26,7 @@ def _path_checked(path: Path, cwd=Path.cwd()) -> Optional[Path]:
     cwd = cwd.resolve()
     return path if path.is_relative_to(cwd) else None
 
+
 def _path_helper(path_str: str, callback, cwd=Path.cwd()):
     path = _path_checked(Path(path_str), cwd=cwd)
     if path is None:
@@ -18,59 +37,77 @@ def _path_helper(path_str: str, callback, cwd=Path.cwd()):
         except Exception as e:
             return {"status": "error", "reason": f"Failure: {e}"}
 
-def ls(args, workdir=Path.cwd()):
+
+def _ls(args: JSON, config: AgentConfig) -> JSON:
     '{"action": "ls", "path": PATH} - list files at PATH'
+
     def inner(path: Path):
         return {"content": list(map(lambda s: str(s), path.iterdir()))}
-    match args:
-        case {"action": "ls", "path": path_str}:
-            return _path_helper(path_str, inner, cwd=workdir)
 
-def read(args, workdir=Path.cwd()):
+    return _path_helper(args["path"], inner, cwd=config.workdir)
+
+
+def _read(args: JSON, config: AgentConfig) -> JSON:
     '{"action": "read", "path": PATH} - read file at PATH'
+
     def inner(path: Path):
         contents = path.read_text(encoding="utf-8")
         return {"content": contents}
-    match args:
-        case {"action": "read", "path": path_str}:
-            return _path_helper(path_str, inner, cwd=workdir)
 
-def write(args, workdir=Path.cwd()):
+    return _path_helper(args["path"], inner, cwd=config.workdir)
+
+
+def _write(args: JSON, config: AgentConfig) -> JSON:
     '{"action": "write", "path": PATH, "content": CONTENT} - write CONTENT into file at PATH'
+
     def inner(path: Path, text: str):
         n_bytes = path.write_text(text)
         return {"bytes": n_bytes}
-    match args:
-        case {"action": "write", "path": path_str, "content": content}:
-            return _path_helper(path_str, lambda path: inner(path, content), cwd=workdir)
 
-def tldr(args, workdir=Path.cwd()):
+    return _path_helper(
+        args["path"], lambda path: inner(path, args["content"]), cwd=config.workdir
+    )
+
+
+def _tldr(args: JSON, config: AgentConfig) -> JSON:
     '{"action": "tldr", "program": PROGRAM} - read tldr docs for PROGRAM'
-    match args:
-        case {"action": "tldr", "program": program_str}:
-            try:
-                result = subprocess.run(
-                    ["tldr", program_str],
-                    capture_output=True,
-                    text=True,
-                    check=True
-                )
-                return {"status": "success", "content": result.stdout}
-            except subprocess.CalledProcessError:
-                return {"status": "error", "reason": f"Tldr for '{program_str}' not found."}
-            except Exception as e:
-                return {"status": "error", "reason": f"An unexpected error occurred: {e}"}
+    try:
+        result = subprocess.run(
+            ["tldr", args["program"]], capture_output=True, text=True, check=True
+        )
+        return {"status": "success", "content": result.stdout}
+    except subprocess.CalledProcessError:
+        return {"status": "error", "reason": f"Tldr for '{args['program']}' not found."}
+    except Exception as e:
+        return {"status": "error", "reason": f"An unexpected error occurred: {e}"}
+
 
-def shell(args, workdir=Path.cwd()):
+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.'
-    match args:
-        case {"action": "shell", "command": cmd}:
-            user_input = input(f"Run cmd `{cmd}`? [y/N]: ")
-            if user_input.lower() == 'y':
-                res = subprocess.run(cmd, shell=True, cwd=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}
-            else:
-                return {"status": "error", "reason": "Permission denied by 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,
+            }
+    else:
+        return {"status": "error", "reason": "Permission denied by user."}
+
+
+ls = Action("ls", _ls.__doc__ or "", _ls)
+read = Action("read", _read.__doc__ or "", _read)
+write = Action("write", _write.__doc__ or "", _write)
+tldr = Action("tldr", _tldr.__doc__ or "", _tldr)
+shell = Action("shell", _shell.__doc__ or "", _shell)
diff --git a/llm.py b/llm.py
index 4ef94d9..b21e81e 100755
--- a/llm.py
+++ b/llm.py
@@ -2,32 +2,51 @@
 import os
 from pathlib import Path
 from openai import OpenAI
-from typing import Optional
+from typing import Optional, List, Dict, TypeVar, TypeAlias
 import argparse
 import re
 import subprocess
 import readline
 import json
+from dataclasses import dataclass
 from actions import *
 
+T = TypeVar("T")
+
 import logging
+
 logger = logging.getLogger(__name__)
-logging.basicConfig(filename="thoughts.log",
-                    level=logging.INFO)
+logging.basicConfig(filename="thoughts.log", level=logging.INFO)
+
+JSON: TypeAlias = Dict[str, "JSON"] | List["JSON"] | str | int | float | bool | None
+
+
+@dataclass
+class AgentConfig:
+    """Configuration for agents."""
+
+    workdir: Path
+    actions: List["Action"]
+
+
+client = OpenAI(
+    api_key=os.environ["AALTO_LLM_KEY"],
+    base_url="https://llm-gateway.k8s.aalto.fi/api/v1",
+)
 
-client = OpenAI(api_key = os.environ["AALTO_LLM_KEY"],
-                base_url = "https://llm-gateway.k8s.aalto.fi/api/v1")
 
 def direct_query(query: str):
     logger.info(f"send: {query}")
-    res = client.responses.create(model = "RedHatAI/gemma-4-31B-it-FP8-Dynamic",
-                                  input = query)
+    res = client.responses.create(
+        model="RedHatAI/gemma-4-31B-it-FP8-Dynamic", input=query
+    )
     logger.info(f"receive: {res.output_text}")
     return res.output_text
 
+
 # def reason_query(input):
 #     # TODO some models may natively support reasoning
-    
+
 #     reasoning = direct_query(input + ({"role": "system",
 #                                        "content": "Reason about the user's query step by step.\n\nAn answer will be generated based on your analysis."},))
 
@@ -39,91 +58,106 @@ def direct_query(query: str):
 
 # # kinda funky
 # # def flexible_query(input):
-   
+
 # #     reasoning_required = direct_query(input[-1:] + ({"role": "system", "content": "Does answering the user's question require reasoning? (y|N)"},))
 # #     if reasoning_required == 'y':
 # #         return reason_query(input)
 # #     else:
 # #         return direct_query(input)
 
-    
-def chat(content, state = (), query_method = direct_query, role = "user"):
-    
+
+def chat(content, state=(), query_method=direct_query, role="user"):
+
     state = state + ({"role": role, "content": content},)
-    
+
     response = query_method(state)
 
     state = state + ({"role": "assistant", "content": response},)
-    
+
     return (response, state)
 
 
-def query_agent(actions: list, query: str, state = (), workdir: Path = Path.cwd()):
+def first(l: List[T], p: Callable[[T], bool]) -> Optional[T]:
+    return next((x for x in l if p(x)), None)
+
+
+def query_agent(query: str, config: AgentConfig, state=()):
 
     state = state + ({"role": "user", "content": query},)
 
-    (res, state) = chat(f"Based on the current state, what should you do next?\n"
-                        "Respond JSON indicating an action, or respond in natural language if you are finished.",
-                        state)
+    (res, state) = chat(
+        f"Based on the current state, what should you do next?\n"
+        "Respond JSON indicating an action, or respond in natural language if you are finished.",
+        state,
+    )
 
     # Replace all fenced JSON blocks
 
     replacements = []
-    result_jsons = []
+    result_jsons: List[JSON] = []
     for m in re.finditer(r"^```json\s*(.*?)\s*^```", res, re.DOTALL | re.MULTILINE):
         json_str = m.group(1)
         try:
             action = json.loads(json_str)
-        except json.JSONDecodeError:
+            action_name = action["action"]
+        except (json.JSONDecodeError, KeyError):
             continue
 
-        for f in actions:
-            result_json = f(action, workdir=workdir)
-            if result_json is not None:
-                replacements.append((m.start(), m.end(), action["action"]))
-                result_jsons.append(result_json)
-                break
+        f = first(config.actions, lambda f: f.name == action_name)
+        if f is not None:
+            result_json = f(action, config)
+            result_jsons.append(result_json)
+            replacements.append((m.start(), m.end(), action_name))
 
     replaced_res = ""
     i = 0
-    for (start, end, action) in replacements:
+    for start, end, action in replacements:
         replaced_res += res[i:start] + f"[performed '{action}']\n"
         i = end
     replaced_res += res[i:]
 
     if len(replacements) > 0:
-        (res, state) = query_agent(actions, json.dumps(result_jsons), state, workdir=workdir)
-        return (replaced_res + "\n\n" + res, state)
+        (res, state) = query_agent(json.dumps(result_jsons), config, state)
+        return (replaced_res + "\n" + res, state)
     else:
         return (replaced_res, state)
 
 
 def main():
     parser = argparse.ArgumentParser(description="Autonomous Coding Agent")
-    parser.add_argument("--workdir", type=Path, default=Path.cwd(), help="The directory the agent has access to (default: current directory)")
+    parser.add_argument(
+        "--workdir",
+        type=Path,
+        default=Path.cwd(),
+        help="The directory the agent has access to (default: current directory)",
+    )
     args = parser.parse_args()
 
-    actions = [ls, read, write, tldr, shell]
-    state = ({"role": "system",
-                      "content": f"""You are an autonomous coding agent that can help with programming tasks.
+    config = AgentConfig(workdir=args.workdir, actions=[ls, read, write, tldr, shell])
+
+    state = (
+        {
+            "role": "system",
+            "content": f"""You are an autonomous coding agent that can help with programming tasks.
 You can explore repositories, read files, search for information, and make changes to code.
 
 You have access to the following actions, which you should embed in JSON markdown blocks "```json ACTION ```" in your responses.
-{'\n'.join(map(lambda f: '* ' + f.__doc__, actions))}
+{"\n".join(map(lambda f: "* " + f.description, config.actions))}
 
 After reading the relevant files, remember to analyze the task at hand and break it into steps.
 
 Once you are finished with your task, give a response with none of the above JSON commands.
-"""},)
-    
+""",
+        },
+    )
+
     while True:
-        
         query = input("> ").strip()
 
-        (response, state) = query_agent(actions, query, state, workdir=args.workdir)
-        
+        (response, state) = query_agent(query, config, state)
+
         print(response)
-        
+
 
 if __name__ == "__main__":
     main()