Coverage for agentos/rag/hybrid

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 collections.abc import Callable

21from dataclasses import dataclass, field

22from typing import Any

24# ── Types ───────────────────────────────────────────────────────────

27@dataclass

28class SearchResult:

29 """A single search result with metadata."""

31 doc_id: str

32 content: str

33 source: str = "" # File path, URL, or source identifier

34 title: str = ""

35 score: float = 0.0

36 dense_score: float = 0.0

37 sparse_score: float = 0.0

38 rerank_score: float = 0.0

39 chunk_index: int = 0

40 metadata: dict[str, Any] = field(default_factory=dict)

41 citations: list[str] = field(default_factory=list) # Specific sentences/quotes

44@dataclass

45class Citation:

46 """A citation from source material."""

48 text: str

49 source: str

50 doc_id: str = ""

51 chunk_index: int = 0

52 start_pos: int = 0

53 end_pos: int = 0

54 confidence: float = 1.0

57# ── BM25 Sparse Retriever ───────────────────────────────────────────

60class BM25Retriever:

61 """Pure Python BM25 implementation for keyword search.

63 No external dependencies. Tokenizes, builds inverted index,

64 and scores documents using Okapi BM25.

65 """

67 def __init__(self, k1: float = 1.5, b: float = 0.75):

68 self.k1 = k1

69 self.b = b

70 self._docs: list[str] = []

71 self._doc_ids: list[str] = []

72 self._doc_lengths: list[int] = []

73 self._avg_dl: float = 0.0

74 self._inverted_index: dict[str, dict[int, int]] = defaultdict(dict)

75 self._idf: dict[str, float] = {}

76 self._N: int = 0

78 def index(self, documents: list[dict[str, str]]):

79 """Build BM25 index from documents.

81 Args:

82 documents: List of {id, content} dicts.

83 """

84 self._docs = [doc.get("content", "") for doc in documents]

85 self._doc_ids = [doc.get("id", f"doc_{i}") for i, doc in enumerate(documents)]

86 self._doc_lengths = [len(self._tokenize(doc)) for doc in self._docs]

87 self._N = len(self._docs)

88 self._avg_dl = sum(self._doc_lengths) / max(self._N, 1)

90 # Build inverted index

91 self._inverted_index.clear()

92 doc_freq: dict[str, int] = defaultdict(int)

94 for doc_id, doc in enumerate(self._docs):

95 tokens = self._tokenize(doc)

96 token_counts = Counter(tokens)

97 for token, count in token_counts.items():

98 self._inverted_index[token][doc_id] = count

99 doc_freq[token] += 1

100

101 # Compute IDF

102 self._idf = {

103 token: math.log(1 + (self._N - freq + 0.5) / (freq + 0.5))

104 for token, freq in doc_freq.items()

105 }

106

107 def search(self, query: str, top_k: int = 10) -> list[SearchResult]:

108 """BM25 keyword search."""

109 if not self._docs:

110 return []

111

112 query_tokens = self._tokenize(query)

113 scores: list[float] = [0.0] * self._N

114

115 for token in query_tokens:

116 if token not in self._inverted_index:

117 continue

118 idf = self._idf.get(token, 0)

119 for doc_id, tf in self._inverted_index[token].items():

120 dl = self._doc_lengths[doc_id]

121 numerator = tf * (self.k1 + 1)

122 denominator = tf + self.k1 * (1 - self.b + self.b * dl / max(self._avg_dl, 1))

123 scores[doc_id] += idf * numerator / max(denominator, 1e-9)

124

125 # Rank and return

126 ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)

127 max_score = ranked[0][1] if ranked else 1.0

128

129 return [

130 SearchResult(

131 doc_id=self._doc_ids[doc_id],

132 content=self._docs[doc_id][:500],

133 sparse_score=score / max(max_score, 1e-9),

134 score=score / max(max_score, 1e-9),

135 metadata={"method": "bm25"},

136 )

137 for doc_id, score in ranked[:top_k]

138 if score > 0

139 ]

140

141 def _tokenize(self, text: str) -> list[str]:

