project/retrieval.py

"""Inspectable local retrieval baseline. No models, APIs, or network calls."""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
import json
import math
from pathlib import Path
import re


STOP_WORDS = frozenset("a an and are as at be can do for from how i in is it may of on or the to what when where which with".split())


def tokens(text: str) -> list[str]:
    """ASCII English words, lowercased; no stemming or semantic understanding."""
    return [word for word in re.findall(r"[a-z0-9]+", text.lower()) if word not in STOP_WORDS]


@dataclass(frozen=True)
class Document:
    id: str
    title: str
    text: str
    source: str
    version: str
    status: str
    audience: str


def load_documents(path: Path) -> list[Document]:
    records = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(records, list):
        raise ValueError("corpus must be a JSON list")
    documents = []
    fields = set(Document.__dataclass_fields__)
    for record in records:
        if not isinstance(record, dict) or set(record) != fields:
            raise ValueError("each document must have exactly the documented fields")
        if any(not isinstance(value, str) or not value.strip() for value in record.values()):
            raise ValueError("document fields must be non-empty strings")
        if record["status"] not in {"current", "archived"}:
            raise ValueError("unknown document status")
        if record["audience"] not in {"public", "staff"}:
            raise ValueError("unknown audience")
        documents.append(Document(**record))
    if len({document.id for document in documents}) != len(documents):
        raise ValueError("document IDs must be unique")
    return documents


@dataclass(frozen=True)
class Hit:
    document: Document
    score: float


class Index:
    """TF-IDF and cosine over current public passages only.

    This metadata policy is a toy corpus filter, not authentication or an
    authorization system. Do not place private data in this learning project.
    """

    def __init__(self, documents: list[Document]):
        if len({document.id for document in documents}) != len(documents):
            raise ValueError("document IDs must be unique")
        self.documents = tuple(document for document in documents
                               if document.status == "current" and document.audience == "public")
        bags = [Counter(tokens(document.text)) for document in self.documents]
        frequency = Counter(term for bag in bags for term in bag)
        self.idf = {term: math.log((len(bags) + 1) / (count + 1)) + 1
                    for term, count in frequency.items()}
        self.vectors = tuple(self._vector(bag) for bag in bags)

    def _vector(self, bag: Counter) -> dict[str, float]:
        weights = {term: count * self.idf[term] for term, count in bag.items() if term in self.idf}
        length = math.sqrt(sum(weight * weight for weight in weights.values()))
        return {term: weight / length for term, weight in weights.items()} if length else {}

    def search(self, question: str, k: int = 2) -> list[Hit]:
        if isinstance(k, bool) or not isinstance(k, int) or k < 1:
            raise ValueError("k must be a positive integer")
        query = self._vector(Counter(tokens(question)))
        hits = []
        for document, vector in zip(self.documents, self.vectors):
            score = sum(weight * vector.get(term, 0.0) for term, weight in query.items())
            if score > 0:
                hits.append(Hit(document, score))
        return sorted(hits, key=lambda hit: (-hit.score, hit.document.id))[:k]


def evidence_packet(hits: list[Hit], max_quote_chars: int = 1200) -> dict:
    """Copy whole passages within a quote-character budget; never summarize."""
    if isinstance(max_quote_chars, bool) or not isinstance(max_quote_chars, int) or max_quote_chars < 1:
        raise ValueError("max_quote_chars must be a positive integer")
    evidence = []
    used = 0
    for hit in hits:
        document = hit.document
        if used + len(document.text) > max_quote_chars:
            continue
        evidence.append({"id": document.id, "title": document.title, "quote": document.text,
                         "source": document.source, "version": document.version,
                         "retrieval_score": round(hit.score, 6)})
        used += len(document.text)
    return {"mode": "quote_only", "model_used": False,
            "status": "evidence_found" if evidence else "no_evidence",
            "warning": "A retrieved passage may not answer the question. A score is not confidence.",
            "quote_characters": used, "evidence": evidence}


def retrieval_metrics(ranked_ids: list[str], relevant_ids: set[str], k: int) -> dict:
    """Recall@k and reciprocal rank@k for a question with labeled evidence."""
    if isinstance(k, bool) or not isinstance(k, int) or k < 1:
        raise ValueError("k must be a positive integer")
    if not relevant_ids:
        raise ValueError("retrieval metrics need at least one relevant document")
    if len(set(ranked_ids)) != len(ranked_ids):
        raise ValueError("ranked IDs must be unique")
    top = ranked_ids[:k]
    first = next((rank for rank, doc_id in enumerate(top, 1) if doc_id in relevant_ids), None)
    return {"recall_at_k": len(set(top) & relevant_ids) / len(relevant_ids),
            "reciprocal_rank_at_k": 1 / first if first else 0.0}


def evaluate(index: Index, cases: list[dict], k: int = 2) -> dict:
    """Evaluate retrieval, separately accounting for questions with no evidence."""
    # Validate k even when cases is empty.
    index.search("", k)
    ids = set()
    eligible = {document.id for document in index.documents}
    rows = []
    for case in cases:
        if (not isinstance(case, dict) or set(case) != {"id", "question", "relevant_ids", "note"}
                or not isinstance(case["id"], str) or not case["id"]
                or not isinstance(case["question"], str) or not case["question"].strip()
                or not isinstance(case["note"], str)
                or not isinstance(case["relevant_ids"], list)
                or any(not isinstance(item, str) for item in case["relevant_ids"])):
            raise ValueError("invalid evaluation case")
        if case["id"] in ids:
            raise ValueError("evaluation case IDs must be unique")
        ids.add(case["id"])
        relevant = set(case["relevant_ids"])
        if len(relevant) != len(case["relevant_ids"]) or not relevant <= eligible:
            raise ValueError("relevance labels must be unique eligible document IDs")
        ranked = [hit.document.id for hit in index.search(case["question"], k)]
        row = {"id": case["id"], "question": case["question"], "relevant_ids": sorted(relevant),
               "retrieved_ids": ranked, "note": case["note"]}
        if relevant:
            row.update(retrieval_metrics(ranked, relevant, k))
        else:
            row["correct_retrieval_abstention"] = not ranked
        rows.append(row)
    answerable = [row for row in rows if row["relevant_ids"]]
    absent = [row for row in rows if not row["relevant_ids"]]
    return {"kind": "synthetic_offline_retrieval_evaluation", "k": k,
            "answerable_count": len(answerable), "no_evidence_count": len(absent),
            "mean_recall_at_k": (sum(row["recall_at_k"] for row in answerable) / len(answerable)
                                 if answerable else None),
            "mean_reciprocal_rank_at_k": (
                sum(row["reciprocal_rank_at_k"] for row in answerable) / len(answerable)
                if answerable else None),
            "correct_retrieval_abstentions": sum(row["correct_retrieval_abstention"] for row in absent),
            "false_evidence_returns": sum(not row["correct_retrieval_abstention"] for row in absent),
            "cases": rows}