rrf.ts · Source

rrf.tspackages/search/src/rrf.tsraw ← src2,310 B← Back to the app
1/**
2 * @factory/search — reciprocal rank fusion (search-architecture §4.5).
3 *
4 * `score(d) = Σ 1 / (k + rank_i(d))` over the ranked lists, with the standard
5 * RRF constant k = 60. Ranks are 1-based. Documents missing from a list
6 * simply do not contribute for that list.
7 *
8 * Pure and deterministic: same lists → same fused order (ties broken by
9 * docId for stability).
10 */
11 
12/** The standard RRF constant (search-architecture §4.5). */
13export const RRF_K = 60;
14 
15export interface RankedList {
16  name: string;
17  /** Doc ids in rank order (1-based ranks derived from index). */
18  docIds: string[];
19}
20 
21export interface FusedResult {
22  docId: string;
23  score: number;
24  /** Which lists contained the doc (sorted list names). */
25  sources: string[];
26}
27 
28/**
29 * Fuse two (or more) ranked lists with RRF.
30 *
31 * @param lists ranked lists (typically `bm25` + `vec`)
32 * @param k RRF constant, default 60
33 * @param limit max results (0 = all)
34 */
35export function reciprocalRankFusion(lists: readonly RankedList[], k: number = RRF_K, limit = 0): FusedResult[] {
36  if (k <= 0) throw new Error("RRF k must be positive");
37  const scores = new Map<string, { score: number; sources: Set<string> }>();
38  for (const list of lists) {
39    list.docIds.forEach((docId, index) => {
40      const rank = index + 1;
41      const contribution = 1 / (k + rank);
42      const existing = scores.get(docId);
43      if (existing) {
44        existing.score += contribution;
45        existing.sources.add(list.name);
46      } else {
47        scores.set(docId, { score: contribution, sources: new Set([list.name]) });
48      }
49    });
50  }
51  const fused: FusedResult[] = [];
52  for (const [docId, { score, sources }] of scores) {
53    fused.push({ docId, score, sources: [...sources].sort() });
54  }
55  fused.sort((a, b) => (b.score !== a.score ? b.score - a.score : a.docId < b.docId ? -1 : 1));
56  return limit > 0 ? fused.slice(0, limit) : fused;
57}
58 
59/**
60 * Convenience: fuse the two hybrid-search lists (BM25 rank list + kNN rank
61 * list) and return doc ids with their fused scores.
62 */
63export function fuseHybrid(bm25: readonly string[], vec: readonly string[], limit = 0): FusedResult[] {
64  return reciprocalRankFusion(
65    [
66      { name: "bm25", docIds: [...bm25] },
67      { name: "vec", docIds: [...vec] },
68    ],
69    RRF_K,
70    limit,
71  );
72}
73