File size: 5,659 Bytes
c55ab5e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d103096
 
 
c55ab5e
 
d103096
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c55ab5e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""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