Coverage for agentos/rag/hybrid_search.py: 23%
267 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +0800
1"""
2Hybrid Search + Re-Ranking for RAG (v1.9.0)
4Production-grade hybrid search combining:
5 - Dense (semantic) retrieval via embeddings
6 - Sparse (keyword) retrieval via BM25
7 - Cross-encoder re-ranking for precision
8 - Citation tracking with source provenance
9 - Multi-modal: text, code, markdown, tables
10 - Fusion algorithms: RRF, weighted sum, cascade
12Compatible with existing ChromaStore + RAGPipeline.
13"""
15from __future__ import annotations
17import math
18import re
19from collections import Counter, defaultdict
20from dataclasses import dataclass, field
21from typing import Any, Optional, Callable
24# ── Types ───────────────────────────────────────────────────────────
26@dataclass
27class SearchResult:
28 """A single search result with metadata."""
29 doc_id: str
30 content: str
31 source: str = "" # File path, URL, or source identifier
32 title: str = ""
33 score: float = 0.0
34 dense_score: float = 0.0
35 sparse_score: float = 0.0
36 rerank_score: float = 0.0
37 chunk_index: int = 0
38 metadata: dict[str, Any] = field(default_factory=dict)
39 citations: list[str] = field(default_factory=list) # Specific sentences/quotes
42@dataclass
43class Citation:
44 """A citation from source material."""
45 text: str
46 source: str
47 doc_id: str = ""
48 chunk_index: int = 0
49 start_pos: int = 0
50 end_pos: int = 0
51 confidence: float = 1.0
54# ── BM25 Sparse Retriever ───────────────────────────────────────────
56class BM25Retriever:
57 """Pure Python BM25 implementation for keyword search.
59 No external dependencies. Tokenizes, builds inverted index,
60 and scores documents using Okapi BM25.
61 """
63 def __init__(self, k1: float = 1.5, b: float = 0.75):
64 self.k1 = k1
65 self.b = b
66 self._docs: list[str] = []
67 self._doc_ids: list[str] = []
68 self._doc_lengths: list[int] = []
69 self._avg_dl: float = 0.0
70 self._inverted_index: dict[str, dict[int, int]] = defaultdict(dict)
71 self._idf: dict[str, float] = {}
72 self._N: int = 0
74 def index(self, documents: list[dict[str, str]]):
75 """Build BM25 index from documents.
77 Args:
78 documents: List of {id, content} dicts.
79 """
80 self._docs = [doc.get("content", "") for doc in documents]
81 self._doc_ids = [doc.get("id", f"doc_{i}") for i, doc in enumerate(documents)]
82 self._doc_lengths = [len(self._tokenize(doc)) for doc in self._docs]
83 self._N = len(self._docs)
84 self._avg_dl = sum(self._doc_lengths) / max(self._N, 1)
86 # Build inverted index
87 self._inverted_index.clear()
88 doc_freq: dict[str, int] = defaultdict(int)
90 for doc_id, doc in enumerate(self._docs):
91 tokens = self._tokenize(doc)
92 token_counts = Counter(tokens)
93 for token, count in token_counts.items():
94 self._inverted_index[token][doc_id] = count
95 doc_freq[token] += 1
97 # Compute IDF
98 self._idf = {
99 token: math.log(1 + (self._N - freq + 0.5) / (freq + 0.5))
100 for token, freq in doc_freq.items()
101 }
103 def search(self, query: str, top_k: int = 10) -> list[SearchResult]:
104 """BM25 keyword search."""
105 if not self._docs:
106 return []
108 query_tokens = self._tokenize(query)
109 scores: list[float] = [0.0] * self._N
111 for token in query_tokens:
112 if token not in self._inverted_index:
113 continue
114 idf = self._idf.get(token, 0)
115 for doc_id, tf in self._inverted_index[token].items():
116 dl = self._doc_lengths[doc_id]
117 numerator = tf * (self.k1 + 1)
118 denominator = tf + self.k1 * (1 - self.b + self.b * dl / max(self._avg_dl, 1))
119 scores[doc_id] += idf * numerator / max(denominator, 1e-9)
121 # Rank and return
122 ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
123 max_score = ranked[0][1] if ranked else 1.0
125 return [
126 SearchResult(
127 doc_id=self._doc_ids[doc_id],
128 content=self._docs[doc_id][:500],
129 sparse_score=score / max(max_score, 1e-9),
130 score=score / max(max_score, 1e-9),
131 metadata={"method": "bm25"},
132 )
133 for doc_id, score in ranked[:top_k] if score > 0
134 ]
136 def _tokenize(self, text: str) -> list[str]:
137 """Simple tokenization: lowercase, split on non-alphanumeric, filter short tokens."""
138 tokens = re.findall(r'[\w\u4e00-\u9fff]+', text.lower())
139 return [t for t in tokens if len(t) > 1]
142# ── Dense Retriever ─────────────────────────────────────────────────
144class DenseRetriever:
145 """Semantic search via embeddings.
147 Wraps an embedding function (e.g., OpenAI embeddings, sentence-transformers)
148 and a vector store (ChromaDB or similar).
149 """
151 def __init__(
152 self,
153 vector_store=None,
154 embed_fn: Optional[Callable[[str], list[float]]] = None,
155 ):
156 self._store = vector_store
157 self._embed = embed_fn
159 async def search(self, query: str, top_k: int = 10) -> list[SearchResult]:
160 """Dense vector search."""
161 if not self._store:
162 return []
164 try:
165 results = await self._store.search(query, top_k=top_k)
167 max_score = results[0].get("score", 1.0) if results else 1.0
169 return [
170 SearchResult(
171 doc_id=result.get("id", ""),
172 content=result.get("content", "")[:500],
173 dense_score=result.get("score", 0) / max(max_score, 1e-9),
174 score=result.get("score", 0) / max(max_score, 1e-9),
175 metadata=result.get("metadata", {}),
176 )
177 for result in results
178 ]
179 except Exception:
180 return []
183# ── Cross-Encoder Re-Ranker ─────────────────────────────────────────
185class CrossEncoderReranker:
186 """Re-rank search results with a cross-encoder model.
188 Instead of embedding query and documents independently (bi-encoder),
189 a cross-encoder processes (query, document) pairs together for higher
190 accuracy — at the cost of more computation.
192 Supports:
193 - HuggingFace cross-encoder models (e.g., ms-marco-MiniLM)
194 - Custom scoring functions
195 - LLM-based re-ranking (use an LLM to judge relevance)
196 """
198 def __init__(
199 self,
200 model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
201 use_llm: bool = False,
202 llm_client=None,
203 ):
204 self._model_name = model_name
205 self._model = None
206 self._use_llm = use_llm
207 self._llm = llm_client
209 async def rerank(
210 self,
211 query: str,
212 candidates: list[SearchResult],
213 top_k: int = 5,
214 ) -> list[SearchResult]:
215 """Re-rank candidates by relevance to query.
217 Args:
218 query: Original search query
219 candidates: Initial retrieval results
220 top_k: Number of results to return after re-ranking
222 Returns:
223 Re-ranked candidates with updated rerank_score.
224 """
225 if not candidates:
226 return []
228 if self._use_llm and self._llm:
229 return await self._llm_rerank(query, candidates, top_k)
230 else:
231 return await self._cross_encoder_rerank(query, candidates, top_k)
233 async def _cross_encoder_rerank(
234 self,
235 query: str,
236 candidates: list[SearchResult],
237 top_k: int,
238 ) -> list[SearchResult]:
239 """Re-rank using HuggingFace cross-encoder."""
240 try:
241 from sentence_transformers import CrossEncoder
242 if self._model is None:
243 self._model = CrossEncoder(self._model_name)
245 pairs = [(query, c.content[:1000]) for c in candidates]
246 scores = self._model.predict(pairs)
248 for candidate, score in zip(candidates, scores):
249 candidate.rerank_score = float(score)
250 # Weighted fusion
251 candidate.score = (
252 candidate.dense_score * 0.3 +
253 candidate.sparse_score * 0.2 +
254 float(score) * 0.5
255 )
257 candidates.sort(key=lambda x: x.rerank_score, reverse=True)
258 return candidates[:top_k]
260 except ImportError:
261 return candidates[:top_k] # Fallback: no re-ranking
263 async def _llm_rerank(
264 self,
265 query: str,
266 candidates: list[SearchResult],
267 top_k: int,
268 ) -> list[SearchResult]:
269 """Re-rank using LLM relevance judgment."""
270 if not self._llm:
271 return candidates[:top_k]
273 prompt = f"Query: {query}\n\nRate each document's relevance on a scale of 0-10:\n\n"
274 for i, c in enumerate(candidates[:20]):
275 prompt += f"[{i}] {c.content[:300]}\n\n"
276 prompt += "Output format: [doc_id] score"
278 try:
279 response = await self._llm.complete(prompt)
280 # Parse scores
281 scores: dict[int, float] = {}
282 for line in response.split("\n"):
283 match = re.match(r'\[(\d+)\]\s*(\d+(?:\.\d+)?)', line.strip())
284 if match:
285 idx = int(match.group(1))
286 score = float(match.group(2)) / 10.0
287 if idx < len(candidates):
288 scores[idx] = score
290 for i, candidate in enumerate(candidates):
291 candidate.rerank_score = scores.get(i, 0.5)
292 candidate.score = (
293 candidate.dense_score * 0.25 +
294 candidate.sparse_score * 0.15 +
295 candidate.rerank_score * 0.6
296 )
298 candidates.sort(key=lambda x: x.rerank_score, reverse=True)
299 return candidates[:top_k]
301 except Exception:
302 return candidates[:top_k]
305# ── Fusion Algorithms ───────────────────────────────────────────────
307class FusionMethod:
308 """Collection of rank fusion algorithms."""
310 @staticmethod
311 def reciprocal_rank_fusion(
312 dense_results: list[SearchResult],
313 sparse_results: list[SearchResult],
314 k: int = 60,
315 ) -> list[SearchResult]:
316 """RRF: Reciprocal Rank Fusion.
318 RRF_score(d) = sum_{ranker} 1 / (k + rank(d))
319 """
320 scores: dict[str, float] = {}
321 docs: dict[str, SearchResult] = {}
323 for rank, result in enumerate(dense_results):
324 scores[result.doc_id] = 1.0 / (k + rank + 1)
325 docs[result.doc_id] = result
327 for rank, result in enumerate(sparse_results):
328 if result.doc_id in scores:
329 scores[result.doc_id] += 1.0 / (k + rank + 1)
330 else:
331 scores[result.doc_id] = 1.0 / (k + rank + 1)
332 docs[result.doc_id] = result
334 fused = sorted(scores.items(), key=lambda x: x[1], reverse=True)
335 results = []
336 for doc_id, score in fused:
337 doc = docs[doc_id]
338 doc.score = score
339 results.append(doc)
341 return results
343 @staticmethod
344 def weighted_sum(
345 dense_results: list[SearchResult],
346 sparse_results: list[SearchResult],
347 dense_weight: float = 0.6,
348 sparse_weight: float = 0.4,
349 ) -> list[SearchResult]:
350 """Weighted score summation."""
351 scores: dict[str, list[float]] = defaultdict(list)
352 docs: dict[str, SearchResult] = {}
354 for result in dense_results:
355 scores[result.doc_id].append(result.dense_score * dense_weight)
356 docs[result.doc_id] = result
358 for result in sparse_results:
359 scores[result.doc_id].append(result.sparse_score * sparse_weight)
360 if result.doc_id not in docs:
361 docs[result.doc_id] = result
363 fused = []
364 for doc_id, wscores in scores.items():
365 doc = docs[doc_id]
366 doc.score = sum(wscores)
367 fused.append(doc)
369 fused.sort(key=lambda x: x.score, reverse=True)
370 return fused
372 @staticmethod
373 def cascade(
374 dense_results: list[SearchResult],
375 sparse_results: list[SearchResult],
376 ) -> list[SearchResult]:
377 """Cascade: dense first, then sparse fills gaps."""
378 seen: set[str] = set()
379 results: list[SearchResult] = []
381 for r in dense_results:
382 results.append(r)
383 seen.add(r.doc_id)
385 for r in sparse_results:
386 if r.doc_id not in seen:
387 results.append(r)
388 seen.add(r.doc_id)
390 return results
393# ── Citation Tracker ────────────────────────────────────────────────
395class CitationTracker:
396 """Track and verify citations from source documents.
398 Key features:
399 - Extract citations from generated text
400 - Verify against source documents
401 - Mark unverifiable (potential hallucination)
402 - Track citation usage statistics
403 """
405 def __init__(self):
406 self._citations: list[Citation] = []
407 self._source_index: dict[str, dict] = {} # doc_id → metadata
409 def add_source(self, doc_id: str, content: str, metadata: dict[str, Any] | None = None):
410 """Register a source document."""
411 self._source_index[doc_id] = {
412 "content": content,
413 "metadata": metadata or {},
414 }
416 def extract_citations(self, text: str, sources: list[SearchResult]) -> list[Citation]:
417 """Extract and verify citations from generated text.
419 Args:
420 text: Generated response text
421 sources: Source documents used for generation
423 Returns:
424 List of verified Citation objects.
425 """
426 citations: list[Citation] = []
428 for source in sources:
429 # Find substrings of generated text that appear in source
430 source_content = source.content.lower()
431 text_lower = text.lower()
433 # Extract sentences from generated text
434 sentences = re.split(r'[.!?]+', text)
435 for sent in sentences:
436 sent = sent.strip()
437 if len(sent) < 15:
438 continue
440 # Check if this sentence appears in source (with fuzzy matching)
441 if self._is_from_source(sent.lower(), source_content):
442 citations.append(Citation(
443 text=sent,
444 source=source.source or source.doc_id,
445 doc_id=source.doc_id,
446 chunk_index=source.chunk_index,
447 confidence=0.9,
448 ))
450 # Deduplicate
451 seen: set[str] = set()
452 unique = []
453 for c in citations:
454 key = c.text[:50]
455 if key not in seen:
456 seen.add(key)
457 unique.append(c)
459 self._citations.extend(unique)
460 return unique
462 def verify(self, text: str, sources: list[SearchResult]) -> dict[str, Any]:
463 """Verify all claims in text against source documents.
465 Returns:
466 Dict with verified/unverified segments and hallucination score.
467 """
468 citations = self.extract_citations(text, sources)
470 sentences = re.split(r'[.!?]+', text)
471 total_sentences = len(sentences)
472 cited_sentences = sum(
473 1 for s in sentences
474 if any(c.text[:30].lower() in s.strip().lower() for c in citations)
475 )
477 uncited = total_sentences - cited_sentences
478 hallucination_risk = uncited / max(total_sentences, 1)
480 return {
481 "total_sentences": total_sentences,
482 "cited_sentences": cited_sentences,
483 "uncited_sentences": uncited,
484 "hallucination_risk": round(hallucination_risk, 3),
485 "citations": [
486 {"text": c.text[:100], "source": c.source, "confidence": c.confidence}
487 for c in citations[:10]
488 ],
489 "status": "clean" if hallucination_risk < 0.3 else "medium_risk" if hallucination_risk < 0.6 else "high_risk",
490 }
492 def _is_from_source(self, text: str, source: str, threshold: float = 0.6) -> bool:
493 """Check if text originated from source using substring and word overlap."""
494 if text in source:
495 return True
497 text_words = set(text.split())
498 source_words = set(source.split())
499 if not text_words:
500 return False
502 overlap = len(text_words & source_words) / len(text_words)
503 return overlap >= threshold
505 def get_stats(self) -> dict[str, Any]:
506 """Get citation statistics."""
507 return {
508 "total_citations": len(self._citations),
509 "by_source": Counter(c.source for c in self._citations),
510 "avg_confidence": (
511 sum(c.confidence for c in self._citations) / len(self._citations)
512 if self._citations else 0
513 ),
514 "sources_indexed": len(self._source_index),
515 }
518# ── Hybrid Search Engine ────────────────────────────────────────────
520class HybridSearchEngine:
521 """Unified hybrid search engine.
523 Combines dense + sparse retrieval with fusion and re-ranking.
525 Usage:
526 engine = HybridSearchEngine(
527 dense_retriever=DenseRetriever(vector_store=chroma_store),
528 sparse_retriever=BM25Retriever(),
529 )
531 # Index documents
532 engine.index_sparse(documents)
534 # Hybrid search
535 results = await engine.search("How to implement retry logic?")
536 for r in results:
537 print(f"{r.score:.3f} | {r.content[:100]}")
538 """
540 def __init__(
541 self,
542 dense_retriever: Optional[DenseRetriever] = None,
543 sparse_retriever: Optional[BM25Retriever] = None,
544 reranker: Optional[CrossEncoderReranker] = None,
545 citation_tracker: Optional[CitationTracker] = None,
546 fusion_method: str = "rrf",
547 dense_weight: float = 0.6,
548 ):
549 self.dense = dense_retriever or DenseRetriever()
550 self.sparse = sparse_retriever or BM25Retriever()
551 self.reranker = reranker or CrossEncoderReranker()
552 self.citations = citation_tracker or CitationTracker()
554 self.fusion_method = fusion_method
555 self.dense_weight = dense_weight
557 def index_sparse(self, documents: list[dict[str, str]]):
558 """Build sparse index from documents."""
559 self.sparse.index(documents)
560 for doc in documents:
561 self.citations.add_source(
562 doc_id=doc.get("id", ""),
563 content=doc.get("content", ""),
564 metadata=doc.get("metadata"),
565 )
567 async def search(
568 self,
569 query: str,
570 top_k: int = 10,
571 rerank: bool = True,
572 return_citations: bool = False,
573 ) -> list[SearchResult]:
574 """Hybrid search: dense + sparse → fusion → rerank.
576 Args:
577 query: Search query
578 top_k: Number of results
579 rerank: Whether to apply re-ranking
580 return_citations: Whether to attach citation info
582 Returns:
583 Ranked SearchResults.
584 """
585 # Step 1: Parallel retrieval
586 dense_results = await self.dense.search(query, top_k=top_k * 2)
587 sparse_results = self.sparse.search(query, top_k=top_k * 2)
589 # Step 2: Fusion
590 if self.fusion_method == "rrf":
591 fused = FusionMethod.reciprocal_rank_fusion(dense_results, sparse_results)
592 elif self.fusion_method == "cascade":
593 fused = FusionMethod.cascade(dense_results, sparse_results)
594 else: # weighted_sum
595 fused = FusionMethod.weighted_sum(
596 dense_results, sparse_results,
597 dense_weight=self.dense_weight,
598 sparse_weight=1.0 - self.dense_weight,
599 )
601 # Step 3: Re-rank (optional)
602 if rerank and len(fused) > top_k:
603 fused = await self.reranker.rerank(query, fused, top_k=top_k)
604 else:
605 fused = fused[:top_k]
607 # Step 4: Attach citations (optional)
608 if return_citations and fused:
609 for result in fused:
610 result.citations = [
611 c.text for c in self.citations.extract_citations(
612 result.content, [result]
613 )
614 ]
616 return fused
618 async def search_with_citations(
619 self,
620 query: str,
621 top_k: int = 10,
622 ) -> dict[str, Any]:
623 """Search and return both results and verified citations."""
624 results = await self.search(query, top_k=top_k, return_citations=True)
626 # Build combined text from top results
627 combined = "\n\n".join(r.content for r in results)
629 # Verify citations
630 verification = self.citations.verify(combined, results)
632 return {
633 "results": results,
634 "verification": verification,
635 "top_result": results[0] if results else None,
636 "citation_stats": self.citations.get_stats(),
637 }
639 def get_stats(self) -> dict[str, Any]:
640 """Get search engine statistics."""
641 return {
642 "bm25_documents": self.sparse._N if self.sparse else 0,
643 "bm25_vocabulary": len(self.sparse._idf) if self.sparse else 0,
644 "citation_stats": self.citations.get_stats(),
645 "fusion_method": self.fusion_method,
646 }