Retrieval-Augmented Generation

"""Retrieval-augmented generation, transcribed from the objective.

No ML library. The marginal written out so the equation and the practice
can be read against each other:

    p(y | x) = sum_z p_eta(z | x) * p_theta(y | x, z),   z in top-k(x)

Two things to read for.

First, the sum ranges over the RETRIEVED set only. Every document outside
it carries probability exactly zero, so a retrieval miss is not merely
unlikely -- it is unrecoverable, and the generator will answer fluently
anyway. That is the whole reason retrieval recall bounds the system.

Second, marginalize_answer() below implements that equation, and
build_prompt() implements what production actually does: concatenate the
top documents into one context and run a single forward pass. They are
not the same computation. The concatenation contains no marginalization
at all -- the retriever chose which chunks appear and then stopped
mattering, and the weighting has quietly become string formatting.
"""

import math
import re


def tokenize(text):
    return re.findall(r"[a-z0-9]+", text.lower())


def stable_hash(token):
    """Deterministic across processes, unlike the built-in hash().

    Not a detail. An index built in one process and queried from another
    must agree about the embedding, and Python salts hash() per process,
    so the built-in silently produces two incompatible vector spaces. The
    production form of this bug is upgrading the encoder without
    re-indexing: same shape, same silence, no error anywhere.
    """
    digest = 2166136261
    for character in token:
        digest = ((digest ^ ord(character)) * 16777619) % (1 << 32)
    return digest


def chunk_document(text, chunk_tokens, overlap_tokens):
    """Split a document into overlapping windows.

    The most consequential function in this file and the one nobody
    tunes. A chunk smaller than the answer leaves that answer retrievable
    only as two poorly-scoring halves. A chunk much larger than the
    answer dilutes its embedding with unrelated text until it stops
    scoring at all. Neither failure appears in the output -- the system
    simply cannot answer a question it looks like it should be able to.

    The overlap exists for answers that straddle a boundary, and it costs
    diversity: two overlapping chunks that both match will crowd a third
    document out of a context window that only holds a handful.
    """
    if chunk_tokens <= 0:
        raise ValueError("chunk_tokens must be positive")
    if overlap_tokens < 0 or overlap_tokens >= chunk_tokens:
        raise ValueError("overlap must be in [0, chunk_tokens)")

    words = tokenize(text)
    stride = chunk_tokens - overlap_tokens
    chunks = []
    start = 0
    while start < len(words):
        window = words[start : start + chunk_tokens]
        if not window:
            break
        chunks.append(window)
        start += stride
    return chunks


def embed(tokens, width):
    """A hashing encoder standing in for a trained one.

    Deliberately crude, because this entry is about the pipeline and
    swapping this for a real sentence encoder changes retrieval quality
    while changing nothing about the reasoning below it. What it does
    reproduce faithfully is the weakness that matters: a hashed bag of
    words has no notion of paraphrase. That is exactly the gap a dense
    encoder closes, and exactly why the lexical index below still earns
    its place after you close it.
    """
    vector = [0.0] * width
    for token in tokens:
        vector[stable_hash(token) % width] += 1.0

    norm = math.sqrt(sum(value * value for value in vector))
    if norm < 1e-12:
        return vector
    return [value / norm for value in vector]


def cosine(left, right):
    return sum(left[idx] * right[idx] for idx in range(len(left)))


def build_index(documents, chunk_tokens, overlap_tokens, width):
    """Chunk, embed, and keep the tokens alongside the vector.

    Keeping the tokens is not redundancy. The lexical retriever needs
    them, the reranker needs them, and the prompt needs them -- a store
    that holds vectors only forces a second round trip to the source of
    truth on every single query, and the two stores then drift.
    """
    index = []
    for document_id, text in documents:
        for ordinal, tokens in enumerate(chunk_document(text, chunk_tokens, overlap_tokens)):
            index.append(
                {
                    "document_id": document_id,
                    "ordinal": ordinal,
                    "tokens": tokens,
                    "vector": embed(tokens, width),
                }
            )
    return index


def dense_search(query_vector, index, k):
    """Exact search over every chunk.

    Exact here so retrieval and the objective provably share a metric. In
    production this becomes an approximate index, and the difference is
    not cosmetic: recall measured by exact search overstates what the
    served system returns, sometimes substantially. Evaluating against
    exact search and serving from an approximate index is how a pipeline
    passes its own tests and still misses passages in production.
    """
    scored = [
        (cosine(query_vector, entry["vector"]), position)
        for position, entry in enumerate(index)
    ]
    scored.sort(reverse=True)
    return scored[:k]


def bm25_statistics(index):
    """Document frequencies and mean length, computed once over the corpus."""
    document_frequency = {}
    total_length = 0
    for entry in index:
        total_length += len(entry["tokens"])
        for token in set(entry["tokens"]):
            document_frequency[token] = document_frequency.get(token, 0) + 1

    return {
        "document_frequency": document_frequency,
        "mean_length": total_length / len(index) if index else 0.0,
        "n_documents": len(index),
    }


