File size: 6,001 Bytes
71b4454
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
"""
Base Agent — single-agent ReAct-style execution loop.

Real execution: think -> act -> observe, backed by persistent memory and
real tool calls. Every step calls the LLM gateway; there is no stubbed or
canned response path in this loop.
"""
from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Awaitable, Callable, Optional

from .llm_gateway import LLMGateway
from .memory import AgentMemoryStore

ToolFunc = Callable[..., Awaitable[Any]]


class StepType(str, Enum):
    ACTION = "action"
    OBSERVATION = "observation"
    FINAL = "final"


@dataclass(slots=True)
class AgentStep:
    type: StepType
    content: str
    tool: Optional[str] = None
    tool_input: Optional[dict[str, Any]] = None


@dataclass(slots=True)
class Tool:
    name: str
    description: str
    func: ToolFunc
    parameters: dict[str, Any] = field(default_factory=dict)


class AgentTimeoutError(RuntimeError):
    """Raised when an agent exhausts its step budget without a final answer."""


class BaseAgent:
    """A single agent with a bounded ReAct loop, tool access, and memory."""

    def __init__(
        self,
        agent_id: str,
        role: str,
        system_prompt: str,
        gateway: LLMGateway,
        memory: AgentMemoryStore,
        tools: Optional[list[Tool]] = None,
        max_steps: int = 12,
    ) -> None:
        self.agent_id = agent_id
        self.role = role
        self.system_prompt = system_prompt
        self.gateway = gateway
        self.memory = memory
        self.tools: dict[str, Tool] = {t.name: t for t in (tools or [])}
        self.max_steps = max_steps

    def register_tool(self, tool: Tool) -> None:
        self.tools[tool.name] = tool

    def _tool_catalog(self) -> str:
        if not self.tools:
            return "No tools available."
        return "\n".join(
            f"- {t.name}: {t.description} args={t.parameters}" for t in self.tools.values()
        )

    def _build_prompt(self, task_id: str, goal: str) -> list[dict[str, str]]:
        history = self.memory.recent_events(self.agent_id, limit=20, task_id=task_id)
        transcript = "\n".join(f"[{e.role}] {e.content}" for e in history)
        instructions = (
            f"{self.system_prompt}\n\n"
            f"You are agent '{self.agent_id}' ({self.role}).\n"
            f"Available tools:\n{self._tool_catalog()}\n\n"
            "Respond with EXACTLY one JSON object per turn, no other text:\n"
            '{"type": "action", "tool": "<name>", "input": {...}} to call a tool, or\n'
            '{"type": "final", "content": "<answer>"} once the goal is complete.'
        )
        messages = [{"role": "system", "content": instructions}]
        if transcript:
            messages.append({"role": "user", "content": f"Prior steps:\n{transcript}"})
        messages.append({"role": "user", "content": f"Goal: {goal}"})
        return messages

    async def run(self, goal: str, task_id: Optional[str] = None) -> str:
        task_id = task_id or str(uuid.uuid4())
        self.memory.record_event(self.agent_id, "goal", goal, task_id=task_id)

        for _ in range(self.max_steps):
            messages = self._build_prompt(task_id, goal)
            raw = await self.gateway.complete(messages)
            step = self._parse_step(raw)

            if step.type is StepType.FINAL:
                self.memory.record_event(self.agent_id, "final", step.content, task_id=task_id)
                return step.content

            if step.type is StepType.ACTION and step.tool:
                self.memory.record_event(
                    self.agent_id, "action",
                    json.dumps({"tool": step.tool, "input": step.tool_input}),
                    task_id=task_id,
                )
                observation = await self._execute_tool(step.tool, step.tool_input or {})
                self.memory.record_event(self.agent_id, "observation", observation, task_id=task_id)
            else:
                self.memory.record_event(
                    self.agent_id, "observation", f"Could not parse step: {raw[:500]}", task_id=task_id
                )

        raise AgentTimeoutError(
            f"Agent '{self.agent_id}' did not reach a final answer within {self.max_steps} steps."
        )

    def _parse_step(self, raw: str) -> AgentStep:
        data = self._extract_json_object(raw)
        if data is None:
            return AgentStep(type=StepType.OBSERVATION, content=raw)

        step_type = data.get("type")
        if step_type == "final":
            return AgentStep(type=StepType.FINAL, content=str(data.get("content", "")))
        if step_type == "action":
            return AgentStep(
                type=StepType.ACTION,
                content=raw,
                tool=data.get("tool"),
                tool_input=data.get("input", {}),
            )
        return AgentStep(type=StepType.OBSERVATION, content=raw)

    @staticmethod
    def _extract_json_object(raw: str) -> Optional[dict[str, Any]]:
        text = raw.strip()
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass
        start, end = text.find("{"), text.rfind("}")
        if start == -1 or end == -1 or end <= start:
            return None
        try:
            return json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            return None

    async def _execute_tool(self, tool_name: str, tool_input: dict[str, Any]) -> str:
        tool = self.tools.get(tool_name)
        if tool is None:
            return f"Error: no such tool '{tool_name}'. Available: {list(self.tools)}"
        try:
            result = await tool.func(**tool_input)
            return result if isinstance(result, str) else json.dumps(result)
        except Exception as exc:  # noqa: BLE001 — tool failures become observations, not crashes
            return f"Error executing '{tool_name}': {exc}"