backend/agents/base.py

"""
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}"