def bm25_score(query_tokens, entry, statistics, k1=1.5, b=0.75):
    """Lexical relevance: the complement to the dense score, not the legacy.

    Dense retrieval is strong on paraphrase and reliably weak on exact
    identifiers -- part numbers, version strings, error codes, names the
    encoder never saw. Those are precisely the queries where being one
    token wrong is being completely wrong, and they are far more common
    in the corpora people actually deploy against than in the paraphrase
    benchmarks embeddings are chosen on. Dropping the lexical index
    because the embeddings look good is the most common self-inflicted
    regression in this area.
    """
    tokens = entry["tokens"]
    length = len(tokens)
    if length == 0:
        return 0.0

    term_frequency = {}
    for token in tokens:
        term_frequency[token] = term_frequency.get(token, 0) + 1

    n_documents = statistics["n_documents"]
    mean_length = statistics["mean_length"] or 1.0

    total = 0.0
    for token in set(query_tokens):
        frequency = term_frequency.get(token, 0)
        if frequency == 0:
            continue
        document_frequency = statistics["document_frequency"].get(token, 0)
        # Standard BM25 IDF, smoothed so a term present in every document
        # contributes a small positive weight rather than a negative one.
        idf = math.log(
            1.0 + (n_documents - document_frequency + 0.5) / (document_frequency + 0.5)
        )
        saturation = (frequency * (k1 + 1.0)) / (
            frequency + k1 * (1.0 - b + b * length / mean_length)
        )
        total += idf * saturation
    return total


def lexical_search(query_tokens, index, statistics, k):
    scored = [
        (bm25_score(query_tokens, entry, statistics), position)
        for position, entry in enumerate(index)
    ]
    scored.sort(reverse=True)
    return scored[:k]


def reciprocal_rank_fusion(dense_ranked, lexical_ranked, weight_dense, rrf_k=60):
    """Merge two rankings without pretending the scores are comparable.

    A cosine and a BM25 score live on different scales and neither is
    calibrated, so a weighted sum of the raw numbers is arbitrary: the
    weight would be doing scale conversion and relevance weighting at the
    same time, and it cannot be tuned for both at once. Fusing on RANK
    sidesteps that, because rank is the only thing the two retrievers
    genuinely agree about.

    What it costs: rank discards magnitude. A document the dense
    retriever scored far above everything else arrives merely first, and
    the strength of that signal is gone.
    """
    if not 0.0 <= weight_dense <= 1.0:
        raise ValueError("weight_dense must be in [0, 1]")

    fused = {}
    for rank, (_, position) in enumerate(dense_ranked, start=1):
        fused[position] = fused.get(position, 0.0) + weight_dense / (rrf_k + rank)
    for rank, (_, position) in enumerate(lexical_ranked, start=1):
        fused[position] = fused.get(position, 0.0) + (1.0 - weight_dense) / (rrf_k + rank)

    merged = [(score, position) for position, score in fused.items()]
    merged.sort(reverse=True)
    return merged


def cross_encoder_score(query_tokens, chunk_tokens):
    """A stand-in for a trained cross-encoder.

    The real thing scores query and passage JOINTLY, which is both why it
    is far more accurate per pair and why it cannot be precomputed --
    there is no passage vector to store when the score depends on the
    query. That is the entire architectural argument for two stages: the
    accurate scorer is affordable only on a short list the cheap scorer
    produced.
    """
    query_set = set(query_tokens)
    if not query_set or not chunk_tokens:
        return 0.0

    chunk_set = set(chunk_tokens)
    coverage = sum(1 for token in query_set if token in chunk_set) / len(query_set)

    # Proximity, which a bag-of-words score cannot see at all: a passage
    # where the query terms cluster is likelier to answer the query than
    # one where they are scattered across unrelated sentences.
    positions = [idx for idx, token in enumerate(chunk_tokens) if token in query_set]
    if len(positions) >= 2:
        spread = (positions[-1] - positions[0]) / len(chunk_tokens)
    else:
        spread = 1.0

    return coverage * (1.0 - 0.5 * spread)


def rerank(query_tokens, candidates, index, m):
    """Cut k candidates down to the m that reach the prompt.

    Reducing rather than passing everything through matters for a reason
    that surprises people: more context is not better. Generators attend
    unevenly over a long context, material in the middle gets measurably
    less attention than material at either end, and a plausible but
    irrelevant passage actively distracts. So m has an interior optimum
    that is usually lower than anyone guesses, and it has to be found by
    sweeping rather than assumed.
    """
    scored = [
        (cross_encoder_score(query_tokens, index[position]["tokens"]), position)
        for _, position in candidates
    ]
    scored.sort(reverse=True)
    return scored[:m]


def search(query_text, index, statistics, k, m, weight_dense, width):
    """The pipeline. Each stage exists because the previous one has a
    specific weakness, and each stage bounds every stage after it.
    """
    query_tokens = tokenize(query_text)
    query_vector = embed(query_tokens, width)

    dense_ranked = dense_search(query_vector, index, k)
    lexical_ranked = lexical_search(query_tokens, index, statistics, k)
    fused = reciprocal_rank_fusion(dense_ranked, lexical_ranked, weight_dense)
    return rerank(query_tokens, fused, index, m)


def build_prompt(query_text, retrieved, index):
    """What production actually does: concatenate, then generate once.

    Read this against marginalize_answer() below -- they are different
    computations. There is no sum over documents here and no weighting at
    all. The retriever's scores decided WHICH chunks appear and then
    stopped mattering, so the generator sees a flat context with no
    indication which passage was retrieved confidently and which barely
    cleared the cut.
    """
    parts = ["Answer using only the sources below. Cite them as [n]."]
    for ordinal, (_, position) in enumerate(retrieved, start=1):
        entry = index[position]
        parts.append(
            "[%d] (%s#%d) %s"
            % (ordinal, entry["document_id"], entry["ordinal"], " ".join(entry["tokens"]))
        )
    parts.append("Question: " + query_text)
    return "\n\n".join(parts)