142 """Simple tokenization: lowercase, split on non-alphanumeric, filter short tokens."""

143 tokens = re.findall(r"[\w\u4e00-\u9fff]+", text.lower())

144 return [t for t in tokens if len(t) > 1]

145

146

147# ── Dense Retriever ─────────────────────────────────────────────────

148

149

150class DenseRetriever:

151 """Semantic search via embeddings.

152

153 Wraps an embedding function (e.g., OpenAI embeddings, sentence-transformers)

154 and a vector store (ChromaDB or similar).

155 """

156

157 def __init__(

158 self,

159 vector_store=None,

160 embed_fn: Callable[[str], list[float]] | None = None,

161 ):

162 self._store = vector_store

163 self._embed = embed_fn

164

165 async def search(self, query: str, top_k: int = 10) -> list[SearchResult]:

166 """Dense vector search."""

167 if not self._store:

168 return []

169

170 try:

171 results = await self._store.search(query, top_k=top_k)

172

173 max_score = results[0].get("score", 1.0) if results else 1.0

174

175 return [

176 SearchResult(

177 doc_id=result.get("id", ""),

178 content=result.get("content", "")[:500],

179 dense_score=result.get("score", 0) / max(max_score, 1e-9),

180 score=result.get("score", 0) / max(max_score, 1e-9),

181 metadata=result.get("metadata", {}),

182 )

183 for result in results

184 ]

185 except Exception:

186 return []

187

188

189# ── Cross-Encoder Re-Ranker ─────────────────────────────────────────

190

191

192class CrossEncoderReranker:

193 """Re-rank search results with a cross-encoder model.

194

195 Instead of embedding query and documents independently (bi-encoder),

196 a cross-encoder processes (query, document) pairs together for higher

197 accuracy — at the cost of more computation.

198

199 Supports:

200 - HuggingFace cross-encoder models (e.g., ms-marco-MiniLM)

201 - Custom scoring functions

202 - LLM-based re-ranking (use an LLM to judge relevance)

203 """

204

205 def __init__(

206 self,

207 model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",

208 use_llm: bool = False,

209 llm_client=None,

210 ):

211 self._model_name = model_name

212 self._model = None

213 self._use_llm = use_llm

214 self._llm = llm_client

215

216 async def rerank(

217 self,

218 query: str,

219 candidates: list[SearchResult],

220 top_k: int = 5,

221 ) -> list[SearchResult]:

222 """Re-rank candidates by relevance to query.

223

224 Args:

225 query: Original search query

226 candidates: Initial retrieval results

227 top_k: Number of results to return after re-ranking

228

229 Returns:

230 Re-ranked candidates with updated rerank_score.

231 """

232 if not candidates:

233 return []

234

235 if self._use_llm and self._llm:

236 return await self._llm_rerank(query, candidates, top_k)

237 else:

238 return await self._cross_encoder_rerank(query, candidates, top_k)

239

240 async def _cross_encoder_rerank(

241 self,

242 query: str,

243 candidates: list[SearchResult],

244 top_k: int,

245 ) -> list[SearchResult]:

246 """Re-rank using HuggingFace cross-encoder."""

247 try:

248 from sentence_transformers import CrossEncoder

249

250 if self._model is None:

251 self._model = CrossEncoder(self._model_name)

252

253 pairs = [(query, c.content[:1000]) for c in candidates]

254 scores = self._model.predict(pairs)

255

256 for candidate, score in zip(candidates, scores):

257 candidate.rerank_score = float(score)

258 # Weighted fusion

259 candidate.score = (

260 candidate.dense_score * 0.3 + candidate.sparse_score * 0.2 + float(score) * 0.5

261 )

262

263 candidates.sort(key=lambda x: x.rerank_score, reverse=True)

264 return candidates[:top_k]

265

266 except ImportError:

267 return candidates[:top_k] # Fallback: no re-ranking

268

269 async def _llm_rerank(

270 self,

271 query: str,

272 candidates: list[SearchResult],

273 top_k: int,

274 ) -> list[SearchResult]:

275 """Re-rank using LLM relevance judgment."""

276 if not self._llm:

277 return candidates[:top_k]

