lean_ctx/core/
qdrant_store.rs

1//! Optional Qdrant backend for dense (embedding) search.
2//!
3//! This module is behind the `qdrant` feature flag. It is intentionally
4//! dependency-light (uses `ureq`) and relies on lean-ctx's existing embedding
5//! pipeline for vector generation.
6
7use std::collections::HashSet;
8use std::path::Path;
9
10use md5::{Digest, Md5};
11use serde::{Deserialize, Serialize};
12
13use crate::core::bm25_index::{BM25Index, CodeChunk};
14
15#[derive(Debug, Clone)]
16pub struct QdrantConfig {
17    pub url: String,
18    pub api_key: Option<String>,
19    pub timeout_secs: u64,
20    pub collection_prefix: String,
21}
22
23impl QdrantConfig {
24    pub fn from_env() -> Result<Self, String> {
25        let url = std::env::var("LEANCTX_QDRANT_URL")
26            .map_err(|_| "LEANCTX_QDRANT_URL is required for qdrant backend".to_string())?;
27        let url = url.trim().trim_end_matches('/').to_string();
28        if url.is_empty() {
29            return Err("LEANCTX_QDRANT_URL is required for qdrant backend".to_string());
30        }
31
32        let api_key = std::env::var("LEANCTX_QDRANT_API_KEY")
33            .ok()
34            .map(|v| v.trim().to_string())
35            .filter(|v| !v.is_empty());
36
37        let timeout_secs = std::env::var("LEANCTX_QDRANT_TIMEOUT_SECS")
38            .ok()
39            .and_then(|v| v.trim().parse::<u64>().ok())
40            .filter(|v| *v > 0)
41            .unwrap_or(10);
42
43        let collection_prefix = std::env::var("LEANCTX_QDRANT_COLLECTION_PREFIX")
44            .ok()
45            .map(|v| v.trim().to_string())
46            .filter(|v| !v.is_empty())
47            .unwrap_or_else(|| "lctx_code_".to_string());
48
49        Ok(Self {
50            url,
51            api_key,
52            timeout_secs,
53            collection_prefix,
54        })
55    }
56}
57
58#[derive(Debug, Clone)]
59pub struct QdrantStore {
60    cfg: QdrantConfig,
61    agent: ureq::Agent,
62}
63
64#[derive(Debug, Clone)]
65pub struct QdrantHit {
66    pub score: f32,
67    pub file_path: String,
68    pub symbol_name: String,
69    pub kind: crate::core::bm25_index::ChunkKind,
70    pub start_line: usize,
71    pub end_line: usize,
72}
73
74impl QdrantStore {
75    pub fn from_env() -> Result<Self, String> {
76        let cfg = QdrantConfig::from_env()?;
77        let agent = crate::core::http_client::ureq_agent(
78            ureq::config::Config::builder()
79                .tls_config(crate::core::http_client::platform_tls_config())
80                .timeout_global(Some(std::time::Duration::from_secs(cfg.timeout_secs)))
81                .http_status_as_error(false)
82                .build(),
83        );
84        Ok(Self { cfg, agent })
85    }
86
87    pub fn collection_name(&self, root: &Path, dimensions: usize) -> Result<String, String> {
88        let ns = crate::core::index_namespace::namespace_hash(root);
89        Ok(format!(
90            "{}{}_d{}",
91            self.cfg.collection_prefix, ns, dimensions
92        ))
93    }
94
95    /// Ensure the collection exists. Returns `true` if it was created.
96    pub fn ensure_collection(&self, collection: &str, dimensions: usize) -> Result<bool, String> {
97        let url = format!("{}/collections/{collection}", self.cfg.url);
98        let payload = serde_json::json!({
99            "vectors": { "size": dimensions, "distance": "Cosine" }
100        });
101        let payload_bytes = serde_json::to_vec(&payload).map_err(|e| e.to_string())?;
102
103        let mut req = self
104            .agent
105            .put(&url)
106            .header("Content-Type", "application/json");
107        if let Some(ref key) = self.cfg.api_key {
108            req = req.header("api-key", key);
109        }
110
111        let resp = req
112            .send(payload_bytes.as_slice())
113            .map_err(|e| format!("qdrant create collection failed: {e}"))?;
114        let status = resp.status().as_u16();
115
116        if (200..300).contains(&status) {
117            return Ok(true);
118        }
119        if status == 409 {
120            return Ok(false);
121        }
122
123        let body = resp.into_body().read_to_string().unwrap_or_default();
124        Err(format!(
125            "qdrant create collection failed ({status}): {body}"
126        ))
127    }
128
129    pub fn sync_index(
130        &self,
131        collection: &str,
132        index: &BM25Index,
133        aligned_embeddings: &[Vec<f32>],
134        changed_files: &[String],
135        created_new: bool,
136    ) -> Result<(), String> {
137        if index.chunks.len() != aligned_embeddings.len() {
138            return Err("embedding alignment length mismatch".to_string());
139        }
140
141        if created_new {
142            // Fresh collection: upsert everything once.
143            return self.upsert_all(collection, index, aligned_embeddings);
144        }
145
146        if changed_files.is_empty() {
147            return Ok(());
148        }
149
150        let mut unique: Vec<String> = changed_files.to_vec();
151        unique.sort();
152        unique.dedup();
153
154        let mut changed_set: HashSet<&str> = HashSet::with_capacity(unique.len());
155        for f in &unique {
156            changed_set.insert(f.as_str());
157        }
158
159        // For each changed file we "replace": delete all points for the file, then upsert current chunks.
160        for file in &unique {
161            self.delete_by_file(collection, file)?;
162        }
163
164        self.upsert_files(collection, index, aligned_embeddings, &changed_set)
165    }
166
167    pub fn search(
168        &self,
169        collection: &str,
170        query_vec: &[f32],
171        limit: usize,
172    ) -> Result<Vec<QdrantHit>, String> {
173        let url = format!("{}/collections/{collection}/points/search", self.cfg.url);
174        let payload = serde_json::json!({
175            "vector": query_vec,
176            "limit": limit,
177            "with_payload": true,
178            "with_vector": false,
179        });
180        let payload_bytes = serde_json::to_vec(&payload).map_err(|e| e.to_string())?;
181
182        let mut req = self
183            .agent
184            .post(&url)
185            .header("Content-Type", "application/json");
186        if let Some(ref key) = self.cfg.api_key {
187            req = req.header("api-key", key);
188        }
189
190        let resp = req
191            .send(payload_bytes.as_slice())
192            .map_err(|e| format!("qdrant search failed: {e}"))?;
193        let status = resp.status().as_u16();
194        let body = resp
195            .into_body()
196            .read_to_string()
197            .map_err(|e| e.to_string())?;
198
199        if status >= 400 {
200            return Err(format!("qdrant search failed ({status}): {body}"));
201        }
202
203        let resp: QdrantResponse<Vec<QdrantSearchHit>> =
204            serde_json::from_str(&body).map_err(|e| format!("invalid qdrant json: {e}"))?;
205
206        let mut out = Vec::with_capacity(resp.result.len());
207        for h in resp.result {
208            let Some(payload) = h.payload else { continue };
209            out.push(QdrantHit {
210                score: h.score,
211                file_path: payload.file_path,
212                symbol_name: payload.symbol_name,
213                kind: crate::core::dense_backend::kind_from_str(&payload.kind),
214                start_line: payload.start_line,
215                end_line: payload.end_line,
216            });
217        }
218        Ok(out)
219    }
220
221    fn upsert_all(
222        &self,
223        collection: &str,
224        index: &BM25Index,
225        aligned_embeddings: &[Vec<f32>],
226    ) -> Result<(), String> {
227        let mut batch: Vec<QdrantPoint<'_>> = Vec::new();
228        for (i, chunk) in index.chunks.iter().enumerate() {
229            let vec = aligned_embeddings
230                .get(i)
231                .ok_or_else(|| "embedding alignment missing".to_string())?;
232            batch.push(point_for_chunk(chunk, vec.as_slice()));
233            if batch.len() >= UPSERT_BATCH_POINTS {
234                self.upsert_points(collection, &batch)?;
235                batch.clear();
236            }
237        }
238        if !batch.is_empty() {
239            self.upsert_points(collection, &batch)?;
240        }
241        Ok(())
242    }
243
244    fn upsert_files(
245        &self,
246        collection: &str,
247        index: &BM25Index,
248        aligned_embeddings: &[Vec<f32>],
249        changed_set: &HashSet<&str>,
250    ) -> Result<(), String> {
251        let mut batch: Vec<QdrantPoint<'_>> = Vec::new();
252        for (i, chunk) in index.chunks.iter().enumerate() {
253            if !changed_set.contains(chunk.file_path.as_str()) {
254                continue;
255            }
256            let vec = aligned_embeddings
257                .get(i)
258                .ok_or_else(|| "embedding alignment missing".to_string())?;
259            batch.push(point_for_chunk(chunk, vec.as_slice()));
260            if batch.len() >= UPSERT_BATCH_POINTS {
261                self.upsert_points(collection, &batch)?;
262                batch.clear();
263            }
264        }
265        if !batch.is_empty() {
266            self.upsert_points(collection, &batch)?;
267        }
268        Ok(())
269    }
270
271    fn upsert_points(&self, collection: &str, points: &[QdrantPoint<'_>]) -> Result<(), String> {
272        let url = format!("{}/collections/{collection}/points?wait=true", self.cfg.url);
273        let payload = QdrantUpsertBody { points };
274        let payload_bytes = serde_json::to_vec(&payload).map_err(|e| e.to_string())?;
275
276        let mut req = self
277            .agent
278            .put(&url)
279            .header("Content-Type", "application/json");
280        if let Some(ref key) = self.cfg.api_key {
281            req = req.header("api-key", key);
282        }
283
284        let resp = req
285            .send(payload_bytes.as_slice())
286            .map_err(|e| format!("qdrant upsert failed: {e}"))?;
287        let status = resp.status().as_u16();
288        if status >= 400 {
289            let body = resp.into_body().read_to_string().unwrap_or_default();
290            return Err(format!("qdrant upsert failed ({status}): {body}"));
291        }
292        Ok(())
293    }
294
295    fn delete_by_file(&self, collection: &str, file_path: &str) -> Result<(), String> {
296        let url = format!(
297            "{}/collections/{collection}/points/delete?wait=true",
298            self.cfg.url
299        );
300        let payload = serde_json::json!({
301            "filter": {
302                "must": [
303                    { "key": "file_path", "match": { "value": file_path } }
304                ]
305            }
306        });
307        let payload_bytes = serde_json::to_vec(&payload).map_err(|e| e.to_string())?;
308
309        let mut req = self
310            .agent
311            .post(&url)
312            .header("Content-Type", "application/json");
313        if let Some(ref key) = self.cfg.api_key {
314            req = req.header("api-key", key);
315        }
316
317        let resp = req
318            .send(payload_bytes.as_slice())
319            .map_err(|e| format!("qdrant delete-by-file failed: {e}"))?;
320        let status = resp.status().as_u16();
321        if status >= 400 {
322            let body = resp.into_body().read_to_string().unwrap_or_default();
323            return Err(format!("qdrant delete-by-file failed ({status}): {body}"));
324        }
325        Ok(())
326    }
327}
328
329const UPSERT_BATCH_POINTS: usize = 256;
330
331#[derive(Debug, Deserialize)]
332struct QdrantResponse<T> {
333    result: T,
334}
335
336#[derive(Debug, Deserialize)]
337struct QdrantSearchHit {
338    score: f32,
339    payload: Option<QdrantPayload>,
340}
341
342#[derive(Debug, Deserialize)]
343struct QdrantPayload {
344    file_path: String,
345    symbol_name: String,
346    kind: String,
347    start_line: usize,
348    end_line: usize,
349}
350
351#[derive(Debug, Serialize)]
352struct QdrantUpsertBody<'a> {
353    points: &'a [QdrantPoint<'a>],
354}
355
356#[derive(Debug, Serialize)]
357struct QdrantPoint<'a> {
358    id: u64,
359    vector: &'a [f32],
360    payload: QdrantPointPayload<'a>,
361}
362
363#[derive(Debug, Serialize)]
364struct QdrantPointPayload<'a> {
365    file_path: &'a str,
366    symbol_name: &'a str,
367    kind: &'a str,
368    start_line: usize,
369    end_line: usize,
370}
371
372fn point_for_chunk<'a>(chunk: &'a CodeChunk, vector: &'a [f32]) -> QdrantPoint<'a> {
373    QdrantPoint {
374        id: point_id_for_chunk(chunk),
375        vector,
376        payload: QdrantPointPayload {
377            file_path: chunk.file_path.as_str(),
378            symbol_name: chunk.symbol_name.as_str(),
379            kind: crate::core::dense_backend::kind_to_str(&chunk.kind),
380            start_line: chunk.start_line,
381            end_line: chunk.end_line,
382        },
383    }
384}
385
386fn point_id_for_chunk(chunk: &CodeChunk) -> u64 {
387    let mut h = Md5::new();
388    h.update(chunk.file_path.as_bytes());
389    h.update(chunk.start_line.to_le_bytes());
390    h.update(chunk.end_line.to_le_bytes());
391    h.update(chunk.symbol_name.as_bytes());
392    // include kind tag to avoid collisions between same-named entities
393    h.update(crate::core::dense_backend::kind_to_str(&chunk.kind).as_bytes());
394    let out = h.finalize();
395    u64::from_le_bytes(out[0..8].try_into().unwrap_or([0u8; 8]))
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401    use crate::core::bm25_index::ChunkKind;
402
403    fn chunk(file: &str, name: &str, start: usize, end: usize, kind: ChunkKind) -> CodeChunk {
404        CodeChunk {
405            file_path: file.to_string(),
406            symbol_name: name.to_string(),
407            kind,
408            start_line: start,
409            end_line: end,
410            content: "fn x() {}".to_string(),
411            tokens: vec![],
412            token_count: 0,
413        }
414    }
415
416    #[test]
417    fn point_id_is_stable() {
418        let c = chunk("src/main.rs", "main", 1, 10, ChunkKind::Function);
419        let a = point_id_for_chunk(&c);
420        let b = point_id_for_chunk(&c);
421        assert_eq!(a, b);
422    }
423
424    #[test]
425    fn point_id_changes_when_location_changes() {
426        let c1 = chunk("src/main.rs", "main", 1, 10, ChunkKind::Function);
427        let c2 = chunk("src/main.rs", "main", 2, 10, ChunkKind::Function);
428        assert_ne!(point_id_for_chunk(&c1), point_id_for_chunk(&c2));
429    }
430}
lean_ctx/core/qdrant_store.rs

lean_ctx/core/
qdrant_store.rs
