searchi/crates/searchi-service/tests/early_text_google_cpu.rs

//! 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(())
}