"""Score fusion and evidence-set selection algorithms. Retrieval is usually presented as a ranked-list problem, while generation consumes a *set* of passages under a token budget. This module makes that boundary explicit. It contains calibrated score fusion, weighted reciprocal rank fusion, budgeted maximum-coverage selection, and diagnostics for evidence lost between retrieval and generation. The implementations are dependency-free and intentionally expose every intermediate score. They are suitable for experiments and small corpora; a production system should vectorize the same objectives and persist its fitted calibration parameters. """ import math from dataclasses import dataclass from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple from .models import Chunk, SearchResult @dataclass(frozen=True) class SelectionCandidate: """A candidate evidence unit with claims it can support and a unit cost.""" result: SearchResult supports: Tuple[str, ...] cost: int def __post_init__(self) -> None: if self.cost <= 0: raise ValueError("candidate cost must be positive") @dataclass(frozen=True) class CoverageSelection: """Result of greedy budgeted maximum coverage.""" selected: Tuple[SelectionCandidate, ...] covered: Tuple[str, ...] uncovered: Tuple[str, ...] spent: int utility: float @dataclass(frozen=True) class EvidenceFlow: """Document-level survival through candidate, rerank, and packing stages.""" relevant: Tuple[str, ...] retrieved: Tuple[str, ...] reranked: Tuple[str, ...] packed: Tuple[str, ...] retrieval_recall: float rerank_survival: float pack_survival: float end_to_end_recall: float lost_at_retrieval: Tuple[str, ...] lost_at_rerank: Tuple[str, ...] lost_at_pack: Tuple[str, ...] def min_max(values: Sequence[float]) -> Tuple[float, ...]: """Scale values to ``[0, 1]`` with a stable all-equal convention.""" if not values: return () low, high = min(values), max(values) if high == low: return tuple(1.0 for _ in values) return tuple((value - low) / (high - low) for value in values) def z_scores(values: Sequence[float]) -> Tuple[float, ...]: """Population z-scores; an all-equal list maps to zero.""" if not values: return () mean = sum(values) / len(values) variance = sum((value - mean) ** 2 for value in values) / len(values) if variance == 0.0: return tuple(0.0 for _ in values) deviation = math.sqrt(variance) return tuple((value - mean) / deviation for value in values) def _candidate_table( rankings: Mapping[str, Sequence[SearchResult]], ) -> Tuple[Dict[str, Chunk], Dict[str, Dict[str, SearchResult]]]: chunks: Dict[str, Chunk] = {} by_system: Dict[str, Dict[str, SearchResult]] = {} for system, results in rankings.items(): system_results: Dict[str, SearchResult] = {} for result in results: chunks[result.chunk.id] = result.chunk system_results[result.chunk.id] = result by_system[system] = system_results return chunks, by_system def calibrated_comb_sum( rankings: Mapping[str, Sequence[SearchResult]], k: int, weights: Optional[Mapping[str, float]] = None, normalization: str = "minmax", missing_score: float = 0.0, ) -> List[SearchResult]: """Fuse scores after per-system calibration. ``CombSUM`` is meaningful only when component scores share a scale. The function therefore fits a deterministic query-local transform per input ranking. Query-local min-max or z-score normalization is useful for a lab, but a deployment should fit calibration on held-out judgments and monitor it after corpus/model changes. """ if k <= 0 or not rankings: return [] if normalization not in {"minmax", "zscore", "none"}: raise ValueError("normalization must be minmax, zscore, or none") chunks, by_system = _candidate_table(rankings) normalized: Dict[str, Dict[str, float]] = {} for system, results in rankings.items(): raw = [result.score for result in results] scaled = ( min_max(raw) if normalization == "minmax" else z_scores(raw) if normalization == "zscore" else tuple(raw) ) normalized[system] = { result.chunk.id: score for result, score in zip(results, scaled) } fused: List[Tuple[float, str, Dict[str, float]]] = [] for chunk_id in chunks: score = 0.0 components: Dict[str, float] = {} for system in rankings: weight = 1.0 if weights is None else weights.get(system, 1.0) value = normalized[system].get(chunk_id, missing_score) contribution = weight * value components[system + "_calibrated"] = value components[system + "_weighted"] = contribution score += contribution raw_result = by_system[system].get(chunk_id) if raw_result is not None: components[system + "_raw"] = raw_result.score fused.append((score, chunk_id, components)) fused.sort(key=lambda item: (-item[0], item[1])) return [ SearchResult( chunk=chunks[chunk_id], score=score, rank=rank, retriever="calibrated-combsum", component_scores=components, ) for rank, (score, chunk_id, components) in enumerate(fused[:k], start=1) ] def reciprocal_rank_fusion( rankings: Mapping[str, Sequence[SearchResult]], k: int, constant: int = 60, weights: Optional[Mapping[str, float]] = None, ) -> List[SearchResult]: """Fuse ordinal ranks without assuming comparable component scores.""" if constant <= 0: raise ValueError("constant must be positive") if k <= 0 or not rankings: return [] chunks, _ = _candidate_table(rankings) totals: Dict[str, float] = {chunk_id: 0.0 for chunk_id in chunks} components: Dict[str, Dict[str, float]] = {chunk_id: {} for chunk_id in chunks} for system, results in rankings.items(): weight = 1.0 if weights is None else weights.get(system, 1.0) for ordinal, result in enumerate(results, start=1): contribution = weight / (constant + ordinal) totals[result.chunk.id] += contribution components[result.chunk.id][system + "_rrf"] = contribution ranked = sorted(totals, key=lambda chunk_id: (-totals[chunk_id], chunk_id))[:k] return [ SearchResult( chunk=chunks[chunk_id], score=totals[chunk_id], rank=rank, retriever="rrf", component_scores=components[chunk_id], ) for rank, chunk_id in enumerate(ranked, start=1) ] def greedy_budgeted_coverage( candidates: Sequence[SelectionCandidate], required: Iterable[str], budget: int, claim_weights: Optional[Mapping[str, float]] = None, relevance_weight: float = 0.05, ) -> CoverageSelection: """Select evidence by marginal supported-claim utility per unit cost. Weighted maximum coverage under a knapsack budget is NP-hard. This greedy approximation repeatedly chooses the candidate with the highest marginal coverage-plus-relevance utility per cost. Stable ID tie-breaking keeps experiments reproducible. """ if budget < 0: raise ValueError("budget must be non-negative") required_set = set(required) weights = {claim: 1.0 for claim in required_set} if claim_weights: for claim, weight in claim_weights.items(): if weight < 0: raise ValueError("claim weights must be non-negative") if claim in required_set: weights[claim] = weight remaining = list(candidates) selected: List[SelectionCandidate] = [] covered: Set[str] = set() spent = 0 utility = 0.0 while remaining: feasible = [candidate for candidate in remaining if spent + candidate.cost <= budget] if not feasible: break scored = [] for candidate in feasible: new_claims = (set(candidate.supports) & required_set) - covered coverage_gain = sum(weights[claim] for claim in new_claims) relevance_gain = relevance_weight * max(0.0, candidate.result.score) gain = coverage_gain + relevance_gain scored.append((gain / candidate.cost, gain, candidate.result.chunk.id, candidate)) scored.sort(key=lambda item: (-item[0], -item[1], item[2])) _, gain, _, chosen = scored[0] if gain <= 0.0: break selected.append(chosen) covered.update(set(chosen.supports) & required_set) spent += chosen.cost utility += gain remaining = [item for item in remaining if item.result.chunk.id != chosen.result.chunk.id] ordered_covered = tuple(sorted(covered)) return CoverageSelection( selected=tuple(selected), covered=ordered_covered, uncovered=tuple(sorted(required_set - covered)), spent=spent, utility=utility, ) def _documents(results: Sequence[SearchResult]) -> Tuple[str, ...]: seen: Set[str] = set() ordered: List[str] = [] for result in results: identifier = result.chunk.document_id if identifier not in seen: ordered.append(identifier) seen.add(identifier) return tuple(ordered) def evidence_flow( relevant_document_ids: Iterable[str], retrieved: Sequence[SearchResult], reranked: Sequence[SearchResult], packed: Sequence[SearchResult], ) -> EvidenceFlow: """Attribute recall loss to retrieval, reranking, or context packing.""" relevant = tuple(sorted(set(relevant_document_ids))) relevant_set = set(relevant) retrieved_ids = _documents(retrieved) reranked_ids = _documents(reranked) packed_ids = _documents(packed) retrieved_relevant = relevant_set & set(retrieved_ids) reranked_relevant = relevant_set & set(reranked_ids) packed_relevant = relevant_set & set(packed_ids) def ratio(numerator: int, denominator: int) -> float: return numerator / denominator if denominator else 1.0 return EvidenceFlow( relevant=relevant, retrieved=retrieved_ids, reranked=reranked_ids, packed=packed_ids, retrieval_recall=ratio(len(retrieved_relevant), len(relevant_set)), rerank_survival=ratio(len(reranked_relevant), len(retrieved_relevant)), pack_survival=ratio(len(packed_relevant), len(reranked_relevant)), end_to_end_recall=ratio(len(packed_relevant), len(relevant_set)), lost_at_retrieval=tuple(sorted(relevant_set - retrieved_relevant)), lost_at_rerank=tuple(sorted(retrieved_relevant - reranked_relevant)), lost_at_pack=tuple(sorted(reranked_relevant - packed_relevant)), ) def selection_regret(oracle_utility: float, observed_utility: float) -> float: """Normalized utility regret, clamped to ``[0, 1]``.""" if oracle_utility < 0 or observed_utility < 0: raise ValueError("utilities must be non-negative") if oracle_utility == 0: return 0.0 if observed_utility == 0 else 0.0 return min(1.0, max(0.0, (oracle_utility - observed_utility) / oracle_utility))