""" Helper functions for hybrid search (semantic + BM25 + graph). """ from typing import Any from .types import MergedCandidate, RetrievalResult def reciprocal_rank_fusion(result_lists: list[list[RetrievalResult]], k: int = 60) -> list[MergedCandidate]: """ Merge multiple ranked result lists using Reciprocal Rank Fusion. RRF formula: score(d) = sum_over_lists(1 / (k + rank(d))) Args: result_lists: List of result lists, each containing RetrievalResult objects k: Constant for RRF formula (default: 60) Returns: Merged list of MergedCandidate objects, sorted by RRF score Example: semantic_results = [RetrievalResult(...), RetrievalResult(...), ...] bm25_results = [RetrievalResult(...), RetrievalResult(...), ...] graph_results = [RetrievalResult(...), RetrievalResult(...), ...] merged = reciprocal_rank_fusion([semantic_results, bm25_results, graph_results]) # Returns: [MergedCandidate(...), MergedCandidate(...), ...] """ # Track scores from each list rrf_scores = {} source_ranks = {} # Track rank from each source for each doc_id all_retrievals = {} # Store the actual RetrievalResult (use first occurrence) source_names = ["semantic", "bm25", "graph", "temporal"] for source_idx, results in enumerate(result_lists): source_name = source_names[source_idx] if source_idx < len(source_names) else f"source_{source_idx}" for rank, retrieval in enumerate(results, start=1): # Type check to catch tuple issues if isinstance(retrieval, tuple): raise TypeError( f"Expected RetrievalResult but got tuple in {source_name} results at rank {rank}. " f"Tuple value: {retrieval[:2] if len(retrieval) >= 2 else retrieval}. " f"This suggests the retrieval function returned tuples instead of RetrievalResult objects." ) if not isinstance(retrieval, RetrievalResult): raise TypeError( f"Expected RetrievalResult but got {type(retrieval).__name__} in {source_name} results at rank {rank}" ) doc_id = retrieval.id # Store retrieval result (use first occurrence) if doc_id not in all_retrievals: all_retrievals[doc_id] = retrieval # Calculate RRF score contribution if doc_id not in rrf_scores: rrf_scores[doc_id] = 0.0 source_ranks[doc_id] = {} rrf_scores[doc_id] += 1.0 / (k + rank) source_ranks[doc_id][f"{source_name}_rank"] = rank # Combine into final results with metadata merged_results = [] for rrf_rank, (doc_id, rrf_score) in enumerate( sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True), start=1 ): merged_candidate = MergedCandidate( retrieval=all_retrievals[doc_id], rrf_score=rrf_score, rrf_rank=rrf_rank, source_ranks=source_ranks[doc_id] ) merged_results.append(merged_candidate) return merged_results def normalize_scores_on_deltas(results: list[dict[str, Any]], score_keys: list[str]) -> list[dict[str, Any]]: """ Normalize scores based on deltas (min-max normalization within result set). This ensures all scores are in [0, 1] range based on the spread in THIS result set. Args: results: List of result dicts score_keys: Keys to normalize (e.g., ["recency", "frequency"]) Returns: Results with normalized scores added as "{key}_normalized" """ for key in score_keys: values = [r.get(key, 0.0) for r in results if key in r] if not values: continue min_val = min(values) max_val = max(values) delta = max_val - min_val if delta > 0: for r in results: if key in r: r[f"{key}_normalized"] = (r[key] - min_val) / delta else: # All values are the same, set to 0.5 for r in results: if key in r: r[f"{key}_normalized"] = 0.5 return results