"""Pipeline orchestrator for hybrid retrieval with cross-retriever score calibration. Coordinates BM25 sparse retrieval, FAISS dense retrieval, score calibration, fusion strategies, and cross-encoder reranking across 7 retrieval modes. Modes: dense - FAISS only (original behavior) sparse - BM25 only rrf - Both → Reciprocal Rank Fusion hybrid_fixed - Both → min-max calibration → α=0.5 linear fusion hybrid_calibrated - Both → CDF calibration → entropy-weighted fusion hybrid_fixed_rerank - hybrid_fixed + cross-encoder reranking hybrid_calibrated_rerank - hybrid_calibrated + cross-encoder reranking """ import contextlib import os import pickle import time from pathlib import Path from typing import Any import numpy as np from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.vectorstores import FAISS from langchain_core.documents import Document from langchain_huggingface import HuggingFaceEmbeddings from pypdf import PdfReader from app.calibration import load_corpus_cdfs, load_corpus_lm from app.fusion import entropy_fuse, linear_fuse, rrf_fuse from app.reranker import CrossEncoderReranker from app.sparse_retriever import BM25Index BASE_DIR = Path(os.getenv("DATA_DIR", "data")) UPLOAD_DIR = BASE_DIR / "uploads" INDEX_DIR = BASE_DIR / "faiss_index" BM25_INDEX_PATH = INDEX_DIR / "bm25_index.pkl" DOC_LOOKUP_PATH = INDEX_DIR / "doc_lookup.pkl" CORPUS_LM_PATH = INDEX_DIR / "corpus_lm.pkl" EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2") # Minimum relevance score for dense-only mode. Benchmark evaluation should not # threshold top-k retrieval, so the default is disabled. MIN_RELEVANCE_SCORE = float(os.getenv("MIN_RELEVANCE_SCORE", "0.0")) RETRIEVAL_MODES = { "dense": "FAISS dense retrieval only (original behavior)", "sparse": "BM25 sparse retrieval only", "rrf": "BM25 + FAISS → Reciprocal Rank Fusion (k=60)", "hybrid_fixed": "BM25 + FAISS → min-max calibration → α=0.5 linear fusion", "hybrid_calibrated": "BM25 + FAISS → CDF calibration → entropy-weighted fusion", "hybrid_fixed_rerank": "hybrid_fixed + cross-encoder reranking", "hybrid_calibrated_rerank": "hybrid_calibrated + cross-encoder reranking (full pipeline)", } # Number of candidates to retrieve from each retriever before fusion CANDIDATE_K = int(os.getenv("CANDIDATE_K", "20")) class DocumentStore: def __init__(self) -> None: self.vector_store: FAISS | None = None self._embeddings: HuggingFaceEmbeddings | None = None self.bm25_index = BM25Index() self._reranker: CrossEncoderReranker | None = None self.documents: set[str] = set() self.corpus_cdfs: tuple[np.ndarray, np.ndarray] | None = None self.corpus_lm: tuple[Any, int] | None = None self.doc_lookup: dict[str, str] = {} @property def embeddings(self) -> HuggingFaceEmbeddings: if self._embeddings is None: self._embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL) return self._embeddings @property def reranker(self) -> CrossEncoderReranker: if self._reranker is None: self._reranker = CrossEncoderReranker() return self._reranker def load(self) -> None: """Load all indices, CDFs, and lookup tables from disk.""" if INDEX_DIR.exists(): with contextlib.suppress(Exception): self.vector_store = FAISS.load_local( str(INDEX_DIR), self.embeddings, allow_dangerous_deserialization=True, ) self.bm25_index.load(BM25_INDEX_PATH) self.documents = {str(m["source"]) for m in self.bm25_index.doc_metadata} # Load corpus-level CDFs self.corpus_cdfs = load_corpus_cdfs(INDEX_DIR) # Load corpus language model (for Clarity Score) self.corpus_lm = load_corpus_lm(CORPUS_LM_PATH) # Load document text lookup if DOC_LOOKUP_PATH.exists(): with open(DOC_LOOKUP_PATH, "rb") as f: self.doc_lookup = pickle.load(f) # noqa: S301 def save(self) -> None: if self.vector_store is not None: INDEX_DIR.mkdir(parents=True, exist_ok=True) self.vector_store.save_local(str(INDEX_DIR)) def add_file(self, file_path: Path) -> dict[str, Any]: text = extract_text(file_path) if not text.strip(): raise ValueError("No readable text found in the uploaded file.") chunks = chunk_text(text, file_path.name) stats = self.add_documents(chunks) return {"filename": file_path.name, **stats, "characters": len(text)} def add_documents(self, documents: list[Document]) -> dict[str, Any]: """Bulk-add already constructed LangChain documents.""" if not documents: return {"chunks": 0} # Build FAISS index if self.vector_store is None: self.vector_store = FAISS.from_documents(documents, self.embeddings) else: self.vector_store.add_documents(documents) # Build BM25 index bm25_docs = [ { "text": doc.page_content, "source": doc.metadata["source"], "chunk": doc.metadata["chunk"], } for doc in documents ] self.bm25_index.add_documents(bm25_docs) self.bm25_index.save(BM25_INDEX_PATH) # Update document text lookup for doc in documents: key = f"{doc.metadata['source']}:{doc.metadata['chunk']}" self.doc_lookup[key] = doc.page_content self.documents.add(str(doc.metadata["source"])) DOC_LOOKUP_PATH.parent.mkdir(parents=True, exist_ok=True) with open(DOC_LOOKUP_PATH, "wb") as f: pickle.dump(self.doc_lookup, f) self.save() return {"chunks": len(documents)} def _search_dense(self, query: str, top_k: int = CANDIDATE_K) -> list[dict[str, Any]]: """Dense retrieval via FAISS.""" if self.vector_store is None: return [] matches = self.vector_store.similarity_search_with_score(query, k=top_k) results = [] for doc, distance in matches: score = 1 / (1 + float(distance)) # L2 to similarity results.append( { "source": doc.metadata.get("source", "unknown"), "chunk": doc.metadata.get("chunk", 0), "text": doc.page_content, "score": round(score, 6), } ) return results def _search_sparse(self, query: str, top_k: int = CANDIDATE_K) -> list[dict[str, Any]]: """Sparse retrieval via BM25.""" return self.bm25_index.search(query, top_k=top_k) def _dense_score_all(self, query: str) -> np.ndarray: """Score ALL documents with dense retriever. For CDF construction. Uses IndexFlatL2 exhaustive search — embeddings are already stored in FAISS, no re-encoding needed. """ if self.vector_store is None: return np.array([]) # Get all docs with scores (large k = all docs) n_docs = self.vector_store.index.ntotal if n_docs == 0: return np.array([]) matches = self.vector_store.similarity_search_with_score(query, k=n_docs) # Convert L2 distances to similarity scores return np.array([1 / (1 + float(dist)) for _, dist in matches]) def search( self, query: str, top_k: int = 3, mode: str = "hybrid_calibrated_rerank" ) -> tuple[list[dict[str, Any]], float, dict[str, Any]]: """Search using the specified retrieval mode. Returns (results, total_latency_ms, telemetry). """ t_start = time.perf_counter() telemetry: dict[str, Any] = {"mode": mode} if mode == "dense": results = self._search_dense(query, top_k=top_k) # Apply legacy score filter for backward compatibility results = [r for r in results if r["score"] >= MIN_RELEVANCE_SCORE] elif mode == "sparse": results = self._search_sparse(query, top_k=top_k) elif mode == "rrf": sparse = self._search_sparse(query, top_k=CANDIDATE_K) dense = self._search_dense(query, top_k=CANDIDATE_K) results, meta = rrf_fuse(sparse, dense, top_k=top_k) telemetry.update(meta) elif mode == "hybrid_fixed": sparse = self._search_sparse(query, top_k=CANDIDATE_K) dense = self._search_dense(query, top_k=CANDIDATE_K) results, meta = linear_fuse(sparse, dense, alpha=0.5, calibration="minmax", top_k=top_k) telemetry.update(meta) elif mode == "hybrid_calibrated": sparse = self._search_sparse(query, top_k=CANDIDATE_K) dense = self._search_dense(query, top_k=CANDIDATE_K) results, meta = entropy_fuse( sparse, dense, calibration="cdf", corpus_cdfs=self.corpus_cdfs, top_k=top_k, ) telemetry.update(meta) elif mode == "hybrid_fixed_rerank": sparse = self._search_sparse(query, top_k=CANDIDATE_K) dense = self._search_dense(query, top_k=CANDIDATE_K) fused, meta = linear_fuse( sparse, dense, alpha=0.5, calibration="minmax", top_k=CANDIDATE_K, ) telemetry.update(meta) results = self.reranker.rerank(query, fused, top_k=top_k) telemetry["reranked"] = True elif mode == "hybrid_calibrated_rerank": sparse = self._search_sparse(query, top_k=CANDIDATE_K) dense = self._search_dense(query, top_k=CANDIDATE_K) fused, meta = entropy_fuse( sparse, dense, calibration="cdf", corpus_cdfs=self.corpus_cdfs, top_k=CANDIDATE_K, ) telemetry.update(meta) results = self.reranker.rerank(query, fused, top_k=top_k) telemetry["reranked"] = True else: raise ValueError(f"Unknown retrieval mode: {mode}. Available: {list(RETRIEVAL_MODES)}") latency_ms = (time.perf_counter() - t_start) * 1000 telemetry["latency_ms"] = round(latency_ms, 2) return results, round(latency_ms, 2), telemetry def answer( self, query: str, top_k: int = 3, mode: str = "hybrid_calibrated_rerank" ) -> dict[str, Any]: results, latency_ms, telemetry = self.search(query, top_k, mode) if not results: return { "answer": "I do not have any indexed documents yet. Upload a PDF or text file first.", "sources": [], "retrieval_latency_ms": latency_ms, "telemetry": telemetry, } # Build answer from top chunks answer_parts: list[str] = [] for r in results[:3]: snippet = r["text"].strip() if len(snippet) > 300: snippet = snippet[:300].rsplit(" ", 1)[0] + "..." answer_parts.append(f"[{r['source']}] {snippet}") return { "answer": "\n\n".join(answer_parts), "sources": results, "retrieval_latency_ms": latency_ms, "telemetry": telemetry, } # --------------------------------------------------------------------------- # Text extraction and chunking (unchanged from original) # --------------------------------------------------------------------------- def extract_text(file_path: Path) -> str: suffix = file_path.suffix.lower() if suffix == ".pdf": reader = PdfReader(str(file_path)) return "\n".join(page.extract_text() or "" for page in reader.pages) if suffix in {".txt", ".md"}: return file_path.read_text(encoding="utf-8", errors="ignore") raise ValueError("Only PDF, TXT, and MD files are supported.") def chunk_text(text: str, source: str) -> list[Document]: splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=120) pieces = splitter.split_text(text) return [ Document(page_content=piece, metadata={"source": source, "chunk": index + 1}) for index, piece in enumerate(pieces) ]