278

279 prompt = f"Query: {query}\n\nRate each document's relevance on a scale of 0-10:\n\n"

280 for i, c in enumerate(candidates[:20]):

281 prompt += f"[{i}] {c.content[:300]}\n\n"

282 prompt += "Output format: [doc_id] score"

283

284 try:

285 response = await self._llm.complete(prompt)

286 # Parse scores

287 scores: dict[int, float] = {}

288 for line in response.split("\n"):

289 match = re.match(r"\[(\d+)\]\s*(\d+(?:\.\d+)?)", line.strip())

290 if match:

291 idx = int(match.group(1))

292 score = float(match.group(2)) / 10.0

293 if idx < len(candidates):

294 scores[idx] = score

295

296 for i, candidate in enumerate(candidates):

297 candidate.rerank_score = scores.get(i, 0.5)

298 candidate.score = (

299 candidate.dense_score * 0.25

300 + candidate.sparse_score * 0.15

301 + candidate.rerank_score * 0.6

302 )

303

304 candidates.sort(key=lambda x: x.rerank_score, reverse=True)

305 return candidates[:top_k]

306

307 except Exception:

308 return candidates[:top_k]

309

310

311# ── Fusion Algorithms ───────────────────────────────────────────────

312

313

314class FusionMethod:

315 """Collection of rank fusion algorithms."""

316

317 @staticmethod

318 def reciprocal_rank_fusion(

319 dense_results: list[SearchResult],

320 sparse_results: list[SearchResult],

321 k: int = 60,

322 ) -> list[SearchResult]:

323 """RRF: Reciprocal Rank Fusion.

324

325 RRF_score(d) = sum_{ranker} 1 / (k + rank(d))

326 """

327 scores: dict[str, float] = {}

328 docs: dict[str, SearchResult] = {}

329

330 for rank, result in enumerate(dense_results):

331 scores[result.doc_id] = 1.0 / (k + rank + 1)

332 docs[result.doc_id] = result

333

334 for rank, result in enumerate(sparse_results):

335 if result.doc_id in scores:

336 scores[result.doc_id] += 1.0 / (k + rank + 1)

337 else:

338 scores[result.doc_id] = 1.0 / (k + rank + 1)

339 docs[result.doc_id] = result

340

341 fused = sorted(scores.items(), key=lambda x: x[1], reverse=True)

342 results = []

343 for doc_id, score in fused:

344 doc = docs[doc_id]

345 doc.score = score

346 results.append(doc)

347

348 return results

349

350 @staticmethod

351 def weighted_sum(

352 dense_results: list[SearchResult],

353 sparse_results: list[SearchResult],

354 dense_weight: float = 0.6,

355 sparse_weight: float = 0.4,

356 ) -> list[SearchResult]:

357 """Weighted score summation."""

358 scores: dict[str, list[float]] = defaultdict(list)

359 docs: dict[str, SearchResult] = {}

360

361 for result in dense_results:

362 scores[result.doc_id].append(result.dense_score * dense_weight)

363 docs[result.doc_id] = result

364

365 for result in sparse_results:

366 scores[result.doc_id].append(result.sparse_score * sparse_weight)

367 if result.doc_id not in docs:

368 docs[result.doc_id] = result

369

370 fused = []

371 for doc_id, wscores in scores.items():

372 doc = docs[doc_id]

373 doc.score = sum(wscores)

374 fused.append(doc)

375

376 fused.sort(key=lambda x: x.score, reverse=True)

377 return fused

378

379 @staticmethod

380 def cascade(

381 dense_results: list[SearchResult],

382 sparse_results: list[SearchResult],

383 ) -> list[SearchResult]:

384 """Cascade: dense first, then sparse fills gaps."""

385 seen: set[str] = set()

386 results: list[SearchResult] = []

387

388 for r in dense_results:

389 results.append(r)

390 seen.add(r.doc_id)

391

392 for r in sparse_results:

393 if r.doc_id not in seen:

394 results.append(r)

395 seen.add(r.doc_id)

396

397 return results

398

399

400# ── Citation Tracker ────────────────────────────────────────────────