def marginalize_answer(candidate_answers, retrieved, generator_log_prob, index):
    """The equation, for once actually computed.

        p(y | x) = sum_z p(z | x) * p(y | x, z)

    One generator call per retrieved document instead of one call over a
    concatenation, then a retriever-weighted average. This is what the
    original formulation meant, and it buys two real things: a document
    that would have been drowned out in a long context gets its own
    forward pass, and the per-document likelihoods are inspectable, so
    "which passage produced this answer" becomes a question with an
    answer rather than a guess.

    Why almost nobody does it: k generator calls instead of one, which is
    linear cost at the most expensive stage in the pipeline.

    The honesty caveat lives in the weights. A softmax over retrieval
    scores is not a posterior over documents -- a similarity is not a
    likelihood, and nothing here was trained to make it one -- so p(z | x)
    is a plausible-looking number with no calibration behind it. The
    arithmetic is exact; the probabilities are not.
    """
    if not retrieved:
        raise ValueError("nothing retrieved: there is no marginal to take")

    scores = [score for score, _ in retrieved]
    top = max(scores)
    exponentiated = [math.exp(score - top) for score in scores]
    total = sum(exponentiated)
    document_weights = [value / total for value in exponentiated]

    marginals = []
    for answer in candidate_answers:
        accumulated = 0.0
        for weight, (_, position) in zip(document_weights, retrieved):
            accumulated += weight * math.exp(
                generator_log_prob(answer, index[position]["tokens"])
            )
        marginals.append((accumulated, answer))

    marginals.sort(reverse=True)
    return marginals


def retrieval_recall_at_m(labelled_queries, retrieve_fn, m):
    """The one metric that has to be measured on its own.

    Every failure in this pipeline presents identically -- a bad answer --
    and the generator is what produced the visible output, so it gets
    blamed by default. This measures whether the answer was even
    reachable. If it is low, nothing done to the generator matters, and
    teams routinely spend a fine-tuning budget discovering that.
    """
    if not labelled_queries:
        raise ValueError("recall needs a labelled query set")

    hits = 0
    for query in labelled_queries:
        retrieved = retrieve_fn(query["text"], m)
        if any(position in query["relevant"] for _, position in retrieved):
            hits += 1
    return hits / len(labelled_queries)


def citation_support(answer_tokens, retrieved, index):
    """A weak check, stated weakly on purpose.

    Overlap between the answer and the cited passages is the only part of
    grounding that can be verified mechanically. It confirms the answer
    reuses the passage's vocabulary. It does not confirm the answer
    follows from the passage, and a citation proves only that the passage
    was retrieved. Reviewers consistently read the stronger claim into it.
    """
    supported = set()
    for _, position in retrieved:
        supported.update(index[position]["tokens"])
    if not answer_tokens:
        return 0.0
    return sum(1 for token in answer_tokens if token in supported) / len(answer_tokens)

"""Retrieval-augmented generation with the pipeline's contracts named.

The design point: this technique is a sequence of stages where each one
bounds every stage after it, so the things worth making explicit are the
stage boundaries and the couplings that cross them. Three of those
couplings cause essentially all of the production pain, and all three are
invisible in a version that passes dicts around:

  1. The encoder version ties the query to the index. Change one and
     every stored vector is meaningless -- with no error.
  2. Access control belongs at retrieval time, per caller. Applied later
     it has already leaked the passage into the ranking, and the
     generated answer launders where it came from.
  3. Retrieval recall bounds answer quality. Measured only end to end it
     is unattributable, and the generator gets blamed by default.

So each is a type here rather than a convention someone remembers.
"""

from __future__ import annotations

import math
import re
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Callable, Iterator, Protocol, Sequence

TOKEN_PATTERN = re.compile(r"[a-z0-9]+")


class RagError(Exception):
    """Base class, so a caller can catch this pipeline specifically."""


class ChunkingError(RagError):
    """A chunking policy that cannot produce the windows it was asked for."""


class EmptyIndexError(RagError):
    """Retrieval against an index with nothing in it."""


class EncoderMismatchError(RagError):
    """Query encoder disagrees with the encoder that built the index.

    The most expensive silent failure in this technique. Embeddings from
    two encoder versions are not comparable, so retrieval returns
    confident nonsense, and a partially re-indexed corpus fails this way
    for only the fraction of chunks that moved.
    """


class CoverageGapError(RagError):
    """Nothing retrieved cleared the relevance floor.

    Deliberately an error rather than an empty list. The alternative is
    generating from irrelevant context, which produces a fluent wrong
    answer -- and a system that never refuses is fabricating instead.
    """


@dataclass(frozen=True)
class ChunkingPolicy:
    """Chunking as a named, versioned decision.

    Frozen and recorded on the index because it determines what is
    answerable at all, and because comparing retrieval numbers across two
    policies is the most useful experiment available here. Left implicit,
    it becomes a default nobody revisits.
    """

    chunk_tokens: int = 400
    overlap_tokens: int = 40

    def __post_init__(self) -> None:
        if self.chunk_tokens <= 0:
            raise ChunkingError("chunk_tokens must be positive")
        if not 0 <= self.overlap_tokens < self.chunk_tokens:
            raise ChunkingError("overlap must be in [0, chunk_tokens)")

    @property
    def stride(self) -> int:
        return self.chunk_tokens - self.overlap_tokens


