1pub mod bm25;
36pub mod chunk;
37pub mod corpus;
38pub mod embed;
39pub mod fuse;
40pub mod hnsw;
41pub mod text;
42
43pub use bm25::Bm25Index;
44pub use chunk::{chunk_text, chunk_text_with_overlap};
45pub use embed::{embed, DEMO_DIM};
46pub use fuse::{rrf_fuse, rrf_fuse_scored, DEFAULT_RRF_K};
47pub use hnsw::{cosine_distance, HnswIndex, HnswParams};
48
49pub struct HybridIndex {
53 lexical: Bm25Index,
54 vector: HnswIndex,
55 rrf_k: f64,
56}
57
58impl HybridIndex {
59 pub fn new(dimensions: usize) -> Self {
61 Self::with_params(dimensions, HnswParams::default(), DEFAULT_RRF_K)
62 }
63
64 pub fn with_params(dimensions: usize, hnsw: HnswParams, rrf_k: f64) -> Self {
65 Self {
66 lexical: Bm25Index::new(),
67 vector: HnswIndex::with_params(dimensions, hnsw),
68 rrf_k,
69 }
70 }
71
72 pub fn add(&mut self, id: &str, text: &str, embedding: &[f32]) {
75 self.lexical.add_document(id, text);
76 self.vector.insert(id, embedding);
77 }
78
79 pub fn remove(&mut self, id: &str) -> bool {
82 self.lexical.remove_document(id);
83 self.vector.remove(id)
84 }
85
86 pub fn len(&self) -> usize {
88 self.vector.len()
89 }
90
91 pub fn is_empty(&self) -> bool {
92 self.vector.is_empty()
93 }
94
95 pub fn search(
102 &self,
103 query_text: &str,
104 query_embedding: &[f32],
105 k: usize,
106 ) -> Vec<(String, f64)> {
107 if k == 0 {
108 return Vec::new();
109 }
110 let fetch = k.saturating_mul(3);
111 let lexical = self.lexical.search(query_text, fetch);
112 let semantic = self.vector.search(query_embedding, fetch);
113
114 let mut fused = rrf_fuse_scored(&lexical, &semantic, self.rrf_k);
115 fused.truncate(k);
116 fused
117 }
118
119 pub fn search_lexical(&self, query: &str, k: usize) -> Vec<(String, f64)> {
121 self.lexical.search(query, k)
122 }
123
124 pub fn search_semantic(&self, query: &[f32], k: usize) -> Vec<(String, f32)> {
126 self.vector.search(query, k)
127 }
128}
129
130#[cfg(test)]
131mod tests {
132 use super::*;
133
134 #[test]
135 fn hybrid_prefers_doc_matching_both_signals() {
136 let mut index = HybridIndex::new(4);
137 index.add(
139 "doc-a",
140 "rust async runtime internals",
141 &[1.0, 0.0, 0.0, 0.0],
142 );
143 index.add("doc-b", "rust cookbook recipes", &[0.0, 1.0, 0.0, 0.0]);
145 index.add("doc-c", "tokio scheduler design", &[0.9, 0.1, 0.0, 0.0]);
147
148 let results = index.search("rust runtime", &[1.0, 0.0, 0.0, 0.0], 3);
149 assert_eq!(results[0].0, "doc-a");
150 }
151
152 #[test]
153 fn remove_drops_from_both_legs() {
154 let mut index = HybridIndex::new(2);
155 index.add("x", "unique pelican words", &[1.0, 0.0]);
156 index.add("y", "other content", &[0.0, 1.0]);
157
158 assert!(index.remove("x"));
159 let results = index.search("pelican", &[1.0, 0.0], 5);
160 assert!(results.iter().all(|(id, _)| id != "x"));
161 assert_eq!(index.len(), 1);
162 }
163
164 #[test]
165 fn k_zero_returns_empty() {
166 let mut index = HybridIndex::new(2);
167 index.add("x", "abc", &[1.0, 0.0]);
168 assert!(index.search("abc", &[1.0, 0.0], 0).is_empty());
169 }
170}