llm_kernel/embedding/
vector_index.rs

1//! Vector index trait for compressed approximate nearest neighbor search.
2//!
3//! Defines the abstract interface that concrete implementations (e.g.,
4//! `llm-kernel-vector-index` with TurboQuant) must satisfy. This module has
5//! **zero external dependencies** — implementations live in separate crates.
6//!
7//! ```
8//! use llm_kernel::embedding::vector_index::SearchHit;
9//!
10//! let hit = SearchHit { id: 42, score: 0.95 };
11//! assert_eq!(hit.id, 42);
12//! ```
13
14use std::path::Path;
15
16use crate::error::Result;
17
18/// A single search hit from vector index lookup.
19///
20/// Sorts by **descending** score (highest similarity first). Ties are broken
21/// by ascending ID for deterministic ordering.
22#[derive(Debug, Clone, Copy, PartialEq)]
23pub struct SearchHit {
24    /// External identifier for the matched vector.
25    pub id: u64,
26    /// Similarity score (higher = more similar).
27    pub score: f32,
28}
29
30impl PartialOrd for SearchHit {
31    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
32        // f32 is not Ord, so we use total_cmp for a total ordering.
33        // Reverse score order: highest score first, then ascending ID.
34        Some(
35            other
36                .score
37                .total_cmp(&self.score)
38                .then_with(|| self.id.cmp(&other.id)),
39        )
40    }
41}
42
43/// Trait for compressed vector indexes.
44///
45/// Implementations provide approximate nearest neighbor search with
46/// quantization-based compression. Follows the same pattern as
47/// [`EmbeddingProvider`](crate::embedding::EmbeddingProvider).
48///
49/// The trait is defined here with zero dependencies. Concrete implementations
50/// live in separate crates (e.g., `llm-kernel-vector-index` with TurboQuant).
51///
52/// This trait is fully object-safe — `load` is intentionally not included
53/// because it requires `Self: Sized`. Concrete types provide their own
54/// `load` inherent methods instead.
55pub trait VectorIndex: Send + Sync {
56    /// Add vectors with auto-assigned sequential IDs.
57    fn add(&mut self, vectors: &[Vec<f32>]) -> Result<()>;
58
59    /// Add vectors with explicit external IDs.
60    fn add_with_ids(&mut self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()>;
61
62    /// Remove vectors by their external IDs.
63    ///
64    /// IDs that do not exist in the index are silently ignored.
65    /// Passing an empty slice is a no-op.
66    fn remove(&mut self, ids: &[u64]) -> Result<()>;
67
68    /// Search for the `k` nearest neighbors of `query`.
69    fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>>;
70
71    /// Search restricted to an allowlist of candidate IDs.
72    ///
73    /// Useful for hybrid retrieval: first narrow candidates via BM25 or
74    /// metadata filter, then dense-rerank within that set.
75    fn search_filtered(&self, query: &[f32], k: usize, allowlist: &[u64])
76    -> Result<Vec<SearchHit>>;
77
78    /// Number of vectors currently indexed.
79    fn len(&self) -> usize;
80
81    /// Whether the index is empty.
82    fn is_empty(&self) -> bool;
83
84    /// Vector dimensionality.
85    fn dim(&self) -> usize;
86
87    /// Persist the index to disk.
88    fn save(&self, path: &Path) -> Result<()>;
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn search_hit_fields() {
97        let hit = SearchHit {
98            id: 42,
99            score: 0.95,
100        };
101        assert_eq!(hit.id, 42);
102        assert!((hit.score - 0.95).abs() < f32::EPSILON);
103    }
104
105    #[test]
106    fn search_hit_copy() {
107        let hit = SearchHit { id: 1, score: 0.5 };
108        let copied = hit; // Copy semantics — no .clone() needed
109        assert_eq!(copied.id, hit.id);
110        assert_eq!(copied.score, hit.score);
111    }
112
113    #[test]
114    fn search_hit_sort_descending_by_score() {
115        let mut hits = [
116            SearchHit { id: 1, score: 0.3 },
117            SearchHit { id: 2, score: 0.9 },
118            SearchHit { id: 3, score: 0.5 },
119        ];
120        hits.sort_by(|a, b| a.partial_cmp(b).unwrap());
121        assert_eq!(hits[0].id, 2); // highest score first
122        assert_eq!(hits[1].id, 3);
123        assert_eq!(hits[2].id, 1);
124    }
125
126    #[test]
127    fn search_hit_tie_break_by_id() {
128        let mut hits = [
129            SearchHit { id: 30, score: 0.5 },
130            SearchHit { id: 10, score: 0.5 },
131            SearchHit { id: 20, score: 0.5 },
132        ];
133        hits.sort_by(|a, b| a.partial_cmp(b).unwrap());
134        assert_eq!(hits[0].id, 10);
135        assert_eq!(hits[1].id, 20);
136        assert_eq!(hits[2].id, 30);
137    }
138}
llm_kernel/embedding/vector_index.rs

llm_kernel/embedding/
vector_index.rs
