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 numpy as np 

6 

7from lilbee.core.config import cfg 

8 

9from .types import SearchChunk 

10 

11 

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

13 """Cosine similarity between two vectors.""" 

14 arr_a = np.asarray(a, dtype=np.float64) 

15 arr_b = np.asarray(b, dtype=np.float64) 

16 norm_a = float(np.linalg.norm(arr_a)) 

17 norm_b = float(np.linalg.norm(arr_b)) 

18 if norm_a == 0.0 or norm_b == 0.0: 

19 return 0.0 

20 return float(np.dot(arr_a, arr_b) / (norm_a * norm_b)) 

21 

22 

23def mmr_rerank( 

24 query_vector: list[float], 

25 results: list[SearchChunk], 

26 top_k: int, 

27 mmr_lambda: float | None = None, 

28) -> list[SearchChunk]: 

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

30 Algorithm: Carbonell & Goldstein 1998, 

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

32 and Producing Summaries." 

33 

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

35 0.0 = maximum diversity, 1.0 = pure relevance. 

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

37 

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

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

40 vector via one matmul rather than recomputing pairs pairwise. 

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

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

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

44 since sub-float32 differences are below retrieval signal. 

45 """ 

46 if mmr_lambda is None: 

47 mmr_lambda = cfg.mmr_lambda 

48 if len(results) <= top_k: 

49 return results 

50 

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

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

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

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

55 cand_norms[cand_norms == 0] = 1.0 

56 cand_unit = candidate_vecs / cand_norms 

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

58 query_unit = query / query_norm 

59 

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

61 

62 n = len(results) 

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

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

65 selected: list[SearchChunk] = [] 

66 

67 for picks in range(top_k): 

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

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

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

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

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

73 selected.append(results[best]) 

74 available[best] = False 

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

76 similarity = cand_unit @ cand_unit[best] 

77 max_redundancy = np.maximum(max_redundancy, similarity) 

78 

79 return selected