scripts/demo_agent.py

"""Drive the full agent loop end-to-end and save a shareable trace.

No weights needed: a ScriptedClient stands in for the fine-tuned model, replaying the
tool calls it would make. This proves the loop — plan → call tools → observe →
answer — and writes the trace to data/sample/trace_example.json for the Hub (📡).

Run from the project root:  python -m scripts.demo_agent
"""

import sys
from pathlib import Path

try:
    sys.stdout.reconfigure(encoding="utf-8")
except (AttributeError, ValueError):
    pass

from src.agent import Agent, AssistantTurn, ScriptedClient, ToolCall, ToolContext, build_tools
from src.ledger import Ledger


def seed_may():
    lg = Ledger(":memory:")
    lg.record_income("2024-05-03", "Branding — Café Luna", 18000, iva_rate="0.16")
    lg.record_expense("2024-05-05", "Adobe Creative Cloud", 1300, iva_rate="0.16")
    lg.record_expense("2024-05-12", "Renta de oficina (parte)", 4000, iva_rate="0.16")
    lg.record_expense("2024-05-22", "Comida no deducible", 600, iva_rate="0.16",
                      deductible=False)
    return lg


def main():
    ctx = ToolContext(seed_may())
    tools = build_tools(ctx)

    user = ("Registra un ingreso de 27000 por un sitio web el 17 de mayo de 2024, "
            "con retención de ISR de 2700 e IVA de 2880. Luego dime qué régimen me "
            "conviene y cuánto IVA debo en mayo.")

    # What the fine-tuned model is expected to do, replayed deterministically:
    client = ScriptedClient([
        AssistantTurn(tool_calls=[ToolCall("record_income", {
            "date": "2024-05-17", "description": "Sitio web — Dental MX",
            "amount": 27000, "iva_rate": 0.16,
            "isr_retenido": 2700, "iva_retenido": 2880})]),
        AssistantTurn(tool_calls=[
            ToolCall("compare_regimes", {"year": 2024, "month": 5}),
            ToolCall("compute_iva", {"year": 2024, "month": 5}),
        ]),
        AssistantTurn(text=(
            "Listo, registré el ingreso de $27,000.\n"
            "• Régimen: te conviene RESICO — pagarías $495.00 de ISR vs $6,994.74 en "
            "el régimen general (fuente: tarifas SAT 2024).\n"
            "• IVA de mayo: $3,472.00 a cargo (IVA cobrado $7,200 − acreditable $848 − "
            "retenido por clientes $2,880; LIVA Art. 5-D).\n"
            "Recuerda: la elección de régimen es anual y tiene requisitos — confírmalo "
            "con tu contador.")),
    ])

    agent = Agent(client, tools)
    trace = agent.run(user)

    print("USER:\n ", user, "\n")
    print("TOOL CALLS:")
    for s in trace.steps:
        flag = "ok" if s.ok else "ERR"
        head = s.result.get("amount") or s.result.get("recommended") or s.result.get("status") or ""
        print(f"  [{flag}] {s.tool}({s.arguments}) -> {head}")
    print("\nFINAL ANSWER:\n", trace.final_answer)

    out = Path("data/sample/trace_example.json")
    out.write_text(trace.to_json(), encoding="utf-8")
    print(f"\nTrace saved → {out}  ({len(trace.steps)} tool calls)")

    ctx.ledger.close()


if __name__ == "__main__":
    main()