"""utils/retriever.py — Embedding, FAISS indexing, and hybrid BM25+dense retrieval. Mirrors src/rag_system/components/retriever/__init__.py's HybridRetriever: RRF fusion (k=60, dense x 0.7 + BM25 x 0.3). """ from __future__ import annotations import re from dataclasses import dataclass from typing import Dict, List, Optional, Tuple import numpy as np from utils.pdf_processor import DocumentChunk @dataclass class RetrievedChunk: chunk: DocumentChunk dense_score: float bm25_score: float rrf_score: float rank: int @property def text(self) -> str: return self.chunk.text @property def source(self) -> str: return f"{self.chunk.source_filename}, page {self.chunk.page_number}" class EmbeddingModel: """Lazy-loaded sentence-transformers embedder with in-memory cache.""" def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5") -> None: self._model_name = model_name self._model = None self._cache: Dict[str, List[float]] = {} def _get_model(self): if self._model is None: from sentence_transformers import SentenceTransformer self._model = SentenceTransformer(self._model_name) return self._model def embed(self, texts: List[str]) -> np.ndarray: model = self._get_model() uncached = [t for t in texts if t not in self._cache] if uncached: vecs = model.encode(uncached, normalize_embeddings=True, show_progress_bar=False) for text, vec in zip(uncached, vecs, strict=True): self._cache[text] = vec.tolist() return np.array([self._cache[t] for t in texts], dtype=np.float32) def embed_query(self, query: str) -> np.ndarray: return self.embed([query])[0] class BM25Index: def __init__(self) -> None: self._bm25 = None self._corpus: List[List[str]] = [] def _tokenize(self, text: str) -> List[str]: return re.sub(r"[^a-zA-Z0-9$%.,]", " ", text.lower()).split() def build(self, chunks: List[DocumentChunk]) -> None: self._corpus = [self._tokenize(c.text) for c in chunks] try: from rank_bm25 import BM25Okapi self._bm25 = BM25Okapi(self._corpus) except ImportError: self._bm25 = None def score(self, query: str) -> np.ndarray: tokens = self._tokenize(query) if not self._corpus: return np.array([]) if self._bm25 is not None: scores = np.array(self._bm25.get_scores(tokens)) else: scores = np.array([ sum(doc.count(t) for t in tokens) / max(len(doc), 1) for doc in self._corpus ], dtype=np.float32) max_score = scores.max() if max_score > 0: scores = scores / max_score return scores def reciprocal_rank_fusion( dense_ranks: List[int], bm25_ranks: List[int], k: int = 60, dense_weight: float = 0.7, bm25_weight: float = 0.3, ) -> List[float]: """RRF(d) = sum(weight_i / (k + rank_i(d))). k=60 per Cormack et al. 2009.""" n = len(dense_ranks) scores = [] for i in range(n): rrf = (dense_weight / (k + dense_ranks[i] + 1)) + (bm25_weight / (k + bm25_ranks[i] + 1)) scores.append(rrf) return scores class VectorIndex: """FAISS-backed dense vector index with BM25 hybrid retrieval.""" def __init__(self, embedding_model: Optional[EmbeddingModel] = None) -> None: self._embedder = embedding_model or EmbeddingModel() self._chunks: List[DocumentChunk] = [] self._index = None self._bm25 = BM25Index() def build(self, chunks: List[DocumentChunk]) -> List[str]: if not chunks: return ["No chunks to index"] steps = [] self._chunks = chunks steps.append(f"Embedding {len(chunks)} chunks with BAAI/bge-small-en-v1.5...") texts = [c.text for c in chunks] embeddings = self._embedder.embed(texts) try: import faiss dim = embeddings.shape[1] self._index = faiss.IndexFlatIP(dim) self._index.add(embeddings) steps.append(f"Built FAISS IndexFlatIP (dim={dim}, {len(chunks)} vectors)") except ImportError: self._index = embeddings steps.append("Built NumPy cosine index (FAISS not installed)") self._bm25.build(chunks) steps.append("Built BM25 keyword index for hybrid retrieval") steps.append(f"Index ready - {len(chunks)} chunks searchable") return steps def search( self, query: str, top_k: int = 5, chunk_type_filter: Optional[str] = None, ) -> Tuple[List[RetrievedChunk], List[str]]: steps = [] if not self._chunks: return [], ["Index is empty - ingest a document first"] steps.append("Embedding query for dense retrieval...") query_vec = self._embedder.embed_query(query).astype(np.float32) candidate_k = min(len(self._chunks), top_k * 4) if hasattr(self._index, "search"): scores, indices = self._index.search(query_vec.reshape(1, -1), candidate_k) dense_scores = scores[0].tolist() dense_indices = indices[0].tolist() else: sims = self._index @ query_vec ranked = np.argsort(-sims)[:candidate_k] dense_indices = ranked.tolist() dense_scores = sims[ranked].tolist() steps.append(f"Dense retrieval: top-{candidate_k} candidates via cosine similarity") bm25_all_scores = self._bm25.score(query) if len(bm25_all_scores) > 0: bm25_ranked = np.argsort(-bm25_all_scores)[:candidate_k].tolist() else: bm25_ranked = list(range(min(candidate_k, len(self._chunks)))) steps.append(f"BM25 keyword retrieval: top-{len(bm25_ranked)} candidates") all_candidate_indices = list(dict.fromkeys(dense_indices + bm25_ranked)) def _dense_rank(idx: int) -> int: try: return dense_indices.index(idx) except ValueError: return candidate_k def _bm25_rank(idx: int) -> int: try: return bm25_ranked.index(idx) except ValueError: return candidate_k rrf_scores = [] for idx in all_candidate_indices: rrf = (0.7 / (60 + _dense_rank(idx) + 1)) + (0.3 / (60 + _bm25_rank(idx) + 1)) rrf_scores.append((idx, rrf)) rrf_scores.sort(key=lambda x: x[1], reverse=True) steps.append(f"RRF fusion (k=60, dense x0.7 + BM25 x0.3): merged {len(rrf_scores)} candidates") results = [] for rank, (idx, rrf_score) in enumerate(rrf_scores): if idx >= len(self._chunks): continue chunk = self._chunks[idx] if chunk_type_filter and chunk.chunk_type != chunk_type_filter: continue d_score = dense_scores[dense_indices.index(idx)] if idx in dense_indices else 0.0 b_score = float(bm25_all_scores[idx]) if len(bm25_all_scores) > idx else 0.0 results.append(RetrievedChunk( chunk=chunk, dense_score=float(d_score), bm25_score=b_score, rrf_score=rrf_score, rank=rank + 1, )) if len(results) >= top_k: break steps.append(f"Retrieved {len(results)} chunks (top-k={top_k})") return results, steps