perf: audit EmbeddingGemma 2 and reuse measured query embeddings

crates/searchi-service/src/lib.rs

@@ -13,10 +13,14 @@ pub use searchi_core::{SearchMetrics, SearchOptions, SearchResponse, SearchResul
 pub use searchi_semantic::WorkerInfo;
 use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
 use serde::Serialize;
-use std::{collections::HashMap, path::Path};
+use std::{
+    collections::{HashMap, VecDeque},
+    path::Path,
+};
 use store::{Store, MAX_SCAN_CHUNKS};
 
 const MAX_BATCH: usize = 8;
+const QUERY_CACHE_ENTRIES: usize = 32;
 const BYTE_SLICE: usize = 64 * 1024;
 const BYTE_OVERLAP: usize = 512;
 const MAX_PASSAGE_WINDOWS: usize = 2048;
@@ -135,6 +139,10 @@ pub struct HybridSearch {
     embedder: Box<dyn EmbeddingProvider>,
     store: Store,
     generation_fingerprint: String,
+    // 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>)>,
 }
 
 impl HybridSearch {
@@ -237,6 +245,8 @@ impl HybridSearch {
             _ => expected,
         };
         Ok(Self {
+            query_cache_fingerprint: embedder.fingerprint().to_owned(),
+            query_cache: VecDeque::new(),
             lexical,
             embedder,
             store,
@@ -386,22 +396,7 @@ impl HybridSearch {
         trace.timed("hybrid.identifier_priority", parent, |_, _| {
             promote_identifiers(query.trim(), &mut lexical)
         });
-        let embedding_span = trace.begin("semantic.query_embedding", parent);
-        let vector = if trace.is_debug() {
-            self.embedder.embed_query_debug(query.trim())
-        } else {
-            self.embedder.embed_query(query.trim())
-        };
-        trace.counter(embedding_span, "failed", u64::from(vector.is_err()));
-        trace.end(
-            embedding_span,
-            if trace.is_debug() {
-                self.embedder.query_diagnostics()
-            } else {
-                Default::default()
-            },
-        );
-        let vector = vector?;
+        let vector = self.query_vector(query, trace, parent)?;
         let semantic =
             self.routed_semantic_traced(query, &vector, candidates, options, trace, parent)?;
         let results = trace.timed("hybrid.rank_fusion", parent, |trace, span| {
@@ -423,9 +418,78 @@ impl HybridSearch {
         if query.trim().is_empty() || limit == 0 {
             return Ok(Vec::new());
         }
-        let vector = self.embedder.embed_query(query.trim())?;
+        let mut trace = TraceCollector::new(MetricsMode::Basic);
+        let vector = self.query_vector(query, &mut trace, None)?;
         self.routed_semantic(query, &vector, limit.min(1000), options)
     }
+
+    fn query_vector(
+        &mut self,
+        query: &str,
+        trace: &mut TraceCollector,
+        parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        let query = query.trim();
+        ensure!(
+            !query.is_empty() && query.len() <= 4096,
+            "Query must be nonempty and at most 4096 UTF-8 bytes"
+        );
+        let span = trace.begin("semantic.query_embedding", parent);
+        if self.query_cache_fingerprint != self.embedder.fingerprint() {
+            self.query_cache.clear();
+            self.query_cache_fingerprint = self.embedder.fingerprint().to_owned();
+        }
+        if let Some(position) = self
+            .query_cache
+            .iter()
+            .position(|(cached, _)| cached == query)
+        {
+            let entry = self
+                .query_cache
+                .remove(position)
+                .expect("Located cache entry");
+            let vector = entry.1.clone();
+            self.query_cache.push_back(entry);
+            trace.counter(span, "query_cache_hit", 1);
+            trace.counter(span, "query_cache_miss", 0);
+            trace.counter(span, "failed", 0);
+            // No worker was called: never attach its previous diagnostics.
+            trace.end(span, Default::default());
+            return Ok(vector);
+        }
+        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)?
+            };
+            ensure!(
+                vector.len() == self.embedder.dimensions()
+                    && vector.iter().all(|value| value.is_finite())
+                    && vector.iter().any(|value| *value != 0.0),
+                "Provider returned an invalid query vector"
+            );
+            Ok(vector)
+        })();
+        trace.counter(span, "failed", u64::from(vector.is_err()));
+        trace.end(
+            span,
+            if trace.is_debug() {
+                self.embedder.query_diagnostics()
+            } else {
+                Default::default()
+            },
+        );
+        let vector = vector?;
+        if self.query_cache.len() == QUERY_CACHE_ENTRIES {
+            self.query_cache.pop_front();
+        }
+        self.query_cache
+            .push_back((query.to_owned(), vector.clone()));
+        Ok(vector)
+    }
 }
 
 fn bounded_title(title: &str) -> &str {