Coverage for src/lilbee/data/store/fusion.py: 100%
58 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""Reciprocal-rank fusion of the vector, BM25, and optional title arms.
3Each arm contributes ``(K + 1) / (K + rank)`` for the rows it retrieved
4(rank is 1-based, K = 60, the standard RRF constant); a row's canonical
5score is the weight-normalized sum of its arm contributions, so it lives
6in [0, 1] and a row ranked first by every arm scores exactly 1. Rows seen
7by only one arm score that arm's share of the total weight, which still
8places an arm's top hit above every row deep in the other arms: the
9property that keeps a lexical-only identifier match visible next to dense
10neighbors. The vector arm weighs 1; the chunk-BM25 arm weighs
11``lexical_weight`` (1.0 = equal voice, lower lets a strong dense arm
12dominate); the optional title arm weighs ``title_weight``. Weights rescale
13the shares without leaving the canonical range.
15Scores normalize against the weights of the arms in the call (the title
16arm counts only when it returned rows), so every call reports the fraction
17of its trusted rank support on one canonical [0, 1] scale and
18``min_relevance_score`` keeps one meaning everywhere. A fixed shared
19denominator was tried and rejected: it capped adaptively quieted calls at
20the vector share and demoted their hits in cross-variant merges.
22Rank fusion is deliberate. A convex combination of normalized raw scores
23(``alpha * vector_similarity + (1 - alpha) * normalized_bm25``) was tried
24here and measurably regressed graded precision: cosine similarities sit
25in a high narrow band, giving every dense neighbor a floor that outranks
26lexically-certain rows, and no blend weight fixes that asymmetry. Ranks
27are scale-free, so neither arm's score distribution can crowd out the
28other. Arm depth matters as much as the formula: rows both arms rank
29mid-pool accumulate two contributions, so deep candidate pools crowd out
30single-arm certainty; the hybrid path therefore feeds fusion arms of
31exactly ``top_k`` rows.
33Per-query fusion serves the history-rewrite path: the condensed rewrite's
34ranking fuses with the typed question's own retrieval at equal weight, so
35a rewrite that absorbed conversation context cannot crowd out the
36question's evidence. Shares normalize over the live lists, so one silent
37query does not demote the other, and a row found by both keeps the
38closest distance and any lexical support.
39"""
41from __future__ import annotations
43from statistics import fmean
45from .types import SearchChunk
47# Standard RRF smoothing constant (Cormack, Clarke & Buettcher 2009).
48_RRF_K = 60
50# Adaptive fusion needs a top hit plus at least one field row to measure a margin.
51_MIN_ROWS_FOR_MARGIN = 2
53# Runners-up window for the margin; a fixed window keeps the signal independent
54# of retrieval depth (a full-pool mean grew with candidate count).
55_MARGIN_WINDOW = 5
57# Adaptive scale lower bound: arms are quieted, never silenced, so BM25
58# provenance and the distance-cut exemption survive. Exact zero hard-dropped
59# lexical-only rows.
60_ADAPTIVE_SCALE_FLOOR = 0.05
63def vector_similarity(distance: float) -> float:
64 """Cosine distance to canonical [0, 1] similarity (distance spans [0, 2])."""
65 return max(0.0, min(1.0, 1.0 - distance))
68def adaptive_weight_scale(vector_rows: list[SearchChunk], margin_scale: float) -> float:
69 """A [0, 1] factor to shrink the lexical arms by when the vector arm is
70 confident about this query.
72 A *peaked* vector ranking -- a top hit standing well clear of the field --
73 means the dense embedder already located the answer and the lexical arms
74 mostly add term-match noise. A *flat* ranking means dense is unsure and
75 BM25's exact-term matching is worth trusting. The confidence signal is the
76 margin between the top similarity and the mean of the next
77 ``_MARGIN_WINDOW`` similarities (a fixed window, so the signal does not
78 change with retrieval depth), divided by *margin_scale*: at or above that
79 margin the factor bottoms out at ``_ADAPTIVE_SCALE_FLOOR`` (arms quieted
80 but their provenance kept), at zero margin it is 1 (arms kept), scaling
81 linearly between. Returns 1.0 when there is nothing to measure (fewer than
82 two scored rows) or when *margin_scale* <= 0 (adaptation off).
83 """
84 if margin_scale <= 0:
85 return 1.0
86 sims = sorted(
87 (vector_similarity(r.distance) for r in vector_rows if r.distance is not None),
88 reverse=True,
89 )
90 if len(sims) < _MIN_ROWS_FOR_MARGIN:
91 return 1.0
92 margin = max(0.0, sims[0] - fmean(sims[1 : 1 + _MARGIN_WINDOW]))
93 confidence = min(1.0, margin / margin_scale)
94 return max(1.0 - confidence, _ADAPTIVE_SCALE_FLOOR)
97def normalized_bm25(scores: list[float]) -> list[float]:
98 """Scale raw BM25 scores against the list maximum, into (0, 1].
100 BM25 has no absolute scale, so the top hit anchors the list; relative
101 strength within one query's results is the meaningful quantity.
102 Non-positive or absent maxima map everything to 0.
103 """
104 top = max(scores, default=0.0)
105 if top <= 0.0:
106 return [0.0] * len(scores)
107 return [max(0.0, s) / top for s in scores]
110def _key(chunk: SearchChunk) -> tuple[str, int]:
111 return (chunk.source, chunk.chunk_index)
114def _rank_weight(rank: int) -> float:
115 """Reciprocal-rank contribution in (0, 1]; 1.0 at rank 1."""
116 return (_RRF_K + 1) / (_RRF_K + rank)
119def _merge_arm(
120 merged: dict[tuple[str, int], SearchChunk],
121 rows: list[SearchChunk],
122 share: float,
123 *,
124 keep_closest_distance: bool = False,
125) -> None:
126 """Fold one arm's ranked rows into *merged*, each contributing *share* of its rank weight.
128 A row seen before keeps the first lexical support; with
129 *keep_closest_distance* it also keeps the closest distance.
130 """
131 for rank, row in enumerate(rows, start=1):
132 key = _key(row)
133 contribution = _rank_weight(rank) * share
134 seen = merged.get(key)
135 if seen is None:
136 merged[key] = row.model_copy(update={"score": contribution})
137 else:
138 update: dict[str, object] = {"score": (seen.score or 0.0) + contribution}
139 if (
140 keep_closest_distance
141 and row.distance is not None
142 and (seen.distance is None or row.distance < seen.distance)
143 ):
144 update["distance"] = row.distance
145 if seen.bm25_score is None and row.bm25_score is not None:
146 update["bm25_score"] = row.bm25_score
147 merged[key] = seen.model_copy(update=update)
150def fuse_arms(
151 vector_rows: list[SearchChunk],
152 fts_rows: list[SearchChunk],
153 title_rows: list[SearchChunk] | None = None,
154 *,
155 lexical_weight: float = 1.0,
156 title_weight: float = 1.0,
157) -> list[SearchChunk]:
158 """Merge the arms into one list scored by reciprocal rank.
160 The vector arm weighs 1; the chunk-FTS (lexical) arm weighs *lexical_weight*
161 relative to it (1.0 = equal voice, lower lets a strong dense arm dominate);
162 a non-empty *title_rows* arm joins at *title_weight*. Rows found by several
163 arms carry every provenance field (``distance`` from the vector arm,
164 ``bm25_score`` from the FTS arms). The result is sorted by ``score``
165 descending and deduplicated on ``(source, chunk_index)``.
167 Scores normalize against the weights of the arms in this call (title only
168 when it returned rows); see the module docstring for why the denominator
169 is per-call rather than shared.
170 """
171 weight_total = 1.0 + lexical_weight + (title_weight if title_rows else 0.0)
172 merged: dict[tuple[str, int], SearchChunk] = {}
173 _merge_arm(merged, vector_rows, 1.0 / weight_total)
174 # A zero-weight arm is configured off (adaptive scaling floors above zero),
175 # so skip it rather than folding in zero-score rows that would still carry
176 # lexical provenance (and its downstream distance/structural exemptions).
177 if lexical_weight > 0:
178 _merge_arm(merged, fts_rows, lexical_weight / weight_total)
179 if title_rows and title_weight > 0:
180 _merge_arm(merged, title_rows, title_weight / weight_total)
181 return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True)
184def fuse_ranked_lists(query_lists: list[list[SearchChunk]]) -> list[SearchChunk]:
185 """Fuse equally-weighted per-query rankings by reciprocal rank.
187 Each live list contributes its rank weights scaled by its share of the
188 lists in the call, so one silent query does not demote the other. A row
189 found by several queries keeps the closest distance and any lexical
190 support. The result is sorted by score descending and deduplicated on
191 (source, chunk_index).
192 """
193 lists = [rows for rows in query_lists if rows]
194 if not lists:
195 return []
196 share = 1.0 / len(lists)
197 merged: dict[tuple[str, int], SearchChunk] = {}
198 for rows in lists:
199 _merge_arm(merged, rows, share, keep_closest_distance=True)
200 return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True)