tests/test_retrieval.py

"""Tests for the regulation retriever and the cite-or-abstain guardrail.

Run from the project root:  python -m unittest discover -s tests
"""

import unittest

from src.agent import Agent, AssistantTurn, ScriptedClient, ToolCall, ToolContext, build_tools
from src.ledger import Ledger
from src.retrieval import Document, Retriever, tokenize


def tiny_corpus():
    return [
        Document("ded-computo", "Deducción de equipo de cómputo", "LISR Art. 34", "MX", 2024,
                 "Puedes deducir tu laptop o computadora (equipo de cómputo) por "
                 "depreciación al 30% anual si es indispensable y tiene CFDI."),
        Document("iva-mensual", "Cálculo mensual del IVA", "LIVA Art. 5-D", "MX", 2024,
                 "Cada mes se resta el IVA acreditable del IVA trasladado para obtener "
                 "el IVA a cargo o el saldo a favor."),
        Document("us-se", "Self-employment tax", "IRS Schedule SE", "US", 2024,
                 "Self-employment tax is 15.3% on 92.35% of net profit for Social "
                 "Security and Medicare."),
    ]


class TestTokenize(unittest.TestCase):
    def test_deaccent_stopwords_and_stemming(self):
        toks = tokenize("¿Cómputo y depreciación de la laptop?")
        self.assertIn("computo", toks)        # accent stripped
        self.assertNotIn("la", toks)          # stopword removed
        # 'depreciación' and 'depreciaciones' must share a stem so plurals/derivations match
        self.assertEqual(tokenize("depreciación"), tokenize("depreciaciones"))

    def test_derivations_share_stem(self):
        # the flagship case: these must collapse to one stem for recall
        stems = {tokenize(w)[0] for w in ["deducir", "deducción", "deducible"]}
        self.assertEqual(len(stems), 1)


class TestRetriever(unittest.TestCase):
    def setUp(self):
        self.r = Retriever(tiny_corpus())

    def test_finds_relevant_doc(self):
        top = self.r.search("puedo deducir mi laptop")[0]
        self.assertEqual(top.source, "LISR Art. 34")

    def test_jurisdiction_filter(self):
        results = self.r.search("tax", jurisdiction="US")
        self.assertTrue(all(p.jurisdiction == "US" for p in results))

    def test_grounded_query(self):
        res = self.r.cite("cómo se calcula el IVA mensual")
        self.assertTrue(res.grounded)
        self.assertIn("5-D", res.passages[0].source)

    def test_abstains_on_unrelated(self):
        res = self.r.cite("xyzzy quux flibberty nonsense")
        self.assertFalse(res.grounded)
        self.assertIn("contador", res.message.lower())


class TestRealCorpusLoads(unittest.TestCase):
    def test_corpus_dir_loads(self):
        r = Retriever.from_corpus_dir("data/regulation")
        self.assertGreaterEqual(len(r.docs), 10)
        # The flagship "can I deduct my laptop?" question must ground correctly.
        res = r.cite("puedo deducir mi laptop o computadora")
        self.assertTrue(res.grounded)
        self.assertEqual(res.passages[0].source, "LISR Art. 34")


class TestCiteToolInAgent(unittest.TestCase):
    def test_tool_present_only_with_retriever(self):
        lg = Ledger(":memory:")
        self.assertNotIn("cite_regulation", build_tools(ToolContext(lg)))
        ctx = ToolContext(lg, retriever=Retriever(tiny_corpus()))
        self.assertIn("cite_regulation", build_tools(ctx))
        lg.close()

    def test_agent_grounds_a_rule_question(self):
        lg = Ledger(":memory:")
        ctx = ToolContext(lg, retriever=Retriever(tiny_corpus()))
        tools = build_tools(ctx)
        client = ScriptedClient([
            AssistantTurn(tool_calls=[ToolCall("cite_regulation",
                                               {"query": "puedo deducir mi laptop"})]),
            AssistantTurn(text="Sí, por depreciación al 30% (LISR Art. 34)."),
        ])
        trace = Agent(client, tools).run("¿Puedo deducir mi laptop?")
        self.assertEqual(trace.steps[0].tool, "cite_regulation")
        self.assertTrue(trace.steps[0].result["grounded"])
        self.assertEqual(trace.steps[0].result["citations"][0]["source"], "LISR Art. 34")
        lg.close()

    def test_agent_abstains_when_ungrounded(self):
        lg = Ledger(":memory:")
        ctx = ToolContext(lg, retriever=Retriever(tiny_corpus()))
        tools = build_tools(ctx)
        r = tools["cite_regulation"].handler(query="zzzqqq unrelated gibberish")
        self.assertFalse(r["grounded"])
        lg.close()


if __name__ == "__main__":
    unittest.main()