Coverage for agentos/rag/hybrid.py: 0%
100 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 01:44 +0800
1"""Hybrid search (dense + sparse) for RAG pipeline.
3Combines dense vector search with BM25 sparse retrieval
4using reciprocal rank fusion (RRF) or weighted score fusion.
5"""
7from __future__ import annotations
9import math
10from dataclasses import dataclass
11from typing import Any
14@dataclass
15class HybridConfig:
16 """Configuration for hybrid search."""
18 dense_weight: float = 0.6 # weight for dense scores
19 sparse_weight: float = 0.4 # weight for BM25 scores
20 fusion_method: str = "weighted" # "weighted" | "rrf"
21 rrf_k: int = 60 # RRF constant
22 bm25_k1: float = 1.5 # BM25 term frequency saturation
23 bm25_b: float = 0.75 # BM25 document length normalization
24 top_k_per_source: int = 20 # candidates from each retriever before fusion
27class BM25Retriever:
28 """BM25 sparse retrieval with Okapi BM25 scoring.
30 Works with pre-chunked documents, builds an in-memory inverted index.
31 """
33 def __init__(
34 self,
35 k1: float = 1.5,
36 b: float = 0.75,
37 stop_words: list[str] | None = None,
38 ):
39 self.k1 = k1
40 self.b = b
41 self.stop_words = set(stop_words or _DEFAULT_STOP_WORDS)
42 self._docs: list[str] = []
43 self._doc_lens: list[int] = []
44 self._avgdl: float = 0.0
45 self._df: dict[str, int] = {} # term -> document frequency
46 self._term_freqs: list[dict[str, int]] = [] # per-doc term freqs
47 self._built = False
49 def _tokenize(self, text: str) -> list[str]:
50 """Simple whitespace + punctuation tokenization."""
51 import re
53 tokens = re.findall(r"\w+", text.lower())
54 return [t for t in tokens if t not in self.stop_words and len(t) > 1]
56 def index(self, documents: list[str]):
57 """Build BM25 index from documents."""
58 self._docs = documents
59 self._doc_lens = [len(d) for d in documents]
60 self._avgdl = sum(self._doc_lens) / max(len(documents), 1)
61 self._df = {}
62 self._term_freqs = []
64 for doc in documents:
65 tokens = self._tokenize(doc)
66 tf = {}
67 for t in tokens:
68 tf[t] = tf.get(t, 0) + 1
69 self._term_freqs.append(tf)
70 for t in tf:
71 self._df[t] = self._df.get(t, 0) + 1
73 self._built = True
75 def search(self, query: str, top_k: int = 10) -> list[tuple[int, float]]:
76 """Search and return (doc_index, score) sorted by score descending."""
77 if not self._built:
78 return []
80 query_tokens = self._tokenize(query)
81 idf_cache = {
82 t: math.log(1 + (len(self._docs) - freq + 0.5) / (freq + 0.5))
83 for t, freq in self._df.items()
84 if t in query_tokens
85 }
87 scores = []
88 for i, tf in enumerate(self._term_freqs):
89 score = 0.0
90 for t in query_tokens:
91 if t not in tf:
92 continue
93 idf = idf_cache.get(t, 0.0)
94 f = tf[t]
95 dl = self._doc_lens[i]
96 numerator = f * (self.k1 + 1)
97 denominator = f + self.k1 * (1 - self.b + self.b * dl / self._avgdl)
98 score += idf * numerator / denominator
99 if score > 0:
100 scores.append((i, score))
102 scores.sort(key=lambda x: x[1], reverse=True)
103 return scores[:top_k]
106class HybridRetriever:
107 """Combined dense + sparse retrieval with score fusion.
109 Usage:
110 retriever = HybridRetriever(
111 dense_fn=your_dense_search_fn,
112 bm25=bm25_retriever,
113 )
114 results = await retriever.search(query="how to train a model", top_k=5)
115 """
117 def __init__(
118 self,
119 dense_fn,
120 bm25: BM25Retriever | None = None,
121 config: HybridConfig | None = None,
122 ):
123 self.dense_fn = dense_fn
124 self.bm25 = bm25 or BM25Retriever()
125 self.config = config or HybridConfig()
127 def index_documents(self, documents: list[str]):
128 """Index documents for BM25 sparse retrieval."""
129 self.bm25.index(documents)
131 async def search(
132 self,
133 query: str,
134 top_k: int = 5,
135 ) -> list[dict[str, Any]]:
136 """Hybrid search combining dense and sparse scores.
138 Returns list of dicts with 'text', 'score', 'dense_score',
139 'sparse_score', 'index'.
140 """
141 # Dense retrieval
142 dense_results = await self.dense_fn(query, self.config.top_k_per_source)
144 # BM25 retrieval
145 bm25_pairs = self.bm25.search(query, self.config.top_k_per_source)
147 # Normalize and fuse scores
148 fused = self._fuse_scores(dense_results, bm25_pairs)
149 fused.sort(key=lambda x: x["score"], reverse=True)
151 return fused[:top_k]
153 def _fuse_scores(
154 self,
155 dense_results: list[dict[str, Any]],
156 bm25_pairs: list[tuple[int, float]],
157 ) -> list[dict[str, Any]]:
158 """Fuse dense and sparse scores using configured method."""
159 # Build lookup: doc_index -> result
160 index_map: dict[int, dict[str, Any]] = {}
161 for i, r in enumerate(dense_results):
162 idx = r.get("index", i)
163 index_map[idx] = {
164 "text": r.get("text", ""),
165 "dense_score": r.get("score", 0.0),
166 "sparse_score": 0.0,
167 "index": idx,
168 "metadata": r.get("metadata", {}),
169 "dense_rank": i + 1, # 1-based rank
170 "sparse_rank": 0,
171 }
173 for rank, (idx, bm25_score) in enumerate(bm25_pairs):
174 rank_p1 = rank + 1
175 if idx in index_map:
176 index_map[idx]["sparse_score"] = bm25_score
177 index_map[idx]["sparse_rank"] = rank_p1
178 else:
179 index_map[idx] = {
180 "text": "",
181 "dense_score": 0.0,
182 "sparse_score": bm25_score,
183 "index": idx,
184 "metadata": {},
185 "dense_rank": 0,
186 "sparse_rank": rank_p1,
187 }
189 # Compute fused score
190 for idx, entry in index_map.items():
191 if self.config.fusion_method == "rrf":
192 dr = entry["dense_rank"] or (self.config.top_k_per_source + 1)
193 sr = entry["sparse_rank"] or (self.config.top_k_per_source + 1)
194 entry["score"] = 1.0 / (self.config.rrf_k + dr) + 1.0 / (self.config.rrf_k + sr)
195 else:
196 entry["score"] = self.config.dense_weight * entry[
197 "dense_score"
198 ] + self.config.sparse_weight * self._normalize_bm25(entry["sparse_score"])
200 return list(index_map.values())
202 def _normalize_bm25(self, score: float) -> float:
203 """Simple sigmoid normalization for BM25 scores."""
204 if score <= 0:
205 return 0.0
206 return 2.0 / (1.0 + math.exp(-score / 3.0)) - 1.0
209_DEFAULT_STOP_WORDS = {
210 "a",
211 "an",
212 "the",
213 "and",
214 "or",
215 "but",
216 "in",
217 "on",
218 "at",
219 "to",
220 "for",
221 "of",
222 "with",
223 "by",
224 "from",
225 "is",
226 "are",
227 "was",
228 "were",
229 "be",
230 "been",
231 "being",
232 "have",
233 "has",
234 "had",
235 "do",
236 "does",
237 "did",
238 "will",
239 "would",
240 "could",
241 "should",
242 "may",
243 "might",
244 "can",
245 "shall",
246 "it",
247 "its",
248 "this",
249 "that",
250 "these",
251 "those",
252 "i",
253 "you",
254 "he",
255 "she",
256 "they",
257 "we",
258 "my",
259 "your",
260 "his",
261 "her",
262 "our",
263 "their",
264 "not",
265 "no",
266 "if",
267 "so",
268 "as",
269 "than",
270 "then",
271 "just",
272 "about",
273 "also",
274 "very",
275 "too",
276 "into",
277 "over",
278 "after",
279 "before",
280 "between",
281 "under",
282 "more",
283 "up",
284 "out",
285 "some",
286 "such",
287 "only",
288 "other",
289 "each",
290 "all",
291 "both",
292 "few",
293 "most",
294}