Metafazer/finrag-backend
0
1"""Hybrid retrieval with Reciprocal Rank Fusion.2 3Combines BM25 sparse search and ChromaDB dense vector search4via Reciprocal Rank Fusion (RRF). Supports multi-query expansion5and HyDE for broader recall on financial queries.6 7Design decisions:8- RRF over linear combination: rank-based fusion is robust to9 score distribution differences between BM25 and cosine similarity.10 BM25 scores are unbounded; cosine distances are [0, 2]. RRF11 sidesteps normalization entirely by using ranks, not scores.12- Pluggable multi-query and HyDE strategies via callables, so the13 LLM-based versions drop in cleanly on Day 7-8.14- Metadata filtering applied at both retrieval layers for efficiency.15 Filtering at the store level avoids fetching irrelevant candidates.16- Deduplication within each source before fusion prevents a document17 appearing multiple times in one ranked list from inflating its score.18 19Debt: DAY-5-001 — Multi-query uses rule-based synonym expansion.20 LLM-based rephrasings will improve recall. Resolve on Day 7.21Debt: DAY-5-002 — HyDE is a no-op stub. Needs LLM to generate22 hypothetical answer documents. Resolve on Day 7-8.23"""24 25from collections import defaultdict26from collections.abc import Callable27 28import structlog29 30from finrag.retrieval.bm25_index import BM25Index31from finrag.vectorstore.chroma_store import ChromaStore32 33logger = structlog.get_logger(__name__)34 35# --------------------------------------------------------------------------- #36# Constants37# --------------------------------------------------------------------------- #38 39# RRF constant k. Controls how much lower-ranked documents are penalized.40# k=60 is the standard value from the original RRF paper (Cormack et al. 2009).41# Higher k = more equal weighting across ranks.42# Lower k = top-ranked documents get disproportionately more weight.43RRF_K = 6044 45# Default number of candidates to fetch from each retriever before fusion.46# Fetching more than the final n_results gives RRF enough signal to work with.47# Too few candidates = fusion has poor recall. Too many = slower retrieval.48DEFAULT_CANDIDATES_PER_RETRIEVER = 2049 50# Financial term expansions for rule-based multi-query.51# Maps common abbreviations and short forms to domain-specific synonyms.52# Limited to 2 synonyms per term to keep query count manageable.53FINANCIAL_SYNONYMS: dict[str, list[str]] = {54 "revenue": ["net sales", "total revenue"],55 "profit": ["net income", "earnings"],56 "eps": ["earnings per share", "diluted eps"],57 "margin": ["gross margin", "operating margin"],58 "debt": ["long-term debt", "total debt"],59 "cash": ["cash and cash equivalents", "free cash flow"],60 "growth": ["year over year growth", "yoy growth"],61 "risk": ["risk factors", "risks and uncertainties"],62 "capex": ["capital expenditure", "capital expenditures"],63 "r&d": ["research and development"],64 "sg&a": ["selling general and administrative"],65 "ebitda": ["earnings before interest taxes depreciation amortization"],66 "roe": ["return on equity"],67 "roa": ["return on assets"],68 "fcf": ["free cash flow"],69 "goodwill": ["goodwill impairment", "intangible assets"],70}71 72 73# --------------------------------------------------------------------------- #74# Reciprocal Rank Fusion75# --------------------------------------------------------------------------- #76 77 78def reciprocal_rank_fusion(79 ranked_lists: list[list[dict]],80 k: int = RRF_K,81) -> list[dict]:82 """Merge multiple ranked result lists using Reciprocal Rank Fusion.83 84 RRF computes a fused score for each document:85 score(d) = sum over r in ranked_lists of: 1 / (k + rank_r(d))86 87 where rank_r(d) is the 1-based rank of document d in list r,88 and k is a constant that controls rank sensitivity.89 90 Documents appearing in multiple lists get higher fused scores,91 which is exactly the behavior we want: a chunk that both BM2592 and vector search agree on is more likely to be relevant.93 94 Args:95 ranked_lists: List of ranked result lists. Each result dict96 must contain a 'chunk_id' key for deduplication.97 k: RRF constant (default 60). Higher = more equal weighting.98 99 Returns:100 Merged and re-ranked list of result dicts, sorted by fused101 score descending. Each result includes:102 - rrf_score: The fused relevance score103 - retrieval_sources: List of sources that found this chunk104 """105 # Accumulate RRF scores by chunk_id106 scores: dict[str, float] = defaultdict(float)107 # Track the best result dict for each chunk_id108 best_result: dict[str, dict] = {}109 # Track which sources contributed to each result110 sources: dict[str, list[str]] = defaultdict(list)111 112 source_names = ["dense", "sparse"]113 114 for list_idx, ranked_list in enumerate(ranked_lists):115 source_name = source_names[list_idx] if list_idx < len(source_names) else f"source_{list_idx}"116 for rank, result in enumerate(ranked_list, start=1):117 chunk_id = result["chunk_id"]118 scores[chunk_id] += 1.0 / (k + rank)119 sources[chunk_id].append(source_name)120 121 # Keep the first occurrence's full result dict122 if chunk_id not in best_result:123 best_result[chunk_id] = result.copy()124 125 # Build fused results126 fused: list[dict] = []127 for chunk_id, rrf_score in scores.items():128 result = best_result[chunk_id]129 result["rrf_score"] = rrf_score130 result["retrieval_sources"] = sources[chunk_id]131 # Remove retriever-specific score fields to avoid confusion132 result.pop("score", None)133 result.pop("distance", None)134 fused.append(result)135 136 # Sort by RRF score descending137 fused.sort(key=lambda x: x["rrf_score"], reverse=True)138 139 return fused140 141 142# --------------------------------------------------------------------------- #143# Multi-Query Expansion144# --------------------------------------------------------------------------- #145 146 147def expand_financial_query(query: str) -> list[str]:148 """Generate query variations using financial term expansion.149 150 Rule-based approach that identifies financial abbreviations and151 terms in the query and generates variations with their synonyms.152 153 This is a Day 5 placeholder. LLM-based rephrasings will replace154 this on Day 7 for much better recall. The LLM version will155 generate semantically diverse rephrasings rather than just156 swapping synonyms.157 158 Args:159 query: Original query string.160 161 Returns:162 List of query variations including the original.163 Always returns at least the original query.164 """165 variations = [query]166 query_lower = query.lower()167 168 for term, synonyms in FINANCIAL_SYNONYMS.items():169 if term in query_lower:170 for synonym in synonyms[:2]:171 variation = query_lower.replace(term, synonym)172 if variation != query_lower and variation not in variations:173 variations.append(variation)174 175 # Deduplicate while preserving order176 seen: set[str] = set()177 unique: list[str] = []178 for v in variations:179 normalized = v.lower().strip()180 if normalized not in seen:181 seen.add(normalized)182 unique.append(v)183 184 logger.debug(185 "query_expanded",186 original=query[:80],187 variation_count=len(unique),188 )189 return unique190 191 192# --------------------------------------------------------------------------- #193# HyDE Stub194# --------------------------------------------------------------------------- #195 196 197def hyde_passthrough(query: str) -> str:198 """No-op HyDE stub that returns the query unchanged.199 200 HyDE (Hypothetical Document Embeddings) generates a hypothetical201 answer document using an LLM, then embeds THAT instead of the202 query. This gives the vector search a document-like embedding203 to match against, bridging the query-document distribution gap.204 205 Example: Query "What was Apple's revenue?" becomes a hypothetical206 response like "Apple reported total revenue of $X billion for207 fiscal year 2024..." which is then embedded for similarity search.208 209 This stub is a placeholder until LLM integration on Day 7-8.210 The strategy pattern allows dropping in the real HyDE with zero211 changes to the HybridRetriever.212 213 Args:214 query: Original query string.215 216 Returns:217 The query unchanged.218 """219 return query220 221 222# --------------------------------------------------------------------------- #223# HybridRetriever224# --------------------------------------------------------------------------- #225 226 227class HybridRetriever:228 """Hybrid retriever combining BM25 and vector search with RRF.229 230 Runs BM25 keyword search and ChromaDB semantic search in sequence,231 then fuses results using Reciprocal Rank Fusion. Supports232 multi-query expansion for broader recall and HyDE for improved233 vector search on question-style queries.234 235 The key insight: BM25 excels at exact financial terms like236 "diluted EPS" or "goodwill impairment" that embeddings miss.237 Vector search excels at semantic similarity like finding revenue238 discussion when the query says "top line performance". RRF239 combines both signals without needing to calibrate score scales.240 241 Args:242 chroma_store: ChromaDB vector store instance.243 bm25_index: BM25 sparse retrieval index.244 rrf_k: RRF fusion constant (default 60).245 candidates_per_retriever: Number of candidates to fetch246 from each retriever before fusion.247 multi_query_fn: Optional callable for query expansion.248 Signature: (query: str) -> list[str].249 Default: rule-based financial term expansion.250 hyde_fn: Optional callable for HyDE transformation.251 Signature: (query: str) -> str.252 Default: passthrough (no-op).253 """254 255 def __init__(256 self,257 chroma_store: ChromaStore,258 bm25_index: BM25Index,259 rrf_k: int = RRF_K,260 candidates_per_retriever: int = DEFAULT_CANDIDATES_PER_RETRIEVER,261 multi_query_fn: Callable[[str], list[str]] | None = None,262 hyde_fn: Callable[[str], str] | None = None,263 ) -> None:264 """Initialize the hybrid retriever.265 266 Args:267 chroma_store: Initialized ChromaStore with embedded chunks.268 bm25_index: Built BM25Index with indexed chunks.269 rrf_k: RRF constant k (default 60).270 candidates_per_retriever: Candidates per source (default 20).271 multi_query_fn: Custom query expansion function.272 hyde_fn: Custom HyDE transformation function.273 """274 self._chroma = chroma_store275 self._bm25 = bm25_index276 self._rrf_k = rrf_k277 self._candidates = candidates_per_retriever278 self._multi_query_fn = multi_query_fn or expand_financial_query279 self._hyde_fn = hyde_fn or hyde_passthrough280 281 logger.info(282 "hybrid_retriever_initialized",283 rrf_k=rrf_k,284 candidates_per_retriever=candidates_per_retriever,285 has_custom_multi_query=multi_query_fn is not None,286 has_custom_hyde=hyde_fn is not None,287 )288 289 def retrieve(290 self,291 query: str,292 n_results: int = 10,293 where: dict | None = None,294 use_multi_query: bool = False,295 use_hyde: bool = False,296 ) -> list[dict]:297 """Run hybrid retrieval with optional multi-query and HyDE.298 299 Pipeline:300 1. Optionally expand query into variations (multi-query)301 2. Optionally transform query for vector search (HyDE)302 3. Run BM25 search for each query variant303 4. Run vector search for each query variant304 5. Deduplicate within each source305 6. Fuse via RRF306 7. Return top n_results307 308 Args:309 query: Natural language query string.310 n_results: Maximum results to return after fusion.311 where: Optional metadata filter dict. Applied to both312 BM25 and vector search.313 Example: {"ticker": "AAPL"} or314 {"$and": [{"ticker": "AAPL"}, {"form_type": "10-K"}]}315 use_multi_query: If True, expand query into variations316 for broader recall.317 use_hyde: If True, generate hypothetical document embedding318 for vector search.319 320 Returns:321 List of result dicts sorted by RRF score descending,322 each containing:323 - chunk_id: The deterministic chunk ID324 - text: The chunk text325 - metadata: Dict of chunk metadata326 - rrf_score: Fused relevance score327 - retrieval_sources: List of sources ("dense", "sparse")328 """329 if not query.strip():330 return []331 332 # Convert flat where filter to ChromaDB $and syntax if needed333 if where and len(where) > 1 and "$and" not in where and "$or" not in where:334 where = {"$and": [{k: v} for k, v in where.items()]}335 336 # Step 1: Determine query variants337 if use_multi_query:338 queries = self._multi_query_fn(query)339 logger.info(340 "multi_query_expanded",341 original=query[:80],342 variations=len(queries),343 )344 else:345 queries = [query]346 347 # Step 2: Apply HyDE transformation for vector search348 if use_hyde:349 hyde_query = self._hyde_fn(query)350 logger.info(351 "hyde_applied",352 original_preview=query[:80],353 hyde_preview=hyde_query[:80],354 )355 else:356 hyde_query = query357 358 # Step 3: Run retrievals for all query variants359 all_dense_results: list[dict] = []360 all_sparse_results: list[dict] = []361 362 for q in queries:363 # Dense search: use HyDE query for the original query,364 # use the expanded variant directly for multi-query expansions.365 dense_q = hyde_query if q == query else q366 367 dense_results = self._chroma.query(368 query_text=dense_q,369 n_results=self._candidates,370 where=where,371 )372 sparse_results = self._bm25.query(373 query_text=q,374 n_results=self._candidates,375 where=where,376 )377 378 all_dense_results.extend(dense_results)379 all_sparse_results.extend(sparse_results)380 381 # Step 4: Deduplicate within each source (keep first = highest ranked)382 dense_deduped = self._deduplicate_results(all_dense_results)383 sparse_deduped = self._deduplicate_results(all_sparse_results)384 385 # Step 5: Fuse via RRF386 fused = reciprocal_rank_fusion(387 ranked_lists=[dense_deduped, sparse_deduped],388 k=self._rrf_k,389 )390 391 # Step 6: Trim to requested count392 results = fused[:n_results]393 394 logger.info(395 "hybrid_retrieval_complete",396 query_preview=query[:80],397 queries_executed=len(queries),398 dense_candidates=len(dense_deduped),399 sparse_candidates=len(sparse_deduped),400 fused_total=len(fused),401 returned=len(results),402 filter=where,403 multi_query=use_multi_query,404 hyde=use_hyde,405 )406 407 return results408 409 def retrieve_dense_only(410 self,411 query: str,412 n_results: int = 10,413 where: dict | None = None,414 ) -> list[dict]:415 """Run vector search only, bypassing BM25 and RRF.416 417 Useful for benchmarking dense vs hybrid retrieval.418 419 Args:420 query: Natural language query string.421 n_results: Maximum results to return.422 where: Optional metadata filter dict.423 424 Returns:425 List of result dicts from vector search only.426 """427 return self._chroma.query(428 query_text=query,429 n_results=n_results,430 where=where,431 )432 433 def retrieve_sparse_only(434 self,435 query: str,436 n_results: int = 10,437 where: dict | None = None,438 ) -> list[dict]:439 """Run BM25 search only, bypassing vector search and RRF.440 441 Useful for benchmarking sparse vs hybrid retrieval.442 443 Args:444 query: Natural language query string.445 n_results: Maximum results to return.446 where: Optional metadata filter dict.447 448 Returns:449 List of result dicts from BM25 search only.450 """451 return self._bm25.query(452 query_text=query,453 n_results=n_results,454 where=where,455 )456 457 def _deduplicate_results(self, results: list[dict]) -> list[dict]:458 """Remove duplicate chunk_ids, keeping the first occurrence.459 460 When multi-query produces overlapping results across query461 variants, we keep only the highest-ranked occurrence. This462 prevents a single chunk from appearing multiple times in463 one ranked list, which would inflate its RRF score unfairly.464 465 Args:466 results: List of result dicts with chunk_id field.467 468 Returns:469 Deduplicated list preserving original order.470 """471 seen: set[str] = set()472 deduped: list[dict] = []473 for result in results:474 chunk_id = result["chunk_id"]475 if chunk_id not in seen:476 seen.add(chunk_id)477 deduped.append(result)478 return deduped479 480 def get_stats(self) -> dict:481 """Return retriever configuration and component stats.482 483 Returns:484 Dict with hybrid retriever config and sub-component stats.485 """486 return {487 "rrf_k": self._rrf_k,488 "candidates_per_retriever": self._candidates,489 "bm25_stats": self._bm25.get_stats(),490 "chroma_stats": self._chroma.get_stats(),491 }492