src/retrieval/index.py

"""Local retriever over the regulation corpus — no cloud API, no downloads.

A pure-Python TF-IDF cosine retriever. It is deterministic, offline, and good enough
to ground tax answers in the right article. The ``Retriever`` interface leaves room
for a heavier embedding backend later (sentence-transformers), but the default
requires nothing beyond the standard library — which keeps the 🔌 Off-the-Grid badge
honest.

The cite-or-abstain rule lives here: if the best match scores below a threshold, the
agent must NOT answer from the corpus — it abstains and recommends a CPA.
"""

from __future__ import annotations

import math
import re
import unicodedata
from collections import Counter
from dataclasses import dataclass
from typing import Dict, List, Optional

from .corpus import Document, load_corpus

# Small Spanish/English stopword set — enough to stop common words dominating.
_STOPWORDS = {
    "de", "la", "el", "los", "las", "un", "una", "y", "o", "en", "que", "del",
    "se", "su", "por", "para", "con", "es", "al", "lo", "como", "mas", "si",
    "no", "the", "a", "an", "of", "to", "in", "is", "for", "on", "and", "or",
    "be", "as", "it", "this", "that", "are", "i", "my", "me", "do", "can",
}


def _deaccent(s: str) -> str:
    return "".join(c for c in unicodedata.normalize("NFD", s)
                    if unicodedata.category(c) != "Mn")


# Derivational suffixes, longest first, so 'deducción' / 'deducible' / 'deducir'
# all reduce to the same stem 'deduc'. Conservative — only stripped when a stem of
# length >= 3 remains.
_SUFFIXES = ("aciones", "iciones", "acion", "icion", "ciones", "cion", "ibles",
             "ables", "ible", "able", "mente", "idades", "idad", "ir", "ar", "er")


def _stem(tok: str) -> str:
    # plurals
    if len(tok) > 4 and tok.endswith("es"):
        tok = tok[:-2]
    elif len(tok) > 3 and tok.endswith("s"):
        tok = tok[:-1]
    # one derivational suffix
    for suf in _SUFFIXES:
        if tok.endswith(suf) and len(tok) - len(suf) >= 3:
            return tok[: -len(suf)]
    return tok


def tokenize(text: str) -> List[str]:
    text = _deaccent(text.lower())
    return [_stem(t) for t in re.findall(r"[a-z0-9]+", text)
            if len(t) >= 2 and t not in _STOPWORDS]


@dataclass
class Passage:
    doc_id: str
    title: str
    source: str
    jurisdiction: str
    year: Optional[int]
    text: str
    score: float
    matched_terms: int = 0  # distinct query terms found in the passage


@dataclass
class CiteResult:
    grounded: bool
    query: str
    passages: List[Passage]
    message: str


# Below this cosine score, we treat the corpus as not supporting the query.
DEFAULT_MIN_SCORE = 0.08


class Retriever:
    def __init__(self, documents: List[Document]):
        self.docs = documents
        self._idf: Dict[str, float] = {}
        self._vectors: List[Dict[str, float]] = []
        self._build()

    @classmethod
    def from_corpus_dir(cls, path) -> "Retriever":
        return cls(load_corpus(path))

    def _build(self) -> None:
        tokenized = [tokenize(f"{d.title} {d.text}") for d in self.docs]
        n = len(self.docs)
        df: Counter = Counter()
        for toks in tokenized:
            df.update(set(toks))
        self._idf = {term: math.log((n + 1) / (freq + 1)) + 1
                     for term, freq in df.items()}
        self._vectors = [self._vectorize(toks) for toks in tokenized]

    def _vectorize(self, tokens: List[str]) -> Dict[str, float]:
        tf = Counter(tokens)
        vec = {t: c * self._idf.get(t, 0.0) for t, c in tf.items()}
        norm = math.sqrt(sum(v * v for v in vec.values())) or 1.0
        return {t: v / norm for t, v in vec.items()}

    def _query_vector(self, query: str) -> Dict[str, float]:
        tf = Counter(tokenize(query))
        vec = {t: c * self._idf.get(t, 0.0) for t, c in tf.items() if t in self._idf}
        norm = math.sqrt(sum(v * v for v in vec.values())) or 1.0
        return {t: v / norm for t, v in vec.items()}

    def search(self, query: str, k: int = 3,
               jurisdiction: Optional[str] = None) -> List[Passage]:
        qv = self._query_vector(query)
        q_terms = set(qv)
        scored: List[Passage] = []
        for doc, dv in zip(self.docs, self._vectors):
            if jurisdiction and doc.jurisdiction != jurisdiction:
                continue
            score = sum(w * dv.get(t, 0.0) for t, w in qv.items())
            if score <= 0:
                continue
            matched = len(q_terms & set(dv))
            scored.append(Passage(doc.id, doc.title, doc.source, doc.jurisdiction,
                                  doc.year, doc.text, round(score, 4), matched))
        scored.sort(key=lambda p: p.score, reverse=True)
        return scored[:k]

    def cite(self, query: str, k: int = 3, min_score: float = DEFAULT_MIN_SCORE,
             jurisdiction: Optional[str] = None) -> CiteResult:
        """Retrieve with the cite-or-abstain guardrail.

        A match must clear the score threshold AND share at least two distinct terms
        with the passage (one term, for single-word queries). The two-term rule kills
        false positives where the query overlaps a corpus doc only on one generic word
        (e.g. 'crédito' in 'crédito fiscal' matching 'tarjeta de crédito').
        """
        passages = self.search(query, k=k, jurisdiction=jurisdiction)
        min_terms = 2 if len(set(tokenize(query))) >= 2 else 1
        top_ok = bool(passages) and (
            passages[0].score >= min_score and passages[0].matched_terms >= min_terms
        )
        if not top_ok:
            return CiteResult(
                grounded=False, query=query, passages=passages,
                message=("No tengo una fuente en mi base de regulación que respalde "
                         "esto con confianza. Te recomiendo confirmarlo con un "
                         "contador público / CPA."),
            )
        cites = ", ".join(f"{p.source}" for p in passages)
        return CiteResult(
            grounded=True, query=query, passages=passages,
            message=f"Respaldado por: {cites}.",
        )