ripvec_core/
hybrid.rs

1//! Hybrid semantic + keyword search with Reciprocal Rank Fusion (RRF).
2//!
3//! [`HybridIndex`] wraps a [`SearchIndex`] (dense vector search) and a
4//! [`Bm25Index`] (BM25 keyword search) and fuses their ranked results via
5//! Reciprocal Rank Fusion so that chunks appearing high in either list
6//! bubble to the top of the combined ranking.
7
8use std::collections::HashMap;
9use std::fmt;
10use std::str::FromStr;
11
12use crate::bm25::Bm25Index;
13use crate::chunk::CodeChunk;
14use crate::index::SearchIndex;
15
16/// Controls which retrieval strategy is used during search.
17#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
18pub enum SearchMode {
19    /// Fuse semantic (vector) and keyword (BM25) results via RRF.
20    #[default]
21    Hybrid,
22    /// Dense vector cosine-similarity ranking only.
23    Semantic,
24    /// BM25 keyword ranking only.
25    Keyword,
26}
27
28impl fmt::Display for SearchMode {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        match self {
31            Self::Hybrid => f.write_str("hybrid"),
32            Self::Semantic => f.write_str("semantic"),
33            Self::Keyword => f.write_str("keyword"),
34        }
35    }
36}
37
38/// Error returned when a `SearchMode` string cannot be parsed.
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub struct ParseSearchModeError(String);
41
42impl fmt::Display for ParseSearchModeError {
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        write!(
45            f,
46            "unknown search mode {:?}; expected hybrid, semantic, or keyword",
47            self.0
48        )
49    }
50}
51
52impl std::error::Error for ParseSearchModeError {}
53
54impl FromStr for SearchMode {
55    type Err = ParseSearchModeError;
56
57    fn from_str(s: &str) -> Result<Self, Self::Err> {
58        match s {
59            "hybrid" => Ok(Self::Hybrid),
60            "semantic" => Ok(Self::Semantic),
61            "keyword" => Ok(Self::Keyword),
62            other => Err(ParseSearchModeError(other.to_string())),
63        }
64    }
65}
66
67/// Combined semantic + keyword search index with RRF fusion.
68///
69/// Build once from chunks and pre-computed embeddings; query repeatedly
70/// via [`search`](Self::search).
71pub struct HybridIndex {
72    /// Semantic (dense vector) search index.
73    pub semantic: SearchIndex,
74    /// BM25 keyword search index.
75    bm25: Bm25Index,
76}
77
78impl HybridIndex {
79    /// Build a `HybridIndex` from raw chunks and their pre-computed embeddings.
80    ///
81    /// Constructs both the [`SearchIndex`] and [`Bm25Index`] in one call.
82    /// `cascade_dim` is forwarded to [`SearchIndex::new`] for optional MRL
83    /// cascade pre-filtering.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error if the BM25 index cannot be built (e.g., tantivy
88    /// schema or writer failure).
89    pub fn new(
90        chunks: Vec<CodeChunk>,
91        embeddings: &[Vec<f32>],
92        cascade_dim: Option<usize>,
93    ) -> crate::Result<Self> {
94        let bm25 = Bm25Index::build(&chunks)?;
95        let semantic = SearchIndex::new(chunks, embeddings, cascade_dim);
96        Ok(Self { semantic, bm25 })
97    }
98
99    /// Assemble a `HybridIndex` from pre-built components.
100    ///
101    /// Useful when the caller has already constructed the sub-indices
102    /// separately (e.g., loaded from a cache).
103    #[must_use]
104    pub fn from_parts(semantic: SearchIndex, bm25: Bm25Index) -> Self {
105        Self { semantic, bm25 }
106    }
107
108    /// Search the index and return `(chunk_index, score)` pairs.
109    ///
110    /// Dispatches based on `mode`:
111    /// - [`SearchMode::Semantic`] — pure dense vector search via
112    ///   [`SearchIndex::rank`].
113    /// - [`SearchMode::Keyword`] — pure BM25 keyword search, truncated to
114    ///   `top_k`.
115    /// - [`SearchMode::Hybrid`] — retrieves both ranked lists, fuses them
116    ///   with [`rrf_fuse`], then truncates to `top_k`.
117    ///
118    /// Scores are min-max normalized to `[0, 1]` regardless of mode, so
119    /// a threshold of 0.5 always means "above midpoint of the score range"
120    /// whether the underlying scores are cosine similarity, BM25, or RRF.
121    #[must_use]
122    pub fn search(
123        &self,
124        query_embedding: &[f32],
125        query_text: &str,
126        top_k: usize,
127        threshold: f32,
128        mode: SearchMode,
129    ) -> Vec<(usize, f32)> {
130        let mut raw = match mode {
131            SearchMode::Semantic => {
132                // Fetch more than top_k so normalization has a meaningful range.
133                self.semantic
134                    .rank_turboquant(query_embedding, top_k.max(100), 0.0)
135            }
136            SearchMode::Keyword => self.bm25.search(query_text, top_k.max(100)),
137            SearchMode::Hybrid => {
138                let sem = self
139                    .semantic
140                    .rank_turboquant(query_embedding, top_k.max(100), 0.0);
141                let kw = self.bm25.search(query_text, top_k.max(100));
142                rrf_fuse(&sem, &kw, 60.0)
143            }
144        };
145
146        // Min-max normalize scores to [0, 1] so threshold is model-agnostic.
147        if let (Some(max), Some(min)) = (raw.first().map(|(_, s)| *s), raw.last().map(|(_, s)| *s))
148        {
149            let range = max - min;
150            if range > f32::EPSILON {
151                for (_, score) in &mut raw {
152                    *score = (*score - min) / range;
153                }
154            } else {
155                // All scores identical — normalize to 1.0
156                for (_, score) in &mut raw {
157                    *score = 1.0;
158                }
159            }
160        }
161
162        // Apply threshold on normalized scores, then truncate
163        raw.retain(|(_, score)| *score >= threshold);
164        raw.truncate(top_k);
165        raw
166    }
167
168    /// All chunks in the index.
169    #[must_use]
170    pub fn chunks(&self) -> &[CodeChunk] {
171        &self.semantic.chunks
172    }
173}
174
175/// Reciprocal Rank Fusion of two ranked lists.
176///
177/// Each entry in `semantic` and `bm25` is `(chunk_index, _score)`.
178/// The fused score for a chunk is the sum of `1 / (k + rank + 1)` across
179/// every list the chunk appears in, where `rank` is 0-based.
180///
181/// Returns all chunks that appear in either list, sorted descending by
182/// fused RRF score.
183///
184/// `k` should typically be 60.0 — a conventional constant that smooths the
185/// ranking boost for the very top results.
186#[must_use]
187pub fn rrf_fuse(semantic: &[(usize, f32)], bm25: &[(usize, f32)], k: f32) -> Vec<(usize, f32)> {
188    let mut scores: HashMap<usize, f32> = HashMap::new();
189
190    for (rank, &(idx, _)) in semantic.iter().enumerate() {
191        *scores.entry(idx).or_insert(0.0) += 1.0 / (k + rank as f32 + 1.0);
192    }
193    for (rank, &(idx, _)) in bm25.iter().enumerate() {
194        *scores.entry(idx).or_insert(0.0) += 1.0 / (k + rank as f32 + 1.0);
195    }
196
197    let mut results: Vec<(usize, f32)> = scores.into_iter().collect();
198    results.sort_unstable_by(|a, b| {
199        b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)) // stable tie-break by chunk index
200    });
201    results
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn rrf_union_semantics() {
210        // sem: [0, 1, 2], bm25: [3, 0, 4]
211        // Chunk 0 appears in both lists → highest RRF score.
212        // Chunks 1, 2, 3, 4 appear in exactly one list → all five appear.
213        let sem = vec![(0, 0.9), (1, 0.8), (2, 0.7)];
214        let bm25 = vec![(3, 10.0), (0, 8.0), (4, 6.0)];
215
216        let fused = rrf_fuse(&sem, &bm25, 60.0);
217
218        let indices: Vec<usize> = fused.iter().map(|&(i, _)| i).collect();
219
220        // All 5 unique chunks must appear
221        for expected in [0, 1, 2, 3, 4] {
222            assert!(
223                indices.contains(&expected),
224                "chunk {expected} missing from fused results"
225            );
226        }
227        assert_eq!(fused.len(), 5);
228
229        // Chunk 0 must rank first (double-list bonus)
230        assert_eq!(indices[0], 0, "chunk 0 should rank first");
231    }
232
233    #[test]
234    fn rrf_single_list() {
235        // Only semantic results; BM25 is empty.
236        let sem = vec![(0, 0.9), (1, 0.8)];
237        let bm25: Vec<(usize, f32)> = vec![];
238
239        let fused = rrf_fuse(&sem, &bm25, 60.0);
240
241        assert_eq!(fused.len(), 2);
242        // Chunk 0 ranked first in sem list → higher RRF score than chunk 1
243        assert_eq!(fused[0].0, 0);
244        assert_eq!(fused[1].0, 1);
245        assert!(fused[0].1 > fused[1].1);
246    }
247
248    #[test]
249    fn search_mode_roundtrip() {
250        assert_eq!("hybrid".parse::<SearchMode>().unwrap(), SearchMode::Hybrid);
251        assert_eq!(
252            "semantic".parse::<SearchMode>().unwrap(),
253            SearchMode::Semantic
254        );
255        assert_eq!(
256            "keyword".parse::<SearchMode>().unwrap(),
257            SearchMode::Keyword
258        );
259
260        let err = "invalid".parse::<SearchMode>();
261        assert!(err.is_err(), "expected parse error for 'invalid'");
262        let msg = err.unwrap_err().to_string();
263        assert!(
264            msg.contains("invalid"),
265            "error message should echo the bad input"
266        );
267    }
268
269    #[test]
270    fn search_mode_display() {
271        assert_eq!(SearchMode::Hybrid.to_string(), "hybrid");
272        assert_eq!(SearchMode::Semantic.to_string(), "semantic");
273        assert_eq!(SearchMode::Keyword.to_string(), "keyword");
274    }
275}
ripvec_core/hybrid.rs

ripvec_core/
hybrid.rs
