Add agent file

agent.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from model import ModelRunner
+from parsing import parse_response
+from tools import ToolRegistry
+import time
+
+
+@dataclass
+class Agent:
+    runner: ModelRunner
+    tools: ToolRegistry
+    system_prompt: str
+
+    def run(
+        self,
+        task: str,
+        max_steps: int = 5,
+        temperature: float = 0.3,
+        verbose: bool = False,
+    ) -> tuple[str | None, list[dict]]:
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": task},
+        ]
+        trace: list[dict] = []
+
+        for step in range(1, max_steps + 1):
+            time.sleep(1.5)
+            response = self.runner.generate(messages, temperature=temperature)
+            if verbose:
+                print(f"[step {step}] {response}")
+            messages.append({"role": "assistant", "content": response})
+
+            kind, payload = parse_response(response)
+
+            if kind == "final":
+                tool_succeeded = any(
+                    entry.get("type") == "action" for entry in trace
+                )
+                if not tool_succeeded:
+                    trace.append({
+                        "step": step,
+                        "type": "blocked_final",
+                        "content": payload,
+                    })
+                    messages.append({
+                        "role": "user",
+                        "content": (
+                            "You cannot give a FINAL answer without first calling a tool. "
+                            "Call the appropriate tool first before answering."
+                        ),
+                    })
+                    continue
+                trace.append({"step": step, "type": "final", "content": payload})
+                return payload, trace
+
+            if kind == "action":
+                name, args = payload
+                try:
+                    result = self.tools.execute(name, *args)
+                    result_text = (
+                        f"${float(result):.2f}"
+                        if isinstance(result, (int, float))
+                        else str(result)
+                    )
+                    observation = f"Tool result: {result_text}"
+                    trace.append({
+                        "step": step,
+                        "type": "action",
+                        "tool": name,
+                        "args": args,
+                        "result": result_text,
+                    })
+                except KeyError as error:
+                    observation = f"Error: {error}"
+                    trace.append({
+                        "step": step,
+                        "type": "error",
+                        "tool": name,
+                        "args": args,
+                        "error": str(error),
+                    })
+                except Exception as error:
+                    observation = f"Error: {error}"
+                    trace.append({
+                        "step": step,
+                        "type": "error",
+                        "tool": name,
+                        "args": args,
+                        "error": str(error),
+                    })
+                messages.append({"role": "user", "content": observation})
+            else:
+                trace.append({"step": step, "type": "unknown", "raw": payload})
+                messages.append({
+                    "role": "user",
+                    "content": (
+                        "Invalid format. Respond with exactly one line starting with "
+                        "ACTION: tool_name(ARGS) or FINAL: <answer>."
+                    ),
+                })
+
+        return None, trace