gemini_memory_rs/retrieval/
fusion.rs

1//! Reciprocal rank fusion across independent lexical queries.
2//!
3//! A plan issues several queries because they fail differently: an entity query
4//! finds everything about a person but ranks it arbitrarily; a topic query
5//! ranks well but misses records that never name the topic. Fusing on *rank*
6//! rather than score means one query's score scale cannot dominate another's.
7
8use std::collections::HashMap;
9
10use crate::bm25::SearchHit;
11use crate::core::MemoryId;
12
13/// The RRF smoothing constant. 60 is the value from the original TREC work and
14/// behaves well when result lists are short.
15pub const RRF_K: f32 = 60.0;
16
17/// A candidate after fusion.
18#[derive(Debug, Clone)]
19pub struct FusedCandidate {
20    /// The best-scoring hit for this record across all queries.
21    pub hit: SearchHit,
22    /// Fused rank score.
23    pub rrf_score: f32,
24    /// How many queries surfaced this record.
25    pub appearances: usize,
26    /// Best rank achieved in any single query (0-based).
27    pub best_rank: usize,
28}
29
30impl FusedCandidate {
31    /// The record's identity.
32    pub fn id(&self) -> &MemoryId {
33        &self.hit.id
34    }
35}
36
37/// Fuse per-query result lists into one ranking.
38pub fn reciprocal_rank_fusion(rankings: &[Vec<SearchHit>]) -> Vec<FusedCandidate> {
39    let mut fused: HashMap<MemoryId, FusedCandidate> = HashMap::new();
40
41    for ranking in rankings {
42        for (rank, hit) in ranking.iter().enumerate() {
43            let contribution = 1.0 / (RRF_K + rank as f32 + 1.0);
44            fused
45                .entry(hit.id.clone())
46                .and_modify(|existing| {
47                    existing.rrf_score += contribution;
48                    existing.appearances += 1;
49                    existing.best_rank = existing.best_rank.min(rank);
50                    if hit.score > existing.hit.score {
51                        existing.hit = hit.clone();
52                    }
53                })
54                .or_insert_with(|| FusedCandidate {
55                    hit: hit.clone(),
56                    rrf_score: contribution,
57                    appearances: 1,
58                    best_rank: rank,
59                });
60        }
61    }
62
63    let mut out: Vec<FusedCandidate> = fused.into_values().collect();
64    out.sort_by(|a, b| {
65        b.rrf_score
66            .partial_cmp(&a.rrf_score)
67            .unwrap_or(std::cmp::Ordering::Equal)
68            .then_with(|| {
69                b.hit
70                    .score
71                    .partial_cmp(&a.hit.score)
72                    .unwrap_or(std::cmp::Ordering::Equal)
73            })
74            .then_with(|| a.hit.id.as_str().cmp(b.hit.id.as_str()))
75    });
76    out
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use crate::bm25::{MemoryOrigin, SearchExplanation};
83    use crate::core::MemoryKind;
84
85    fn hit(id: &str, score: f32) -> SearchHit {
86        SearchHit {
87            id: MemoryId::new(id),
88            score,
89            statement: format!("statement for {id}"),
90            kind: MemoryKind::Preference,
91            origin: MemoryOrigin::Canonical,
92            explanation: SearchExplanation {
93                memory_id: MemoryId::new(id),
94                components: Vec::new(),
95                boosts: Vec::new(),
96                lexical_score: score,
97                final_score: score,
98            },
99        }
100    }
101
102    #[test]
103    fn a_record_found_by_several_queries_outranks_one_found_once() {
104        // `mem_b` tops one list but appears nowhere else; `mem_a` is second in
105        // both, and agreement across queries wins.
106        let fused = reciprocal_rank_fusion(&[
107            vec![hit("mem_b", 9.0), hit("mem_a", 5.0)],
108            vec![hit("mem_c", 8.0), hit("mem_a", 5.0)],
109        ]);
110        assert_eq!(fused[0].id().as_str(), "mem_a");
111        assert_eq!(fused[0].appearances, 2);
112    }
113
114    #[test]
115    fn fusion_keeps_the_best_scoring_variant_of_a_record() {
116        let fused = reciprocal_rank_fusion(&[vec![hit("mem_a", 2.0)], vec![hit("mem_a", 7.0)]]);
117        assert_eq!(fused.len(), 1);
118        assert_eq!(fused[0].hit.score, 7.0);
119        assert_eq!(fused[0].best_rank, 0);
120    }
121
122    #[test]
123    fn fusion_of_nothing_is_nothing() {
124        assert!(reciprocal_rank_fusion(&[]).is_empty());
125        assert!(reciprocal_rank_fusion(&[vec![], vec![]]).is_empty());
126    }
127
128    #[test]
129    fn ordering_is_deterministic_for_equal_scores() {
130        let first = reciprocal_rank_fusion(&[vec![hit("mem_b", 1.0), hit("mem_a", 1.0)]]);
131        let second = reciprocal_rank_fusion(&[vec![hit("mem_b", 1.0), hit("mem_a", 1.0)]]);
132        let ids: Vec<_> = first.iter().map(|c| c.id().to_string()).collect();
133        let ids2: Vec<_> = second.iter().map(|c| c.id().to_string()).collect();
134        assert_eq!(ids, ids2);
135    }
136}
gemini_memory_rs/retrieval/fusion.rs

gemini_memory_rs/retrieval/
fusion.rs
