rerankgate 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
rerankgate/judge.py ADDED
@@ -0,0 +1,88 @@
1
+ """Faithfulness judges: heuristic default, LLM-backed opt-in."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import re
6
+ from collections.abc import Callable
7
+ from dataclasses import dataclass, field
8
+ from typing import Protocol
9
+
10
+ from rerankgate.retrieval import ScoredChunk
11
+ from rerankgate.tokenize import content_tokens
12
+
13
+ _SENT = re.compile(r"(?<=[.!?])\s+")
14
+
15
+
16
+ @dataclass
17
+ class FaithfulnessReport:
18
+ score: float # mean per-sentence grounding, [0,1]
19
+ unsupported: list[str] = field(default_factory=list) # sentences below the floor
20
+ per_sentence: list[dict] = field(default_factory=dict)
21
+
22
+ def to_dict(self) -> dict:
23
+ return {
24
+ "score": round(self.score, 4),
25
+ "unsupported": self.unsupported,
26
+ "per_sentence": self.per_sentence,
27
+ }
28
+
29
+
30
+ class Judge(Protocol):
31
+ name: str
32
+
33
+ def check(self, answer: str, evidence: list[ScoredChunk]) -> FaithfulnessReport: ...
34
+
35
+
36
+ class HeuristicJudge:
37
+ """Ground every answer sentence in the evidence via token-overlap F1.
38
+
39
+ A sentence is 'supported' if some evidence chunk shares enough content
40
+ tokens (F1 >= 0.35 by default). Heuristic — documented as such — but
41
+ free, instant, and dependency-free.
42
+ """
43
+
44
+ name = "heuristic"
45
+
46
+ def __init__(self, support_floor: float = 0.35) -> None:
47
+ self.support_floor = support_floor
48
+
49
+ def check(self, answer: str, evidence: list[ScoredChunk]) -> FaithfulnessReport:
50
+ sentences = [s.strip() for s in _SENT.split(answer.strip()) if s.strip()]
51
+ ev_sets = [set(content_tokens(c.chunk.text)) for c in evidence]
52
+ per_sentence: list[dict] = []
53
+ unsupported: list[str] = []
54
+ for s in sentences:
55
+ s_toks = set(content_tokens(s))
56
+ best = 0.0
57
+ for ev in ev_sets:
58
+ if not s_toks or not ev:
59
+ continue
60
+ inter = s_toks & ev
61
+ prec = len(inter) / len(s_toks)
62
+ rec = len(inter) / len(ev)
63
+ f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0
64
+ best = max(best, f1)
65
+ per_sentence.append({"sentence": s, "grounding": round(best, 4)})
66
+ if best < self.support_floor:
67
+ unsupported.append(s)
68
+ if per_sentence:
69
+ score = sum(p["grounding"] for p in per_sentence) / len(per_sentence)
70
+ else:
71
+ score = 0.0
72
+ return FaithfulnessReport(
73
+ score=round(score, 4), unsupported=unsupported, per_sentence=per_sentence
74
+ )
75
+
76
+
77
+ class LLMJudge:
78
+ """Wrap any callable as a model-based judge: fn(answer, evidence_texts) -> float in [0,1]."""
79
+
80
+ name = "llm"
81
+
82
+ def __init__(self, fn: Callable[[str, list[str]], float]) -> None:
83
+ self._fn = fn
84
+
85
+ def check(self, answer: str, evidence: list[ScoredChunk]) -> FaithfulnessReport:
86
+ score = float(self._fn(answer, [c.chunk.text for c in evidence]))
87
+ score = max(0.0, min(1.0, score))
88
+ return FaithfulnessReport(score=round(score, 4))
rerankgate/metrics.py ADDED
@@ -0,0 +1,106 @@
1
+ """Evaluation aggregates: retrieval metrics + gate confusion matrix."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass, field
6
+
7
+
8
+ @dataclass
9
+ class EvalItem:
10
+ query: str
11
+ relevant_ids: list[str] # chunk or doc ids considered relevant
12
+ answerable: bool = True
13
+
14
+
15
+ def _match(ids: list[str], relevant: set[str]) -> bool:
16
+ # A retrieved chunk matches if its id or its doc id is in the relevant set.
17
+ for i in ids:
18
+ if i in relevant or i.split("#c")[0] in relevant:
19
+ return True
20
+ return False
21
+
22
+
23
+ @dataclass
24
+ class EvalReport:
25
+ n: int = 0
26
+ recall_at_k: dict[int, float] = field(default_factory=dict)
27
+ mrr: float = 0.0
28
+ gate_precision: float = 0.0 # of gate-PASS, fraction actually answerable
29
+ gate_recall: float = 0.0 # of answerable, fraction gate passed
30
+ abstention_rate: float = 0.0 # fraction verdict == insufficient
31
+ answerable_rate: float = 0.0 # fraction verdict == answerable
32
+ per_query: list[dict] = field(default_factory=list)
33
+
34
+ def to_dict(self) -> dict:
35
+ return {
36
+ "n": self.n,
37
+ "recall_at_k": {str(k): round(v, 4) for k, v in self.recall_at_k.items()},
38
+ "mrr": round(self.mrr, 4),
39
+ "gate_precision": round(self.gate_precision, 4),
40
+ "gate_recall": round(self.gate_recall, 4),
41
+ "abstention_rate": round(self.abstention_rate, 4),
42
+ "answerable_rate": round(self.answerable_rate, 4),
43
+ "per_query": self.per_query,
44
+ }
45
+
46
+
47
+ def aggregate(
48
+ items: list[EvalItem],
49
+ results: list[dict],
50
+ ks: tuple[int, ...] = (1, 3, 5, 10),
51
+ ) -> EvalReport:
52
+ """Aggregate per-query pipeline results into retrieval + gate metrics.
53
+
54
+ Each result dict must have: `retrieved_ids` (ranked chunk ids),
55
+ `verdict` ("answerable"|"insufficient"|"needs_review").
56
+ """
57
+ rep = EvalReport(n=len(items))
58
+ if not items:
59
+ return rep
60
+ # Retrieval metrics are computed over answerable items only (those with
61
+ # labeled relevant ids); the gate confusion matrix covers all items.
62
+ retrievable = [(it, rs) for it, rs in zip(items, results, strict=True) if it.relevant_ids]
63
+ hits = dict.fromkeys(ks, 0)
64
+ rr_sum = 0.0
65
+ tp = fp = fn = 0 # gate PASS vs labeled answerable
66
+ abstentions = 0
67
+ answerable_out = 0
68
+ for item, res in zip(items, results, strict=True):
69
+ verdict = res.get("verdict")
70
+ passed = verdict == "answerable"
71
+ if verdict == "insufficient":
72
+ abstentions += 1
73
+ if passed:
74
+ answerable_out += 1
75
+ if passed and item.answerable:
76
+ tp += 1
77
+ elif passed and not item.answerable:
78
+ fp += 1
79
+ elif not passed and item.answerable:
80
+ fn += 1
81
+ rep.per_query.append(
82
+ {
83
+ "query": item.query,
84
+ "verdict": verdict,
85
+ "expected_answerable": item.answerable,
86
+ "retrieved_ids": res.get("retrieved_ids", [])[: max(ks)],
87
+ }
88
+ )
89
+ for item, res in retrievable:
90
+ rel = set(item.relevant_ids)
91
+ ranked = res.get("retrieved_ids", [])
92
+ for k in ks:
93
+ if _match(ranked[:k], rel):
94
+ hits[k] += 1
95
+ for rank, cid in enumerate(ranked, start=1):
96
+ if cid in rel or cid.split("#c")[0] in rel:
97
+ rr_sum += 1.0 / rank
98
+ break
99
+ denom = len(retrievable) or 1
100
+ rep.recall_at_k = {k: hits[k] / denom for k in ks}
101
+ rep.mrr = rr_sum / denom
102
+ rep.gate_precision = tp / (tp + fp) if (tp + fp) else 0.0
103
+ rep.gate_recall = tp / (tp + fn) if (tp + fn) else 0.0
104
+ rep.abstention_rate = abstentions / len(items)
105
+ rep.answerable_rate = answerable_out / len(items)
106
+ return rep
rerankgate/pipeline.py ADDED
@@ -0,0 +1,301 @@
1
+ """RAGPipeline: retrieve -> rerank -> gate -> rewrite/retry -> verdict."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import warnings
6
+ from collections.abc import Callable
7
+ from dataclasses import dataclass, field
8
+ from time import perf_counter
9
+ from typing import Literal
10
+
11
+ from rerankgate.bm25 import BM25Index
12
+ from rerankgate.config import PipelineConfig
13
+ from rerankgate.documents import Chunk, Document, chunk_text
14
+ from rerankgate.embeddings import EmbeddingBackend, SentenceTransformerBackend
15
+ from rerankgate.gate import GateDecision, SufficiencyGate
16
+ from rerankgate.judge import FaithfulnessReport, HeuristicJudge, Judge
17
+ from rerankgate.metrics import EvalItem, aggregate
18
+ from rerankgate.rerank import FeatureReranker, Reranker
19
+ from rerankgate.retrieval import HybridRetriever, RetrievalResult, ScoredChunk
20
+ from rerankgate.rewrite import RM3Rewriter
21
+ from rerankgate.store import DocumentStore
22
+ from rerankgate.vectorstore import VectorStore
23
+
24
+ ANSWERABLE = "answerable"
25
+ INSUFFICIENT = "insufficient"
26
+ NEEDS_REVIEW = "needs_review"
27
+
28
+ Verdict = Literal[ANSWERABLE, INSUFFICIENT, NEEDS_REVIEW] # type: ignore[valid-type]
29
+
30
+
31
+ @dataclass
32
+ class AttemptRecord:
33
+ query: str
34
+ retrieval_metrics: dict
35
+ gate: dict
36
+ evidence_ids: list[str]
37
+ rewrite: str | None = None
38
+
39
+
40
+ @dataclass
41
+ class AnswerResult:
42
+ verdict: Verdict
43
+ query: str
44
+ evidence: list[ScoredChunk]
45
+ gate: GateDecision
46
+ attempts: list[AttemptRecord] = field(default_factory=list)
47
+ answer: str | None = None
48
+ faithfulness: FaithfulnessReport | None = None
49
+ trace: dict = field(default_factory=dict)
50
+
51
+ def to_dict(self) -> dict:
52
+ return {
53
+ "verdict": self.verdict,
54
+ "query": self.query,
55
+ "evidence": [
56
+ {
57
+ "chunk_id": c.chunk.id,
58
+ "doc_id": c.chunk.doc_id,
59
+ "score": round(c.score, 4),
60
+ "text": c.chunk.text,
61
+ "sources": c.sources,
62
+ }
63
+ for c in self.evidence
64
+ ],
65
+ "gate": self.gate.to_dict(),
66
+ "attempts": [
67
+ {
68
+ "query": a.query,
69
+ "retrieval_metrics": a.retrieval_metrics,
70
+ "gate": a.gate,
71
+ "evidence_ids": a.evidence_ids,
72
+ "rewrite": a.rewrite,
73
+ }
74
+ for a in self.attempts
75
+ ],
76
+ "answer": self.answer,
77
+ "faithfulness": self.faithfulness.to_dict() if self.faithfulness else None,
78
+ "trace": self.trace,
79
+ }
80
+
81
+
82
+ class RAGPipeline:
83
+ """Drop-in retrieval pipeline with a sufficiency gate.
84
+
85
+ If the retrieval wasn't good enough, the answer never gets generated:
86
+ `ask()` returns verdict="insufficient" and abstains instead of
87
+ handing weak evidence to an LLM.
88
+ """
89
+
90
+ def __init__(
91
+ self,
92
+ config: PipelineConfig | None = None,
93
+ embeddings: EmbeddingBackend | None = None,
94
+ reranker: Reranker | None = None,
95
+ judge: Judge | None = None,
96
+ ) -> None:
97
+ self.config = config or PipelineConfig()
98
+ self.bm25 = BM25Index()
99
+ self.chunks: list[Chunk] = []
100
+ self.reranker: Reranker = reranker or FeatureReranker()
101
+ self.judge: Judge = judge or HeuristicJudge()
102
+ self.gate = SufficiencyGate(self.config.effective_gate())
103
+ self.rewriter = RM3Rewriter(expansion_terms=self.config.retry.expansion_terms)
104
+ self.faithfulness_floor = 0.35
105
+
106
+ self.embeddings: EmbeddingBackend | None = None
107
+ self.vector_store: VectorStore | None = None
108
+ self.degraded_reason: str | None = None
109
+ if self.config.use_vectors:
110
+ if embeddings is not None:
111
+ self._enable_vectors(embeddings)
112
+ else:
113
+ try:
114
+ self._enable_vectors(SentenceTransformerBackend())
115
+ except RuntimeError as exc:
116
+ self.degraded_reason = str(exc)
117
+ warnings.warn(
118
+ "rerankgate: embeddings unavailable, running BM25-only. "
119
+ "Install with `pip install rerankgate[local]` for hybrid mode.",
120
+ stacklevel=2,
121
+ )
122
+
123
+ def _enable_vectors(self, backend: EmbeddingBackend) -> None:
124
+ self.embeddings = backend
125
+ self.vector_store = VectorStore(dim=backend.dim)
126
+
127
+ # -- ingestion -----------------------------------------------------
128
+ def ingest(self, documents: list[Document]) -> int:
129
+ """Chunk, index, and embed documents. Returns chunk count added."""
130
+ new_chunks: list[Chunk] = []
131
+ for doc in documents:
132
+ new_chunks.extend(
133
+ chunk_text(
134
+ doc.text,
135
+ doc_id=doc.id,
136
+ size=self.config.chunk_size,
137
+ overlap=self.config.chunk_overlap,
138
+ metadata=doc.metadata,
139
+ )
140
+ )
141
+ self.chunks.extend(new_chunks)
142
+ self.bm25.add(new_chunks)
143
+ if self.embeddings is not None and self.vector_store is not None:
144
+ vecs = self.embeddings.embed([c.text for c in new_chunks])
145
+ self.vector_store.add(new_chunks, vecs)
146
+ return len(new_chunks)
147
+
148
+ def save(self, path: str) -> None:
149
+ vectors: list[list[float]] | None = None
150
+ if self.vector_store is not None:
151
+ vectors = self.vector_store.vectors
152
+ DocumentStore(path).save(
153
+ self.chunks,
154
+ vectors,
155
+ meta={"embedding": self.embeddings.name if self.embeddings else None},
156
+ )
157
+
158
+ @classmethod
159
+ def load(
160
+ cls,
161
+ path: str,
162
+ config: PipelineConfig | None = None,
163
+ embeddings: EmbeddingBackend | None = None,
164
+ ) -> RAGPipeline:
165
+ """Load a persisted store and rebuild indexes."""
166
+ store = DocumentStore(path)
167
+ chunks, vectors, meta = store.load()
168
+ pipe = cls(config=config, embeddings=embeddings)
169
+ # If the store has vectors but this process has no embedding backend
170
+ # (e.g. sentence-transformers missing), stay BM25-only: retrieval
171
+ # still works, quality just degrades gracefully.
172
+ if vectors and pipe.vector_store is not None and pipe.embeddings is not None:
173
+ if len(vectors) == len(chunks) and len(vectors[0]) == pipe.embeddings.dim:
174
+ pipe.vector_store.add(chunks, vectors)
175
+ # else: dim mismatch (different model) -> fall back to re-embed
176
+ # lazily on next ingest; existing vectors are not trusted.
177
+ pipe.chunks = chunks
178
+ pipe.bm25.add(chunks)
179
+ if vectors and (pipe.vector_store is None or len(pipe.vector_store) == 0):
180
+ # Re-embed with the current backend if dims allow.
181
+ if pipe.embeddings is not None:
182
+ pipe.vector_store = VectorStore(dim=pipe.embeddings.dim)
183
+ pipe.vector_store.add(chunks, pipe.embeddings.embed([c.text for c in chunks]))
184
+ return pipe
185
+
186
+ # -- querying ------------------------------------------------------
187
+ def _retriever(self) -> HybridRetriever:
188
+ retr = HybridRetriever(
189
+ self.bm25,
190
+ vector_store=self.vector_store,
191
+ embeddings=self.embeddings,
192
+ rrf_k=self.config.rrf_k,
193
+ bm25_weight=self.config.bm25_weight,
194
+ vector_weight=self.config.vector_weight,
195
+ )
196
+ retr.degraded = self.degraded_reason
197
+ return retr
198
+
199
+ def retrieve(self, query: str, k: int | None = None) -> RetrievalResult:
200
+ return self._retriever().retrieve(query, k=k or self.config.retrieve_k)
201
+
202
+ def ask(
203
+ self,
204
+ query: str,
205
+ generator: Callable[[str, list[ScoredChunk]], str] | None = None,
206
+ gate_override: bool = False,
207
+ ) -> AnswerResult:
208
+ """Run retrieve -> rerank -> gate, with rewrite retries, then verdict.
209
+
210
+ If `generator` is given and the gate passes, the generated answer is
211
+ faithfulness-checked; a failing check flips the verdict to needs_review.
212
+
213
+ `gate_override=True` is the escape hatch for suspected gate false
214
+ negatives: the best evidence is returned as answerable regardless of
215
+ the gate, and the override is recorded in the trace.
216
+ """
217
+ t0 = perf_counter()
218
+ cfg = self.config
219
+ attempts: list[AttemptRecord] = []
220
+ best: tuple[float, list[ScoredChunk], GateDecision, str] | None = None
221
+ current_query = query
222
+
223
+ for attempt in range(cfg.retry.max_retries + 1):
224
+ retrieval = self._retriever().retrieve(current_query, k=cfg.retrieve_k)
225
+ reranked = self.reranker.rerank(current_query, retrieval.chunks, k=cfg.rerank_k)
226
+ decision = self.gate.evaluate(current_query, reranked)
227
+ attempts.append(
228
+ AttemptRecord(
229
+ query=current_query,
230
+ retrieval_metrics=retrieval.metrics,
231
+ gate=decision.to_dict(),
232
+ evidence_ids=[c.chunk.id for c in reranked],
233
+ )
234
+ )
235
+ prev_best = best[0] if best is not None else -1.0
236
+ if best is None or decision.score > best[0]:
237
+ best = (decision.score, reranked, decision, current_query)
238
+ if decision.passed:
239
+ break
240
+ # Diminishing-returns early stop: don't burn retries going nowhere.
241
+ improved = decision.score > prev_best + cfg.retry.min_improvement
242
+ if len(attempts) >= 2 and not improved:
243
+ break
244
+ if attempt < cfg.retry.max_retries:
245
+ rewritten = self.rewriter.rewrite(current_query, reranked)
246
+ if rewritten == current_query:
247
+ break # rewriter found nothing new; further retries are pointless
248
+ attempts[-1].rewrite = rewritten
249
+ current_query = rewritten
250
+
251
+ assert best is not None
252
+ _, evidence, decision, _ = best
253
+ verdict: Verdict = ANSWERABLE if decision.passed else INSUFFICIENT
254
+ overridden = False
255
+ if gate_override and verdict == INSUFFICIENT:
256
+ verdict = ANSWERABLE
257
+ overridden = True
258
+
259
+ answer: str | None = None
260
+ faithfulness: FaithfulnessReport | None = None
261
+ if verdict == ANSWERABLE and generator is not None:
262
+ answer = generator(query, evidence)
263
+ faithfulness = self.judge.check(answer, evidence)
264
+ if faithfulness.score < self.faithfulness_floor:
265
+ verdict = NEEDS_REVIEW
266
+
267
+ total_ms = round((perf_counter() - t0) * 1000, 2)
268
+ trace = {
269
+ "modes": self._retriever().modes,
270
+ "degraded": self.degraded_reason,
271
+ "attempts": len(attempts),
272
+ "retries_used": len(attempts) - 1,
273
+ "reranker": self.reranker.name,
274
+ "judge": self.judge.name,
275
+ "gate_overridden": overridden,
276
+ "total_ms": total_ms,
277
+ "chunk_count": len(self.chunks),
278
+ }
279
+ return AnswerResult(
280
+ verdict=verdict,
281
+ query=query,
282
+ evidence=evidence,
283
+ gate=decision,
284
+ attempts=attempts,
285
+ answer=answer,
286
+ faithfulness=faithfulness,
287
+ trace=trace,
288
+ )
289
+
290
+ def evaluate(self, items: list[EvalItem]) -> dict:
291
+ """Run the eval harness: retrieval metrics + gate confusion matrix."""
292
+ results = []
293
+ for item in items:
294
+ res = self.ask(item.query)
295
+ results.append(
296
+ {
297
+ "retrieved_ids": [c.chunk.id for c in res.evidence],
298
+ "verdict": res.verdict,
299
+ }
300
+ )
301
+ return aggregate(items, results).to_dict()
rerankgate/rerank.py ADDED
@@ -0,0 +1,94 @@
1
+ """Rerankers: zero-dependency feature reranker (default) + cross-encoder (opt-in)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from time import perf_counter
6
+ from typing import Protocol
7
+
8
+ from rerankgate.retrieval import ScoredChunk
9
+ from rerankgate.tokenize import content_tokens, tokenize
10
+
11
+
12
+ class Reranker(Protocol):
13
+ name: str
14
+
15
+ def rerank(self, query: str, candidates: list[ScoredChunk], k: int) -> list[ScoredChunk]: ...
16
+
17
+
18
+ class FeatureReranker:
19
+ """Deterministic, explainable reranker. No model required.
20
+
21
+ score = 0.5 * norm(fusion) + 0.3 * term_coverage + 0.2 * phrase_bonus
22
+ """
23
+
24
+ name = "feature"
25
+
26
+ def rerank(self, query: str, candidates: list[ScoredChunk], k: int) -> list[ScoredChunk]:
27
+ t0 = perf_counter()
28
+ q_terms = set(content_tokens(query))
29
+ q_phrases = _phrases(query)
30
+ max_fused = max((c.score for c in candidates), default=1.0) or 1.0
31
+
32
+ rescored: list[ScoredChunk] = []
33
+ for c in candidates:
34
+ toks = set(tokenize(c.chunk.text))
35
+ coverage = len(q_terms & toks) / len(q_terms) if q_terms else 0.0
36
+ text_l = c.chunk.text.lower()
37
+ phrase_bonus = 1.0 if q_phrases and any(p in text_l for p in q_phrases) else 0.0
38
+ norm_fused = c.score / max_fused
39
+ score = 0.5 * norm_fused + 0.3 * coverage + 0.2 * phrase_bonus
40
+ rescored.append(
41
+ ScoredChunk(
42
+ chunk=c.chunk,
43
+ score=score,
44
+ sources={
45
+ **c.sources,
46
+ "rerank": round(score, 4),
47
+ "coverage": round(coverage, 4),
48
+ "phrase": phrase_bonus,
49
+ },
50
+ )
51
+ )
52
+ rescored.sort(key=lambda c: c.score, reverse=True)
53
+ self.last_ms = round((perf_counter() - t0) * 1000, 2)
54
+ return rescored[:k]
55
+
56
+
57
+ def _phrases(query: str) -> list[str]:
58
+ toks = content_tokens(query)
59
+ return [" ".join(toks[i : i + 2]) for i in range(len(toks) - 1)]
60
+
61
+
62
+ class CrossEncoderReranker:
63
+ """Cross-encoder reranker (ms-marco MiniLM). Requires `rerankgate[local]`."""
64
+
65
+ name = "cross-encoder"
66
+
67
+ def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2") -> None:
68
+ try:
69
+ from sentence_transformers import CrossEncoder # type: ignore
70
+ except ImportError as exc:
71
+ raise RuntimeError(
72
+ "sentence-transformers is not installed. "
73
+ "Install it with `pip install rerankgate[local]`."
74
+ ) from exc
75
+ self._model = CrossEncoder(model_name)
76
+ self.last_ms = 0.0
77
+
78
+ def rerank(self, query: str, candidates: list[ScoredChunk], k: int) -> list[ScoredChunk]:
79
+ import math
80
+
81
+ if not candidates:
82
+ return []
83
+ t0 = perf_counter()
84
+ pairs = [(query, c.chunk.text) for c in candidates]
85
+ logits = self._model.predict(pairs)
86
+ rescored = []
87
+ for c, logit in zip(candidates, logits, strict=True):
88
+ score = 1 / (1 + math.exp(-float(logit))) # sigmoid to [0,1]
89
+ rescored.append(
90
+ ScoredChunk(chunk=c.chunk, score=score, sources={**c.sources, "rerank": score})
91
+ )
92
+ rescored.sort(key=lambda c: c.score, reverse=True)
93
+ self.last_ms = round((perf_counter() - t0) * 1000, 2)
94
+ return rescored[:k]
@@ -0,0 +1,99 @@
1
+ """Hybrid retrieval: BM25 + vector search fused with Reciprocal Rank Fusion."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass, field
6
+ from time import perf_counter
7
+ from typing import Any
8
+
9
+ from rerankgate.bm25 import BM25Index
10
+ from rerankgate.documents import Chunk
11
+ from rerankgate.embeddings import EmbeddingBackend
12
+ from rerankgate.tokenize import tokenize
13
+ from rerankgate.vectorstore import VectorStore
14
+
15
+
16
+ @dataclass
17
+ class ScoredChunk:
18
+ chunk: Chunk
19
+ score: float
20
+ sources: dict[str, float] = field(default_factory=dict) # per-system raw scores
21
+
22
+
23
+ @dataclass
24
+ class RetrievalResult:
25
+ chunks: list[ScoredChunk]
26
+ metrics: dict[str, Any]
27
+
28
+
29
+ class HybridRetriever:
30
+ """Fuses BM25 and vector rankings with RRF (rank-based, no score normalization)."""
31
+
32
+ def __init__(
33
+ self,
34
+ bm25: BM25Index,
35
+ vector_store: VectorStore | None = None,
36
+ embeddings: EmbeddingBackend | None = None,
37
+ rrf_k: int = 60,
38
+ bm25_weight: float = 1.0,
39
+ vector_weight: float = 1.0,
40
+ ) -> None:
41
+ self.bm25 = bm25
42
+ self.vector_store = vector_store
43
+ self.embeddings = embeddings
44
+ self.rrf_k = rrf_k
45
+ self.bm25_weight = bm25_weight
46
+ self.vector_weight = vector_weight
47
+ self.degraded: str | None = None # reason if running BM25-only
48
+
49
+ @property
50
+ def modes(self) -> list[str]:
51
+ modes = ["bm25"]
52
+ if self.vector_store is not None and self.embeddings is not None:
53
+ modes.append("vector")
54
+ return modes
55
+
56
+ def retrieve(self, query: str, k: int = 100) -> RetrievalResult:
57
+ t0 = perf_counter()
58
+ rankings: list[tuple[str, list[tuple[Chunk, float]], float]] = []
59
+
60
+ t_bm25 = perf_counter()
61
+ bm25_hits = self.bm25.search(query, k=k)
62
+ bm25_ms = (perf_counter() - t_bm25) * 1000
63
+ rankings.append(("bm25", bm25_hits, self.bm25_weight))
64
+
65
+ vector_ms = 0.0
66
+ if self.vector_store is not None and self.embeddings is not None:
67
+ t_vec = perf_counter()
68
+ qvec = self.embeddings.embed([query])[0]
69
+ vec_hits = self.vector_store.search(qvec, k=k)
70
+ vector_ms = (perf_counter() - t_vec) * 1000
71
+ rankings.append(("vector", vec_hits, self.vector_weight))
72
+
73
+ # Reciprocal Rank Fusion over ranks (1-based).
74
+ fused: dict[str, float] = {}
75
+ by_id: dict[str, Chunk] = {}
76
+ sources: dict[str, dict[str, float]] = {}
77
+ for name, hits, weight in rankings:
78
+ for rank, (ch, raw) in enumerate(hits, start=1):
79
+ fused[ch.id] = fused.get(ch.id, 0.0) + weight / (self.rrf_k + rank)
80
+ by_id[ch.id] = ch
81
+ sources.setdefault(ch.id, {})[name] = raw
82
+
83
+ ordered = sorted(fused.items(), key=lambda kv: kv[1], reverse=True)[:k]
84
+ chunks = [
85
+ ScoredChunk(chunk=by_id[cid], score=s, sources=dict(sources[cid])) for cid, s in ordered
86
+ ]
87
+ total_ms = (perf_counter() - t0) * 1000
88
+ metrics = {
89
+ "modes": self.modes,
90
+ "degraded": self.degraded,
91
+ "k": k,
92
+ "bm25_hits": len(bm25_hits),
93
+ "bm25_ms": round(bm25_ms, 2),
94
+ "vector_ms": round(vector_ms, 2),
95
+ "total_ms": round(total_ms, 2),
96
+ "fused_count": len(chunks),
97
+ "query_terms": len(tokenize(query)),
98
+ }
99
+ return RetrievalResult(chunks=chunks, metrics=metrics)