llm_kernel/embedding/
turbovec.rs

1//! TurboQuant-backed vector index implementation.
2
3use std::path::Path;
4
5use super::vector_index::{SearchHit, VectorIndex};
6use crate::error::{KernelError, Result};
7
8/// Compressed vector index backed by TurboQuant.
9///
10/// Wraps `turbovec::IdMapIndex` with dimension validation and a consistent
11/// error-handling layer. Supports online ingest (no training step),
12/// filtered search with allowlists, and persistence via `save`/`load`.
13pub struct TurbovecIndex {
14    inner: turbovec::IdMapIndex,
15    dim: usize,
16    bit_width: u8,
17}
18
19impl std::fmt::Debug for TurbovecIndex {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        f.debug_struct("TurbovecIndex")
22            .field("dim", &self.dim)
23            .field("bit_width", &self.bit_width)
24            .field("len", &self.inner.len())
25            .finish()
26    }
27}
28
29impl TurbovecIndex {
30    /// Create a new index for vectors of the given dimension.
31    ///
32    /// `bit_width` must be 2 or 4, controlling the quantization level:
33    /// - **2-bit**: 16x compression, lower recall at low k
34    /// - **4-bit**: 8x compression, higher recall (recommended default)
35    pub fn new(dim: usize, bit_width: u8) -> Result<Self> {
36        if bit_width != 2 && bit_width != 4 {
37            return Err(KernelError::Embedding(format!(
38                "bit_width must be 2 or 4, got {bit_width}"
39            )));
40        }
41        let inner = turbovec::IdMapIndex::new(dim, bit_width as usize)
42            .map_err(|e| KernelError::Embedding(format!("failed to create index: {e}")))?;
43        Ok(Self {
44            inner,
45            dim,
46            bit_width,
47        })
48    }
49
50    /// Quantization bit width (2 or 4).
51    pub fn bit_width(&self) -> u8 {
52        self.bit_width
53    }
54
55    /// Load a previously saved index from disk.
56    ///
57    /// This is an inherent method (not on the `VectorIndex` trait) so that the
58    /// trait remains fully object-safe. Callers must use the concrete type:
59    /// `TurbovecIndex::load(path)`.
60    pub fn load(path: &Path) -> Result<Self> {
61        let inner = turbovec::IdMapIndex::load(path)
62            .map_err(|e| KernelError::Embedding(format!("failed to load vector index: {e}")))?;
63        let meta_path = path.with_extension("meta.json");
64        let meta: IndexMeta = serde_json::from_str(&std::fs::read_to_string(&meta_path)?)
65            .map_err(KernelError::embedding)?;
66        if meta.bit_width != 2 && meta.bit_width != 4 {
67            return Err(KernelError::Embedding(format!(
68                "corrupted index meta: bit_width must be 2 or 4, got {}",
69                meta.bit_width
70            )));
71        }
72        if meta.dim == 0 {
73            return Err(KernelError::Embedding(
74                "corrupted index meta: dim must be positive, got 0".into(),
75            ));
76        }
77
78        // Cross-validate: loaded index vs sidecar metadata.
79        let inner_dim = inner.dim();
80        if inner_dim != 0 && inner_dim != meta.dim {
81            return Err(KernelError::Embedding(format!(
82                "index-meta mismatch: index dim={inner_dim}, meta dim={}",
83                meta.dim
84            )));
85        }
86        let inner_bw = inner.bit_width();
87        if inner_bw != meta.bit_width as usize {
88            return Err(KernelError::Embedding(format!(
89                "index-meta mismatch: index bit_width={inner_bw}, meta bit_width={}",
90                meta.bit_width
91            )));
92        }
93
94        Ok(Self {
95            inner,
96            dim: meta.dim,
97            bit_width: meta.bit_width,
98        })
99    }
100
101    fn validate_dim(&self, v: &[f32]) -> Result<()> {
102        if v.len() != self.dim {
103            return Err(KernelError::Embedding(format!(
104                "vector dimension mismatch: expected {}, got {}",
105                self.dim,
106                v.len()
107            )));
108        }
109        Ok(())
110    }
111
112    fn validate_dims(&self, vectors: &[Vec<f32>]) -> Result<()> {
113        for v in vectors {
114            self.validate_dim(v)?;
115        }
116        Ok(())
117    }
118}
119
120impl VectorIndex for TurbovecIndex {
121    fn add(&mut self, vectors: &[Vec<f32>]) -> Result<()> {
122        if vectors.is_empty() {
123            return Ok(());
124        }
125        self.validate_dims(vectors)?;
126        let start_id = self.inner.len() as u64;
127        let ids: Vec<u64> = (start_id..start_id + vectors.len() as u64).collect();
128        // Skip validation — already checked above.
129        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
130        self.inner
131            .add_with_ids_2d(&flat, self.dim, &ids)
132            .map_err(|e| KernelError::Embedding(format!("add failed: {e}")))?;
133        Ok(())
134    }
135
136    fn add_with_ids(&mut self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
137        if vectors.len() != ids.len() {
138            return Err(KernelError::Embedding(format!(
139                "vectors ({} entries) and ids ({} entries) must have the same length",
140                vectors.len(),
141                ids.len()
142            )));
143        }
144        self.validate_dims(vectors)?;
145        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
146        self.inner
147            .add_with_ids_2d(&flat, self.dim, ids)
148            .map_err(|e| KernelError::Embedding(format!("add failed: {e}")))?;
149        Ok(())
150    }
151
152    fn remove(&mut self, ids: &[u64]) -> Result<()> {
153        for &id in ids {
154            self.inner.remove(id);
155        }
156        Ok(())
157    }
158
159    fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>> {
160        self.validate_dim(query)?;
161        if self.inner.is_empty() {
162            return Ok(vec![]);
163        }
164        let (scores, ids) = self.inner.search(query, k);
165        Ok(scores
166            .into_iter()
167            .zip(ids)
168            .map(|(score, id)| SearchHit { id, score })
169            .collect())
170    }
171
172    fn search_filtered(
173        &self,
174        query: &[f32],
175        k: usize,
176        allowlist: &[u64],
177    ) -> Result<Vec<SearchHit>> {
178        self.validate_dim(query)?;
179        if self.inner.is_empty() || allowlist.is_empty() {
180            return Ok(vec![]);
181        }
182        let (scores, ids) = self.inner.search_with_allowlist(query, k, Some(allowlist));
183        Ok(scores
184            .into_iter()
185            .zip(ids)
186            .map(|(score, id)| SearchHit { id, score })
187            .collect())
188    }
189
190    fn len(&self) -> usize {
191        self.inner.len()
192    }
193
194    fn is_empty(&self) -> bool {
195        self.inner.is_empty()
196    }
197
198    fn dim(&self) -> usize {
199        self.dim
200    }
201
202    fn save(&self, path: &Path) -> Result<()> {
203        // Atomic save: write to temp files, fsync, then rename.
204        let tmp_index = path.with_extension("tvim.tmp");
205        let tmp_meta = path.with_extension("meta.tmp");
206
207        self.inner
208            .write(&tmp_index)
209            .map_err(|e| KernelError::Embedding(format!("failed to write vector index: {e}")))?;
210
211        let meta = IndexMeta {
212            dim: self.dim,
213            bit_width: self.bit_width,
214        };
215        let json = serde_json::to_string_pretty(&meta).map_err(KernelError::embedding)?;
216        std::fs::write(&tmp_meta, &json)?;
217
218        // Fsync temp files to ensure data is on disk.
219        if let Ok(f) = std::fs::File::open(&tmp_index) {
220            let _ = f.sync_all();
221        }
222        if let Ok(f) = std::fs::File::open(&tmp_meta) {
223            let _ = f.sync_all();
224        }
225
226        // Atomic rename — POSIX guarantees rename is atomic.
227        std::fs::rename(&tmp_meta, path.with_extension("meta.json"))?;
228        std::fs::rename(&tmp_index, path)?;
229
230        Ok(())
231    }
232}
233
234#[derive(serde::Serialize, serde::Deserialize)]
235struct IndexMeta {
236    dim: usize,
237    bit_width: u8,
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use tempfile::TempDir;
244
245    fn make_index(dim: usize, bit_width: u8) -> TurbovecIndex {
246        TurbovecIndex::new(dim, bit_width).unwrap()
247    }
248
249    fn random_vector(dim: usize, seed: f32) -> Vec<f32> {
250        (0..dim).map(|i| (seed + i as f32 * 0.001).sin()).collect()
251    }
252
253    #[test]
254    fn new_valid_bit_widths() {
255        assert!(TurbovecIndex::new(128, 2).is_ok());
256        assert!(TurbovecIndex::new(128, 4).is_ok());
257    }
258
259    #[test]
260    fn new_invalid_bit_width() {
261        assert!(TurbovecIndex::new(128, 3).is_err());
262        assert!(TurbovecIndex::new(128, 8).is_err());
263        assert!(TurbovecIndex::new(128, 1).is_err());
264    }
265
266    #[test]
267    fn add_and_len() {
268        let mut idx = make_index(64, 4);
269        assert!(idx.is_empty());
270        idx.add(&[random_vector(64, 1.0), random_vector(64, 2.0)])
271            .unwrap();
272        assert_eq!(idx.len(), 2);
273    }
274
275    #[test]
276    fn add_empty() {
277        let mut idx = make_index(64, 4);
278        idx.add(&[]).unwrap();
279        assert!(idx.is_empty());
280    }
281
282    #[test]
283    fn add_with_explicit_ids() {
284        let mut idx = make_index(64, 4);
285        idx.add_with_ids(&[random_vector(64, 1.0)], &[42u64])
286            .unwrap();
287        assert_eq!(idx.len(), 1);
288    }
289
290    #[test]
291    fn add_dimension_mismatch() {
292        let mut idx = make_index(64, 4);
293        let result = idx.add(&[vec![0.0; 32]]);
294        assert!(result.is_err());
295        assert!(
296            result
297                .unwrap_err()
298                .to_string()
299                .contains("dimension mismatch")
300        );
301    }
302
303    #[test]
304    fn add_with_ids_length_mismatch() {
305        let mut idx = make_index(64, 4);
306        let result = idx.add_with_ids(&[random_vector(64, 1.0), random_vector(64, 2.0)], &[1u64]);
307        assert!(result.is_err());
308        assert!(result.unwrap_err().to_string().contains("same length"));
309    }
310
311    #[test]
312    fn search_empty_index() {
313        let idx = make_index(64, 4);
314        let hits = idx.search(&random_vector(64, 1.0), 5).unwrap();
315        assert!(hits.is_empty());
316    }
317
318    #[test]
319    fn search_returns_nearest() {
320        let mut idx = make_index(64, 4);
321        let target = random_vector(64, 3.0);
322        idx.add_with_ids(
323            &[
324                random_vector(64, 100.0),
325                target.clone(),
326                random_vector(64, 200.0),
327            ],
328            &[0u64, 1u64, 2u64],
329        )
330        .unwrap();
331        let hits = idx.search(&target, 1).unwrap();
332        assert_eq!(hits.len(), 1);
333        assert_eq!(hits[0].id, 1);
334    }
335
336    #[test]
337    fn search_dimension_mismatch() {
338        let mut idx = make_index(64, 4);
339        idx.add(&[random_vector(64, 1.0)]).unwrap();
340        let result = idx.search(&[0.0; 32], 1);
341        assert!(result.is_err());
342    }
343
344    #[test]
345    fn search_filtered_with_allowlist() {
346        let mut idx = make_index(64, 4);
347        idx.add_with_ids(
348            &[
349                random_vector(64, 1.0),
350                random_vector(64, 2.0),
351                random_vector(64, 3.0),
352            ],
353            &[10u64, 20u64, 30u64],
354        )
355        .unwrap();
356        let hits = idx
357            .search_filtered(&random_vector(64, 1.0), 10, &[20u64, 30u64])
358            .unwrap();
359        let ids: Vec<u64> = hits.iter().map(|h| h.id).collect();
360        assert!(ids.contains(&20));
361        assert!(ids.contains(&30));
362        assert!(!ids.contains(&10));
363    }
364
365    #[test]
366    fn search_filtered_empty_allowlist() {
367        let mut idx = make_index(64, 4);
368        idx.add(&[random_vector(64, 1.0)]).unwrap();
369        let hits = idx
370            .search_filtered(&random_vector(64, 1.0), 5, &[])
371            .unwrap();
372        assert!(hits.is_empty());
373    }
374
375    #[test]
376    fn save_load_roundtrip() {
377        let dir = TempDir::new().unwrap();
378        let path = dir.path().join("test.tvim");
379        let mut idx = make_index(64, 4);
380        idx.add_with_ids(
381            &[random_vector(64, 1.0), random_vector(64, 2.0)],
382            &[100u64, 200u64],
383        )
384        .unwrap();
385        idx.save(&path).unwrap();
386        let loaded = TurbovecIndex::load(&path).unwrap();
387        assert_eq!(loaded.dim(), 64);
388        assert_eq!(loaded.bit_width(), 4);
389        assert_eq!(loaded.len(), 2);
390    }
391
392    #[test]
393    fn load_rejects_corrupted_meta() {
394        let dir = TempDir::new().unwrap();
395        let path = dir.path().join("corrupt.tvim");
396        let mut idx = make_index(64, 4);
397        idx.add(&[random_vector(64, 1.0)]).unwrap();
398        idx.save(&path).unwrap();
399        let meta_path = path.with_extension("meta.json");
400        std::fs::write(&meta_path, r#"{"dim": 64, "bit_width": 7}"#).unwrap();
401        let result = TurbovecIndex::load(&path);
402        assert!(result.is_err());
403        assert!(result.unwrap_err().to_string().contains("bit_width"));
404    }
405
406    #[test]
407    fn load_rejects_zero_dim() {
408        let dir = TempDir::new().unwrap();
409        let path = dir.path().join("zero.tvim");
410        let mut idx = make_index(64, 4);
411        idx.add(&[random_vector(64, 1.0)]).unwrap();
412        idx.save(&path).unwrap();
413        let meta_path = path.with_extension("meta.json");
414        std::fs::write(&meta_path, r#"{"dim": 0, "bit_width": 4}"#).unwrap();
415        let result = TurbovecIndex::load(&path);
416        assert!(result.is_err());
417        assert!(result.unwrap_err().to_string().contains("dim"));
418    }
419
420    #[test]
421    fn dim_and_bit_width_accessors() {
422        let idx = make_index(128, 2);
423        assert_eq!(idx.dim(), 128);
424        assert_eq!(idx.bit_width(), 2);
425    }
426
427    #[test]
428    fn trait_object_compatibility() {
429        let mut idx: Box<dyn VectorIndex> = Box::new(make_index(64, 4));
430        idx.add(&[random_vector(64, 1.0)]).unwrap();
431        assert_eq!(idx.len(), 1);
432        assert!(!idx.is_empty());
433    }
434
435    #[test]
436    fn remove_existing_id() {
437        let mut idx = make_index(64, 4);
438        idx.add_with_ids(
439            &[
440                random_vector(64, 1.0),
441                random_vector(64, 2.0),
442                random_vector(64, 3.0),
443            ],
444            &[10u64, 20u64, 30u64],
445        )
446        .unwrap();
447        assert_eq!(idx.len(), 3);
448        idx.remove(&[20u64]).unwrap();
449        assert_eq!(idx.len(), 2);
450        let hits = idx.search(&random_vector(64, 2.0), 10).unwrap();
451        let ids: Vec<u64> = hits.iter().map(|h| h.id).collect();
452        assert!(!ids.contains(&20));
453    }
454
455    #[test]
456    fn remove_nonexistent_id() {
457        let mut idx = make_index(64, 4);
458        idx.add_with_ids(&[random_vector(64, 1.0)], &[1u64])
459            .unwrap();
460        idx.remove(&[999u64]).unwrap();
461        assert_eq!(idx.len(), 1);
462    }
463
464    #[test]
465    fn remove_empty_ids() {
466        let mut idx = make_index(64, 4);
467        idx.add(&[random_vector(64, 1.0)]).unwrap();
468        idx.remove(&[]).unwrap();
469        assert_eq!(idx.len(), 1);
470    }
471
472    #[test]
473    fn remove_via_trait_object() {
474        let mut idx: Box<dyn VectorIndex> = Box::new(make_index(64, 4));
475        idx.add_with_ids(&[random_vector(64, 1.0)], &[42u64])
476            .unwrap();
477        idx.remove(&[42u64]).unwrap();
478        assert!(idx.is_empty());
479    }
480
481    #[test]
482    fn load_detects_dim_mismatch() {
483        let dir = TempDir::new().unwrap();
484        let path = dir.path().join("mismatch.tvim");
485        let mut idx = make_index(64, 4);
486        idx.add(&[random_vector(64, 1.0)]).unwrap();
487        idx.save(&path).unwrap();
488        let meta_path = path.with_extension("meta.json");
489        std::fs::write(&meta_path, r#"{"dim": 128, "bit_width": 4}"#).unwrap();
490        let result = TurbovecIndex::load(&path);
491        assert!(result.is_err());
492        let msg = result.unwrap_err().to_string();
493        assert!(
494            msg.contains("mismatch") || msg.contains("dim"),
495            "expected mismatch error, got: {msg}"
496        );
497    }
498}
llm_kernel/embedding/turbovec.rs

llm_kernel/embedding/
turbovec.rs
