gobby_code/search/
semantic.rs1pub 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
59fn 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}