Prepare text search during media setup and retain query sessions through promotion

crates/searchi-service/tests/early_text_google_cpu.rs

@@ -0,0 +1,197 @@
+//! Genuine Google CPU proof: retained query session sees early text and full promotion.
+use anyhow::Result;
+use searchi_core::SearchIndex;
+use searchi_semantic::{Device, GemmaEmbedder, TextWindow, WorkerInfo};
+use searchi_service::{
+    EmbeddingProvider, FfmpegTools, HybridSearch, MediaLimits, PreparedGraphs, SearchOptions,
+    SemanticCoverage,
+};
+use std::{
+    collections::BTreeMap,
+    fs,
+    path::{Path, PathBuf},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+
+struct CountedQuery {
+    google: GemmaEmbedder,
+    calls: Arc<AtomicUsize>,
+}
+impl EmbeddingProvider for CountedQuery {
+    fn fingerprint(&self) -> &str {
+        EmbeddingProvider::fingerprint(&self.google)
+    }
+    fn dimensions(&self) -> usize {
+        EmbeddingProvider::dimensions(&self.google)
+    }
+    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
+        self.google.split_windows(text, title)
+    }
+    fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
+        self.google.embed_documents(texts, title)
+    }
+    fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.calls.fetch_add(1, Ordering::Relaxed);
+        self.google.embed_query(query)
+    }
+    fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.calls.fetch_add(1, Ordering::Relaxed);
+        self.google.embed_query_debug(query)
+    }
+    fn query_diagnostics(&self) -> BTreeMap<String, u64> {
+        self.google.query_diagnostics()
+    }
+    fn worker_info(&self) -> Option<WorkerInfo> {
+        Some(self.google.info.clone())
+    }
+}
+
+fn text_vectors(db: &Path) -> Result<Vec<(String, Vec<u8>)>> {
+    let connection = rusqlite::Connection::open(db)?;
+    let mut statement=connection.prepare("SELECT path,vector FROM semantic_chunks WHERE generation=(SELECT active_generation FROM semantic_meta WHERE singleton=1) AND source_type='text' ORDER BY path,ordinal")?;
+    let rows = statement
+        .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
+        .collect::<rusqlite::Result<_>>()?;
+    Ok(rows)
+}
+
+fn assert_retained_trace(response: &searchi_service::SearchResponse<searchi_core::SearchResult>) {
+    let trace = response.metrics.trace.as_ref().expect("Debug trace");
+    assert!(trace
+        .iter()
+        .any(|span| span.counters.get("query_cache_hit") == Some(&1)));
+    assert!(
+        trace
+            .iter()
+            .any(|span| span.name == "semantic.ann_load"
+                && span.counters.get("cache_hit") == Some(&1))
+    );
+    assert!(!trace.iter().any(|span| matches!(
+        span.name.as_str(),
+        "semantic.ann_blob_read" | "semantic.ann_checksum" | "semantic.ann_deserialize"
+    )));
+}
+
+#[test]
+#[ignore = "Requires verified Google pair, normal worker and CPU FFmpeg helpers"]
+fn retained_query_reader_follows_early_text_and_full_media_without_reencoding_documents(
+) -> Result<()> {
+    assert_eq!(std::env::var("SEARCHI_FORCE_CPU")?, "1");
+    let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Verified models");
+    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Normal worker");
+    let tools = FfmpegTools {
+        ffmpeg: PathBuf::from(std::env::var_os("SEARCHI_TEST_FFMPEG").expect("CPU ffmpeg")),
+        ffprobe: PathBuf::from(std::env::var_os("SEARCHI_TEST_FFPROBE").expect("CPU ffprobe")),
+    };
+    let temp = tempfile::tempdir()?;
+    let root = temp.path().join("corpus");
+    fs::create_dir(&root)?;
+    let supplier = root.join("supplier.txt");
+    fs::write(
+        &supplier,
+        "The supplier agreement may be terminated with sixty days written notice.",
+    )?;
+    fs::write(
+        root.join("garden.txt"),
+        "A vegetable bed and rainwater barrel are planned for the garden.",
+    )?;
+    fs::copy(
+        Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("../searchi-semantic/tests/fixtures/digup-photo.jpg"),
+        root.join("photo.jpg"),
+    )?;
+    let db = temp.path().join("index.db");
+    let mut lexical = SearchIndex::open(&db)?;
+    lexical.index_folder(&root)?;
+    let graphs = PreparedGraphs::default();
+    let calls = Arc::new(AtomicUsize::new(0));
+    let google = GemmaEmbedder::load_query_with_device(&models, &worker, 1, Device::Cpu)?;
+    let mut reader = HybridSearch::with_provider(
+        &db,
+        Box::new(CountedQuery {
+            google,
+            calls: calls.clone(),
+        }),
+    )?;
+    reader.share_prepared_graphs(&graphs)?;
+    reader.warmup()?; // No active generation yet.
+    let query = "Wie beenden wir die Zusammenarbeit mit dem Lieferanten?";
+    assert!(reader
+        .search(query, 5, &SearchOptions::default())?
+        .is_empty());
+    assert_eq!(calls.load(Ordering::Relaxed), 2);
+    let mut publisher = HybridSearch::open(&db, &models, &worker, 1)?;
+    publisher.share_prepared_graphs(&graphs)?;
+    let text = publisher
+        .index_text_embeddings_preserving_media()?
+        .expect("Initial text build");
+    assert_eq!(text.document_inference_calls, 2);
+    assert_eq!(text.media_files, 0);
+    let text_status = HybridSearch::embedding_status_for_coverage(&db, SemanticCoverage::TextOnly)?;
+    assert!(!text_status.stale);
+    assert_eq!(text_status.pending, 0);
+    assert_eq!(text_status.coverage, Some(SemanticCoverage::TextOnly));
+    assert!(HybridSearch::embedding_status(&db)?.stale);
+    let early = reader.search_measured(query, 5, &SearchOptions::default(), true)?;
+    assert_eq!(early.results[0].filename, "supplier.txt");
+    assert!(early.results[0].snippet.contains("sixty days"));
+    assert_eq!(early.results[0].source.as_ref().unwrap().kind, "line");
+    assert_retained_trace(&early);
+    assert_eq!(calls.load(Ordering::Relaxed), 2);
+    let before = text_vectors(&db)?;
+    let full = publisher.index_embeddings_with_media(&tools, &MediaLimits::default())?;
+    assert_eq!(
+        full.document_inference_calls, 0,
+        "Promotion reuses text's identical encoder/window policy"
+    );
+    assert_eq!(full.image_inference_calls, 1);
+    assert_eq!(
+        text_vectors(&db)?,
+        before,
+        "Document vector bytes/citations survive promotion"
+    );
+    let full_status = HybridSearch::embedding_status(&db)?;
+    assert_eq!(full_status.coverage, Some(SemanticCoverage::FullMedia));
+    assert!(!full_status.stale);
+    assert!(!full_status.text_indexing_allowed);
+    let promoted = reader.search_measured(query, 5, &SearchOptions::default(), true)?;
+    assert_eq!(promoted.results[0].filename, "supplier.txt");
+    assert_retained_trace(&promoted);
+    assert_eq!(
+        calls.load(Ordering::Relaxed),
+        2,
+        "No repeated warmup, restart or cached query inference"
+    );
+    assert_eq!(reader.worker_info().unwrap().role, "query");
+    assert!(publisher
+        .index_text_embeddings_preserving_media()?
+        .is_none());
+    assert_eq!(
+        HybridSearch::embedding_status(&db)?.generation,
+        Some(full.generation)
+    );
+    fs::write(
+        &supplier,
+        "An updated supplier agreement has ninety days written notice.",
+    )?;
+    lexical.index_folder(&root)?;
+    let stale = reader.search_measured(query, 5, &SearchOptions::default(), true)?;
+    assert!(
+        stale
+            .results
+            .iter()
+            .all(|result| result.filename != "supplier.txt"),
+        "Old extraction/semantic chunks stay suppressed"
+    );
+    assert_eq!(calls.load(Ordering::Relaxed), 2);
+    assert!(
+        publisher
+            .index_text_embeddings_preserving_media()?
+            .is_none(),
+        "Projector repair must not downgrade stale full generation"
+    );
+    Ok(())
+}