lean_ctx/core/
dense_backend.rs

1use std::path::Path;
2
3use crate::core::bm25_index::BM25Index;
4#[cfg(feature = "qdrant")]
5use crate::core::bm25_index::ChunkKind;
6use crate::core::hnsw::FlatEmbeddings;
7use crate::core::hybrid_search::{DenseSearchResult, HybridConfig, HybridResult};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum DenseBackendKind {
11    Local,
12    #[cfg(feature = "qdrant")]
13    Qdrant,
14}
15
16impl DenseBackendKind {
17    pub fn try_from_env() -> Result<Self, String> {
18        let explicit = std::env::var("LEANCTX_DENSE_BACKEND")
19            .ok()
20            .map(|v| v.trim().to_ascii_lowercase())
21            .filter(|v| !v.is_empty());
22
23        let inferred_qdrant =
24            std::env::var("LEANCTX_QDRANT_URL").is_ok_and(|v| !v.trim().is_empty());
25
26        let requested = explicit.or_else(|| inferred_qdrant.then_some("qdrant".to_string()));
27
28        match requested.as_deref() {
29            None | Some("local") => Ok(Self::Local),
30            Some("qdrant") => {
31                #[cfg(feature = "qdrant")]
32                {
33                    Ok(Self::Qdrant)
34                }
35                #[cfg(not(feature = "qdrant"))]
36                {
37                    Err("Dense backend 'qdrant' requested, but feature 'qdrant' is not enabled. Rebuild with --features qdrant.".to_string())
38                }
39            }
40            Some(other) => Err(format!(
41                "Unknown LEANCTX_DENSE_BACKEND={other:?} (expected 'local' or 'qdrant')"
42            )),
43        }
44    }
45
46    pub fn label(&self) -> &'static str {
47        match self {
48            Self::Local => "local",
49            #[cfg(feature = "qdrant")]
50            Self::Qdrant => "qdrant",
51        }
52    }
53}
54
55#[cfg(feature = "embeddings")]
56#[allow(clippy::too_many_arguments)]
57pub fn dense_results_as_hybrid(
58    backend: DenseBackendKind,
59    root: &Path,
60    index: &BM25Index,
61    engine: &crate::core::embeddings::EmbeddingEngine,
62    aligned_embeddings: &FlatEmbeddings,
63    changed_files: &[String],
64    query: &str,
65    top_k: usize,
66    filter: Option<&dyn Fn(&str) -> bool>,
67) -> Result<Vec<HybridResult>, String> {
68    let dense = dense_results(
69        backend,
70        root,
71        index,
72        engine,
73        aligned_embeddings,
74        changed_files,
75        query,
76        top_k,
77        filter,
78    )?;
79
80    Ok(dense
81        .into_iter()
82        .map(|d| HybridResult {
83            file_path: d.file_path,
84            symbol_name: d.symbol_name,
85            kind: d.kind,
86            start_line: d.start_line,
87            end_line: d.end_line,
88            snippet: d.snippet,
89            rrf_score: d.similarity as f64,
90            bm25_score: None,
91            dense_score: Some(d.similarity),
92            bm25_rank: None,
93            dense_rank: None,
94        })
95        .collect())
96}
97
98#[cfg(feature = "embeddings")]
99#[allow(clippy::too_many_arguments)]
100pub fn hybrid_results(
101    backend: DenseBackendKind,
102    root: &Path,
103    index: &BM25Index,
104    engine: &crate::core::embeddings::EmbeddingEngine,
105    aligned_embeddings: &FlatEmbeddings,
106    changed_files: &[String],
107    query: &str,
108    top_k: usize,
109    config: &HybridConfig,
110    filter: Option<&dyn Fn(&str) -> bool>,
111    graph_file_ranks: Option<&std::collections::HashMap<String, usize>>,
112) -> Result<Vec<HybridResult>, String> {
113    match backend {
114        DenseBackendKind::Local => {
115            let _ = (root, changed_files);
116            let mut results = crate::core::hybrid_search::hybrid_search(
117                query,
118                index,
119                Some(engine),
120                Some(aligned_embeddings),
121                top_k,
122                config,
123                graph_file_ranks,
124            );
125            if let Some(pred) = filter {
126                results.retain(|r| pred(&r.file_path));
127            }
128            results.truncate(top_k);
129            Ok(results)
130        }
131        #[cfg(feature = "qdrant")]
132        DenseBackendKind::Qdrant => {
133            let bm25_k = config.bm25_candidates.max(top_k);
134            let dense_k = config.dense_candidates.max(top_k);
135
136            let mut bm25 = index.search(query, bm25_k);
137            if let Some(pred) = filter {
138                bm25.retain(|r| pred(&r.file_path));
139            }
140
141            let dense = dense_results(
142                backend,
143                root,
144                index,
145                engine,
146                aligned_embeddings,
147                changed_files,
148                query,
149                dense_k,
150                filter,
151            )?;
152
153            let mut fused = crate::core::hybrid_search::reciprocal_rank_fusion(
154                &bm25,
155                &dense,
156                config,
157                top_k,
158                graph_file_ranks,
159            );
160            if let Some(pred) = filter {
161                fused.retain(|r| pred(&r.file_path));
162            }
163            fused.truncate(top_k);
164            Ok(fused)
165        }
166    }
167}
168
169#[cfg(feature = "embeddings")]
170#[allow(clippy::too_many_arguments)]
171fn dense_results(
172    backend: DenseBackendKind,
173    root: &Path,
174    index: &BM25Index,
175    engine: &crate::core::embeddings::EmbeddingEngine,
176    aligned_embeddings: &FlatEmbeddings,
177    changed_files: &[String],
178    query: &str,
179    top_k: usize,
180    filter: Option<&dyn Fn(&str) -> bool>,
181) -> Result<Vec<DenseSearchResult>, String> {
182    match backend {
183        DenseBackendKind::Local => {
184            let _ = (root, changed_files);
185            dense_results_local(index, engine, aligned_embeddings, query, top_k, filter)
186        }
187        #[cfg(feature = "qdrant")]
188        DenseBackendKind::Qdrant => {
189            let vecs: Vec<Vec<f32>> = (0..aligned_embeddings.n_vectors())
190                .map(|i| aligned_embeddings.get_vec(i))
191                .collect();
192            dense_results_qdrant(
193                root,
194                index,
195                engine,
196                &vecs,
197                changed_files,
198                query,
199                top_k,
200                filter,
201            )
202        }
203    }
204}
205
206#[cfg(feature = "embeddings")]
207fn dense_results_local(
208    index: &BM25Index,
209    engine: &crate::core::embeddings::EmbeddingEngine,
210    aligned_embeddings: &FlatEmbeddings,
211    query: &str,
212    top_k: usize,
213    filter: Option<&dyn Fn(&str) -> bool>,
214) -> Result<Vec<DenseSearchResult>, String> {
215    let query_embedding = engine
216        .embed_query(query)
217        .map_err(|e| format!("embedding failed: {e}"))?;
218
219    let top = top_k_by_similarity(&query_embedding, aligned_embeddings, top_k, |i| {
220        let Some(pred) = filter else { return true };
221        index.chunks.get(i).is_some_and(|c| pred(&c.file_path))
222    });
223
224    Ok(top
225        .into_iter()
226        .filter_map(|(idx, sim)| {
227            let chunk = index.chunks.get(idx)?;
228            let snippet = chunk.content.lines().take(5).collect::<Vec<_>>().join("\n");
229            Some(DenseSearchResult {
230                chunk_idx: idx,
231                similarity: sim,
232                file_path: chunk.file_path.clone(),
233                symbol_name: chunk.symbol_name.clone(),
234                kind: chunk.kind.clone(),
235                start_line: chunk.start_line,
236                end_line: chunk.end_line,
237                snippet,
238            })
239        })
240        .collect())
241}
242
243/// Min-heap based Top-K selection over a flat embedding buffer.
244/// O(n log k) instead of O(n log n) full sort. The filter is applied inline
245/// during the scan so only matching chunks are considered — post-filtering
246/// cannot drop below `top_k` results regardless of filter selectivity.
247///
248/// Uses sequential memory access (one dereference) via `FlatEmbeddings::get`,
249/// unlike the old `Arc<[Vec<f32>]>` layout which had two-level indirection.
250#[cfg(feature = "embeddings")]
251fn top_k_by_similarity(
252    query: &[f32],
253    embeddings: &FlatEmbeddings,
254    k: usize,
255    filter: impl Fn(usize) -> bool,
256) -> Vec<(usize, f32)> {
257    use std::cmp::Ordering;
258    use std::collections::BinaryHeap;
259
260    #[derive(PartialEq)]
261    struct MinEntry(f32, usize);
262
263    impl Eq for MinEntry {}
264    impl PartialOrd for MinEntry {
265        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
266            Some(self.cmp(other))
267        }
268    }
269    impl Ord for MinEntry {
270        fn cmp(&self, other: &Self) -> Ordering {
271            other
272                .0
273                .partial_cmp(&self.0)
274                .unwrap_or(Ordering::Equal)
275                .then_with(|| self.1.cmp(&other.1))
276        }
277    }
278
279    let n = embeddings.n_vectors();
280    let mut heap: BinaryHeap<MinEntry> = BinaryHeap::with_capacity(k + 1);
281
282    for i in 0..n {
283        if !filter(i) {
284            continue;
285        }
286        let emb = embeddings.get(i);
287        let sim = crate::core::embeddings::cosine_similarity(query, emb);
288        if heap.len() < k {
289            heap.push(MinEntry(sim, i));
290        } else if let Some(min) = heap.peek()
291            && sim > min.0
292        {
293            heap.pop();
294            heap.push(MinEntry(sim, i));
295        }
296    }
297
298    let mut result: Vec<(usize, f32)> = heap.into_iter().map(|e| (e.1, e.0)).collect();
299    result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
300    result
301}
302
303#[cfg(feature = "qdrant")]
304#[cfg(feature = "embeddings")]
305fn dense_results_qdrant(
306    root: &Path,
307    index: &BM25Index,
308    engine: &crate::core::embeddings::EmbeddingEngine,
309    aligned_embeddings: &[Vec<f32>],
310    changed_files: &[String],
311    query: &str,
312    top_k: usize,
313    filter: Option<&dyn Fn(&str) -> bool>,
314) -> Result<Vec<DenseSearchResult>, String> {
315    let store = crate::core::qdrant_store::QdrantStore::from_env()?;
316    let collection = store.collection_name(root, engine.dimensions())?;
317    let created_new = store.ensure_collection(&collection, engine.dimensions())?;
318    store.sync_index(
319        &collection,
320        index,
321        aligned_embeddings,
322        changed_files,
323        created_new,
324    )?;
325
326    let query_vec = engine
327        .embed_query(query)
328        .map_err(|e| format!("embedding failed: {e}"))?;
329
330    let hits = store.search(&collection, &query_vec, top_k)?;
331    let mut out = Vec::with_capacity(hits.len());
332    for hit in hits {
333        if let Some(pred) = filter
334            && !pred(&hit.file_path)
335        {
336            continue;
337        }
338        let snippet = snippet_from_disk(root, &hit.file_path, hit.start_line, hit.end_line, 5);
339        out.push(DenseSearchResult {
340            chunk_idx: 0,
341            similarity: hit.score,
342            file_path: hit.file_path,
343            symbol_name: hit.symbol_name,
344            kind: hit.kind,
345            start_line: hit.start_line,
346            end_line: hit.end_line,
347            snippet,
348        });
349    }
350    Ok(out)
351}
352
353#[cfg(feature = "qdrant")]
354fn snippet_from_disk(
355    root: &Path,
356    rel_path: &str,
357    start_line: usize,
358    end_line: usize,
359    max_lines: usize,
360) -> String {
361    let Ok(path) = crate::core::pathjail::jail_path(&root.join(rel_path), root) else {
362        return String::new();
363    };
364    let Ok(content) = std::fs::read_to_string(path) else {
365        return String::new();
366    };
367    let lines: Vec<&str> = content.lines().collect();
368    if lines.is_empty() {
369        return String::new();
370    }
371    let start = start_line.saturating_sub(1).min(lines.len());
372    let end = end_line.max(start_line).min(lines.len());
373    let mut slice = &lines[start..end];
374    if slice.len() > max_lines {
375        slice = &slice[..max_lines];
376    }
377    slice.join("\n")
378}
379
380#[cfg(feature = "qdrant")]
381pub(crate) fn kind_to_str(kind: &ChunkKind) -> &'static str {
382    match kind {
383        ChunkKind::Function => "Function",
384        ChunkKind::Struct => "Struct",
385        ChunkKind::Impl => "Impl",
386        ChunkKind::Module => "Module",
387        ChunkKind::Class => "Class",
388        ChunkKind::Method => "Method",
389        ChunkKind::Issue => "Issue",
390        ChunkKind::PullRequest => "PullRequest",
391        ChunkKind::WikiPage => "WikiPage",
392        ChunkKind::DbSchema => "DbSchema",
393        ChunkKind::ApiEndpoint => "ApiEndpoint",
394        ChunkKind::Ticket => "Ticket",
395        ChunkKind::ExternalOther => "ExternalOther",
396        ChunkKind::Other => "Other",
397    }
398}
399
400#[cfg(feature = "qdrant")]
401pub(crate) fn kind_from_str(s: &str) -> ChunkKind {
402    match s {
403        "Function" => ChunkKind::Function,
404        "Struct" => ChunkKind::Struct,
405        "Impl" => ChunkKind::Impl,
406        "Module" => ChunkKind::Module,
407        "Class" => ChunkKind::Class,
408        "Method" => ChunkKind::Method,
409        "Issue" => ChunkKind::Issue,
410        "PullRequest" => ChunkKind::PullRequest,
411        "WikiPage" => ChunkKind::WikiPage,
412        "DbSchema" => ChunkKind::DbSchema,
413        "ApiEndpoint" => ChunkKind::ApiEndpoint,
414        "Ticket" => ChunkKind::Ticket,
415        "ExternalOther" => ChunkKind::ExternalOther,
416        _ => ChunkKind::Other,
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423    use std::sync::Mutex;
424
425    static ENV_LOCK: Mutex<()> = Mutex::new(());
426
427    fn set_env(key: &str, value: Option<&str>) -> Option<String> {
428        let old = std::env::var(key).ok();
429        match value {
430            Some(v) => crate::test_env::set_var(key, v),
431            None => crate::test_env::remove_var(key),
432        }
433        old
434    }
435
436    fn restore_env(key: &str, old: Option<String>) {
437        match old {
438            Some(v) => crate::test_env::set_var(key, v),
439            None => crate::test_env::remove_var(key),
440        }
441    }
442
443    #[test]
444    fn dense_backend_defaults_to_local() {
445        let _g = ENV_LOCK.lock().unwrap();
446        let old_backend = set_env("LEANCTX_DENSE_BACKEND", None);
447        let old_url = set_env("LEANCTX_QDRANT_URL", None);
448
449        let got = DenseBackendKind::try_from_env().unwrap();
450        assert_eq!(got, DenseBackendKind::Local);
451
452        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
453        restore_env("LEANCTX_QDRANT_URL", old_url);
454    }
455
456    #[test]
457    fn dense_backend_unknown_value_errors() {
458        let _g = ENV_LOCK.lock().unwrap();
459        let old_backend = set_env("LEANCTX_DENSE_BACKEND", Some("wat"));
460        let old_url = set_env("LEANCTX_QDRANT_URL", None);
461
462        let err = DenseBackendKind::try_from_env().unwrap_err();
463        assert!(err.contains("Unknown LEANCTX_DENSE_BACKEND"));
464
465        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
466        restore_env("LEANCTX_QDRANT_URL", old_url);
467    }
468
469    #[cfg(feature = "qdrant")]
470    #[test]
471    fn dense_backend_infers_qdrant_from_url() {
472        let _g = ENV_LOCK.lock().unwrap();
473        let old_backend = set_env("LEANCTX_DENSE_BACKEND", None);
474        let old_url = set_env("LEANCTX_QDRANT_URL", Some("http://127.0.0.1:6333"));
475
476        let got = DenseBackendKind::try_from_env().unwrap();
477        assert_eq!(got, DenseBackendKind::Qdrant);
478
479        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
480        restore_env("LEANCTX_QDRANT_URL", old_url);
481    }
482
483    #[cfg(not(feature = "qdrant"))]
484    #[test]
485    fn dense_backend_qdrant_requires_feature() {
486        let _g = ENV_LOCK.lock().unwrap();
487        let old_backend = set_env("LEANCTX_DENSE_BACKEND", Some("qdrant"));
488        let old_url = set_env("LEANCTX_QDRANT_URL", None);
489
490        let err = DenseBackendKind::try_from_env().unwrap_err();
491        assert!(err.contains("feature 'qdrant' is not enabled"));
492
493        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
494        restore_env("LEANCTX_QDRANT_URL", old_url);
495    }
496}
lean_ctx/core/dense_backend.rs

lean_ctx/core/
dense_backend.rs
