llm_kernel/embedding/
async_vector_index.rs

1//! Async vector index trait for remote/shared backends.
2//!
3//! The existing [`VectorIndex`](crate::embedding::VectorIndex) is a synchronous,
4//! in-process trait (`&mut self`, blocking `search`). That fits compressed
5//! in-memory indexes (TurboQuant) but not remote vector services such as
6//! Qdrant or Elasticsearch, whose clients are **async-only** and naturally
7//! shared (`&self`) rather than exclusively borrowed.
8//!
9//! [`AsyncVectorIndex`] is the async, object-safe counterpart for those
10//! backends. It mirrors the useful subset of [`VectorIndex`] — add, remove,
11//! search, filtered search, length, dimensionality — and omits `save` because
12//! remote backends persist server-side (just as [`VectorIndex`] omits `load`
13//! to stay object-safe). Concrete implementations live in this crate behind
14//! feature flags: the `qdrant` feature (`src/embedding/qdrant.rs`) and the
15//! `elastic` feature (`src/embedding/elastic.rs`).
16//!
17//! The trait has no concrete dependencies beyond `async_trait`. It is defined
18//! behind the `embedding` feature so the shared contract stays in the kernel
19//! while the heavy client crates remain opt-in.
20
21use crate::embedding::vector_index::SearchHit;
22use crate::error::Result;
23
24/// Async, object-safe vector index for remote/shared backends.
25///
26/// Implementations are remote vector services (Qdrant, Elasticsearch, …) whose
27/// clients are async and shareable. Use `dyn AsyncVectorIndex` to abstract over
28/// concrete backends.
29///
30/// IDs are always explicit — remote indexes do not auto-assign sequential IDs
31/// the way an in-memory index can, so callers supply the `u64` external IDs.
32#[async_trait::async_trait]
33pub trait AsyncVectorIndex: Send + Sync {
34    /// Upsert vectors keyed by their explicit external IDs.
35    ///
36    /// Re-upserting an existing ID replaces its vector. `vectors.len()` must
37    /// equal `ids.len()`.
38    async fn add(&self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()>;
39
40    /// Remove vectors by their external IDs.
41    ///
42    /// IDs that do not exist are silently ignored. An empty slice is a no-op.
43    async fn remove(&self, ids: &[u64]) -> Result<()>;
44
45    /// Search for the `k` nearest neighbors of `query`.
46    async fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>>;
47
48    /// Search restricted to an allowlist of candidate IDs.
49    ///
50    /// An **empty** allowlist yields no candidates, so an empty `Vec` is
51    /// returned (it does not fall back to an unfiltered search).
52    ///
53    /// Mirrors [`VectorIndex::search_filtered`](crate::embedding::VectorIndex::search_filtered):
54    /// narrow candidates (e.g. by metadata or BM25), then dense-rerank.
55    async fn search_filtered(
56        &self,
57        query: &[f32],
58        k: usize,
59        allowlist: &[u64],
60    ) -> Result<Vec<SearchHit>>;
61
62    /// Number of vectors currently indexed.
63    async fn len(&self) -> Result<usize>;
64
65    /// Whether the index is empty.
66    async fn is_empty(&self) -> Result<bool> {
67        Ok(self.len().await? == 0)
68    }
69
70    /// Vector dimensionality.
71    fn dim(&self) -> usize;
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    /// The trait is object-safe: a blanket stub demonstrates `dyn
79    /// AsyncVectorIndex` compiles. The real backends live behind feature flags
80    /// (e.g. the `qdrant` feature at `src/embedding/qdrant.rs`).
81    struct StubIndex {
82        d: usize,
83    }
84
85    #[async_trait::async_trait]
86    impl AsyncVectorIndex for StubIndex {
87        async fn add(&self, _vectors: &[Vec<f32>], _ids: &[u64]) -> Result<()> {
88            Ok(())
89        }
90        async fn remove(&self, _ids: &[u64]) -> Result<()> {
91            Ok(())
92        }
93        async fn search(&self, _query: &[f32], _k: usize) -> Result<Vec<SearchHit>> {
94            Ok(Vec::new())
95        }
96        async fn search_filtered(
97            &self,
98            _query: &[f32],
99            _k: usize,
100            _allowlist: &[u64],
101        ) -> Result<Vec<SearchHit>> {
102            Ok(Vec::new())
103        }
104        async fn len(&self) -> Result<usize> {
105            Ok(0)
106        }
107        fn dim(&self) -> usize {
108            self.d
109        }
110    }
111
112    /// AC2: `dyn AsyncVectorIndex` is usable (object-safety) and the default
113    /// `is_empty` method composes over `len`.
114    #[tokio::test]
115    async fn dyn_async_vector_index_object_safe() {
116        let idx: Box<dyn AsyncVectorIndex> = Box::new(StubIndex { d: 4 });
117        assert_eq!(idx.dim(), 4);
118        idx.add(&[vec![0.0; 4]], &[1]).await.unwrap();
119        assert!(idx.is_empty().await.unwrap());
120    }
121}
llm_kernel/embedding/async_vector_index.rs

llm_kernel/embedding/
async_vector_index.rs
