apps/api/app/rag/rerank.py

"""Cross-encoder reranking, and the refusal gate that rides on its scores.

Bi-encoder vs cross-encoder — the distinction the whole stage rests on
---------------------------------------------------------------------
The retriever is a **bi-encoder**: query and passage are embedded *separately* and
compared by cosine. That is what makes it fast — every passage vector is precomputed —
and it is also its ceiling, because the passage was encoded without ever having seen the
query. "Can my employer keep my passport?" and a passage about custody of documents get
one shared summary vector each, and nuance is averaged away.

A **cross-encoder** concatenates the pair and runs them through the transformer together,
so attention operates across both. It cannot precompute anything — cost is O(candidates)
per query, not O(1) — which is precisely why it is used on 20 candidates and not on 181
chunks. Retrieve wide with the cheap model, judge narrow with the expensive one.

The output is an unbounded relevance logit, not a probability. That is useful twice: it
orders the candidates, and its absolute value is a calibrated signal of whether the
corpus contains an answer at all. :func:`apply_refusal_gate` uses the second property —
knowing when *not* to answer is the feature this project exists to demonstrate.
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass

from tokenizers import Tokenizer

from app.core.models import ScoredChunk
from app.core.settings import Settings, get_settings

logger = logging.getLogger(__name__)

_LOCK = threading.Lock()
_ENCODER: object | None = None


class RerankerUnavailableError(RuntimeError):
    """The cross-encoder could not be loaded."""


def get_cross_encoder(settings: Settings | None = None) -> object:
    """Process-wide cross-encoder.

    Two interchangeable backends. ``fastembed`` runs the cross-encoder as ONNX on CPU
    with no PyTorch in the image; ``sentence-transformers`` offers the more familiar API
    at the cost of a multi-gigabyte torch dependency. ONNX is the default because image
    size is what decides whether this fits a free Cloud Run or Hugging Face Space at all.
    Either backend serves whichever checkpoint ``reranker_model`` names.
    """
    global _ENCODER  # noqa: PLW0603 - deliberate process-wide singleton
    if _ENCODER is not None:
        return _ENCODER
    cfg = settings or get_settings()
    with _LOCK:
        if _ENCODER is None:
            cfg.models_cache_dir.mkdir(parents=True, exist_ok=True)
            logger.info("loading reranker %s via %s", cfg.reranker_model, cfg.reranker_backend)
            if cfg.reranker_backend == "fastembed":
                from fastembed.rerank.cross_encoder import TextCrossEncoder

                _ENCODER = TextCrossEncoder(
                    model_name=cfg.reranker_model,
                    cache_dir=str(cfg.models_cache_dir),
                )
            else:  # pragma: no cover - optional heavyweight backend
                try:
                    from sentence_transformers import CrossEncoder
                except ImportError as exc:
                    raise RerankerUnavailableError(
                        "reranker_backend='sentence-transformers' requires the optional "
                        "`sentence-transformers` extra; the default 'fastembed' backend "
                        "runs the same weights as ONNX with no torch dependency."
                    ) from exc
                _ENCODER = CrossEncoder(cfg.reranker_model)
    return _ENCODER


def _prepare(documents: list[str], settings: Settings) -> list[tuple[str, int]]:
    """Cap each passage at ``rerank_max_tokens`` and return it with its token length.

    The length is returned rather than recomputed later so bucketing sorts on the real
    padded cost instead of a character-count proxy.
    """
    tokenizer = get_reranker_tokenizer(settings)
    if tokenizer is None:
        return [(document, len(document)) for document in documents]
    limit = settings.rerank_max_tokens
    prepared: list[tuple[str, int]] = []
    for document in documents:
        ids = tokenizer.encode(document, add_special_tokens=False).ids
        if 0 < limit < len(ids):
            prepared.append((tokenizer.decode(ids[:limit]), limit))
        else:
            prepared.append((document, len(ids)))
    return prepared


def _score_batch(query: str, documents: list[str], settings: Settings) -> list[float]:
    encoder = get_cross_encoder(settings)
    if settings.reranker_backend == "fastembed":
        scores = encoder.rerank(query, documents, batch_size=len(documents))  # type: ignore[attr-defined]
        return [float(value) for value in scores]
    pairs = [(query, document) for document in documents]  # pragma: no cover
    return [float(value) for value in encoder.predict(pairs)]  # type: ignore[attr-defined]


def _score(query: str, documents: list[str], settings: Settings) -> list[float]:
    """Score every (query, passage) pair, length-bucketed for throughput.

    A transformer batch is padded to its longest member, so scoring one 506-token
    passage alongside nineteen 90-token ones costs as much as twenty long ones. Sorting
    by length and scoring in small batches removes that waste. The transformation is
    purely a reordering — scores are bit-identical to a single large batch, which
    ``tests/test_rerank.py`` asserts — so it buys latency and changes nothing else.
    """
    if not documents:
        return []
    prepared = _prepare(documents, settings)
    batch_size = max(1, settings.rerank_batch_size)
    order = sorted(range(len(prepared)), key=lambda i: prepared[i][1])
    scores = [0.0] * len(prepared)
    for start in range(0, len(order), batch_size):
        indices = order[start : start + batch_size]
        batch = [prepared[i][0] for i in indices]
        for index, value in zip(indices, _score_batch(query, batch, settings), strict=True):
            scores[index] = value
    return scores


def rerank(
    query: str,
    candidates: list[ScoredChunk],
    settings: Settings | None = None,
    top_k: int | None = None,
) -> list[ScoredChunk]:
    """Rescore candidates with the cross-encoder and keep the best ``top_k``."""
    cfg = settings or get_settings()
    limit = top_k if top_k is not None else cfg.rerank_top_k
    if not candidates:
        return []

    scores = _score(query, [candidate.chunk.text for candidate in candidates], cfg)
    if len(scores) != len(candidates):  # pragma: no cover - defensive
        raise RerankerUnavailableError(
            f"reranker returned {len(scores)} scores for {len(candidates)} candidates"
        )

    scored = [
        candidate.model_copy(update={"rerank_score": score})
        for candidate, score in zip(candidates, scores, strict=True)
    ]
    # Ties broken by chunk_id so the same input always yields the same output order.
    scored.sort(key=lambda item: (-(item.rerank_score or 0.0), item.chunk.chunk_id))
    return [
        item.model_copy(update={"final_rank": position + 1})
        for position, item in enumerate(scored[:limit])
    ]


def passthrough(
    candidates: list[ScoredChunk],
    settings: Settings | None = None,
    top_k: int | None = None,
) -> list[ScoredChunk]:
    """Take the top fused candidates without reranking.

    This is the ``--no-rerank`` arm of the evaluation. It exists so the reranker's
    contribution is a measured delta rather than an assertion.
    """
    cfg = settings or get_settings()
    limit = top_k if top_k is not None else cfg.rerank_top_k
    return [
        candidate.model_copy(update={"final_rank": position + 1})
        for position, candidate in enumerate(candidates[:limit])
    ]


@dataclass(frozen=True, slots=True)
class GateOutcome:
    """Whether the corpus covers the question, and the evidence either way."""

    covered: bool
    evidence: tuple[ScoredChunk, ...]
    near_misses: tuple[ScoredChunk, ...]
    best_score: float | None
    floor: float
    best_dense: float | None = None
    dense_floor: float | None = None
    reason: str = ""
    signal: str = ""


def apply_refusal_gate(
    reranked: list[ScoredChunk],
    settings: Settings | None = None,
    *,
    best_dense: float | None = None,
    scope: object | None = None,
) -> GateOutcome:
    """Decide whether the retrieved evidence is good enough to answer from.

    Retrieval always returns *something*: nearest-neighbour search over a non-empty index
    cannot return nothing, and a question the corpus has never heard of still comes back
    with five confidently-ranked passages. Generating from them is exactly how a RAG
    system produces a fluent, well-cited, wrong answer.

    Three independent signals must all pass, because each catches what the others miss —
    all three thresholds were fitted on the labelled eval set, not chosen by intuition:

    1. **Scope.** Does the question name a legal system the corpus does not contain?
       No similarity score can answer this; see ``app.rag.scope``.
    2. **Domain floor** on the best dense similarity. Answers "is this question even
       about the corpus's subject matter?"
    3. **Relevance floor** on the best cross-encoder score. Answers "is the single best
       passage actually responsive?"

    A refusal hands back the near misses, so it is auditable: the user sees what was
    considered and can judge the call themselves.

    With reranking disabled the cross-encoder signal does not exist, so only the scope
    and domain checks apply. AUDIT.md reports refusal accuracy per configuration rather
    than implying the un-reranked path is equally protected — measured, it is not.
    """
    cfg = settings or get_settings()
    if scope is not None:
        reason = getattr(scope, "reason", "")
        signal = getattr(scope, "signal", "out-of-scope")
        return GateOutcome(
            covered=False,
            evidence=(),
            near_misses=tuple(reranked[: cfg.refusal_near_miss_count]),
            best_score=reranked[0].rerank_score if reranked else None,
            floor=cfg.refusal_score_floor,
            best_dense=best_dense,
            dense_floor=cfg.refusal_dense_floor,
            reason=reason,
            signal=signal,
        )

    if not reranked:
        return GateOutcome(
            covered=False,
            evidence=(),
            near_misses=(),
            best_score=None,
            floor=cfg.refusal_score_floor,
            best_dense=best_dense,
            dense_floor=cfg.refusal_dense_floor,
            reason="Retrieval returned no candidates.",
            signal="empty-retrieval",
        )

    if best_dense is not None and best_dense < cfg.refusal_dense_floor:
        return GateOutcome(
            covered=False,
            evidence=(),
            near_misses=tuple(reranked[: cfg.refusal_near_miss_count]),
            best_score=reranked[0].rerank_score,
            floor=cfg.refusal_score_floor,
            best_dense=best_dense,
            dense_floor=cfg.refusal_dense_floor,
            reason="No indexed passage is close enough to this question's subject matter.",
            signal="below-domain-floor",
        )

    best = reranked[0].rerank_score
    if best is None or best >= cfg.refusal_score_floor:
        return GateOutcome(
            covered=True,
            evidence=tuple(reranked),
            near_misses=(),
            best_score=best,
            floor=cfg.refusal_score_floor,
            best_dense=best_dense,
            dense_floor=cfg.refusal_dense_floor,
        )
    return GateOutcome(
        covered=False,
        evidence=(),
        near_misses=tuple(reranked[: cfg.refusal_near_miss_count]),
        best_score=best,
        floor=cfg.refusal_score_floor,
        best_dense=best_dense,
        dense_floor=cfg.refusal_dense_floor,
        reason="The closest passage is not responsive enough to answer from.",
        signal="below-relevance-floor",
    )


_TOKENIZERS: dict[str, Tokenizer] = {}


def get_reranker_tokenizer(settings: Settings | None = None) -> Tokenizer | None:
    """The reranker's own tokenizer, used to cap passage length before scoring.

    Two details here are load-bearing, and both were found by a CI failure that could
    not be reproduced locally:

    1. **The cross-encoder is loaded first.** The tokenizer is located by globbing the
       model cache, so on a cold cache — a fresh container, a CI runner with no restored
       cache — the file does not exist yet because nothing has downloaded it. Forcing the
       encoder to load first guarantees the files are on disk before the search.

    2. **A miss is never cached.** With ``lru_cache`` a single early miss was memoised for
       the lifetime of the process, so truncation stayed silently disabled long after the
       model had arrived. Only successful lookups are cached.

    Why it matters that truncation actually happens: without it, passages exceed the
    model's window and the runtime truncates per batch, which makes a score depend on
    which other documents happen to share its batch. Length bucketing then stops being a
    pure reordering. With truncation on, scores are identical across batch sizes —
    measured at max |delta| 0.0000 over the corpus.
    """
    cfg = settings or get_settings()
    cached = _TOKENIZERS.get(cfg.reranker_model)
    if cached is not None:
        return cached

    # Ensure the model — and therefore its tokenizer.json — is on disk before looking.
    get_cross_encoder(cfg)

    stem = cfg.reranker_model.split("/")[-1]
    candidates = [
        path for path in cfg.models_cache_dir.rglob("tokenizer.json") if stem in str(path)
    ]
    if not candidates:
        logger.warning(
            "no cached tokenizer for %s under %s; reranking without truncation, which "
            "makes scores batch-dependent",
            cfg.reranker_model,
            cfg.models_cache_dir,
        )
        return None

    tokenizer = Tokenizer.from_file(str(max(candidates, key=lambda path: path.stat().st_mtime)))
    _TOKENIZERS[cfg.reranker_model] = tokenizer
    return tokenizer


def reset_cross_encoder() -> None:
    """Drop the cached encoder. Used by tests."""
    global _ENCODER  # noqa: PLW0603 - mirrors get_cross_encoder
    with _LOCK:
        _ENCODER = None
    _TOKENIZERS.clear()