Coverage for src/lilbee/retrieval/query/dedup.py: 100%
95 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Result filtering, sorting, deduplication, and source-diversity helpers."""
3from __future__ import annotations
5import re
7from lilbee.core.config import active_config
8from lilbee.data.store import SearchChunk
10_DEFAULT_RELEVANCE_WEIGHT = 0.5
12_WHITESPACE_RE = re.compile(r"\s+")
14# Neutral point of the [0, 1] fusion scale: used for a candidate with no usable
15# score and for a cohort whose scores are all equal (no spread to rank on).
16_NEUTRAL_SCORE = 0.5
19def _relevance_weight(result: SearchChunk) -> float:
20 """Return a [0, 1] relevance weight for distance-aware selection.
22 Every store search path stamps the canonical ``score``; a scoreless row
23 (constructed by hand, never by retrieval) weighs neutrally rather than
24 resurrecting the pre-score per-arm arithmetic.
25 """
26 if result.score is not None:
27 return min(1.0, max(0.0, result.score))
28 return _DEFAULT_RELEVANCE_WEIGHT
31def normalize_scores(scores: list[float]) -> list[float]:
32 """Min-max normalize scores to [0, 1]; an all-equal set maps to the midpoint."""
33 min_score = min(scores)
34 max_score = max(scores)
35 score_range = max_score - min_score
36 if score_range > 0:
37 return [(s - min_score) / score_range for s in scores]
38 return [_NEUTRAL_SCORE] * len(scores)
41def fusion_norms(results: list[SearchChunk]) -> list[float]:
42 """Normalize each chunk's canonical score to [0, 1] across the result set.
44 Every retrieval path scores in the same [0, 1] family, so one min-max
45 pass suffices; a scoreless hand-built row sits at the neutral midpoint.
46 """
47 scored = [i for i, r in enumerate(results) if r.score is not None]
48 norms = [_NEUTRAL_SCORE] * len(results)
49 if scored:
50 scaled = normalize_scores([results[i].score or 0.0 for i in scored])
51 for i, value in zip(scored, scaled, strict=True):
52 norms[i] = value
53 return norms
56def order_by_fusion(results: list[SearchChunk]) -> list[SearchChunk]:
57 """Sort results best-first by canonical score (stable on ties)."""
58 norms = fusion_norms(results)
59 order = sorted(range(len(results)), key=lambda i: norms[i], reverse=True)
60 return [results[i] for i in order]
63def _greedy_cover(
64 chunk_tokens: list[set[str]],
65 question_terms: set[str],
66 term_weights: dict[str, float],
67 budget: int,
68 relevance_weights: list[float] | None = None,
69) -> list[int]:
70 """Greedy weighted set cover: pick chunks that add the most uncovered weight.
72 Standard (1 - 1/e) approximation for weighted set cover. Budget is
73 always filled, falling back to retrieval order once no chunk can
74 contribute any new weight. When *relevance_weights* is provided,
75 each chunk's IDF gain is scaled by its relevance so that far-away
76 chunks are penalised even when they share query terms.
77 """
78 selected: list[int] = []
79 covered: set[str] = set()
80 remaining = list(range(len(chunk_tokens)))
81 while remaining and len(selected) < budget:
82 best_pos = -1
83 best_gain = 0.0
84 for pos, idx in enumerate(remaining):
85 new_terms = (chunk_tokens[idx] & question_terms) - covered
86 gain = sum(term_weights[t] for t in new_terms)
87 if relevance_weights is not None:
88 gain *= relevance_weights[idx]
89 if gain > best_gain:
90 best_gain = gain
91 best_pos = pos
92 if best_pos < 0:
93 break
94 chosen = remaining.pop(best_pos)
95 selected.append(chosen)
96 covered |= chunk_tokens[chosen] & question_terms
98 for idx in remaining:
99 if len(selected) >= budget:
100 break
101 selected.append(idx)
102 return selected
105def filter_results(
106 results: list[SearchChunk],
107 max_distance: float,
108 min_relevance_score: float = 0.0,
109) -> list[SearchChunk]:
110 """Drop results below min_relevance_score or above max_distance.
112 ``min_relevance_score`` gates on the [0, 1] fused score, which normalizes
113 against the arms in play, so the threshold means the same thing across
114 queries. ``max_distance`` additionally drops rows whose only
115 signal is a far vector match (a row with lexical support keeps its standing
116 regardless of distance). Pass max_distance=0 to disable distance filtering.
117 """
118 if max_distance <= 0 and min_relevance_score <= 0:
119 return results
120 filtered: list[SearchChunk] = []
121 for r in results:
122 if min_relevance_score > 0 and r.score is not None and r.score < min_relevance_score:
123 continue
124 if (
125 max_distance > 0
126 and r.bm25_score is None
127 and r.distance is not None
128 and r.distance > max_distance
129 ):
130 continue
131 filtered.append(r)
132 return filtered
135def _sort_key(r: SearchChunk) -> float:
136 """Sort key: lower = more relevant; a scoreless hand-built row sorts last."""
137 if r.score is not None:
138 return -r.score
139 return float("inf")
142def sort_by_relevance(results: list[SearchChunk]) -> list[SearchChunk]:
143 """Sort search results by relevance (works for both hybrid and vector results)."""
144 return sorted(results, key=_sort_key)
147def diversify_sources(
148 results: list[SearchChunk], max_per_source: int | None = None
149) -> list[SearchChunk]:
150 """Cap results per source document to ensure diversity.
151 Source diversity filtering: Zhai 2008, "Statistical Language Models for
152 Information Retrieval" -- caps per-source representation to prevent
153 any single document from dominating results.
155 Callers holding a Config pass its ``diversity_max_per_source`` so the
156 library API's scoped config is honored; the active-config default only
157 covers direct ad-hoc calls.
158 """
159 if max_per_source is None:
160 max_per_source = active_config().diversity_max_per_source
161 counts: dict[str, int] = {}
162 diverse: list[SearchChunk] = []
163 for r in results:
164 count = counts.get(r.source, 0)
165 if count < max_per_source:
166 diverse.append(r)
167 counts[r.source] = count + 1
168 return diverse
171def dedup_near_identical(results: list[SearchChunk]) -> list[SearchChunk]:
172 """Keep the first (best-ranked) copy of passages with identical normalized text.
174 The per-source cap cannot catch the same file ingested under two paths or
175 boilerplate repeated across documents; those copies add no information and
176 crowd real passages out of the context.
177 """
178 seen: set[str] = set()
179 kept: list[SearchChunk] = []
180 for r in results:
181 key = _WHITESPACE_RE.sub(" ", r.chunk).strip().lower()
182 if key in seen:
183 continue
184 seen.add(key)
185 kept.append(r)
186 return kept
189def prepare_results(
190 results: list[SearchChunk], max_per_source: int | None = None
191) -> list[SearchChunk]:
192 """Sort by relevance, drop near-identical copies, apply the source diversity cap."""
193 return diversify_sources(dedup_near_identical(sort_by_relevance(results)), max_per_source)