refactor: add AgentConfig and more type annotations

1 files changed, 74 insertions, 40 deletions

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()