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

1"""Vector ranking primitives: cosine similarity and Maximal Marginal Relevance reranking.""" 

2 

3from __future__ import annotations 

4 

5import math 

6 

7import numpy as np 

8 

9from lilbee.core.config import cfg 

10 

11from .types import SearchChunk 

12 

13 

14def cosine_sim(a: list[float], b: list[float]) -> float: 

15 """Cosine similarity between two vectors.""" 

16 dot = sum(x * y for x, y in zip(a, b, strict=True)) 

17 norm_a = math.sqrt(sum(x * x for x in a)) 

18 norm_b = math.sqrt(sum(x * x for x in b)) 

19 if norm_a == 0 or norm_b == 0: 

20 return 0.0 

21 return dot / (norm_a * norm_b) 

22 

23 

24def mmr_rerank( 

25 query_vector: list[float], 

26 results: list[SearchChunk], 

27 top_k: int, 

28 mmr_lambda: float | None = None, 

29) -> list[SearchChunk]: 

30 """Maximal Marginal Relevance: select diverse results. 

31 Algorithm: Carbonell & Goldstein 1998, 

32 "The Use of MMR, Diversity-Based Reranking for Reordering Documents 

33 and Producing Summaries." 

34 

35 ``mmr_lambda`` controls the relevance/diversity tradeoff: 

36 0.0 = maximum diversity, 1.0 = pure relevance. 

37 Defaults to ``cfg.mmr_lambda`` (0.5). 

38 

39 Complexity: O(top_k · N · D) time, O(N · D) space for N candidates 

40 of dimension D. Each outer iteration updates a running max-redundancy 

41 vector via one matmul rather than recomputing pairs pairwise. 

42 Candidate vectors run through numpy in ``float32``, which can pick a 

43 different candidate than the pure-Python ``float64`` loop on 

44 ties within ~1e-7; distinct in principle, unobservable in practice 

45 since sub-float32 differences are below retrieval signal. 

46 """ 

47 if mmr_lambda is None: 

48 mmr_lambda = cfg.mmr_lambda 

49 if len(results) <= top_k: 

50 return results 

51 

52 candidate_vecs = np.asarray([r.vector for r in results], dtype=np.float32) 

53 query = np.asarray(query_vector, dtype=np.float32) 

54 # L2-normalize once so cosine becomes a plain dot product. 

55 cand_norms = np.linalg.norm(candidate_vecs, axis=1, keepdims=True) 

56 cand_norms[cand_norms == 0] = 1.0 

57 cand_unit = candidate_vecs / cand_norms 

58 query_norm = float(np.linalg.norm(query)) or 1.0 

59 query_unit = query / query_norm 

60 

61 relevance = cand_unit @ query_unit # shape (N,) 

62 

63 n = len(results) 

64 max_redundancy = np.zeros(n, dtype=np.float32) 

65 available = np.ones(n, dtype=bool) 

66 selected: list[SearchChunk] = [] 

67 

68 for picks in range(top_k): 

69 redundancy_term = max_redundancy if picks > 0 else np.zeros(n, dtype=np.float32) 

70 score = mmr_lambda * relevance - (1.0 - mmr_lambda) * redundancy_term 

71 # Mask already-picked candidates so argmax skips them. 

72 score = np.where(available, score, -np.inf) 

73 best = int(np.argmax(score)) 

74 selected.append(results[best]) 

75 available[best] = False 

76 # Update running max redundancy against the newly-selected vector. 

77 similarity = cand_unit @ cand_unit[best] 

78 max_redundancy = np.maximum(max_redundancy, similarity) 

79 

80 return selected