Skip to main content

gobby_code/search/
semantic.rs

1//! Compatibility wrapper for Qdrant vector search.
2//!
3//! Reusable vector projection behavior lives in `crate::vector::code_symbols`.
4
5pub use crate::vector::code_symbols::{embed_query_with_source, vector_search};
6
7use crate::config::{CODE_SYMBOL_COLLECTION_PREFIX, Context};
8use crate::visibility;
9use gobby_core::qdrant::{CollectionScope, SearchRequest};
10use rayon::prelude::*;
11
12pub fn semantic_search(ctx: &Context, query: &str, limit: usize) -> Vec<(String, f64)> {
13    let project_ids = visibility::visible_project_ids(ctx);
14    let Some(per_project_limit) = per_project_semantic_limit(limit, project_ids.len()) else {
15        return vec![];
16    };
17    let Some(embedding_source) = crate::vector::code_symbols::embedding_source_from_context(ctx)
18    else {
19        return vec![];
20    };
21    let Some(query_vector) = embed_query_with_source(&embedding_source, query) else {
22        return vec![];
23    };
24
25    let mut results = project_ids
26        .par_iter()
27        .flat_map(|project_id| {
28            let collection = gobby_core::qdrant::collection_name(
29                "gcode",
30                CollectionScope::Custom(&format!("{CODE_SYMBOL_COLLECTION_PREFIX}{project_id}")),
31            );
32            let request = SearchRequest {
33                vector: query_vector.clone(),
34                limit: per_project_limit,
35                filter: None,
36            };
37
38            match gobby_core::qdrant::with_qdrant(ctx.qdrant.as_ref(), Vec::new(), |config| {
39                gobby_core::qdrant::search(config, &collection, request)
40            }) {
41                Ok((hits, _state)) => hits
42                    .into_iter()
43                    .map(|hit| (hit.id, f64::from(hit.score)))
44                    .collect::<Vec<_>>(),
45                Err(error) => {
46                    log::warn!(
47                        "semantic Qdrant search failed for collection {collection}: {error}"
48                    );
49                    Vec::new()
50                }
51            }
52        })
53        .collect::<Vec<_>>();
54    results.sort_by(|a, b| b.1.total_cmp(&a.1));
55    results.truncate(limit);
56    results
57}
58
59/// Fetch the requested limit from each visible project, then globally truncate.
60///
61/// Keeping the exact per-project limit preserves single-project behavior while
62/// giving overlay searches enough candidates to merge across visible projects.
63fn per_project_semantic_limit(limit: usize, project_count: usize) -> Option<usize> {
64    if limit == 0 || project_count == 0 {
65        return None;
66    }
67    Some(limit)
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use crate::config::QdrantConfig;
74    use std::path::PathBuf;
75
76    fn make_ctx_no_qdrant() -> Context {
77        Context {
78            database_url: "postgresql://localhost/nonexistent".to_string(),
79            project_root: PathBuf::from("/nonexistent"),
80            project_id: "test".to_string(),
81            quiet: true,
82            falkordb: None,
83            qdrant: None,
84            embedding: None,
85            code_vectors: crate::config::CodeVectorSettings::default(),
86            daemon_url: None,
87            index_scope: crate::config::ProjectIndexScope::Single,
88        }
89    }
90
91    #[test]
92    fn test_semantic_search_no_qdrant() {
93        let ctx = make_ctx_no_qdrant();
94        let result = semantic_search(&ctx, "test query", 10);
95        assert!(result.is_empty());
96    }
97
98    #[test]
99    fn test_semantic_search_no_embedding_config() {
100        let ctx = Context {
101            qdrant: Some(QdrantConfig {
102                url: Some("http://localhost:6333".to_string()),
103                api_key: None,
104            }),
105            ..make_ctx_no_qdrant()
106        };
107        let result = semantic_search(&ctx, "test query", 10);
108        assert!(result.is_empty());
109    }
110
111    #[test]
112    fn per_project_limit_over_fetches_each_visible_project() {
113        assert_eq!(per_project_semantic_limit(10, 2), Some(10));
114        assert_eq!(per_project_semantic_limit(11, 2), Some(11));
115        assert_eq!(per_project_semantic_limit(1, 2), Some(1));
116    }
117
118    #[test]
119    fn per_project_limit_handles_empty_work() {
120        assert_eq!(per_project_semantic_limit(0, 2), None);
121        assert_eq!(per_project_semantic_limit(10, 0), None);
122    }
123}