@dataclass(frozen=True)
class Chunk:
    """A passage plus the provenance a citation needs.

    document_id and ordinal are not bookkeeping: a citation that cannot
    be resolved back to a source location is decoration, and acl_tags
    travel with the chunk because filtering by them after ranking is
    already too late.
    """

    document_id: str
    ordinal: int
    tokens: tuple[str, ...]
    vector: tuple[float, ...]
    acl_tags: frozenset[str] = frozenset()
    effective_from: str | None = None


@dataclass(frozen=True)
class Retrieved:
    chunk: Chunk
    score: float
    stage: str


@dataclass(frozen=True)
class RetrievalConfig:
    candidate_k: int = 60
    context_m: int = 5
    weight_dense: float = 0.5
    rrf_k: int = 60
    relevance_floor: float = 0.05

    def __post_init__(self) -> None:
        if self.candidate_k < self.context_m:
            raise ValueError("candidate_k must be at least context_m")
        if not 0.0 <= self.weight_dense <= 1.0:
            raise ValueError("weight_dense must be in [0, 1]")


@dataclass
class AnswerReport:
    """The answer and the numbers needed to attribute a failure.

    Returned together on purpose. Faithfulness and retrieval quality
    computed by a caller who remembers to are computed by nobody, and
    then every failure in a five-stage pipeline looks like a generator
    problem because the generator produced the visible output.
    """

    answer: str
    cited: list[Retrieved]
    citation_support: float
    best_retrieval_score: float
    strategy: str
    document_posterior: list[float] = field(default_factory=list)


class Encoder(Protocol):
    """Named version, because the version is a compatibility contract."""

    @property
    def version(self) -> str: ...

    def encode(self, tokens: Sequence[str]) -> tuple[float, ...]: ...


class Generator(Protocol):
    def generate(self, prompt: str) -> str: ...

    def log_prob(self, answer: str, context: Sequence[str]) -> float: ...


class ConsumptionStrategy(Protocol):
    """How retrieved documents reach the generator.

    Two strategies, one protocol, because the difference between them is
    the difference between the published equation and what production
    does -- and that difference deserves to be a visible choice rather
    than an implementation detail buried in a prompt builder.
    """

    @property
    def name(self) -> str: ...

    def answer(
        self,
        query: str,
        retrieved: Sequence[Retrieved],
        generator: Generator,
    ) -> AnswerReport: ...


def tokenize(text: str) -> tuple[str, ...]:
    return tuple(TOKEN_PATTERN.findall(text.lower()))


class HashingEncoder:
    """Deterministic hashing encoder, version-stamped.

    The stamp is the point of the class. A real encoder is a neural
    network; what it shares with this one is that changing it invalidates
    every vector in the index, and the only defence is that the index
    knows which version wrote it.
    """

    def __init__(self, width: int = 512, *, version: str = "hash-v1") -> None:
        if width <= 0:
            raise ValueError("width must be positive")
        self._width = width
        self._version = version

    @property
    def version(self) -> str:
        return self._version

    def encode(self, tokens: Sequence[str]) -> tuple[float, ...]:
        vector = [0.0] * self._width
        for token in tokens:
            digest = 2166136261
            for character in token:
                digest = ((digest ^ ord(character)) * 16777619) % (1 << 32)
            vector[digest % self._width] += 1.0

        norm = math.sqrt(sum(value * value for value in vector))
        if norm < 1e-12:
            return tuple(vector)
        return tuple(value / norm for value in vector)


class HybridIndex:
    """Dense vectors, lexical statistics and ACL tags in one store.

    One store rather than two because a vector store and a text store
    drift, and a chunk present in one but not the other is a retrieval
    miss with no error attached. The encoder version is stored here so a
    mismatched query is refused rather than answered.
    """

    def __init__(self, encoder_version: str, policy: ChunkingPolicy) -> None:
        self._encoder_version = encoder_version
        self._policy = policy
        self._chunks: list[Chunk] = []
        self._document_frequency: dict[str, int] = {}
        self._total_tokens = 0

    @property
    def encoder_version(self) -> str:
        return self._encoder_version

    @property
    def policy(self) -> ChunkingPolicy:
        return self._policy

    def __len__(self) -> int:
        return len(self._chunks)

    def add(self, chunk: Chunk) -> None:
        self._chunks.append(chunk)
        self._total_tokens += len(chunk.tokens)
        for token in set(chunk.tokens):
            self._document_frequency[token] = self._document_frequency.get(token, 0) + 1

    def visible_positions(self, caller_tags: frozenset[str]) -> list[int]:
        """ACL filtering BEFORE scoring, which is the only correct order.

        Filtering after ranking means an entitled-to-nothing caller still
        influenced which passages were considered, and any cache keyed on
        the query rather than on the caller will serve them the result.
        """
        return [
            position
            for position, chunk in enumerate(self._chunks)
            if not chunk.acl_tags or chunk.acl_tags & caller_tags
        ]

    def chunk_at(self, position: int) -> Chunk:
        return self._chunks[position]

    def bm25(self, query_tokens: Sequence[str], position: int, k1: float = 1.5, b: float = 0.75) -> float:
        chunk = self._chunks[position]
        length = len(chunk.tokens)
        if length == 0:
            return 0.0

        mean_length = (self._total_tokens / len(self._chunks)) or 1.0
        term_frequency: dict[str, int] = {}
        for token in chunk.tokens:
            term_frequency[token] = term_frequency.get(token, 0) + 1

        total = 0.0
        for token in set(query_tokens):
            frequency = term_frequency.get(token, 0)
            if frequency == 0:
                continue
            document_frequency = self._document_frequency.get(token, 0)
            idf = math.log(
                1.0 + (len(self._chunks) - document_frequency + 0.5) / (document_frequency + 0.5)
            )
            total += idf * (frequency * (k1 + 1.0)) / (
                frequency + k1 * (1.0 - b + b * length / mean_length)
            )
        return total

    def cosine(self, query_vector: Sequence[float], position: int) -> float:
        stored = self._chunks[position].vector
        return sum(query_vector[idx] * stored[idx] for idx in range(len(stored)))