401

402

403class CitationTracker:

404 """Track and verify citations from source documents.

405

406 Key features:

407 - Extract citations from generated text

408 - Verify against source documents

409 - Mark unverifiable (potential hallucination)

410 - Track citation usage statistics

411 """

412

413 def __init__(self):

414 self._citations: list[Citation] = []

415 self._source_index: dict[str, dict] = {} # doc_id → metadata

416

417 def add_source(self, doc_id: str, content: str, metadata: dict[str, Any] | None = None):

418 """Register a source document."""

419 self._source_index[doc_id] = {

420 "content": content,

421 "metadata": metadata or {},

422 }

423

424 def extract_citations(self, text: str, sources: list[SearchResult]) -> list[Citation]:

425 """Extract and verify citations from generated text.

426

427 Args:

428 text: Generated response text

429 sources: Source documents used for generation

430

431 Returns:

432 List of verified Citation objects.

433 """

434 citations: list[Citation] = []

435

436 for source in sources:

437 # Find substrings of generated text that appear in source

438 source_content = source.content.lower()

439 text.lower()

440

441 # Extract sentences from generated text

442 sentences = re.split(r"[.!?]+", text)

443 for sent in sentences:

444 sent = sent.strip()

445 if len(sent) < 15:

446 continue

447

448 # Check if this sentence appears in source (with fuzzy matching)

449 if self._is_from_source(sent.lower(), source_content):

450 citations.append(

451 Citation(

452 text=sent,

453 source=source.source or source.doc_id,

454 doc_id=source.doc_id,

455 chunk_index=source.chunk_index,

456 confidence=0.9,

457 )

458 )

459

460 # Deduplicate

461 seen: set[str] = set()

462 unique = []

463 for c in citations:

464 key = c.text[:50]

465 if key not in seen:

466 seen.add(key)

467 unique.append(c)

468

469 self._citations.extend(unique)

470 return unique

471

472 def verify(self, text: str, sources: list[SearchResult]) -> dict[str, Any]:

473 """Verify all claims in text against source documents.

474

475 Returns:

476 Dict with verified/unverified segments and hallucination score.

477 """

478 citations = self.extract_citations(text, sources)

479

480 sentences = re.split(r"[.!?]+", text)

481 total_sentences = len(sentences)

482 cited_sentences = sum(

483 1 for s in sentences if any(c.text[:30].lower() in s.strip().lower() for c in citations)

484 )

485

486 uncited = total_sentences - cited_sentences

487 hallucination_risk = uncited / max(total_sentences, 1)

488

489 return {

490 "total_sentences": total_sentences,

491 "cited_sentences": cited_sentences,

492 "uncited_sentences": uncited,

493 "hallucination_risk": round(hallucination_risk, 3),

494 "citations": [

495 {"text": c.text[:100], "source": c.source, "confidence": c.confidence}

496 for c in citations[:10]

497 ],

498 "status": (

499 "clean"

500 if hallucination_risk < 0.3

501 else "medium_risk" if hallucination_risk < 0.6 else "high_risk"

502 ),

503 }

504

505 def _is_from_source(self, text: str, source: str, threshold: float = 0.6) -> bool:

506 """Check if text originated from source using substring and word overlap."""

507 if text in source:

508 return True

509

510 text_words = set(text.split())

511 source_words = set(source.split())

512 if not text_words:

513 return False

514

515 overlap = len(text_words & source_words) / len(text_words)

516 return overlap >= threshold

517

518 def get_stats(self) -> dict[str, Any]:

519 """Get citation statistics."""

520 return {

521 "total_citations": len(self._citations),

522 "by_source": Counter(c.source for c in self._citations),

523 "avg_confidence": (

524 sum(c.confidence for c in self._citations) / len(self._citations)

525 if self._citations

526 else 0

527 ),

528 "sources_indexed": len(self._source_index),

529 }

530

531

532# ── Hybrid Search Engine ────────────────────────────────────────────

533

534

535class HybridSearchEngine:

