122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
"""
|
|
Helper functions for hybrid search (semantic + BM25 + graph).
|
|
"""
|
|
|
|
from typing import List, Dict, Any, Tuple
|
|
import asyncio
|
|
from .types import RetrievalResult, MergedCandidate
|
|
|
|
|
|
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
|