Financial_Agent / agent.py
luisejdm's picture
Add agent file
9113b4c
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