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 

5from collections.abc import Sequence 

6 

7import numpy as np 

8 

9from lilbee.core.config import active_config 

10from lilbee.core.vectors import Vector 

11 

12from .types import SearchChunk 

13 

14 

15def cosine_sim(a: Vector | Sequence[float], b: Vector | Sequence[float]) -> float: 

16 """Cosine similarity between two vectors.""" 

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

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

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

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

21 if norm_a == 0.0 or norm_b == 0.0: 

22 return 0.0 

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

24 

25 

26def mmr_rerank( 

27 query_vector: Vector, 

28 results: list[SearchChunk], 

29 top_k: int, 

30 mmr_lambda: float | None = None, 

31) -> list[SearchChunk]: 

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

33 Algorithm: Carbonell & Goldstein 1998, 

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

35 and Producing Summaries." 

36 

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

38 0.0 = maximum diversity, 1.0 = pure relevance. 

39 Defaults to the active config's ``mmr_lambda`` (0.5). 

40 

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

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

43 vector via one matmul rather than recomputing pairs pairwise. 

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

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

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

47 since sub-float32 differences are below retrieval signal. 

48 """ 

49 if mmr_lambda is None: 

50 # active_config(), not the process-global cfg: under the library API 

51 # a config_scope binding is the caller's config, and every sibling in 

52 # the data path resolves through the scope. 

53 mmr_lambda = active_config().mmr_lambda 

54 if len(results) <= top_k: 

55 return results 

56 

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

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

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

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

61 cand_norms[cand_norms == 0] = 1.0 

62 cand_unit = candidate_vecs / cand_norms 

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

64 query_unit = query / query_norm 

65 

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

67 

68 n = len(results) 

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

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

71 selected: list[SearchChunk] = [] 

72 

73 for _ in range(top_k): 

74 # max_redundancy starts at zeros, so the first pass is pure relevance 

75 # without needing a special case for it. 

76 score = mmr_lambda * relevance - (1.0 - mmr_lambda) * max_redundancy 

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

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

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

80 selected.append(results[best]) 

81 available[best] = False 

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

83 similarity = cand_unit @ cand_unit[best] 

84 max_redundancy = np.maximum(max_redundancy, similarity) 

85 

86 return selected