@contextmanager
def index_session(index: HybridIndex, encoder: Encoder) -> Iterator[HybridIndex]:
    """Bind an encoder to an index for the duration of a query batch.

    A context manager rather than a check at each call site because the
    compatibility question is asked once per session and answered wrong
    exactly when nobody is asking it. Real index handles hold sockets and
    memory maps too, and those need releasing on the failure path as much
    as on the happy one.
    """
    if len(index) == 0:
        raise EmptyIndexError("index is empty; nothing is retrievable")
    if index.encoder_version != encoder.version:
        raise EncoderMismatchError(
            f"index built with {index.encoder_version}, querying with {encoder.version}: "
            "re-index before serving"
        )
    try:
        yield index
    finally:
        # Released here rather than by the caller, so an exception raised
        # mid-retrieval cannot leak the handle.
        pass


def build_index(
    documents: Sequence[tuple[str, str, frozenset[str]]],
    encoder: Encoder,
    policy: ChunkingPolicy,
) -> HybridIndex:
    index = HybridIndex(encoder.version, policy)
    for document_id, text, acl_tags in documents:
        words = tokenize(text)
        for ordinal, start in enumerate(range(0, max(len(words), 1), policy.stride)):
            window = words[start : start + policy.chunk_tokens]
            if not window:
                break
            index.add(
                Chunk(
                    document_id=document_id,
                    ordinal=ordinal,
                    tokens=window,
                    vector=encoder.encode(window),
                    acl_tags=acl_tags,
                )
            )
    return index


def retrieve(
    query: str,
    index: HybridIndex,
    encoder: Encoder,
    config: RetrievalConfig,
    caller_tags: frozenset[str],
    rerank_fn: Callable[[Sequence[str], Chunk], float],
) -> list[Retrieved]:
    """Hybrid retrieval, rank fusion, rerank. Guard clauses first.

    caller_tags is required rather than optional because a default of
    "everything" is the wrong default exactly once and then it is a
    breach. Fusion is on rank, not score: a cosine and a BM25 value are
    not comparable, so a weighted sum of the raw numbers would make the
    weight do scale conversion and relevance weighting at once.
    """
    query_tokens = tokenize(query)
    if not query_tokens:
        raise ValueError("query contains no indexable tokens")

    visible = index.visible_positions(caller_tags)
    if not visible:
        raise CoverageGapError("caller is entitled to no chunks in this index")

    query_vector = encoder.encode(query_tokens)

    dense = sorted(
        ((index.cosine(query_vector, position), position) for position in visible),
        reverse=True,
    )[: config.candidate_k]
    lexical = sorted(
        ((index.bm25(query_tokens, position), position) for position in visible),
        reverse=True,
    )[: config.candidate_k]

    fused: dict[int, float] = {}
    for rank, (_, position) in enumerate(dense, start=1):
        fused[position] = fused.get(position, 0.0) + config.weight_dense / (config.rrf_k + rank)
    for rank, (_, position) in enumerate(lexical, start=1):
        fused[position] = fused.get(position, 0.0) + (1.0 - config.weight_dense) / (
            config.rrf_k + rank
        )

    candidates = sorted(fused.items(), key=lambda item: item[1], reverse=True)
    reranked = sorted(
        (
            Retrieved(
                chunk=index.chunk_at(position),
                score=rerank_fn(query_tokens, index.chunk_at(position)),
                stage="rerank",
            )
            for position, _ in candidates
        ),
        key=lambda item: item.score,
        reverse=True,
    )[: config.context_m]

    if not reranked or reranked[0].score < config.relevance_floor:
        raise CoverageGapError(
            f"best reranked score {reranked[0].score if reranked else 0.0:.4f} "
            f"below floor {config.relevance_floor}: refuse rather than generate"
        )
    return reranked


class ConcatenationStrategy:
    """What production does: one prompt, one forward pass.

    Cheap and the default for good reason. Its cost is that the
    retriever's confidence is discarded -- the scores chose which chunks
    appear and then stopped mattering -- and that quality turns over as
    context_m grows, with passages in the middle attended to less than
    passages at either end.
    """

    @property
    def name(self) -> str:
        return "concatenate"

    def answer(
        self,
        query: str,
        retrieved: Sequence[Retrieved],
        generator: Generator,
    ) -> AnswerReport:
        lines = ["Answer using only the sources below. Cite them as [n]."]
        for ordinal, item in enumerate(retrieved, start=1):
            lines.append(
                f"[{ordinal}] ({item.chunk.document_id}#{item.chunk.ordinal}) "
                + " ".join(item.chunk.tokens)
            )
        lines.append(f"Question: {query}")

        text = generator.generate("\n\n".join(lines))
        return AnswerReport(
            answer=text,
            cited=list(retrieved),
            citation_support=_citation_support(tokenize(text), retrieved),
            best_retrieval_score=retrieved[0].score if retrieved else 0.0,
            strategy=self.name,
        )


