Coverage for src/lilbee/data/store/fusion.py: 100%

1"""Reciprocal-rank fusion of the vector, BM25, and optional title arms. 

2 

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. 

14 

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. 

21 

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. 

32""" 

33 

34from __future__ import annotations 

35 

36from statistics import fmean 

37 

38from .types import SearchChunk 

39 

40# Standard RRF smoothing constant (Cormack, Clarke & Buettcher 2009). 

41_RRF_K = 60 

42 

43# Adaptive fusion needs a top hit plus at least one field row to measure a margin. 

44_MIN_ROWS_FOR_MARGIN = 2 

45 

46# Runners-up window for the margin; a fixed window keeps the signal independent 

47# of retrieval depth (a full-pool mean grew with candidate count). 

48_MARGIN_WINDOW = 5 

49 

50# Adaptive scale lower bound: arms are quieted, never silenced, so BM25 

51# provenance and the distance-cut exemption survive. Exact zero hard-dropped 

52# lexical-only rows. 

53_ADAPTIVE_SCALE_FLOOR = 0.05 

54 

55 

56def vector_similarity(distance: float) -> float: 

57 """Cosine distance to canonical [0, 1] similarity (distance spans [0, 2]).""" 

58 return max(0.0, min(1.0, 1.0 - distance)) 

59 

60 

61def adaptive_weight_scale(vector_rows: list[SearchChunk], margin_scale: float) -> float: 

62 """A [0, 1] factor to shrink the lexical arms by when the vector arm is 

63 confident about this query. 

64 

65 A *peaked* vector ranking -- a top hit standing well clear of the field -- 

66 means the dense embedder already located the answer and the lexical arms 

67 mostly add term-match noise. A *flat* ranking means dense is unsure and 

68 BM25's exact-term matching is worth trusting. The confidence signal is the 

69 margin between the top similarity and the mean of the next 

70 ``_MARGIN_WINDOW`` similarities (a fixed window, so the signal does not 

71 change with retrieval depth), divided by *margin_scale*: at or above that 

72 margin the factor bottoms out at ``_ADAPTIVE_SCALE_FLOOR`` (arms quieted 

73 but their provenance kept), at zero margin it is 1 (arms kept), scaling 

74 linearly between. Returns 1.0 when there is nothing to measure (fewer than 

75 two scored rows) or when *margin_scale* <= 0 (adaptation off). 

76 """ 

77 if margin_scale <= 0: 

78 return 1.0 

79 sims = sorted( 

80 (vector_similarity(r.distance) for r in vector_rows if r.distance is not None), 

81 reverse=True, 

82 ) 

83 if len(sims) < _MIN_ROWS_FOR_MARGIN: 

84 return 1.0 

85 margin = max(0.0, sims[0] - fmean(sims[1 : 1 + _MARGIN_WINDOW])) 

86 confidence = min(1.0, margin / margin_scale) 

87 return max(1.0 - confidence, _ADAPTIVE_SCALE_FLOOR) 

88 

89 

90def normalized_bm25(scores: list[float]) -> list[float]: 

91 """Scale raw BM25 scores against the list maximum, into (0, 1]. 

92 

93 BM25 has no absolute scale, so the top hit anchors the list; relative 

94 strength within one query's results is the meaningful quantity. 

95 Non-positive or absent maxima map everything to 0. 

96 """ 

97 top = max(scores, default=0.0) 

98 if top <= 0.0: 

99 return [0.0] * len(scores) 

100 return [max(0.0, s) / top for s in scores] 

101 

102 

103def _key(chunk: SearchChunk) -> tuple[str, int]: 

104 return (chunk.source, chunk.chunk_index) 

105 

106 

107def _rank_weight(rank: int) -> float: 

108 """Reciprocal-rank contribution in (0, 1]; 1.0 at rank 1.""" 

109 return (_RRF_K + 1) / (_RRF_K + rank) 

110 

111 

112def _merge_arm( 

113 merged: dict[tuple[str, int], SearchChunk], 

114 rows: list[SearchChunk], 

115 share: float, 

116) -> None: 

117 """Fold one arm's ranked rows into *merged*, each contributing *share* of its rank weight. 

118 

119 A lexical row (title or chunk FTS) carries ``bm25_score``; when the row was 

120 already seen, that provenance is kept from whichever lexical arm set it 

121 first, so the lexical-support exemption applies either way. 

122 """ 

123 for rank, row in enumerate(rows, start=1): 

124 key = _key(row) 

125 contribution = _rank_weight(rank) * share 

126 seen = merged.get(key) 

127 if seen is None: 

128 merged[key] = row.model_copy(update={"score": contribution}) 

129 else: 

130 update: dict[str, object] = {"score": (seen.score or 0.0) + contribution} 

131 if seen.bm25_score is None and row.bm25_score is not None: 

132 update["bm25_score"] = row.bm25_score 

133 merged[key] = seen.model_copy(update=update) 

134 

135 

136def fuse_arms( 

137 vector_rows: list[SearchChunk], 

138 fts_rows: list[SearchChunk], 

139 title_rows: list[SearchChunk] | None = None, 

140 *, 

141 lexical_weight: float = 1.0, 

142 title_weight: float = 1.0, 

143) -> list[SearchChunk]: 

144 """Merge the arms into one list scored by reciprocal rank. 

145 

146 The vector arm weighs 1; the chunk-FTS (lexical) arm weighs *lexical_weight* 

147 relative to it (1.0 = equal voice, lower lets a strong dense arm dominate); 

148 a non-empty *title_rows* arm joins at *title_weight*. Rows found by several 

149 arms carry every provenance field (``distance`` from the vector arm, 

150 ``bm25_score`` from the FTS arms). The result is sorted by ``score`` 

151 descending and deduplicated on ``(source, chunk_index)``. 

152 

153 Scores normalize against the weights of the arms in this call (title only 

154 when it returned rows); see the module docstring for why the denominator 

155 is per-call rather than shared. 

156 """ 

157 weight_total = 1.0 + lexical_weight + (title_weight if title_rows else 0.0) 

158 merged: dict[tuple[str, int], SearchChunk] = {} 

159 _merge_arm(merged, vector_rows, 1.0 / weight_total) 

160 # A zero-weight arm is configured off (adaptive scaling floors above zero), 

161 # so skip it rather than folding in zero-score rows that would still carry 

162 # lexical provenance (and its downstream distance/structural exemptions). 

163 if lexical_weight > 0: 

164 _merge_arm(merged, fts_rows, lexical_weight / weight_total) 

165 if title_rows and title_weight > 0: 

166 _merge_arm(merged, title_rows, title_weight / weight_total) 

167 return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True)