scripts/demo_grounding.py

"""Cite-or-abstain in action: the agent grounds rule questions in real articles,
and refuses to answer when the corpus doesn't support a claim.

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

import sys

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
from src.retrieval import Retriever


def rule(title):
    print("\n" + "═" * 64 + f"\n{title}\n" + "═" * 64)


def main():
    ctx = ToolContext(Ledger(":memory:"),
                      retriever=Retriever.from_corpus_dir("data/regulation"))
    tools = build_tools(ctx)

    # 1) A groundable question — the agent cites the article.
    rule("Q: ¿Puedo deducir mi laptop?")
    agent = Agent(build_client_grounded(), tools)
    trace = agent.run("¿Puedo deducir mi laptop nueva?")
    cites = trace.steps[0].result["citations"]
    print(f"grounded = {trace.steps[0].result['grounded']}")
    print(f"top citation = {cites[0]['source']} — {cites[0]['title']}")
    print(f"excerpt: {cites[0]['excerpt'][:160]}...")
    print("\nFINAL:\n", trace.final_answer)

    # 2) An out-of-scope question — the agent abstains, no invented law.
    rule("Q: ¿Cuál es el crédito fiscal por instalar paneles solares en 2024?")
    res = tools["cite_regulation"].handler(
        query="crédito fiscal por instalar paneles solares incentivo energía")
    print(f"grounded = {res['grounded']}")
    print("agent says:", res["message"])

    ctx.ledger.close()
    print("\nThe agent only speaks tax law it can cite — everything else it defers to a CPA.")


def build_client_grounded():
    return ScriptedClient([
        AssistantTurn(tool_calls=[ToolCall("cite_regulation",
                                           {"query": "deducir laptop computadora equipo de cómputo",
                                            "jurisdiction": "MX"})]),
        AssistantTurn(text=(
            "Sí. Una laptop usada en tu actividad es deducible, pero como INVERSIÓN: "
            "se deduce por depreciación a una tasa máxima de 30% anual, no como gasto "
            "del mes. Requiere CFDI y ser estrictamente indispensable (LISR Art. 34). "
            "Confirma el tratamiento con tu contador.")),
    ])


if __name__ == "__main__":
    main()