src/agent/tools.py

"""Typed tools — the only things the model is allowed to do.

Each tool is a thin, schema-described wrapper over the deterministic engine or the
ledger. The model's job is to pick a tool and fill its arguments; the tool does the
work and returns a JSON-serializable result that *includes the breakdown and the
source*, so the model can explain without recomputing.

This module is the product surface. Adding a capability = adding a tool here.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Dict, List, Optional

from ..engine import (
    break_even_units,
    cash_runway_months,
    current_ratio,
    federal_income_tax,
    isr_provisional_monthly,
    iva_monthly,
    profit_margin,
    quarterly_estimated_tax,
    resico_isr_monthly,
    schedule_c_net_profit,
    self_employment_tax,
    taxable_base,
    us_annual_estimate,
)
from ..engine.result import CalcResult
from ..finetune import RuleClassifier
from ..finetune.classifier import Classification
from ..ledger import Ledger, Line
from ..retrieval import Retriever


# --- serialization --------------------------------------------------------
def calc_to_dict(r: CalcResult) -> dict:
    """Render a CalcResult as JSON-safe data (Decimals → strings)."""
    return {
        "amount": str(r.amount),
        "label": r.label,
        "breakdown": [[desc, str(val)] for desc, val in r.breakdown],
        "source": r.source,
        "effective_year": r.effective_year,
        "notes": list(r.notes),
    }


def _us_estimate_dict(est: dict) -> dict:
    """Serialize a us_annual_estimate result (mix of CalcResults and Decimals)."""
    out = {}
    for k, v in est.items():
        out[k] = calc_to_dict(v) if isinstance(v, CalcResult) else str(v)
    return out


@dataclass
class ToolContext:
    """Per-session state the tools operate on."""

    ledger: Ledger
    country: str = "MX"
    regime: str = "RESICO"
    retriever: Optional[Retriever] = None  # enables the cite_regulation tool
    classifier: Optional[object] = None    # SAT classifier; defaults to RuleClassifier


@dataclass
class Tool:
    name: str
    description: str
    parameters: dict  # JSON schema for the arguments
    handler: Callable[..., dict]

    def schema(self) -> dict:
        """OpenAI / llama.cpp function-calling schema."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


# --- JSON-schema shorthands ----------------------------------------------
def _obj(properties: dict, required: List[str]) -> dict:
    return {"type": "object", "properties": properties, "required": required}


_NUM = lambda d: {"type": "number", "description": d}  # noqa: E731
_INT = lambda d: {"type": "integer", "description": d}  # noqa: E731
_STR = lambda d: {"type": "string", "description": d}  # noqa: E731
_BOOL = lambda d: {"type": "boolean", "description": d}  # noqa: E731