class MarginalStrategy:
    """The published equation: one generator call per document, then average.

    Buys attributability -- the per-document posterior says which passage
    carried the answer -- and costs a factor of m at the most expensive
    stage in the pipeline, which is why it stays rare. The posterior it
    reports is honest about its own weakness: a softmax over reranker
    scores is not a posterior over documents, because a similarity is not
    a likelihood and nothing trained it to be one.
    """

    def __init__(self, candidates: Sequence[str], *, temperature: float = 1.0) -> None:
        if not candidates:
            raise ValueError("the marginal needs candidate answers to score")
        if temperature <= 0.0:
            raise ValueError("temperature must be positive")
        self._candidates = tuple(candidates)
        self._temperature = temperature

    @property
    def name(self) -> str:
        return "marginalize"

    def answer(
        self,
        query: str,
        retrieved: Sequence[Retrieved],
        generator: Generator,
    ) -> AnswerReport:
        scores = [item.score / self._temperature for item in retrieved]
        top = max(scores)
        exponentiated = [math.exp(score - top) for score in scores]
        total = sum(exponentiated)
        posterior = [value / total for value in exponentiated]

        best_answer = ""
        best_marginal = -1.0
        for candidate in self._candidates:
            marginal = sum(
                weight * math.exp(generator.log_prob(candidate, item.chunk.tokens))
                for weight, item in zip(posterior, retrieved)
            )
            if marginal > best_marginal:
                best_marginal, best_answer = marginal, candidate

        return AnswerReport(
            answer=best_answer,
            cited=list(retrieved),
            citation_support=_citation_support(tokenize(best_answer), retrieved),
            best_retrieval_score=retrieved[0].score if retrieved else 0.0,
            strategy=self.name,
            document_posterior=posterior,
        )


def _citation_support(answer_tokens: Sequence[str], retrieved: Sequence[Retrieved]) -> float:
    """Overlap with the cited passages. A weak check, named weakly.

    It confirms the answer reuses the passages' vocabulary. It does not
    confirm the answer follows from them, and a citation only ever proves
    the passage was retrieved. Reviewers read the stronger claim into it
    anyway, which is why the number is reported rather than thresholded.
    """
    if not answer_tokens:
        return 0.0
    supported: set[str] = set()
    for item in retrieved:
        supported.update(item.chunk.tokens)
    return sum(1 for token in answer_tokens if token in supported) / len(answer_tokens)


def evaluate_retrieval(
    labelled: Sequence[tuple[str, frozenset[str]]],
    retrieve_fn: Callable[[str], Sequence[Retrieved]],
) -> dict[str, float]:
    """Recall and refusal rate, reported apart from answer quality.

    Two numbers rather than one, because an end-to-end score cannot say
    which stage to fix and that is the only question this evaluation
    exists to answer. Refusal rate belongs here: a system that never
    refuses on unanswerable queries is fabricating on them, and recall
    alone will not show it.
    """
    if not labelled:
        raise ValueError("retrieval evaluation needs a labelled query set")

    hits = 0
    refusals = 0
    for query, relevant_documents in labelled:
        try:
            retrieved = retrieve_fn(query)
        except CoverageGapError:
            refusals += 1
            continue
        if any(item.chunk.document_id in relevant_documents for item in retrieved):
            hits += 1

    return {
        "recall_at_m": hits / len(labelled),
        "refusal_rate": refusals / len(labelled),
    }

