rag-wright 0.1.0__py3-none-any.whl

rag_wright/capabilities/reranking.py

@@ -0,0 +1,123 @@
+"""Reranking (FR-C.4, FR-Q.2): a cross-encoder precision gate over the retrieved candidates.
+
+A BGE-reranker cross-encoder re-scores each retrieved candidate against the query and cuts the list to
+a top-k set, the precision gate before any expensive downstream work (synthesis). Unlike the bi-encoder
+retrieval legs (T19/T21), a cross-encoder reads the query and the passage together, so it is more
+precise but too costly to run over the whole index — it runs only over the already-fused candidate list
+(T21).
+
+This capability is a pure function of `(query, passage)` pairs: it is given the candidates *with their
+passage text* and returns them reranked and cut. It deliberately does not fetch the text itself — the
+passage a candidate is scored on (the chunk summary, or the full chunk text from the parse manifest,
+FR-I.1) is a wiring choice the compiled query graph makes, not this capability's, so reranking stays
+decoupled from the store and the manifest. Grounded on `FlagEmbedding.FlagAutoReranker` (the public
+auto-reranker; `.compute_score` over sentence pairs).
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel
+
+
+DEFAULT_TOP_K = 5  # the precision-gate cut before synthesis; the eval and caller can override
+DEFAULT_RERANKER_MODEL = "BAAI/bge-reranker-v2-m3"  # the BGE-M3 companion cross-encoder
+
+
+@runtime_checkable
+class Reranker(Protocol):
+    """The cross-encoder seam: score each passage against the query (higher is more relevant)."""
+
+    def score(self, query: str, passages: list[str]) -> list[float]: ...
+
+
+class Passage(BaseModel):
+    """A retrieved candidate with the passage text to score (the reranker's input unit)."""
+
+    model_config = {"frozen": True}
+
+    chunk_id: str
+    source_doc_id: str
+    text: str
+
+
+class ScoredCandidate(BaseModel):
+    """One reranked candidate: the relevance score the cross-encoder gave it."""
+
+    model_config = {"frozen": True}
+
+    chunk_id: str
+    source_doc_id: str
+    score: float
+
+
+class RerankResult(BaseModel):
+    """The reranking capability's output: the top-k candidates by cross-encoder score (FR-C.4)."""
+
+    model_config = {"frozen": True}
+
+    query: str
+    candidates: list[ScoredCandidate]  # cross-encoder order, best first, cut to top_k
+
+
+class BGEReranker:
+    """The real reranker: `FlagEmbedding.FlagAutoReranker` (model loaded lazily).
+
+    THREAD-SAFE (engine issue 0016): FlagEmbedding mutates the model in place on every call (the same in-place
+    `.to()`/`.eval()` conversions that segfault the shared BGE-M3 embedder under thread concurrency), so a shared
+    reranker is guarded by an instance lock too -- defense-in-depth for any concurrent `score` caller."""
+
+    def __init__(self, model_name: str = DEFAULT_RERANKER_MODEL, *, use_fp16: bool = False,
+                 model: object | None = None) -> None:
+        import threading
+
+        self._lock = threading.Lock()  # issue 0016: serialize the in-place-mutating compute_score across threads
+        if model is not None:  # injected (hermetic tests)
+            self._model = model
+        else:
+            from FlagEmbedding import FlagAutoReranker
+
+            self._model = FlagAutoReranker.from_finetuned(model_name, use_fp16=use_fp16)
+
+    def score(self, query: str, passages: list[str]) -> list[float]:
+        if not passages:
+            return []
+        with self._lock:  # issue 0016
+            scores = self._model.compute_score([(query, passage) for passage in passages])
+        # compute_score returns a scalar for a single pair; normalize to a list of floats.
+        if not isinstance(scores, (list, tuple)):
+            scores = [scores]
+        return [float(s) for s in scores]
+
+
+def rerank(
+    query: str,
+    passages: list[Passage],
+    *,
+    reranker: Reranker,
+    top_k: int = DEFAULT_TOP_K,
+) -> RerankResult:
+    """Rerank the candidate passages by cross-encoder relevance and cut to the top-`k` set.
+
+    The cross-encoder scores each `(query, passage.text)` pair; the passages are sorted by score
+    descending (stable, so equal scores keep their incoming fused order) and cut to `top_k`. An empty
+    candidate list reranks to an empty result.
+    """
+    if not passages:
+        return RerankResult(query=query, candidates=[])
+
+    scores = reranker.score(query, [p.text for p in passages])
+    if len(scores) != len(passages):
+        raise ValueError(
+            f"reranker returned {len(scores)} scores for {len(passages)} passages (must be 1:1)"
+        )
+
+    scored = [
+        ScoredCandidate(chunk_id=p.chunk_id, source_doc_id=p.source_doc_id, score=score)
+        for p, score in zip(passages, scores)
+    ]
+    scored.sort(key=lambda c: c.score, reverse=True)  # stable: ties keep the incoming fused order
+    return RerankResult(query=query, candidates=scored[:top_k])
+
+

rag_wright/capabilities/retrieval_core.py

@@ -0,0 +1,126 @@
+"""CAP-REG-3 (FR-Q, ADR-0033): the KG-primary retrieval core, packaged out of `eval/kg_primary.py`.
+
+The eval script ranked ACORD candidates through inline closures (`pool_of`, `_match`, `_tiebreak`) wrapped in
+MODE/VARIANT/MATCH ablation scaffolding. This module lifts the ADOPTED operating point out of that scaffolding
+as three pure, deterministic `function` capabilities, each registered under its FR-C slug so the query graph
+can bind them (`typed_constraint_match_rank` / `dense_rank_tiebreak` back the adopted Leg B via
+`property_boosted_retrieval`; `candidate_routing` is a general union-combiner utility, unused since the redundant
+`cross_corpus_retrieval` subgraph was retired):
+
+  - **candidate_routing** -- the union combiner: several routing signals (LLM / LegalBERT classifier /
+    dimension-prior) each propose ranked functions; union them (first-wins, recall-safe) and fetch the
+    candidate clause pool for that function set (KG-5e). The store pool lookup is an injected seam.
+  - **typed_constraint_match_rank** -- grade each candidate by how many query (dimension, value) constraints
+    its grounded typed props satisfy, under KG-5a canonicalization + subsumption (`constraint_match_count`).
+    Recall-safe: a zero-match candidate keeps its place (stable sort), never dropped.
+  - **dense_rank_tiebreak** -- order candidates by descending cosine to the query vector: the embedding signal
+    that breaks constraint-match ties meaningfully (KG-6 / V4). Pure -- vectors come from `embedding`.
+
+The three carry NO LLM call (the LLM front door is `query_constraint_extraction` + query function
+classification, upstream); they are exact, deterministic compute -> `function`, not `subgraph`.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Sequence
+
+from pydantic import BaseModel
+
+# The store seam for candidate_routing: given the routed functions, return the candidate clause ids. Bound in
+# production (LG-3c) to the store's function->pool query; injected as a fake in tests.
+PoolFn = Callable[[list[str]], list[str]]
+# EP-CORE-1b: the DOMAIN (dimension, value) match-counter, INJECTED so this retrieval mechanism imports no domain
+# vocab (the contract subsumption/canonicalization lives in `contracts.value_match.constraint_match_count`, provided
+# by the domain caller). `(query_constraints, clause_props) -> int`.
+MatchCountFn = Callable[[set, set], int]
+
+_NONE_FUNCTION = "NONE"  # the classifier's no-function sentinel; never routes a pool
+
+
+class CandidatePool(BaseModel):
+    """candidate_routing output: the unioned routing functions and the candidate clause ids they select."""
+
+    functions: list[str]
+    candidate_ids: list[str]
+
+
+class RankedClause(BaseModel):
+    """One graded candidate: its clause id and how many query constraints its props satisfy (KG-5a)."""
+
+    clause_id: str
+    match_score: float
+
+
+class MatchRanking(BaseModel):
+    """typed_constraint_match_rank output: candidates in descending graded (constraint-match) order."""
+
+    ranked: list[RankedClause]
+
+
+class DenseScoredClause(BaseModel):
+    """One candidate scored by cosine similarity to the query vector."""
+
+    clause_id: str
+    cosine: float
+
+
+class DenseRanking(BaseModel):
+    """dense_rank_tiebreak output: candidates in descending cosine order."""
+
+    ranked: list[DenseScoredClause]
+
+
+def _cosine(a: Sequence[float], b: Sequence[float]) -> float:
+    """Cosine similarity; 0.0 when either vector has zero norm (mirrors eval/kg_primary._cosine)."""
+    dot = sum(x * y for x, y in zip(a, b))
+    na = sum(x * x for x in a) ** 0.5
+    nb = sum(y * y for y in b) ** 0.5
+    return dot / (na * nb) if na and nb else 0.0
+
+
+def candidate_routing(
+    function_predictions: Sequence[Sequence[str]], *, pool_fn: PoolFn
+) -> CandidatePool:
+    """Union the ranked function predictions from each router (first-wins, order-preserving; empty and the
+    `NONE` sentinel dropped), then fetch the candidate clause pool for that function set. No functions -> an
+    empty pool with no store lookup. The recall-safe union combiner (KG-5e)."""
+    functions: list[str] = []
+    for predictions in function_predictions:
+        for function in predictions:
+            if function and function != _NONE_FUNCTION and function not in functions:
+                functions.append(function)
+    candidate_ids = pool_fn(functions) if functions else []
+    return CandidatePool(functions=functions, candidate_ids=candidate_ids)
+
+
+def typed_constraint_match_rank(
+    query_constraints: set, candidate_props: Sequence[tuple[str, set]], *, match_count_fn: MatchCountFn
+) -> MatchRanking:
+    """Grade each candidate by how many of the query's (dimension, value) constraints its grounded typed props
+    satisfy -- the match semantics are the INJECTED `match_count_fn` (the contract pack supplies KG-5a
+    canonicalization + subsumption via `constraint_match_count`; this mechanism stays domain-free). Returns
+    descending graded order. Recall-safe: the sort is stable, so a zero-match candidate keeps its input position."""
+    scored = [
+        RankedClause(clause_id=clause_id, match_score=float(match_count_fn(query_constraints, props)))
+        for clause_id, props in candidate_props
+    ]
+    ranked = sorted(scored, key=lambda r: -r.match_score)  # stable -> ties keep input order
+    return MatchRanking(ranked=ranked)
+
+
+def dense_rank_tiebreak(
+    query_vector: Sequence[float], candidate_vectors: Sequence[tuple[str, Sequence[float]]]
+) -> DenseRanking:
+    """Order candidates by descending cosine similarity to the query vector -- the embedding tiebreak signal
+    (KG-6 / V4). Pure: the vectors come from the `embedding` capability, not from a store or a model here."""
+    scored = [
+        DenseScoredClause(clause_id=clause_id, cosine=_cosine(query_vector, vector))
+        for clause_id, vector in candidate_vectors
+    ]
+    ranked = sorted(scored, key=lambda r: -r.cosine)
+    return DenseRanking(ranked=ranked)
+
+
+# (EP-CORE-1b/ADR-0118: candidate_routing / typed_constraint_match_rank / dense_rank_tiebreak are de-registered
+# from ARD -- generic retrieval primitives now, composed by direct import (the contract leg injects the domain
+# match-counter). Their register_* functions were removed.)