llm_kernel/embedding/
qdrant.rs

1//! Qdrant `AsyncVectorIndex` (`qdrant` feature).
2//!
3//! `QdrantVectorIndex` implements `AsyncVectorIndex` over the official
4//! `qdrant-client`. It is the async counterpart to the in-memory `VectorIndex`
5//! — remote vector services are async-only and naturally shared, so they
6//! cannot implement the synchronous `VectorIndex`.
7
8use crate::error::{KernelError, Result};
9use qdrant_client::qdrant::point_id::PointIdOptions;
10use qdrant_client::qdrant::{
11    Condition, CountPointsBuilder, CreateCollectionBuilder, DeletePointsBuilder, Distance, Filter,
12    PointStruct, PointsIdsList, QueryPointsBuilder, ScoredPoint, UpsertPointsBuilder,
13    VectorParamsBuilder,
14};
15use qdrant_client::{Payload, Qdrant};
16
17use super::{AsyncVectorIndex, SearchHit};
18
19/// Async vector index backed by a Qdrant collection.
20///
21/// The collection is created on construction (Cosine distance) if it does not
22/// already exist. All operations are async over the official `qdrant-client`.
23pub struct QdrantVectorIndex {
24    client: Qdrant,
25    collection: String,
26    dim: usize,
27}
28
29impl QdrantVectorIndex {
30    /// Connect to `url` (e.g. `http://localhost:6334`) and ensure `collection`
31    /// exists with a Cosine-distance vector config of `dim` dimensions.
32    pub async fn new(url: &str, collection: &str, dim: usize) -> Result<Self> {
33        let client = Qdrant::from_url(url)
34            .build()
35            .map_err(KernelError::embedding)?;
36        let idx = Self {
37            client,
38            collection: collection.to_string(),
39            dim,
40        };
41        idx.ensure_collection().await?;
42        Ok(idx)
43    }
44
45    /// Create the collection if it does not already exist.
46    async fn ensure_collection(&self) -> Result<()> {
47        if !self
48            .client
49            .collection_exists(&self.collection)
50            .await
51            .map_err(KernelError::embedding)?
52        {
53            self.client
54                .create_collection(
55                    CreateCollectionBuilder::new(&self.collection).vectors_config(
56                        VectorParamsBuilder::new(self.dim as u64, Distance::Cosine),
57                    ),
58                )
59                .await
60                .map_err(KernelError::embedding)?;
61        }
62        Ok(())
63    }
64
65    /// Drop the collection (useful for test cleanup or full reset).
66    pub async fn delete_collection(&self) -> Result<()> {
67        self.client
68            .delete_collection(&self.collection)
69            .await
70            .map_err(KernelError::embedding)?;
71        Ok(())
72    }
73
74    /// Extract a numeric `u64` id from a Qdrant `PointId`, returning `None` for
75    /// UUID (non-numeric) ids. Pure — unit-testable offline.
76    fn extract_numeric_id(pid: &qdrant_client::qdrant::PointId) -> Option<u64> {
77        match &pid.point_id_options {
78            Some(PointIdOptions::Num(n)) => Some(*n),
79            _ => None,
80        }
81    }
82
83    /// Extract a `u64` `SearchHit` from a Qdrant `ScoredPoint`.
84    ///
85    /// Points with non-numeric IDs (UUIDs) are dropped — this index keys on
86    /// `u64` external IDs, matching `super::VectorIndex`.
87    fn scored_to_hit(point: &ScoredPoint) -> Option<SearchHit> {
88        let id = point.id.as_ref().and_then(Self::extract_numeric_id)?;
89        Some(SearchHit {
90            id,
91            score: point.score,
92        })
93    }
94}
95
96#[async_trait::async_trait]
97impl AsyncVectorIndex for QdrantVectorIndex {
98    async fn add(&self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
99        if vectors.len() != ids.len() {
100            return Err(KernelError::Embedding(format!(
101                "vectors.len() ({}) must equal ids.len() ({})",
102                vectors.len(),
103                ids.len()
104            )));
105        }
106        if vectors.is_empty() {
107            return Ok(());
108        }
109        let payload = Payload::try_from(serde_json::json!({}))
110            .map_err(|e| KernelError::Embedding(format!("invalid empty payload: {e}")))?;
111        let points: Vec<PointStruct> = vectors
112            .iter()
113            .zip(ids.iter())
114            .map(|(v, &id)| PointStruct::new(id, v.clone(), payload.clone()))
115            .collect();
116        self.client
117            .upsert_points(UpsertPointsBuilder::new(&self.collection, points).wait(true))
118            .await
119            .map_err(KernelError::embedding)?;
120        Ok(())
121    }
122
123    async fn remove(&self, ids: &[u64]) -> Result<()> {
124        if ids.is_empty() {
125            return Ok(());
126        }
127        let id_list = PointsIdsList {
128            ids: ids.iter().map(|&id| id.into()).collect(),
129        };
130        self.client
131            .delete_points(
132                DeletePointsBuilder::new(&self.collection)
133                    .points(id_list)
134                    .wait(true),
135            )
136            .await
137            .map_err(KernelError::embedding)?;
138        Ok(())
139    }
140
141    async fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>> {
142        let res = self
143            .client
144            .query(
145                QueryPointsBuilder::new(&self.collection)
146                    .query(query.to_vec())
147                    .limit(k as u64)
148                    .with_payload(false),
149            )
150            .await
151            .map_err(KernelError::embedding)?;
152        Ok(res.result.iter().filter_map(Self::scored_to_hit).collect())
153    }
154
155    async fn search_filtered(
156        &self,
157        query: &[f32],
158        k: usize,
159        allowlist: &[u64],
160    ) -> Result<Vec<SearchHit>> {
161        // An empty allowlist excludes every point (no candidates) → empty.
162        if allowlist.is_empty() {
163            return Ok(vec![]);
164        }
165        let filter = Filter::must([Condition::has_id(allowlist.iter().copied())]);
166        let res = self
167            .client
168            .query(
169                QueryPointsBuilder::new(&self.collection)
170                    .query(query.to_vec())
171                    .limit(k as u64)
172                    .with_payload(false)
173                    .filter(filter),
174            )
175            .await
176            .map_err(KernelError::embedding)?;
177        Ok(res.result.iter().filter_map(Self::scored_to_hit).collect())
178    }
179
180    async fn len(&self) -> Result<usize> {
181        let res = self
182            .client
183            .count(CountPointsBuilder::new(&self.collection).exact(true))
184            .await
185            .map_err(KernelError::embedding)?;
186        Ok(res.result.map(|c| c.count as usize).unwrap_or(0))
187    }
188
189    fn dim(&self) -> usize {
190        self.dim
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use crate::embedding::AsyncVectorIndex;
198
199    const DIM: usize = 4;
200
201    fn unique_collection() -> String {
202        format!("llm_kernel_test_{}", std::process::id())
203    }
204
205    /// Offline (no server): numeric `PointId` extracts; UUIDs/empty are dropped.
206    #[test]
207    fn extract_numeric_id_handles_num_and_uuid() {
208        use qdrant_client::qdrant::{PointId, point_id::PointIdOptions};
209        let num = PointId {
210            point_id_options: Some(PointIdOptions::Num(42)),
211        };
212        assert_eq!(QdrantVectorIndex::extract_numeric_id(&num), Some(42));
213        let uuid = PointId {
214            point_id_options: Some(PointIdOptions::Uuid("x".into())),
215        };
216        assert_eq!(QdrantVectorIndex::extract_numeric_id(&uuid), None);
217        let none = PointId {
218            point_id_options: None,
219        };
220        assert_eq!(QdrantVectorIndex::extract_numeric_id(&none), None);
221    }
222
223    /// Conformance body returning `Result` so failures are errors (not panics),
224    /// letting the caller clean up the throwaway collection in every case.
225    async fn run_live_conformance(idx: &QdrantVectorIndex) -> Result<()> {
226        if idx.dim() != DIM {
227            return Err(KernelError::Embedding("dim mismatch".into()));
228        }
229        if !idx.is_empty().await? {
230            return Err(KernelError::Embedding("not empty at start".into()));
231        }
232        idx.add(
233            &[vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
234            &[1, 2],
235        )
236        .await?;
237        if idx.len().await? != 2 {
238            return Err(KernelError::Embedding("len != 2 after add".into()));
239        }
240
241        let hits = idx.search(&[1.0, 0.0, 0.0, 0.0], 1).await?;
242        if hits.len() != 1 || hits[0].id != 1 {
243            return Err(KernelError::Embedding("nearest neighbor != id 1".into()));
244        }
245
246        let filtered = idx.search_filtered(&[1.0, 0.0, 0.0, 0.0], 2, &[2]).await?;
247        if filtered.len() != 1 || filtered[0].id != 2 {
248            return Err(KernelError::Embedding("filtered search != id 2".into()));
249        }
250
251        idx.add(&[vec![0.9, 0.1, 0.0, 0.0]], &[1]).await?;
252        if idx.len().await? != 2 {
253            return Err(KernelError::Embedding("len != 2 after re-add".into()));
254        }
255
256        idx.remove(&[1]).await?;
257        if idx.len().await? != 1 {
258            return Err(KernelError::Embedding("len != 1 after remove".into()));
259        }
260        let after = idx.search(&[1.0, 0.0, 0.0, 0.0], 5).await?;
261        if after.iter().any(|h| h.id == 1) {
262            return Err(KernelError::Embedding(
263                "id 1 still present after remove".into(),
264            ));
265        }
266        Ok(())
267    }
268
269    /// Live Qdrant conformance (skips without `LLMKERNEL_QDRANT_URL`). The
270    /// throwaway collection is deleted on EVERY exit path (pass or fail) so a
271    /// mid-test failure cannot leak it.
272    #[tokio::test]
273    async fn live_qdrant_conformance() {
274        let url = match std::env::var("LLMKERNEL_QDRANT_URL") {
275            Ok(u) => u,
276            Err(_) => {
277                eprintln!("skipped: LLMKERNEL_QDRANT_URL unset (no live Qdrant)");
278                return;
279            }
280        };
281
282        let coll = unique_collection();
283        let idx = match QdrantVectorIndex::new(&url, &coll, DIM).await {
284            Ok(i) => i,
285            Err(e) => panic!("connect + create collection: {e:?}"),
286        };
287        // Run the body, then ALWAYS drop the throwaway collection before
288        // propagating any failure — panic-safe cleanup.
289        let result = run_live_conformance(&idx).await;
290        let _ = idx.delete_collection().await;
291        result.expect("qdrant conformance failed");
292    }
293}
llm_kernel/embedding/qdrant.rs

llm_kernel/embedding/
qdrant.rs
