Use measured query-only context across desktop CLI and MCP

crates/searchi-semantic/tests/query_context_cpu.rs

@@ -0,0 +1,119 @@
+use anyhow::Result;
+use searchi_semantic::{Device, GemmaEmbedder};
+
+fn cosine(a: &[f32], b: &[f32]) -> f64 {
+    a.iter()
+        .zip(b)
+        .map(|(a, b)| f64::from(*a) * f64::from(*b))
+        .sum()
+}
+
+#[test]
+#[ignore = "Requires verified Google model and current CPU-capable native worker"]
+fn query_context_preserves_real_vectors_ranking_and_token_boundary() -> Result<()> {
+    let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
+    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
+    let mut full = GemmaEmbedder::load_with_device(&model, &worker, 1, Device::Cpu)?;
+    let mut query = GemmaEmbedder::load_query_with_device(&model, &worker, 1, Device::Cpu)?;
+    assert_eq!(full.info.context_tokens, 2048);
+    assert_eq!(query.info.context_tokens, 512);
+    assert_eq!(query.info.role, "query");
+    // This describes the binary's media capability, not this session's role.
+    assert!(query.info.media_supported);
+    let documents = full.embed_documents(
+        &[
+            "The supplier agreement may be terminated with sixty days written notice.",
+            "Business train tickets and hotel stays are reimbursed after submitting receipts.",
+            "Die Gartenplanung sieht ein Hochbeet und eine Regentonne vor.",
+        ],
+        "none",
+    )?;
+    for (text, expected) in [
+        ("Wie beenden wir die Zusammenarbeit mit dem Lieferanten?", 0),
+        ("How do I claim my accommodation and railway expenses?", 1),
+        (
+            "Where is the plan for the vegetable bed and rainwater barrel?",
+            2,
+        ),
+    ] {
+        let before = full.embed_query(text)?;
+        let after = query.embed_query(text)?;
+        let parity = cosine(&before, &after);
+        assert!(parity >= 0.99999, "Context changed vector: {parity}");
+        let rank = |vector: &[f32]| {
+            documents
+                .iter()
+                .enumerate()
+                .max_by(|a, b| cosine(a.1, vector).total_cmp(&cosine(b.1, vector)))
+                .unwrap()
+                .0
+        };
+        assert_eq!(rank(&before), expected);
+        assert_eq!(rank(&after), expected);
+        println!("Query role parity {parity:.8}; expected document {expected}");
+    }
+    for text in [
+        "retry failed requests with exponential backoff",
+        "Zugriffstoken bei Ablauf erneuern",
+        "متى تغلق المكتبة",
+        "ইলিশ মাছ রান্না",
+    ] {
+        assert!(
+            cosine(
+                &full.embed_code_query(text)?,
+                &query.embed_code_query(text)?
+            ) >= 0.99999
+        );
+    }
+    // Actual token counts determine the exact boundary, rather than assuming
+    // words equal tokens. Each repeated plain token contributes one token.
+    let mut repetitions = 490usize;
+    let mut boundary = String::new();
+    for _ in 0..8 {
+        let text = " cat".repeat(repetitions);
+        full.embed_query_debug(&text)?;
+        let tokens = full.query_diagnostics()["worker_token_count"] as usize;
+        if tokens == 512 {
+            boundary = text;
+            break;
+        }
+        assert!(tokens < 512, "Unexpected fixture token count {tokens}");
+        repetitions += 512 - tokens;
+    }
+    assert!(
+        !boundary.is_empty(),
+        "Could not construct exact512token query"
+    );
+    let before = full.embed_query(&boundary)?;
+    let after = query.embed_query_debug(&boundary)?;
+    assert_eq!(query.query_diagnostics()["worker_token_count"], 512);
+    assert!(cosine(&before, &after) >= 0.99999);
+    assert!(query.embed_query(&(boundary + " cat")).is_err());
+    assert!(query
+        .embed_query("supplier cancellation")?
+        .iter()
+        .all(|v| v.is_finite()));
+    assert!(query
+        .embed_documents(&["document"], "none")
+        .unwrap_err()
+        .to_string()
+        .contains("Foreground query worker"));
+    assert!(query.split_windows("document", "none").is_err());
+    assert!(query
+        .embed_audio_pcm(&[0.0; 16])
+        .unwrap_err()
+        .to_string()
+        .contains("Foreground query worker"));
+    assert!(query
+        .embed_images(&[std::path::Path::new("missing.png")])
+        .unwrap_err()
+        .to_string()
+        .contains("Foreground query worker"));
+    // Full2048-token indexing survives, including long document extraction.
+    assert!(
+        full.split_windows(&"ordinary maintenance log ".repeat(1000), "log")?
+            .len()
+            > 1
+    );
+    Ok(())
+}