"""Retrieval as two matrix operations, plus the asymmetry between them.

Dense retrieval over the whole corpus is one GEMM: a contiguous
(n_chunks, dim) float32 matrix times a (dim, n_queries) block. That is
the natural shape of exact search, and it is worth knowing that brute
force stays competitive to surprisingly large corpora -- an approximate
index buys sublinear lookup and pays in recall, and the recall it costs
is invisible unless retrieval is measured through the served index rather
than through exact search.

Lexical retrieval has completely different scaling. A query touches only
the postings for its own terms, so cost is proportional to query length
and posting length, not to corpus size. That asymmetry is the practical
argument for keeping both: the lexical side is nearly free, and it covers
the exact-identifier queries where the dense side is reliably weak.

What every optimization here shares: nothing changes which documents are
retrievable. Chunking already decided that, upstream of all of it.
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

TOKEN_PATTERN = re.compile(r"[a-z0-9]+")

# Additive, and f32-safe, unlike -inf which poisons a row max.
MASK_VALUE = np.float32(-1.0e4)


@dataclass(frozen=True)
class Bm25Params:
    k1: float = 1.5
    b: float = 0.75


class PackedIndex:
    """Dense matrix plus a CSR inverted index, in one object.

    Both retrievers read the same chunk ordering, which is not a
    convenience -- it is what makes rank fusion meaningful. Two stores
    with two orderings produce a fusion that silently mixes up documents,
    and the symptom is a ranking that looks merely mediocre.
    """

    def __init__(
        self,
        vectors: npt.NDArray[np.float32],
        lengths: npt.NDArray[np.int32],
        term_ids: npt.NDArray[np.int32],
        postings: npt.NDArray[np.int32],
        posting_offsets: npt.NDArray[np.int64],
        posting_frequencies: npt.NDArray[np.float32],
        acl_bits: npt.NDArray[np.uint64],
        encoder_version: str,
    ) -> None:
        if vectors.dtype != np.float32 or not vectors.flags.c_contiguous:
            raise TypeError("vectors must be a C-contiguous float32 matrix")
        if vectors.shape[0] != lengths.shape[0]:
            raise ValueError("one length per chunk is required")

        self.vectors = vectors
        self.lengths = lengths
        self.mean_length = float(lengths.mean()) if lengths.size else 1.0
        self.term_ids = term_ids
        self.postings = postings
        self.posting_offsets = posting_offsets
        self.posting_frequencies = posting_frequencies
        self.acl_bits = acl_bits
        self.encoder_version = encoder_version

        # IDF precomputed per term: it depends only on the corpus, so
        # recomputing it per query was pure waste in the earlier stages.
        document_frequency = np.diff(posting_offsets).astype(np.float32)
        n_chunks = np.float32(vectors.shape[0])
        self.idf = np.log(
            1.0 + (n_chunks - document_frequency + 0.5) / (document_frequency + 0.5)
        ).astype(np.float32)

    @property
    def n_chunks(self) -> int:
        return int(self.vectors.shape[0])


class RetrievalBuffers:
    """Allocated once at the peak batch shape, then overwritten.

    The trade is explicit: a serving loop performs no allocation, and the
    memory is never returned. The zeroing in dense_scores() is not
    optional -- a shared accumulator that is not reset leaks scores from
    the previous query into this one, which produces a ranking that is
    wrong in a way no assertion catches.
    """

    def __init__(self, max_queries: int, n_chunks: int) -> None:
        self.dense = np.empty((max_queries, n_chunks), dtype=np.float32)
        self.lexical = np.empty((max_queries, n_chunks), dtype=np.float32)
        self.fused = np.empty((max_queries, n_chunks), dtype=np.float32)
        self.mask = np.empty(n_chunks, dtype=np.float32)
        self.max_queries = max_queries

    def view(self, n_queries: int) -> tuple[npt.NDArray[np.float32], ...]:
        if n_queries > self.max_queries:
            raise ValueError("batch exceeds the shape these buffers were sized for")
        return (
            self.dense[:n_queries],
            self.lexical[:n_queries],
            self.fused[:n_queries],
        )


def encode_batch(
    token_lists: list[list[str]], width: int, out: npt.NDArray[np.float32] | None = None
) -> npt.NDArray[np.float32]:
    """Encode a batch into one contiguous block.

    The loop over tokens survives because hashing is inherently scattered;
    what is removed is the per-query allocation and the per-query
    normalization pass. A real encoder replaces this whole function with
    one batched forward pass, and the shape contract is the same.
    """
    n_queries = len(token_lists)
    matrix = np.zeros((n_queries, width), dtype=np.float32) if out is None else out[:n_queries]
    if out is not None:
        matrix[:] = 0.0

    for row, tokens in enumerate(token_lists):
        for token in tokens:
            digest = 2166136261
            for character in token:
                digest = ((digest ^ ord(character)) * 16777619) % (1 << 32)
            matrix[row, digest % width] += 1.0

    # One fused normalization over the batch rather than per query.
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    np.divide(matrix, np.maximum(norms, 1e-12), out=matrix)
    return matrix


def dense_scores(
    index: PackedIndex,
    query_vectors: npt.NDArray[np.float32],
    out: npt.NDArray[np.float32],
) -> npt.NDArray[np.float32]:
    """Exact cosine over the entire corpus in one GEMM.

    Both operands are unit-normalized, so the inner product IS the
    cosine and there is no per-row division. Written into a pre-allocated
    buffer, which is why the shapes are asserted rather than inferred:
    a silently broadcast mismatch here produces a full score matrix of
    plausible garbage.
    """
    if query_vectors.shape[1] != index.vectors.shape[1]:
        raise ValueError("query dimension does not match the index")

    np.matmul(query_vectors, index.vectors.T, out=out)
    return out


def lexical_scores(
    index: PackedIndex,
    query_term_ids: list[npt.NDArray[np.int32]],
    out: npt.NDArray[np.float32],
    params: Bm25Params = Bm25Params(),
) -> npt.NDArray[np.float32]:
    """BM25 through the inverted index, one posting slice at a time.

    This is where the scaling asymmetry lives. The dense GEMM touches
    every chunk; this touches only chunks containing a query term, so a
    three-term query over a million chunks reads three short arrays. The
    saturation is applied to a whole slice with one expression rather
    than per chunk, and the result accumulates into the pre-zeroed row by
    fancy-index addition.
    """
    out[:] = 0.0
    length_ratio = index.lengths.astype(np.float32) / np.float32(index.mean_length)
    denominator_base = np.float32(params.k1) * (
        np.float32(1.0 - params.b) + np.float32(params.b) * length_ratio
    )

    for row, term_ids in enumerate(query_term_ids):
        for term_id in term_ids:
            start = index.posting_offsets[term_id]
            stop = index.posting_offsets[term_id + 1]
            chunk_ids = index.postings[start:stop]
            frequencies = index.posting_frequencies[start:stop]

            saturation = (frequencies * np.float32(params.k1 + 1.0)) / (
                frequencies + denominator_base[chunk_ids]
            )
            np.add.at(out[row], chunk_ids, index.idf[term_id] * saturation)
    return out


def build_acl_mask(
    index: PackedIndex, caller_bits: np.uint64, out: npt.NDArray[np.float32]
) -> npt.NDArray[np.float32]:
    """Entitlement as an additive mask, applied before selection.

    A bitmask rather than a set intersection so this is one vectorized
    AND over the corpus. The ordering is the security property: masking
    before top-k means an unentitled chunk was never a candidate, where
    a post-filter has already let it influence the ranking and, worse,
    lets a cache keyed on the query alone serve it to the next caller.
    """
    visible = (index.acl_bits & caller_bits) != 0
    unrestricted = index.acl_bits == 0
    np.where(visible | unrestricted, np.float32(0.0), MASK_VALUE, out=out)
    return out


def fuse_and_select(
    dense: npt.NDArray[np.float32],
    lexical: npt.NDArray[np.float32],
    fused: npt.NDArray[np.float32],
    mask: npt.NDArray[np.float32],
    candidate_k: int,
    weight_dense: float,
    rrf_k: int = 60,
) -> npt.NDArray[np.int64]:
    """Rank-fuse both retrievers, then take top-k without sorting the tail.

    Fusion is on rank because a cosine and a BM25 value are not
    comparable. Getting ranks vectorized costs two argsorts, which is the
    honest price of not inventing a score calibration: double-argsort is
    O(n log n) per row where the scoring itself was O(n), so on a large
    corpus the fusion, not the retrieval, becomes the cost. Weighted
    score fusion avoids that and needs a calibration nobody has.

    Selection is argpartition, not sort: the ordering of the discarded
    tail is never read, and only the k survivors get sorted.
    """
    dense_ranks = np.argsort(np.argsort(-dense, axis=1), axis=1) + 1
    lexical_ranks = np.argsort(np.argsort(-lexical, axis=1), axis=1) + 1

    # Fused in place, then masked, so masking cannot be forgotten between
    # the fusion and the selection.
    np.divide(np.float32(weight_dense), (rrf_k + dense_ranks).astype(np.float32), out=fused)
    fused += np.float32(1.0 - weight_dense) / (rrf_k + lexical_ranks).astype(np.float32)
    fused += mask

    top_unsorted = np.argpartition(-fused, candidate_k - 1, axis=1)[:, :candidate_k]
    row_index = np.arange(fused.shape[0])[:, None]
    order = np.argsort(-fused[row_index, top_unsorted], axis=1)
    return np.take_along_axis(top_unsorted, order, axis=1)


def rerank_batch(
    cross_encoder: object,
    query_texts: list[str],
    candidates: npt.NDArray[np.int64],
    chunk_texts: list[str],
    context_m: int,
) -> npt.NDArray[np.int64]:
    """One batched cross-encoder call over every (query, candidate) pair.

    The pairs are flattened into a single batch because the cross-encoder
    is the one stage whose cost is entirely under our control, and it
    scales with candidate_k times the batch size. This is the knob that
    trades latency for accuracy, and it is worth restating that the stage
    after it wants FEWER documents, not more: answer quality turns over
    as context_m grows, so reranking deeper and passing less is usually
    the better allocation of the same budget.
    """
    n_queries, candidate_k = candidates.shape
    pairs = [
        (query_texts[row], chunk_texts[candidates[row, column]])
        for row in range(n_queries)
        for column in range(candidate_k)
    ]
    flat = np.asarray(cross_encoder.score(pairs), dtype=np.float32)  # type: ignore[attr-defined]
    scores = flat.reshape(n_queries, candidate_k)

    keep = np.argsort(-scores, axis=1)[:, :context_m]
    return np.take_along_axis(candidates, keep, axis=1)


def recall_at_m(
    selected: npt.NDArray[np.int64], relevant_mask: npt.NDArray[np.bool_]
) -> float:
    """Recall over a whole evaluation batch, vectorized.

    Batched because this metric is swept, not sampled: chunk size, k, m
    and the fusion weight all move it, and the sweep is the actual tuning
    loop for this technique. A per-query Python loop makes the sweep slow
    enough that people run it once and stop, which is how the highest-
    leverage parameter in the system ends up left at its default.
    """
    if selected.shape[0] == 0:
        raise ValueError("recall needs at least one query")
    hits = np.take_along_axis(relevant_mask, selected, axis=1).any(axis=1)
    return float(hits.mean())


def context_sweep(
    fused_selection: npt.NDArray[np.int64],
    relevant_mask: npt.NDArray[np.bool_],
    candidate_m: tuple[int, ...] = (1, 3, 5, 10, 20),
) -> dict[int, float]:
    """Recall as a function of m, which is the curve worth having.

    Retrieval recall rises monotonically with m -- more candidates can
    only help -- while ANSWER quality turns over, because irrelevant
    context distracts and material in the middle of a long context is
    attended to less than material at the ends. Plotting the two against
    each other is how the interior optimum gets found, and it is usually
    at a smaller m than anyone guesses.
    """
    return {
        m: recall_at_m(fused_selection[:, :m], relevant_mask)
        for m in candidate_m
        if m <= fused_selection.shape[1]
    }