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 {

crates/searchi-service/tests/query_cache.rs

@@ -0,0 +1,208 @@
+//! Cache-control tests count provider calls; they do not measure model relevance.
+//! Production still loads Google and root runs separate real CPU benchmarks.
+use anyhow::Result;
+use searchi_core::{SearchIndex, SearchOptions};
+use searchi_semantic::TextWindow;
+use searchi_service::{EmbeddingProvider, HybridSearch};
+use std::{
+    collections::BTreeMap,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc, Mutex,
+    },
+};
+
+#[derive(Default)]
+struct Calls {
+    queries: Mutex<Vec<String>>,
+    fail: AtomicBool,
+    invalid: AtomicBool,
+    new_fingerprint: AtomicBool,
+}
+struct CountingProvider(Arc<Calls>);
+impl EmbeddingProvider for CountingProvider {
+    fn fingerprint(&self) -> &str {
+        if self.0.new_fingerprint.load(Ordering::Relaxed) {
+            "cache-control-v2"
+        } else {
+            "cache-control-v1"
+        }
+    }
+    fn dimensions(&self) -> usize {
+        768
+    }
+    fn split_windows(&mut self, _: &str, _: &str) -> Result<Vec<TextWindow>> {
+        anyhow::bail!("Not used by cache-control tests")
+    }
+    fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
+        anyhow::bail!("Not used by cache-control tests")
+    }
+    fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.0.queries.lock().unwrap().push(query.into());
+        if self.0.fail.load(Ordering::Relaxed) {
+            anyhow::bail!("Injected provider failure");
+        }
+        let mut vector = vec![0.0; 768];
+        if !self.0.invalid.load(Ordering::Relaxed) {
+            vector[0] = 1.0;
+        }
+        Ok(vector)
+    }
+    fn query_diagnostics(&self) -> BTreeMap<String, u64> {
+        BTreeMap::from([("worker_inference_us".into(), 123)])
+    }
+}
+fn fixture() -> Result<(tempfile::TempDir, HybridSearch, Arc<Calls>)> {
+    let directory = tempfile::tempdir()?;
+    let database = directory.path().join("index.db");
+    let calls = Arc::new(Calls::default());
+    let service = HybridSearch::with_provider(database, Box::new(CountingProvider(calls.clone())))?;
+    Ok((directory, service, calls))
+}
+fn count(calls: &Calls) -> usize {
+    calls.queries.lock().unwrap().len()
+}
+
+#[test]
+fn exact_trimmed_queries_reuse_vectors_but_case_and_model_fingerprint_do_not() -> Result<()> {
+    let (_directory, mut service, calls) = fixture()?;
+    let options = SearchOptions::default();
+    service.search("  supplier  ", 5, &options)?;
+    service.search("supplier", 5, &options)?;
+    assert_eq!(*calls.queries.lock().unwrap(), vec!["supplier"]);
+    service.search("Supplier", 5, &options)?;
+    assert_eq!(count(&calls), 2);
+    calls.new_fingerprint.store(true, Ordering::Relaxed);
+    service.search("supplier", 5, &options)?;
+    assert_eq!(count(&calls), 3);
+    Ok(())
+}
+
+#[test]
+fn filters_and_changed_index_are_recomputed_even_when_vector_is_reused() -> Result<()> {
+    let (directory, mut service, calls) = fixture()?;
+    std::fs::write(directory.path().join("first.txt"), "supplier notice")?;
+    std::fs::write(directory.path().join("second.md"), "supplier notice")?;
+    let mut index = SearchIndex::open(directory.path().join("index.db"))?;
+    index.index_folder(directory.path())?;
+    let txt = SearchOptions {
+        extension: Some("txt".into()),
+        ..Default::default()
+    };
+    let md = SearchOptions {
+        extension: Some("md".into()),
+        ..Default::default()
+    };
+    assert_eq!(
+        service.search("supplier", 5, &txt)?[0].filename,
+        "first.txt"
+    );
+    assert_eq!(service.search("supplier", 5, &md)?[0].filename, "second.md");
+    assert_eq!(count(&calls), 1);
+    std::fs::remove_file(directory.path().join("first.txt"))?;
+    index.index_folder(directory.path())?;
+    assert!(service.search("supplier", 5, &txt)?.is_empty());
+    assert_eq!(count(&calls), 1);
+    Ok(())
+}
+
+#[test]
+fn cache_hit_trace_never_replays_prior_worker_diagnostics() -> Result<()> {
+    let (_directory, mut service, _calls) = fixture()?;
+    let miss = service.search_measured("supplier", 5, &SearchOptions::default(), true)?;
+    let hit = service.search_measured("supplier", 5, &SearchOptions::default(), true)?;
+    let miss = miss
+        .metrics
+        .trace
+        .unwrap()
+        .into_iter()
+        .find(|span| span.name == "semantic.query_embedding")
+        .unwrap();
+    let hit = hit
+        .metrics
+        .trace
+        .unwrap()
+        .into_iter()
+        .find(|span| span.name == "semantic.query_embedding")
+        .unwrap();
+    assert_eq!(miss.counters["query_cache_miss"], 1);
+    assert_eq!(miss.counters["worker_inference_us"], 123);
+    assert_eq!(hit.counters["query_cache_hit"], 1);
+    assert_eq!(hit.counters["query_cache_miss"], 0);
+    assert!(!hit.counters.contains_key("worker_inference_us"));
+    assert_eq!(hit.counters["failed"], 0);
+    Ok(())
+}
+
+#[test]
+fn failures_and_invalid_vectors_are_not_cached() -> Result<()> {
+    let (_directory, mut service, calls) = fixture()?;
+    assert!(service
+        .search_semantic(&"界".repeat(2000), 5, &SearchOptions::default())
+        .is_err());
+    assert_eq!(
+        count(&calls),
+        0,
+        "Oversized query must not reach the provider or cache"
+    );
+    calls.fail.store(true, Ordering::Relaxed);
+    for _ in 0..2 {
+        assert!(service
+            .search("supplier", 5, &SearchOptions::default())
+            .is_err());
+    }
+    calls.fail.store(false, Ordering::Relaxed);
+    calls.invalid.store(true, Ordering::Relaxed);
+    for _ in 0..2 {
+        assert!(service
+            .search("supplier", 5, &SearchOptions::default())
+            .is_err());
+    }
+    calls.invalid.store(false, Ordering::Relaxed);
+    service.search("supplier", 5, &SearchOptions::default())?;
+    service.search("supplier", 5, &SearchOptions::default())?;
+    assert_eq!(count(&calls), 5);
+    Ok(())
+}
+
+#[test]
+fn cache_evicts_least_recently_used_query_at_32_entries() -> Result<()> {
+    let (_directory, mut service, calls) = fixture()?;
+    for number in 0..32 {
+        service.search(&format!("query{number}"), 5, &SearchOptions::default())?;
+    }
+    service.search("query0", 5, &SearchOptions::default())?;
+    service.search("query32", 5, &SearchOptions::default())?;
+    service.search("query0", 5, &SearchOptions::default())?;
+    assert_eq!(count(&calls), 33);
+    service.search("query1", 5, &SearchOptions::default())?;
+    assert_eq!(count(&calls), 34);
+    Ok(())
+}
+
+#[test]
+fn warmup_bypasses_cache_and_never_seeds_an_unseen_measured_query() -> Result<()> {
+    let (_directory, mut service, calls) = fixture()?;
+    service.warmup()?;
+    assert_eq!(count(&calls), 1);
+    let first = service.search_measured(
+        "warm up local file search",
+        5,
+        &SearchOptions::default(),
+        true,
+    )?;
+    assert_eq!(count(&calls), 2);
+    let span = first
+        .metrics
+        .trace
+        .unwrap()
+        .into_iter()
+        .find(|span| span.name == "semantic.query_embedding")
+        .unwrap();
+    assert_eq!(span.counters["query_cache_miss"], 1);
+    service.warmup()?;
+    assert_eq!(count(&calls), 3);
+    service.search("warm up local file search", 5, &SearchOptions::default())?;
+    assert_eq!(count(&calls), 3);
+    Ok(())
+}