src/agent/loop.py

"""The agent loop: plan → call tools → observe → answer.

The model proposes tool calls; this loop executes them against the real engine and
ledger, feeds the results back, and repeats until the model returns a final answer
(or the step budget is exhausted). Every run produces a :class:`Trace` that can be
serialized and shared on the Hub (📡 Sharing is Caring).

The loop never lets the model compute — it only routes the model's chosen tool
calls to deterministic handlers.
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from typing import Dict, List, Optional

from .llm import AssistantTurn, LLMClient, ToolCall
from .tools import Tool

DEFAULT_SYSTEM_PROMPT = """\
You are PocketAccountant, a meticulous accountant agent that helps freelancers and \
small businesses with their taxes in Mexico (🇲🇽) and the USA (🇺🇸). Reason about the \
question, then use tools. Follow these rules without exception:

1. NEVER do arithmetic yourself. To compute any tax, total, ratio or statement, call \
the appropriate tool — the tools are the only valid source of numbers. For Mexico use \
compute_iva / compute_isr_resico / compare_regimes; for the USA use us_tax_summary.
2. To answer questions about the user's finances, call a data/compute tool (e.g. \
month_summary, income_statement, compute_iva, us_tax_summary) — do not guess numbers.
3. For any tax-RULE question (deductibility, rates, obligations, deadlines, "what is …", \
"do I need …"), call cite_regulation. If it returns grounded=false, say plainly you \
cannot confirm it and recommend a licensed accountant / CPA. Never invent tax law.
4. When you state a number, also give its source (included in the tool result) and the \
year of the tax tables.
5. The user message may begin with context tags: [mx] or [us] = the country/tax system \
to use; [en] or [es] = the language to answer in; [YYYY-MM] = the current accounting \
period (pass that year/month to tools). Strip the tags from what you show the user.
6. Be concise and answer in the requested language.

You are an assistant, not a substitute for a professional accountant; tax tables are \
dated and should be verified.
"""


@dataclass
class Step:
    tool: str
    arguments: dict
    result: dict
    ok: bool = True


@dataclass
class Trace:
    user: str
    steps: List[Step] = field(default_factory=list)
    final_answer: str = ""
    messages: List[dict] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "user": self.user,
            "steps": [asdict(s) for s in self.steps],
            "final_answer": self.final_answer,
            "messages": self.messages,
        }

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, indent=indent)


class Agent:
    def __init__(
        self,
        client: LLMClient,
        tools: Dict[str, Tool],
        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
        max_steps: int = 6,
    ):
        self.client = client
        self.tools = tools
        self.system_prompt = system_prompt
        self.max_steps = max_steps

    def _execute(self, call: ToolCall) -> Step:
        tool = self.tools.get(call.name)
        if tool is None:
            return Step(call.name, call.arguments,
                        {"error": f"unknown tool {call.name!r}"}, ok=False)
        try:
            result = tool.handler(**call.arguments)
            return Step(call.name, call.arguments, result, ok=True)
        except Exception as e:  # surface the error to the model so it can recover
            return Step(call.name, call.arguments,
                        {"error": f"{type(e).__name__}: {e}"}, ok=False)

    def run(self, user_message: str, history: Optional[List[dict]] = None) -> Trace:
        messages: List[dict] = [{"role": "system", "content": self.system_prompt}]
        if history:
            messages.extend(history)
        messages.append({"role": "user", "content": user_message})

        trace = Trace(user=user_message)
        tool_schemas = [t.schema() for t in self.tools.values()]

        for _ in range(self.max_steps):
            turn: AssistantTurn = self.client.chat(messages, tool_schemas)

            if turn.is_final:
                trace.final_answer = turn.text or ""
                messages.append({"role": "assistant", "content": trace.final_answer})
                break

            # Record the assistant's tool-call intent.
            messages.append({
                "role": "assistant",
                "content": turn.text or "",
                "tool_calls": [
                    {"id": c.id or f"call_{i}", "type": "function",
                     "function": {"name": c.name, "arguments": json.dumps(c.arguments)}}
                    for i, c in enumerate(turn.tool_calls)
                ],
            })

            # Execute each call and feed results back.
            for i, call in enumerate(turn.tool_calls):
                step = self._execute(call)
                trace.steps.append(step)
                messages.append({
                    "role": "tool",
                    "tool_call_id": call.id or f"call_{i}",
                    "name": call.name,
                    "content": json.dumps(step.result, ensure_ascii=False),
                })
        else:
            trace.final_answer = (
                "I reached the step limit before finishing. Here is what I gathered; "
                "please narrow the question or consult a CPA for the rest."
            )

        trace.messages = messages
        return trace