app/benchmarks/rag_observability/quality.py

"""Labelled-evaluation quality scoring for the RAG observability benchmark.

Implements exactly the "Labelled evaluation quality measures" from the plan
preamble against this benchmark's real candidate shape (`document_id` +
`chunk_index`, as `ChromaDBClient.query_observed` / `_get_chunks` actually
return them -- see `app/websockets/handlers.py::_get_chunks`). `page_hit` from
the preamble's canonical list becomes `section_hit` here: this corpus is plain
text, not paginated PDFs, and the plan preamble explicitly allows substituting
sections for pages in that case.

Two entry points, matching the task-4 brief's two example tests:

* `score_retrieval(label, candidates)` -- pure ranking-quality scoring over an
  already-successful candidate list. Knows nothing about *why* a retrieval
  produced its candidates, only whether they're relevant.
* `score_outcome(label, outcome)` -- the full labelled-quality result for one
  `RetrievalOutcome`, preserving the empty-vs-error distinction the whole
  telemetry contract exists to keep visible: a fetch that legitimately found
  nothing (`success_empty`) is never scored the same as a fetch that failed
  outright (`error_fallback`). Delegates to `score_retrieval` only on the
  success path.

Emits nothing -- pure functions over inputs already in memory, called by
`runner.py` once per measured query invocation.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from app.benchmarks.rag_observability.corpus import QueryLabel, SectionIndex
from app.observability.contracts import RetrievalOutcome


@dataclass(frozen=True)
class RetrievalQualityResult:
    """Ranking-quality result over a candidate list alone (no outcome context)."""

    document_hit: bool
    # None when the label has no expected_sections (not applicable). When not
    # None, this is an ANY-match: True if at least one candidate's section is
    # in `label.expected_sections`, not a requirement that every expected
    # section be matched.
    section_hit: bool | None
    reciprocal_rank: float


@dataclass(frozen=True)
class QualityResult:
    """The full canonical labelled-quality row for one scored query invocation."""

    document_hit: bool
    section_hit: bool | None  # any-expected-section-matched, not all -- see RetrievalQualityResult
    reciprocal_rank: float
    citation_validity: bool | None  # None when the label doesn't require a citation
    empty_result: bool
    retrieval_error: bool


def score_retrieval(label: QueryLabel, candidates: list[dict[str, Any]]) -> RetrievalQualityResult:
    """Score a ranked candidate list against `label`'s expected documents/sections.

    `reciprocal_rank` is standard MRR: 1 / rank of the first candidate whose
    `document_id` is in `label.expected_documents` (1-indexed), or 0.0 if none
    match. `section_hit` only considers candidates that *also* matched an
    expected document -- a section match on an irrelevant document is not
    meaningful (see `test_section_hit_only_counts_among_expected_document_candidates`).
    """
    document_hit = False
    reciprocal_rank = 0.0
    section_hit: bool | None = None if not label.expected_sections else False

    section_index = _section_index() if label.expected_sections else None

    for rank, candidate in enumerate(candidates, start=1):
        doc_id = candidate.get("document_id")
        if doc_id not in label.expected_documents:
            continue
        if not document_hit:
            document_hit = True
            reciprocal_rank = 1.0 / rank
        if section_index is not None:
            section = section_index.section_for_chunk(doc_id, candidate.get("chunk_index"))
            if section in label.expected_sections:
                # Any-match semantics: one hit anywhere in the ranked list is
                # enough. This never resets to False and never requires every
                # entry in `label.expected_sections` to be covered.
                section_hit = True

    return RetrievalQualityResult(
        document_hit=document_hit,
        section_hit=section_hit,
        reciprocal_rank=reciprocal_rank,
    )


def score_outcome(label: QueryLabel, outcome: RetrievalOutcome) -> QualityResult:
    """Score one full `RetrievalOutcome`, preserving empty-vs-error distinctly.

    A failed retrieval (`retrieval_error=True`) is scored as a total miss on
    every ranking measure without ever being conflated with a legitimate empty
    result -- an error and an empty result are different rows in the canonical
    contract and must stay that way here too.
    """
    if outcome.retrieval_error:
        return QualityResult(
            document_hit=False,
            section_hit=None,
            reciprocal_rank=0.0,
            citation_validity=None,
            empty_result=False,
            retrieval_error=True,
        )

    ranking = score_retrieval(label, outcome.candidates)
    citation_validity = _score_citation_validity(label, outcome, ranking)

    return QualityResult(
        document_hit=ranking.document_hit,
        section_hit=ranking.section_hit,
        reciprocal_rank=ranking.reciprocal_rank,
        citation_validity=citation_validity,
        empty_result=outcome.empty_result,
        retrieval_error=False,
    )


def _score_citation_validity(
    label: QueryLabel, outcome: RetrievalOutcome, ranking: RetrievalQualityResult
) -> bool | None:
    """Whether a citation-required query has (or correctly lacks) grounding to cite.

    Not a claim about any LLM-generated citation text -- this benchmark only
    measures retrieval, never synthesis. It is a retrieval-observable proxy:
    for an answerable query, valid means the expected document was actually
    retrieved (`document_hit`); for a deliberately unanswerable query, valid
    means retrieval came back empty -- any non-empty result would be exactly
    the retrieval-side precondition for a downstream fabricated citation.
    """
    if not label.citation_required:
        return None
    if label.answerability == "unanswerable":
        return len(outcome.candidates) == 0
    return ranking.document_hit


_section_index_cache: SectionIndex | None = None


def _section_index() -> SectionIndex:
    """Lazily build (and cache) the corpus section index on first use.

    Scoring is called many times (once per measured repetition) against the
    same fixed corpus, so this avoids re-chunking every document on every
    query. Module-level cache is safe here: the benchmark process only ever
    loads one corpus version per run.
    """
    global _section_index_cache
    if _section_index_cache is None:
        from app.benchmarks.rag_observability.corpus import load_corpus

        _section_index_cache = SectionIndex(load_corpus())
    return _section_index_cache