lean_ctx/core/
dense_backend.rs

1use std::path::Path;
2
3use crate::core::bm25_index::BM25Index;
4#[cfg(any(feature = "qdrant", feature = "pgvector"))]
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    #[cfg(feature = "pgvector")]
15    Pgvector,
16}
17
18impl DenseBackendKind {
19    pub fn try_from_env() -> Result<Self, String> {
20        let explicit = std::env::var("LEANCTX_DENSE_BACKEND")
21            .ok()
22            .map(|v| v.trim().to_ascii_lowercase())
23            .filter(|v| !v.is_empty());
24
25        let inferred_qdrant =
26            std::env::var("LEANCTX_QDRANT_URL").is_ok_and(|v| !v.trim().is_empty());
27        let inferred_pgvector =
28            std::env::var("LEANCTX_PGVECTOR_URL").is_ok_and(|v| !v.trim().is_empty());
29
30        let requested = explicit
31            .or_else(|| inferred_qdrant.then_some("qdrant".to_string()))
32            .or_else(|| inferred_pgvector.then_some("pgvector".to_string()));
33
34        match requested.as_deref() {
35            None | Some("local") => Ok(Self::Local),
36            Some("qdrant") => {
37                #[cfg(feature = "qdrant")]
38                {
39                    Ok(Self::Qdrant)
40                }
41                #[cfg(not(feature = "qdrant"))]
42                {
43                    Err("Dense backend 'qdrant' requested, but feature 'qdrant' is not enabled. Rebuild with --features qdrant.".to_string())
44                }
45            }
46            Some("pgvector") => {
47                #[cfg(feature = "pgvector")]
48                {
49                    Ok(Self::Pgvector)
50                }
51                #[cfg(not(feature = "pgvector"))]
52                {
53                    Err("Dense backend 'pgvector' requested, but feature 'pgvector' is not enabled. Rebuild with --features pgvector.".to_string())
54                }
55            }
56            Some(other) => Err(format!(
57                "Unknown LEANCTX_DENSE_BACKEND={other:?} (expected 'local', 'qdrant' or 'pgvector')"
58            )),
59        }
60    }
61
62    pub fn label(&self) -> &'static str {
63        match self {
64            Self::Local => "local",
65            #[cfg(feature = "qdrant")]
66            Self::Qdrant => "qdrant",
67            #[cfg(feature = "pgvector")]
68            Self::Pgvector => "pgvector",
69        }
70    }
71}
72
73#[cfg(feature = "embeddings")]
74#[allow(clippy::too_many_arguments)]
75pub fn dense_results_as_hybrid(
76    backend: DenseBackendKind,
77    root: &Path,
78    index: &BM25Index,
79    engine: &crate::core::embeddings::EmbeddingEngine,
80    aligned_embeddings: &FlatEmbeddings,
81    changed_files: &[String],
82    query: &str,
83    top_k: usize,
84    filter: Option<&dyn Fn(&str) -> bool>,
85) -> Result<Vec<HybridResult>, String> {
86    let dense = dense_results(
87        backend,
88        root,
89        index,
90        engine,
91        aligned_embeddings,
92        changed_files,
93        query,
94        top_k,
95        filter,
96    )?;
97
98    Ok(dense
99        .into_iter()
100        .map(|d| HybridResult {
101            file_path: d.file_path,
102            symbol_name: d.symbol_name,
103            kind: d.kind,
104            start_line: d.start_line,
105            end_line: d.end_line,
106            snippet: d.snippet,
107            rrf_score: d.similarity as f64,
108            bm25_score: None,
109            dense_score: Some(d.similarity),
110            bm25_rank: None,
111            dense_rank: None,
112        })
113        .collect())
114}
115
116#[cfg(feature = "embeddings")]
117#[allow(clippy::too_many_arguments)]
118pub fn hybrid_results(
119    backend: DenseBackendKind,
120    root: &Path,
121    index: &BM25Index,
122    engine: &crate::core::embeddings::EmbeddingEngine,
123    aligned_embeddings: &FlatEmbeddings,
124    changed_files: &[String],
125    query: &str,
126    top_k: usize,
127    config: &HybridConfig,
128    filter: Option<&dyn Fn(&str) -> bool>,
129    graph_file_ranks: Option<&std::collections::HashMap<String, usize>>,
130) -> Result<Vec<HybridResult>, String> {
131    match backend {
132        DenseBackendKind::Local => {
133            let _ = (root, changed_files);
134            let mut results = crate::core::hybrid_search::hybrid_search(
135                query,
136                index,
137                Some(engine),
138                Some(aligned_embeddings),
139                top_k,
140                config,
141                graph_file_ranks,
142            );
143            if let Some(pred) = filter {
144                results.retain(|r| pred(&r.file_path));
145            }
146            results.truncate(top_k);
147            Ok(results)
148        }
149        #[cfg(feature = "qdrant")]
150        DenseBackendKind::Qdrant => remote_hybrid_results(
151            backend,
152            root,
153            index,
154            engine,
155            aligned_embeddings,
156            changed_files,
157            query,
158            top_k,
159            config,
160            filter,
161            graph_file_ranks,
162        ),
163        #[cfg(feature = "pgvector")]
164        DenseBackendKind::Pgvector => remote_hybrid_results(
165            backend,
166            root,
167            index,
168            engine,
169            aligned_embeddings,
170            changed_files,
171            query,
172            top_k,
173            config,
174            filter,
175            graph_file_ranks,
176        ),
177    }
178}
179
180/// Shared BM25+dense RRF pipeline for remote vector backends (qdrant, pgvector):
181/// the backends differ only in where `dense_results` fetches from.
182#[cfg(all(feature = "embeddings", any(feature = "qdrant", feature = "pgvector")))]
183#[allow(clippy::too_many_arguments)]
184fn remote_hybrid_results(
185    backend: DenseBackendKind,
186    root: &Path,
187    index: &BM25Index,
188    engine: &crate::core::embeddings::EmbeddingEngine,
189    aligned_embeddings: &FlatEmbeddings,
190    changed_files: &[String],
191    query: &str,
192    top_k: usize,
193    config: &HybridConfig,
194    filter: Option<&dyn Fn(&str) -> bool>,
195    graph_file_ranks: Option<&std::collections::HashMap<String, usize>>,
196) -> Result<Vec<HybridResult>, String> {
197    let bm25_k = config.bm25_candidates.max(top_k);
198    let dense_k = config.dense_candidates.max(top_k);
199
200    let mut bm25 = index.search(query, bm25_k);
201    if let Some(pred) = filter {
202        bm25.retain(|r| pred(&r.file_path));
203    }
204
205    let dense = dense_results(
206        backend,
207        root,
208        index,
209        engine,
210        aligned_embeddings,
211        changed_files,
212        query,
213        dense_k,
214        filter,
215    )?;
216
217    let mut fused = crate::core::hybrid_search::reciprocal_rank_fusion(
218        &bm25,
219        &dense,
220        config,
221        top_k,
222        graph_file_ranks,
223    );
224    if let Some(pred) = filter {
225        fused.retain(|r| pred(&r.file_path));
226    }
227    fused.truncate(top_k);
228    Ok(fused)
229}
230
231#[cfg(feature = "embeddings")]
232#[allow(clippy::too_many_arguments)]
233fn dense_results(
234    backend: DenseBackendKind,
235    root: &Path,
236    index: &BM25Index,
237    engine: &crate::core::embeddings::EmbeddingEngine,
238    aligned_embeddings: &FlatEmbeddings,
239    changed_files: &[String],
240    query: &str,
241    top_k: usize,
242    filter: Option<&dyn Fn(&str) -> bool>,
243) -> Result<Vec<DenseSearchResult>, String> {
244    match backend {
245        DenseBackendKind::Local => {
246            let _ = (root, changed_files);
247            dense_results_local(index, engine, aligned_embeddings, query, top_k, filter)
248        }
249        #[cfg(feature = "qdrant")]
250        DenseBackendKind::Qdrant => {
251            let vecs: Vec<Vec<f32>> = (0..aligned_embeddings.n_vectors())
252                .map(|i| aligned_embeddings.get_vec(i))
253                .collect();
254            dense_results_qdrant(
255                root,
256                index,
257                engine,
258                &vecs,
259                changed_files,
260                query,
261                top_k,
262                filter,
263            )
264        }
265        #[cfg(feature = "pgvector")]
266        DenseBackendKind::Pgvector => {
267            let vecs: Vec<Vec<f32>> = (0..aligned_embeddings.n_vectors())
268                .map(|i| aligned_embeddings.get_vec(i))
269                .collect();
270            dense_results_pgvector(
271                root,
272                index,
273                engine,
274                &vecs,
275                changed_files,
276                query,
277                top_k,
278                filter,
279            )
280        }
281    }
282}
283
284#[cfg(feature = "embeddings")]
285fn dense_results_local(
286    index: &BM25Index,
287    engine: &crate::core::embeddings::EmbeddingEngine,
288    aligned_embeddings: &FlatEmbeddings,
289    query: &str,
290    top_k: usize,
291    filter: Option<&dyn Fn(&str) -> bool>,
292) -> Result<Vec<DenseSearchResult>, String> {
293    let query_embedding = engine
294        .embed_query(query)
295        .map_err(|e| format!("embedding failed: {e}"))?;
296
297    let top = top_k_by_similarity(&query_embedding, aligned_embeddings, top_k, |i| {
298        let Some(pred) = filter else { return true };
299        index.chunks.get(i).is_some_and(|c| pred(&c.file_path))
300    });
301
302    Ok(top
303        .into_iter()
304        .filter_map(|(idx, sim)| {
305            let chunk = index.chunks.get(idx)?;
306            let snippet = chunk.content.lines().take(5).collect::<Vec<_>>().join("\n");
307            Some(DenseSearchResult {
308                chunk_idx: idx,
309                similarity: sim,
310                file_path: chunk.file_path.clone(),
311                symbol_name: chunk.symbol_name.clone(),
312                kind: chunk.kind.clone(),
313                start_line: chunk.start_line,
314                end_line: chunk.end_line,
315                snippet,
316            })
317        })
318        .collect())
319}
320
321/// Min-heap based Top-K selection over a flat embedding buffer.
322/// O(n log k) instead of O(n log n) full sort. The filter is applied inline
323/// during the scan so only matching chunks are considered — post-filtering
324/// cannot drop below `top_k` results regardless of filter selectivity.
325///
326/// Uses sequential memory access (one dereference) via `FlatEmbeddings::get`,
327/// unlike the old `Arc<[Vec<f32>]>` layout which had two-level indirection.
328#[cfg(feature = "embeddings")]
329fn top_k_by_similarity(
330    query: &[f32],
331    embeddings: &FlatEmbeddings,
332    k: usize,
333    filter: impl Fn(usize) -> bool,
334) -> Vec<(usize, f32)> {
335    use std::cmp::Ordering;
336    use std::collections::BinaryHeap;
337
338    #[derive(PartialEq)]
339    struct MinEntry(f32, usize);
340
341    impl Eq for MinEntry {}
342    impl PartialOrd for MinEntry {
343        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
344            Some(self.cmp(other))
345        }
346    }
347    impl Ord for MinEntry {
348        fn cmp(&self, other: &Self) -> Ordering {
349            other
350                .0
351                .partial_cmp(&self.0)
352                .unwrap_or(Ordering::Equal)
353                .then_with(|| self.1.cmp(&other.1))
354        }
355    }
356
357    let n = embeddings.n_vectors();
358    let mut heap: BinaryHeap<MinEntry> = BinaryHeap::with_capacity(k + 1);
359
360    for i in 0..n {
361        if !filter(i) {
362            continue;
363        }
364        let emb = embeddings.get(i);
365        let sim = crate::core::embeddings::cosine_similarity(query, emb);
366        if heap.len() < k {
367            heap.push(MinEntry(sim, i));
368        } else if let Some(min) = heap.peek()
369            && sim > min.0
370        {
371            heap.pop();
372            heap.push(MinEntry(sim, i));
373        }
374    }
375
376    let mut result: Vec<(usize, f32)> = heap.into_iter().map(|e| (e.1, e.0)).collect();
377    result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
378    result
379}
380
381#[cfg(feature = "qdrant")]
382#[cfg(feature = "embeddings")]
383fn dense_results_qdrant(
384    root: &Path,
385    index: &BM25Index,
386    engine: &crate::core::embeddings::EmbeddingEngine,
387    aligned_embeddings: &[Vec<f32>],
388    changed_files: &[String],
389    query: &str,
390    top_k: usize,
391    filter: Option<&dyn Fn(&str) -> bool>,
392) -> Result<Vec<DenseSearchResult>, String> {
393    let store = crate::core::qdrant_store::QdrantStore::from_env()?;
394    let collection = store.collection_name(root, engine.dimensions())?;
395    let created_new = store.ensure_collection(&collection, engine.dimensions())?;
396    store.sync_index(
397        &collection,
398        index,
399        aligned_embeddings,
400        changed_files,
401        created_new,
402    )?;
403
404    let query_vec = engine
405        .embed_query(query)
406        .map_err(|e| format!("embedding failed: {e}"))?;
407
408    let hits = store.search(&collection, &query_vec, top_k)?;
409    let mut out = Vec::with_capacity(hits.len());
410    for hit in hits {
411        if let Some(pred) = filter
412            && !pred(&hit.file_path)
413        {
414            continue;
415        }
416        let snippet = snippet_from_disk(root, &hit.file_path, hit.start_line, hit.end_line, 5);
417        out.push(DenseSearchResult {
418            chunk_idx: 0,
419            similarity: hit.score,
420            file_path: hit.file_path,
421            symbol_name: hit.symbol_name,
422            kind: hit.kind,
423            start_line: hit.start_line,
424            end_line: hit.end_line,
425            snippet,
426        });
427    }
428    Ok(out)
429}
430
431#[cfg(feature = "pgvector")]
432#[cfg(feature = "embeddings")]
433#[allow(clippy::too_many_arguments)]
434fn dense_results_pgvector(
435    root: &Path,
436    index: &BM25Index,
437    engine: &crate::core::embeddings::EmbeddingEngine,
438    aligned_embeddings: &[Vec<f32>],
439    changed_files: &[String],
440    query: &str,
441    top_k: usize,
442    filter: Option<&dyn Fn(&str) -> bool>,
443) -> Result<Vec<DenseSearchResult>, String> {
444    let store = crate::core::pgvector_store::PgvectorStore::from_env()?;
445    let table = store.table_name(root, engine.dimensions())?;
446    let created_new = store.ensure_table(&table, engine.dimensions())?;
447    store.sync_index(
448        &table,
449        index,
450        aligned_embeddings,
451        changed_files,
452        created_new,
453    )?;
454
455    let query_vec = engine
456        .embed_query(query)
457        .map_err(|e| format!("embedding failed: {e}"))?;
458
459    let hits = store.search(&table, &query_vec, top_k)?;
460    let mut out = Vec::with_capacity(hits.len());
461    for hit in hits {
462        if let Some(pred) = filter
463            && !pred(&hit.file_path)
464        {
465            continue;
466        }
467        let snippet = snippet_from_disk(root, &hit.file_path, hit.start_line, hit.end_line, 5);
468        out.push(DenseSearchResult {
469            chunk_idx: 0,
470            similarity: hit.score,
471            file_path: hit.file_path,
472            symbol_name: hit.symbol_name,
473            kind: hit.kind,
474            start_line: hit.start_line,
475            end_line: hit.end_line,
476            snippet,
477        });
478    }
479    Ok(out)
480}
481
482#[cfg(any(feature = "qdrant", feature = "pgvector"))]
483fn snippet_from_disk(
484    root: &Path,
485    rel_path: &str,
486    start_line: usize,
487    end_line: usize,
488    max_lines: usize,
489) -> String {
490    let Ok(path) = crate::core::pathjail::jail_path(&root.join(rel_path), root) else {
491        return String::new();
492    };
493    let Ok(content) = std::fs::read_to_string(path) else {
494        return String::new();
495    };
496    let lines: Vec<&str> = content.lines().collect();
497    if lines.is_empty() {
498        return String::new();
499    }
500    let start = start_line.saturating_sub(1).min(lines.len());
501    let end = end_line.max(start_line).min(lines.len());
502    let mut slice = &lines[start..end];
503    if slice.len() > max_lines {
504        slice = &slice[..max_lines];
505    }
506    slice.join("\n")
507}
508
509#[cfg(any(feature = "qdrant", feature = "pgvector"))]
510pub(crate) fn kind_to_str(kind: &ChunkKind) -> &'static str {
511    match kind {
512        ChunkKind::Function => "Function",
513        ChunkKind::Struct => "Struct",
514        ChunkKind::Impl => "Impl",
515        ChunkKind::Module => "Module",
516        ChunkKind::Class => "Class",
517        ChunkKind::Method => "Method",
518        ChunkKind::Issue => "Issue",
519        ChunkKind::PullRequest => "PullRequest",
520        ChunkKind::WikiPage => "WikiPage",
521        ChunkKind::DbSchema => "DbSchema",
522        ChunkKind::ApiEndpoint => "ApiEndpoint",
523        ChunkKind::Ticket => "Ticket",
524        ChunkKind::ExternalOther => "ExternalOther",
525        ChunkKind::Other => "Other",
526    }
527}
528
529#[cfg(any(feature = "qdrant", feature = "pgvector"))]
530pub(crate) fn kind_from_str(s: &str) -> ChunkKind {
531    match s {
532        "Function" => ChunkKind::Function,
533        "Struct" => ChunkKind::Struct,
534        "Impl" => ChunkKind::Impl,
535        "Module" => ChunkKind::Module,
536        "Class" => ChunkKind::Class,
537        "Method" => ChunkKind::Method,
538        "Issue" => ChunkKind::Issue,
539        "PullRequest" => ChunkKind::PullRequest,
540        "WikiPage" => ChunkKind::WikiPage,
541        "DbSchema" => ChunkKind::DbSchema,
542        "ApiEndpoint" => ChunkKind::ApiEndpoint,
543        "Ticket" => ChunkKind::Ticket,
544        "ExternalOther" => ChunkKind::ExternalOther,
545        _ => ChunkKind::Other,
546    }
547}
548
549#[cfg(test)]
550mod tests {
551    use super::*;
552    fn set_env(key: &str, value: Option<&str>) -> Option<String> {
553        let old = std::env::var(key).ok();
554        match value {
555            Some(v) => crate::test_env::set_var(key, v),
556            None => crate::test_env::remove_var(key),
557        }
558        old
559    }
560
561    fn restore_env(key: &str, old: Option<String>) {
562        match old {
563            Some(v) => crate::test_env::set_var(key, v),
564            None => crate::test_env::remove_var(key),
565        }
566    }
567
568    #[test]
569    fn dense_backend_defaults_to_local() {
570        let _g = crate::core::data_dir::test_env_lock();
571        let old_backend = set_env("LEANCTX_DENSE_BACKEND", None);
572        let old_url = set_env("LEANCTX_QDRANT_URL", None);
573        let old_pg = set_env("LEANCTX_PGVECTOR_URL", None);
574
575        let got = DenseBackendKind::try_from_env().unwrap();
576        assert_eq!(got, DenseBackendKind::Local);
577
578        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
579        restore_env("LEANCTX_QDRANT_URL", old_url);
580        restore_env("LEANCTX_PGVECTOR_URL", old_pg);
581    }
582
583    #[test]
584    fn dense_backend_unknown_value_errors() {
585        let _g = crate::core::data_dir::test_env_lock();
586        let old_backend = set_env("LEANCTX_DENSE_BACKEND", Some("wat"));
587        let old_url = set_env("LEANCTX_QDRANT_URL", None);
588
589        let err = DenseBackendKind::try_from_env().unwrap_err();
590        assert!(err.contains("Unknown LEANCTX_DENSE_BACKEND"));
591
592        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
593        restore_env("LEANCTX_QDRANT_URL", old_url);
594    }
595
596    #[cfg(feature = "qdrant")]
597    #[test]
598    fn dense_backend_infers_qdrant_from_url() {
599        let _g = crate::core::data_dir::test_env_lock();
600        let old_backend = set_env("LEANCTX_DENSE_BACKEND", None);
601        let old_url = set_env("LEANCTX_QDRANT_URL", Some("http://127.0.0.1:6333"));
602        let old_pg = set_env("LEANCTX_PGVECTOR_URL", None);
603
604        let got = DenseBackendKind::try_from_env().unwrap();
605        assert_eq!(got, DenseBackendKind::Qdrant);
606
607        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
608        restore_env("LEANCTX_QDRANT_URL", old_url);
609        restore_env("LEANCTX_PGVECTOR_URL", old_pg);
610    }
611
612    #[cfg(feature = "pgvector")]
613    #[test]
614    fn dense_backend_infers_pgvector_from_url() {
615        let _g = crate::core::data_dir::test_env_lock();
616        let old_backend = set_env("LEANCTX_DENSE_BACKEND", None);
617        let old_url = set_env("LEANCTX_QDRANT_URL", None);
618        let old_pg = set_env("LEANCTX_PGVECTOR_URL", Some("postgres://localhost/lctx"));
619
620        let got = DenseBackendKind::try_from_env().unwrap();
621        assert_eq!(got, DenseBackendKind::Pgvector);
622
623        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
624        restore_env("LEANCTX_QDRANT_URL", old_url);
625        restore_env("LEANCTX_PGVECTOR_URL", old_pg);
626    }
627
628    #[cfg(all(feature = "qdrant", feature = "pgvector"))]
629    #[test]
630    fn dense_backend_qdrant_wins_when_both_urls_set() {
631        let _g = crate::core::data_dir::test_env_lock();
632        let old_backend = set_env("LEANCTX_DENSE_BACKEND", None);
633        let old_url = set_env("LEANCTX_QDRANT_URL", Some("http://127.0.0.1:6333"));
634        let old_pg = set_env("LEANCTX_PGVECTOR_URL", Some("postgres://localhost/lctx"));
635
636        let got = DenseBackendKind::try_from_env().unwrap();
637        assert_eq!(got, DenseBackendKind::Qdrant);
638
639        // Explicit selection overrides inference order.
640        crate::test_env::set_var("LEANCTX_DENSE_BACKEND", "pgvector");
641        let got = DenseBackendKind::try_from_env().unwrap();
642        assert_eq!(got, DenseBackendKind::Pgvector);
643
644        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
645        restore_env("LEANCTX_QDRANT_URL", old_url);
646        restore_env("LEANCTX_PGVECTOR_URL", old_pg);
647    }
648
649    #[cfg(not(feature = "qdrant"))]
650    #[test]
651    fn dense_backend_qdrant_requires_feature() {
652        let _g = crate::core::data_dir::test_env_lock();
653        let old_backend = set_env("LEANCTX_DENSE_BACKEND", Some("qdrant"));
654        let old_url = set_env("LEANCTX_QDRANT_URL", None);
655
656        let err = DenseBackendKind::try_from_env().unwrap_err();
657        assert!(err.contains("feature 'qdrant' is not enabled"));
658
659        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
660        restore_env("LEANCTX_QDRANT_URL", old_url);
661    }
662
663    #[cfg(not(feature = "pgvector"))]
664    #[test]
665    fn dense_backend_pgvector_requires_feature() {
666        let _g = crate::core::data_dir::test_env_lock();
667        let old_backend = set_env("LEANCTX_DENSE_BACKEND", Some("pgvector"));
668        let old_url = set_env("LEANCTX_QDRANT_URL", None);
669
670        let err = DenseBackendKind::try_from_env().unwrap_err();
671        assert!(err.contains("feature 'pgvector' is not enabled"));
672
673        restore_env("LEANCTX_DENSE_BACKEND", old_backend);
674        restore_env("LEANCTX_QDRANT_URL", old_url);
675    }
676}
lean_ctx/core/dense_backend.rs

lean_ctx/core/
dense_backend.rs
