mentedb/
reranker.rs

1//! Pluggable reranking: an optional second pass that reorders recall candidates
2//! by a signal beyond the first-pass vector/BM25 score.
3//!
4//! Off by default. The engine ships a dependency-free [`LexicalReranker`], and the
5//! [`Reranker`] trait lets an embedder plug a stronger model (a cross-encoder, or
6//! an LLM judge) without the engine taking on that dependency. Apply it with
7//! [`crate::MenteDb::recall_reranked`].
8
9use std::collections::HashSet;
10
11use mentedb_core::types::MemoryId;
12
13/// One candidate handed to a reranker: its id, text, and first-pass score.
14pub struct RerankCandidate<'a> {
15    pub id: MemoryId,
16    pub content: &'a str,
17    pub score: f32,
18}
19
20/// Reorders recall candidates for a query by returning a new score per id (higher
21/// is more relevant). The caller sorts descending by the returned score.
22pub trait Reranker: Send + Sync {
23    fn rerank(&self, query: &str, candidates: &[RerankCandidate<'_>]) -> Vec<(MemoryId, f32)>;
24}
25
26/// A dependency-free lexical reranker: blends the first-pass score with the
27/// fraction of query terms that appear in the candidate. Cheap and deterministic,
28/// a reasonable default when no model is available, and it lifts exact-term matches
29/// that pure vector similarity can bury.
30pub struct LexicalReranker {
31    /// Weight on lexical overlap versus the first-pass score, clamped to `[0, 1]`.
32    overlap_weight: f32,
33}
34
35impl Default for LexicalReranker {
36    fn default() -> Self {
37        Self {
38            overlap_weight: 0.5,
39        }
40    }
41}
42
43impl LexicalReranker {
44    pub fn new(overlap_weight: f32) -> Self {
45        Self {
46            overlap_weight: overlap_weight.clamp(0.0, 1.0),
47        }
48    }
49}
50
51fn tokens(s: &str) -> HashSet<String> {
52    s.to_lowercase()
53        .split(|c: char| !c.is_alphanumeric())
54        .filter(|t| !t.is_empty())
55        .map(|t| t.to_string())
56        .collect()
57}
58
59impl Reranker for LexicalReranker {
60    fn rerank(&self, query: &str, candidates: &[RerankCandidate<'_>]) -> Vec<(MemoryId, f32)> {
61        let q = tokens(query);
62        candidates
63            .iter()
64            .map(|c| {
65                let ct = tokens(c.content);
66                let overlap = if q.is_empty() {
67                    0.0
68                } else {
69                    // Fraction of query terms present in the candidate.
70                    q.intersection(&ct).count() as f32 / q.len() as f32
71                };
72                let blended = (1.0 - self.overlap_weight) * c.score + self.overlap_weight * overlap;
73                (c.id, blended)
74            })
75            .collect()
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    fn id(n: u128) -> MemoryId {
84        uuid::Uuid::from_u128(n).into()
85    }
86
87    #[test]
88    fn lexical_reranker_lifts_exact_term_matches() {
89        // Candidate B has the lower first-pass score but contains both query terms;
90        // with enough overlap weight it should overtake A.
91        let a = id(1);
92        let b = id(2);
93        let cands = vec![
94            RerankCandidate {
95                id: a,
96                content: "a note about gardening tools",
97                score: 0.9,
98            },
99            RerankCandidate {
100                id: b,
101                content: "the espresso machine broke",
102                score: 0.6,
103            },
104        ];
105        let rr = LexicalReranker::new(0.8);
106        let mut scored = rr.rerank("espresso machine", &cands);
107        scored.sort_by(|x, y| y.1.total_cmp(&x.1));
108        assert_eq!(scored[0].0, b, "exact-term candidate should rank first");
109    }
110
111    #[test]
112    fn empty_query_leaves_first_pass_order() {
113        let a = id(1);
114        let b = id(2);
115        let cands = vec![
116            RerankCandidate {
117                id: a,
118                content: "alpha",
119                score: 0.8,
120            },
121            RerankCandidate {
122                id: b,
123                content: "beta",
124                score: 0.5,
125            },
126        ];
127        let rr = LexicalReranker::default();
128        let scored = rr.rerank("", &cands);
129        // With no query terms, overlap is 0 for all, so scores scale the first pass
130        // uniformly and preserve order.
131        assert!(scored[0].1 > scored[1].1);
132    }
133}
mentedb/reranker.rs

mentedb/
reranker.rs
