feat: add Google code retrieval across desktop CLI and MCP

crates/searchi-service/src/lib.rs

@@ -54,6 +54,12 @@ pub trait EmbeddingProvider: Send {
     fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
     fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>>;
     fn embed_query(&mut self, query: &str) -> Result<Vec<f32>>;
+    fn embed_code_query(&mut self, _query: &str) -> Result<Vec<f32>> {
+        anyhow::bail!("Provider does not support Google code retrieval")
+    }
+    fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.embed_code_query(query)
+    }
     fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
         self.embed_query(query)
     }
@@ -90,6 +96,12 @@ impl EmbeddingProvider for GemmaEmbedder {
     fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
         GemmaEmbedder::embed_query_debug(self, query)
     }
+    fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        GemmaEmbedder::embed_code_query(self, query)
+    }
+    fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
+        GemmaEmbedder::embed_code_query_debug(self, query)
+    }
     fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
         GemmaEmbedder::query_diagnostics(self)
     }
@@ -132,6 +144,11 @@ pub struct EmbeddingStatus {
     pub pending: usize,
 }
 
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum QueryTask {
+    Search,
+    Code,
+}
 /// Shared local hybrid service. Google EmbeddingGemma 2 is the default and
 /// only shipped provider; no alternative model is silently substituted.
 pub struct HybridSearch {
@@ -142,7 +159,7 @@ pub struct HybridSearch {
     // Real provider outputs only, scoped to this model session. Queries never
     // persist and results/filters are never cached here.
     query_cache_fingerprint: String,
-    query_cache: VecDeque<(String, Vec<f32>)>,
+    query_cache: VecDeque<(QueryTask, String, Vec<f32>)>,
 }
 
 impl HybridSearch {
@@ -396,7 +413,7 @@ impl HybridSearch {
         trace.timed("hybrid.identifier_priority", parent, |_, _| {
             promote_identifiers(query.trim(), &mut lexical)
         });
-        let vector = self.query_vector(query, trace, parent)?;
+        let vector = self.query_vector(query, options.code_only, trace, parent)?;
         let semantic =
             self.routed_semantic_traced(query, &vector, candidates, options, trace, parent)?;
         let results = trace.timed("hybrid.rank_fusion", parent, |trace, span| {
@@ -419,13 +436,14 @@ impl HybridSearch {
             return Ok(Vec::new());
         }
         let mut trace = TraceCollector::new(MetricsMode::Basic);
-        let vector = self.query_vector(query, &mut trace, None)?;
+        let vector = self.query_vector(query, options.code_only, &mut trace, None)?;
         self.routed_semantic(query, &vector, limit.min(1000), options)
     }
 
     fn query_vector(
         &mut self,
         query: &str,
+        code: bool,
         trace: &mut TraceCollector,
         parent: Option<SpanToken>,
     ) -> Result<Vec<f32>> {
@@ -435,6 +453,12 @@ impl HybridSearch {
             "Query must be nonempty and at most 4096 UTF-8 bytes"
         );
         let span = trace.begin("semantic.query_embedding", parent);
+        let task = if code {
+            QueryTask::Code
+        } else {
+            QueryTask::Search
+        };
+        trace.counter(span, "code_task", u64::from(code));
         if self.query_cache_fingerprint != self.embedder.fingerprint() {
             self.query_cache.clear();
             self.query_cache_fingerprint = self.embedder.fingerprint().to_owned();
@@ -442,13 +466,13 @@ impl HybridSearch {
         if let Some(position) = self
             .query_cache
             .iter()
-            .position(|(cached, _)| cached == query)
+            .position(|(cached_task, cached, _)| *cached_task == task && cached == query)
         {
             let entry = self
                 .query_cache
                 .remove(position)
                 .expect("Located cache entry");
-            let vector = entry.1.clone();
+            let vector = entry.2.clone();
             self.query_cache.push_back(entry);
             trace.counter(span, "query_cache_hit", 1);
             trace.counter(span, "query_cache_miss", 0);
@@ -460,10 +484,11 @@ impl HybridSearch {
         trace.counter(span, "query_cache_hit", 0);
         trace.counter(span, "query_cache_miss", 1);
         let vector = (|| {
-            let vector = if trace.is_debug() {
-                self.embedder.embed_query_debug(query)?
-            } else {
-                self.embedder.embed_query(query)?
+            let vector = match (code, trace.is_debug()) {
+                (true, true) => self.embedder.embed_code_query_debug(query)?,
+                (true, false) => self.embedder.embed_code_query(query)?,
+                (false, true) => self.embedder.embed_query_debug(query)?,
+                (false, false) => self.embedder.embed_query(query)?,
             };
             ensure!(
                 vector.len() == self.embedder.dimensions()
@@ -487,7 +512,7 @@ impl HybridSearch {
             self.query_cache.pop_front();
         }
         self.query_cache
-            .push_back((query.to_owned(), vector.clone()));
+            .push_back((task, query.to_owned(), vector.clone()));
         Ok(vector)
     }
 }