536 """Unified hybrid search engine.

537

538 Combines dense + sparse retrieval with fusion and re-ranking.

539

540 Usage:

541 engine = HybridSearchEngine(

542 dense_retriever=DenseRetriever(vector_store=chroma_store),

543 sparse_retriever=BM25Retriever(),

544 )

545

546 # Index documents

547 engine.index_sparse(documents)

548

549 # Hybrid search

550 results = await engine.search("How to implement retry logic?")

551 for r in results:

552 print(f"{r.score:.3f} | {r.content[:100]}")

553 """

554

555 def __init__(

556 self,

557 dense_retriever: DenseRetriever | None = None,

558 sparse_retriever: BM25Retriever | None = None,

559 reranker: CrossEncoderReranker | None = None,

560 citation_tracker: CitationTracker | None = None,

561 fusion_method: str = "rrf",

562 dense_weight: float = 0.6,

563 ):

564 self.dense = dense_retriever or DenseRetriever()

565 self.sparse = sparse_retriever or BM25Retriever()

566 self.reranker = reranker or CrossEncoderReranker()

567 self.citations = citation_tracker or CitationTracker()

568

569 self.fusion_method = fusion_method

570 self.dense_weight = dense_weight

571

572 def index_sparse(self, documents: list[dict[str, str]]):

573 """Build sparse index from documents."""

574 self.sparse.index(documents)

575 for doc in documents:

576 self.citations.add_source(

577 doc_id=doc.get("id", ""),

578 content=doc.get("content", ""),

579 metadata=doc.get("metadata"),

580 )

581

582 async def search(

583 self,

584 query: str,

585 top_k: int = 10,

586 rerank: bool = True,

587 return_citations: bool = False,

588 ) -> list[SearchResult]:

589 """Hybrid search: dense + sparse → fusion → rerank.

590

591 Args:

592 query: Search query

593 top_k: Number of results

594 rerank: Whether to apply re-ranking

595 return_citations: Whether to attach citation info

596

597 Returns:

598 Ranked SearchResults.

599 """

600 # Step 1: Parallel retrieval

601 dense_results = await self.dense.search(query, top_k=top_k * 2)

602 sparse_results = self.sparse.search(query, top_k=top_k * 2)

603

604 # Step 2: Fusion

605 if self.fusion_method == "rrf":

606 fused = FusionMethod.reciprocal_rank_fusion(dense_results, sparse_results)

607 elif self.fusion_method == "cascade":

608 fused = FusionMethod.cascade(dense_results, sparse_results)

609 else: # weighted_sum

610 fused = FusionMethod.weighted_sum(

611 dense_results,

612 sparse_results,

613 dense_weight=self.dense_weight,

614 sparse_weight=1.0 - self.dense_weight,

615 )

616

617 # Step 3: Re-rank (optional)

618 if rerank and len(fused) > top_k:

619 fused = await self.reranker.rerank(query, fused, top_k=top_k)

620 else:

621 fused = fused[:top_k]

622

623 # Step 4: Attach citations (optional)

624 if return_citations and fused:

625 for result in fused:

626 result.citations = [

627 c.text for c in self.citations.extract_citations(result.content, [result])

628 ]

629

630 return fused

631

632 async def search_with_citations(

633 self,

634 query: str,

635 top_k: int = 10,

636 ) -> dict[str, Any]:

637 """Search and return both results and verified citations."""

638 results = await self.search(query, top_k=top_k, return_citations=True)

639

640 # Build combined text from top results

641 combined = "\n\n".join(r.content for r in results)

642

643 # Verify citations

644 verification = self.citations.verify(combined, results)

645

646 return {

647 "results": results,

648 "verification": verification,

649 "top_result": results[0] if results else None,

650 "citation_stats": self.citations.get_stats(),

651 }

652

653 def get_stats(self) -> dict[str, Any]:

654 """Get search engine statistics."""

655 return {

656 "bm25_documents": self.sparse._N if self.sparse else 0,

657 "bm25_vocabulary": len(self.sparse._idf) if self.sparse else 0,

658 "citation_stats": self.citations.get_stats(),

659 "fusion_method": self.fusion_method,

660 }

Coverage for agentos/rag/hybrid_search.py: 23%

268 statements