semtree_rag/
hybrid.rs

1use semtree_store::Hit;
2
3use crate::{LexicalIndex, RagError, SearchEngine};
4
5// Reciprocal Rank Fusion constant. 60 is the value from the original
6// Cormack et al. paper and the de-facto standard.
7const RRF_K: f64 = 60.0;
8
9/// How a query is matched against the index.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum SearchMode {
12    /// Fuse semantic (vector) and lexical (BM25) rankings via RRF. Default.
13    Hybrid,
14    /// Vector similarity only.
15    Semantic,
16    /// BM25 keyword matching only.
17    Lexical,
18}
19
20impl SearchMode {
21    /// Parses a mode name; returns `None` for unknown values.
22    pub fn parse(s: &str) -> Option<Self> {
23        match s.to_lowercase().as_str() {
24            "hybrid" => Some(Self::Hybrid),
25            "semantic" | "vector" => Some(Self::Semantic),
26            "lexical" | "keyword" | "bm25" => Some(Self::Lexical),
27            _ => None,
28        }
29    }
30}
31
32/// Combines a vector [`SearchEngine`] with a [`LexicalIndex`] so queries can be
33/// matched by meaning, by keyword, or both (fused). Hybrid is what lets semtree
34/// catch concepts a grep misses *and* keep the exact-identifier precision a
35/// pure vector search loses.
36pub struct HybridSearcher {
37    engine: SearchEngine,
38    lexical: LexicalIndex,
39}
40
41impl HybridSearcher {
42    pub fn new(engine: SearchEngine, lexical: LexicalIndex) -> Self {
43        Self { engine, lexical }
44    }
45
46    /// Returns up to `limit` hits for `query` under the given `mode`.
47    pub async fn search(
48        &self,
49        query: &str,
50        limit: usize,
51        mode: SearchMode,
52    ) -> Result<Vec<Hit>, RagError> {
53        match mode {
54            SearchMode::Semantic => self.engine.search(query, limit).await,
55            SearchMode::Lexical => Ok(self
56                .lexical
57                .search(query, limit)
58                .into_iter()
59                .map(|(id, score)| Hit { id, score })
60                .collect()),
61            SearchMode::Hybrid => {
62                let vector = self.engine.search(query, limit).await?;
63                let lexical = self.lexical.search(query, limit);
64                Ok(fuse_rrf(&vector, &lexical, limit))
65            }
66        }
67    }
68}
69
70/// Reciprocal Rank Fusion: each ranker contributes `1 / (k + rank)` per
71/// document, summed across rankers. Rank-based, so it's robust to the fact
72/// that cosine scores and BM25 scores live on different scales.
73fn fuse_rrf(vector: &[Hit], lexical: &[(String, f32)], limit: usize) -> Vec<Hit> {
74    use std::collections::HashMap;
75
76    let mut fused: HashMap<String, f64> = HashMap::new();
77    for (rank, hit) in vector.iter().enumerate() {
78        *fused.entry(hit.id.clone()).or_insert(0.0) += 1.0 / (RRF_K + (rank + 1) as f64);
79    }
80    for (rank, (id, _)) in lexical.iter().enumerate() {
81        *fused.entry(id.clone()).or_insert(0.0) += 1.0 / (RRF_K + (rank + 1) as f64);
82    }
83
84    let mut ranked: Vec<(String, f64)> = fused.into_iter().collect();
85    ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
86    ranked.truncate(limit);
87    ranked
88        .into_iter()
89        .map(|(id, score)| Hit {
90            id,
91            score: score as f32,
92        })
93        .collect()
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    fn hit(id: &str) -> Hit {
101        Hit {
102            id: id.to_string(),
103            score: 0.0,
104        }
105    }
106
107    #[test]
108    fn rrf_rewards_agreement() {
109        // "b" is ranked highly by both rankers; it should win even though it is
110        // not #1 in either list.
111        let vector = vec![hit("a"), hit("b"), hit("c")];
112        let lexical = vec![("d".to_string(), 9.0), ("b".to_string(), 8.0)];
113        let fused = fuse_rrf(&vector, &lexical, 10);
114        assert_eq!(fused.first().map(|h| h.id.as_str()), Some("b"));
115    }
116
117    #[test]
118    fn rrf_keeps_unique_hits_from_each_ranker() {
119        let vector = vec![hit("a")];
120        let lexical = vec![("b".to_string(), 1.0)];
121        let fused = fuse_rrf(&vector, &lexical, 10);
122        let ids: Vec<&str> = fused.iter().map(|h| h.id.as_str()).collect();
123        assert!(ids.contains(&"a"));
124        assert!(ids.contains(&"b"));
125    }
126
127    #[test]
128    fn mode_parsing() {
129        assert_eq!(SearchMode::parse("hybrid"), Some(SearchMode::Hybrid));
130        assert_eq!(SearchMode::parse("SEMANTIC"), Some(SearchMode::Semantic));
131        assert_eq!(SearchMode::parse("bm25"), Some(SearchMode::Lexical));
132        assert_eq!(SearchMode::parse("nope"), None);
133    }
134}
semtree_rag/hybrid.rs

semtree_rag/
hybrid.rs
