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