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