# --- handlers -------------------------------------------------------------
def build_tools(ctx: ToolContext) -> Dict[str, Tool]:
    """Construct the tool set bound to a session context."""

    lg = ctx.ledger
    classifier = ctx.classifier or RuleClassifier()

    def classify_transaction(description):
        c: Classification = classifier.classify(description)
        return c.to_dict()

    def record_expense_auto(date, description, amount, cfdi_uuid=None):
        """Classify a description and book it to the correct SAT account."""
        c: Classification = classifier.classify(description)
        iva_rate = c.iva_tasa
        if c.kind == "investment":
            # Investments are assets, deducted via depreciation — not a monthly expense.
            lg.ensure_account(c.sat_code, c.cuenta, "asset", "debit")
            from ..engine.money import money as _m
            base = _m(amount)
            iva = _m(base * __import__("decimal").Decimal(iva_rate))
            lines = [Line(c.sat_code, debit=base)]
            if iva > 0:
                lines.append(Line(lg.code_for_role("iva_acreditable"), debit=iva))
            lines.append(Line(lg.code_for_role("bank"), credit=_m(base + iva)))
            tx = lg.post(date, description, lines, kind="investment",
                         iva_treatment=c.iva_tratamiento, iva_rate=iva_rate,
                         cfdi_uuid=cfdi_uuid)
            note = "Booked as an asset — deduct via depreciation (e.g. 30%/yr for cómputo)."
        else:
            lg.ensure_account(c.sat_code, c.cuenta, "expense", "debit")
            tx = lg.record_expense(date, description, amount, iva_rate=iva_rate,
                                   expense_code=c.sat_code, deductible=c.deducible,
                                   cfdi_uuid=cfdi_uuid)
            note = ("Deducible." if c.deducible else "No deducible.")
            if c.deducible_ratio < 1.0:
                note = f"Deducible solo al {c.deducible_ratio*100:.1f}%."
        return {"status": "booked", "transaction_id": tx,
                "classification": c.to_dict(), "note": note}

    def record_income(date, description, amount, iva_rate=0.16,
                      isr_retenido=0, iva_retenido=0, cfdi_uuid=None):
        tx = lg.record_income(date, description, amount, iva_rate=iva_rate,
                              isr_retenido=isr_retenido, iva_retenido=iva_retenido,
                              cfdi_uuid=cfdi_uuid)
        return {"status": "booked", "transaction_id": tx,
                "message": f"Income of {amount} booked on {date}."}

    def record_expense(date, description, amount, iva_rate=0.16,
                       deductible=True, cfdi_uuid=None):
        tx = lg.record_expense(date, description, amount, iva_rate=iva_rate,
                               deductible=deductible, cfdi_uuid=cfdi_uuid)
        return {"status": "booked", "transaction_id": tx,
                "deductible": deductible,
                "message": f"Expense of {amount} booked on {date}."}

    def month_summary(year, month):
        m = lg.month_totals(year, month)
        return {
            "year": year, "month": month,
            "income": str(m.income),
            "deductible_expenses": str(m.deductible_expenses),
            "nondeductible_expenses": str(m.nondeductible_expenses),
            "iva_trasladado": str(m.iva_trasladado),
            "iva_acreditable": str(m.iva_acreditable),
            "iva_retenido": str(m.iva_retenido),
            "isr_retenido": str(m.isr_retenido),
        }

    def compute_isr_resico(year, month):
        m = lg.month_totals(year, month)
        return {"income": str(m.income), **calc_to_dict(resico_isr_monthly(m.income))}

    def compute_isr_general(year, month):
        m = lg.month_totals(year, month)
        base = taxable_base(m.income, m.deductible_expenses)
        return {"taxable_base": str(base), **calc_to_dict(isr_provisional_monthly(base))}

    def compute_iva(year, month):
        m = lg.month_totals(year, month)
        return calc_to_dict(iva_monthly(m.iva_trasladado, m.iva_acreditable, m.iva_retenido))

    def compare_regimes(year, month):
        m = lg.month_totals(year, month)
        resico = resico_isr_monthly(m.income)
        base = taxable_base(m.income, m.deductible_expenses)
        general = isr_provisional_monthly(base)
        cheaper = "RESICO" if resico.amount <= general.amount else "GENERAL"
        savings = abs(resico.amount - general.amount)
        return {
            "resico_isr": str(resico.amount),
            "general_isr": str(general.amount),
            "general_taxable_base": str(base),
            "recommended": cheaper,
            "monthly_savings": str(savings),
            "note": f"{cheaper} is cheaper this month by {savings}. "
                    "Regime choice has eligibility rules and is annual — confirm with a CPA.",
        }

    def income_statement(year, month=None):
        s = lg.income_statement(year, month)
        return {"period": s.period, "revenue": str(s.revenue),
                "expenses": str(s.expenses), "net_profit": str(s.net_profit)}

    def balance_sheet(as_of=None):
        bs = lg.balance_sheet(as_of)
        return {"as_of": bs.as_of, "assets": str(bs.assets),
                "liabilities": str(bs.liabilities), "equity": str(bs.equity)}

    def margin(year, month=None):
        s = lg.income_statement(year, month)
        return calc_to_dict(profit_margin(s.net_profit, s.revenue))

    def break_even(fixed_costs, price_per_unit, variable_cost_per_unit):
        return calc_to_dict(break_even_units(fixed_costs, price_per_unit, variable_cost_per_unit))

    def runway(cash_on_hand, monthly_net_burn):
        return calc_to_dict(cash_runway_months(cash_on_hand, monthly_net_burn))

    def liquidity(current_assets, current_liabilities):
        return calc_to_dict(current_ratio(current_assets, current_liabilities))

    def cite_regulation(query, jurisdiction=None):
        res = ctx.retriever.cite(query, jurisdiction=jurisdiction)
        return {
            "grounded": res.grounded,
            "message": res.message,
            "citations": [
                {"source": p.source, "title": p.title, "year": p.year,
                 "excerpt": p.text[:400], "score": p.score}
                for p in res.passages
            ],
        }

    def us_tax_estimate(gross_income, expenses):
        est = us_annual_estimate(gross_income, expenses)
        return _us_estimate_dict(est)

    def us_tax_summary(year):
        s = lg.income_statement(int(year))  # annual revenue & expenses from the ledger
        est = us_annual_estimate(s.revenue, s.expenses)
        out = _us_estimate_dict(est)
        out["year"] = int(year)
        out["gross_income"] = str(s.revenue)
        out["business_expenses"] = str(s.expenses)
        return out

    tools = [
        Tool("record_income", "Book an income / sale into the ledger (amount is the "
             "pre-IVA subtotal). Use for any money the user earned or invoiced.",
             _obj({
                 "date": _STR("ISO date YYYY-MM-DD"),
                 "description": _STR("what the income was for"),
                 "amount": _NUM("taxable base, pre-IVA subtotal"),
                 "iva_rate": _NUM("IVA rate as a fraction, e.g. 0.16; use 0 for zero-rated"),
                 "isr_retenido": _NUM("ISR withheld by the client, if any"),
                 "iva_retenido": _NUM("IVA withheld by the client, if any"),
                 "cfdi_uuid": _STR("CFDI folio fiscal, if available"),
             }, ["date", "description", "amount"]),
             record_income),
        Tool("record_expense", "Book an expense / purchase into the ledger (amount is "
             "the pre-IVA subtotal). Set deductible=false for non-deductible outlays.",
             _obj({
                 "date": _STR("ISO date YYYY-MM-DD"),
                 "description": _STR("what the expense was for"),
                 "amount": _NUM("taxable base, pre-IVA subtotal"),
                 "iva_rate": _NUM("IVA rate as a fraction, e.g. 0.16; 0 for zero-rated"),
                 "deductible": _BOOL("whether the expense is tax-deductible"),
                 "cfdi_uuid": _STR("CFDI folio fiscal, if available"),
             }, ["date", "description", "amount"]),
             record_expense),
        Tool("classify_transaction", "Classify a free-text transaction description into "
             "its SAT account, deductibility and IVA treatment. Use before booking when "
             "the category is unclear, or to explain how something would be categorized.",
             _obj({"description": _STR("the transaction description / receipt text")},
                  ["description"]),
             classify_transaction),
        Tool("record_expense_auto", "Classify a transaction description AND book it to the "
             "correct SAT account in one step (handles deductibility and IVA, routes "
             "investments to an asset account). Prefer this for quick expense capture.",
             _obj({"date": _STR("ISO date YYYY-MM-DD"),
                   "description": _STR("what the expense was for"),
                   "amount": _NUM("taxable base, pre-IVA subtotal"),
                   "cfdi_uuid": _STR("CFDI folio fiscal, if available")},
                  ["date", "description", "amount"]),
             record_expense_auto),
        Tool("month_summary", "Get the aggregated income, expenses and IVA figures for "
             "a given month from the ledger.",
             _obj({"year": _INT("year"), "month": _INT("month 1-12")}, ["year", "month"]),
             month_summary),
        Tool("compute_isr_resico", "Compute the monthly ISR under the RESICO regime from "
             "the ledger.",
             _obj({"year": _INT("year"), "month": _INT("month 1-12")}, ["year", "month"]),
             compute_isr_resico),
        Tool("compute_isr_general", "Compute the monthly provisional ISR under the general "
             "regime (income minus deductions) from the ledger.",
             _obj({"year": _INT("year"), "month": _INT("month 1-12")}, ["year", "month"]),
             compute_isr_general),
        Tool("compute_iva", "Compute the monthly net IVA position (a cargo or a favor) "
             "from the ledger.",
             _obj({"year": _INT("year"), "month": _INT("month 1-12")}, ["year", "month"]),
             compute_iva),
        Tool("compare_regimes", "Compare RESICO vs general-regime ISR for a month and "
             "recommend the cheaper one.",
             _obj({"year": _INT("year"), "month": _INT("month 1-12")}, ["year", "month"]),
             compare_regimes),
        Tool("income_statement", "Profit & loss for a month (omit month for the whole year).",
             _obj({"year": _INT("year"), "month": _INT("optional month 1-12")}, ["year"]),
             income_statement),
        Tool("balance_sheet", "Assets, liabilities and equity as of a date (omit for latest).",
             _obj({"as_of": _STR("optional ISO date YYYY-MM-DD")}, []),
             balance_sheet),
        Tool("profit_margin", "Net profit margin for a period from the ledger.",
             _obj({"year": _INT("year"), "month": _INT("optional month 1-12")}, ["year"]),
             margin),
        Tool("break_even", "Units to sell to break even given fixed costs, price and "
             "variable cost per unit.",
             _obj({"fixed_costs": _NUM("total fixed costs"),
                   "price_per_unit": _NUM("selling price per unit"),
                   "variable_cost_per_unit": _NUM("variable cost per unit")},
                  ["fixed_costs", "price_per_unit", "variable_cost_per_unit"]),
             break_even),
        Tool("cash_runway", "How many months the cash lasts at the current burn rate.",
             _obj({"cash_on_hand": _NUM("cash available"),
                   "monthly_net_burn": _NUM("net cash burned per month")},
                  ["cash_on_hand", "monthly_net_burn"]),
             runway),
        Tool("current_ratio", "Liquidity ratio = current assets / current liabilities.",
             _obj({"current_assets": _NUM("current assets"),
                   "current_liabilities": _NUM("current liabilities")},
                  ["current_assets", "current_liabilities"]),
             liquidity),
        Tool("us_tax_estimate", "USA self-employed estimate from explicit numbers: "
             "Schedule-C net profit, SE tax, federal income tax (after standard deduction "
             "and half-SE), and quarterly estimated payment.",
             _obj({"gross_income": _NUM("gross US income"),
                   "expenses": _NUM("business expenses")},
                  ["gross_income", "expenses"]),
             us_tax_estimate),
        Tool("us_tax_summary", "USA self-employed tax estimate for a year, computed from "
             "the ledger's annual income and expenses. Use for US users' tax questions.",
             _obj({"year": _INT("year")}, ["year"]),
             us_tax_summary),
    ]

    # Grounding tool — only available when a regulation retriever is configured.
    if ctx.retriever is not None:
        tools.append(Tool(
            "cite_regulation",
            "Look up the SAT / IRS regulation that supports a tax question and return "
            "cited passages. Use this for any rule question (deductibility, rates, "
            "obligations). If it returns grounded=false, do NOT answer the rule from "
            "memory — tell the user to confirm with a CPA.",
            _obj({"query": _STR("the tax-rule question, in the user's words"),
                  "jurisdiction": _STR("optional filter: 'MX' or 'US'")},
                 ["query"]),
            cite_regulation))

    return {t.name: t for t in tools}