diff --git a/AGENTS.md b/AGENTS.md index bc0d7d4..606a91e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -51,6 +51,13 @@ readiness without completed release gates and recorded evidence. thread; yield between document encodes. Trace admission separately from native inference. This does not bound decoder, extraction or ANN CPU work. Do not claim preemption of an already executing media encode on a one-thread budget. +- Foreground desktop, CLI and MCP sessions use the normal worker's query role + with a 512-token context; document/media indexing retains 2048 tokens. Preserve + role on CPU restart and reject indexing operations before projector work. + `media_supported` describes compiled capability, not a loaded projector or + permission to index in a query session. Legacy workers need an explicit update, + never a silent larger-context fallback. Keep private commitment, resident memory, + model verification and cached startup measurements distinct. - Treat native filesystem events as bounded hints; retain work across busy scans and verify periodically. Use targeted, cancellable reconciliation for events. Media hints must advance durable change revisions even with identical stats; diff --git a/apps/desktop/src-tauri/src/semantic_commands.rs b/apps/desktop/src-tauri/src/semantic_commands.rs index 21dec5e..37e4ec2 100644 --- a/apps/desktop/src-tauri/src/semantic_commands.rs +++ b/apps/desktop/src-tauri/src/semantic_commands.rs @@ -109,7 +109,7 @@ impl SemanticState { .acquire(InferenceRole::Foreground, threads) .map_err(|error| error.to_string())?; let mut service = - HybridSearch::open(database, &self.model_directory, &self.worker, threads) + HybridSearch::open_query(database, &self.model_directory, &self.worker, threads) .map_err(|error| error.to_string())?; drop(admission); service = service.with_budget(&self.inference_budget, InferenceRole::Foreground, threads); @@ -396,7 +396,7 @@ pub async fn hybrid_search( .inference_budget .acquire(InferenceRole::Foreground, threads) .map_err(|error| error.to_string())?; - let session = HybridSearch::open( + let session = HybridSearch::open_query( database, &semantic.model_directory, &semantic.worker, @@ -480,7 +480,7 @@ pub async fn hybrid_search_measured( .acquire(InferenceRole::Foreground, threads) .map_err(|error| error.to_string())?; let session = trace.timed("desktop.session_initialize", None, |_, _| { - HybridSearch::open( + HybridSearch::open_query( database, &semantic.model_directory, &semantic.worker, diff --git a/crates/searchi-cli/src/main.rs b/crates/searchi-cli/src/main.rs index 25023e7..bd87f65 100644 --- a/crates/searchi-cli/src/main.rs +++ b/crates/searchi-cli/src/main.rs @@ -203,7 +203,7 @@ fn main() -> Result<()> { }); let results = if let (Some(model), Some(runtime)) = (model_dir, embedding_worker) { let mut service = trace.timed("cli.session_initialize", None, |_, _| { - HybridSearch::open(&db, model, runtime, usize::from(cpu_threads)) + HybridSearch::open_query(&db, model, runtime, usize::from(cpu_threads)) })?; service .search_traced(&query, usize::from(limit), &options, &mut trace, None) diff --git a/crates/searchi-cli/src/mcp.rs b/crates/searchi-cli/src/mcp.rs index d215992..97b5db5 100644 --- a/crates/searchi-cli/src/mcp.rs +++ b/crates/searchi-cli/src/mcp.rs @@ -450,7 +450,7 @@ pub fn serve( "No semantic generation; run explicit indexing/embedding outside MCP first" ); - let mut hybrid = HybridSearch::open(&db, models, worker, threads)?; + let mut hybrid = HybridSearch::open_query(&db, models, worker, threads)?; hybrid.warmup()?; diff --git a/crates/searchi-cli/tests/mcp.rs b/crates/searchi-cli/tests/mcp.rs index 7744c13..59d7314 100644 --- a/crates/searchi-cli/tests/mcp.rs +++ b/crates/searchi-cli/tests/mcp.rs @@ -19,9 +19,22 @@ impl Drop for Server { } impl Server { fn start(db: &std::path::Path) -> Self { - let mut child = Command::new(env!("CARGO_BIN_EXE_searchi")) - .args(["mcp", "--db"]) - .arg(db) + Self::start_with_assets(db, None) + } + fn start_with_assets( + db: &std::path::Path, + assets: Option<(&std::ffi::OsStr, &std::ffi::OsStr)>, + ) -> Self { + let mut command = Command::new(env!("CARGO_BIN_EXE_searchi")); + command.args(["mcp", "--db"]).arg(db); + if let Some((model, worker)) = assets { + command + .arg("--model-dir") + .arg(model) + .arg("--embedding-worker") + .arg(worker); + } + let mut child = command .env("SEARCHI_FORCE_CPU", "1") .stdin(Stdio::piped()) .stdout(Stdio::piped()) @@ -233,3 +246,110 @@ fn missing_database_never_creates_file_or_protocol_noise() { assert!(output.stdout.is_empty()); assert!(!db.exists()); } + +#[test] +#[ignore = "Requires verified Google model and current CPU-capable native worker"] +fn actual_google_cli_and_mcp_use_query_role_with_cited_results() { + let models = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set models"); + let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set worker"); + let fixture = tempfile::tempdir().unwrap(); + let folder = fixture.path().join("documents"); + std::fs::create_dir(&folder).unwrap(); + std::fs::write( + folder.join("supplier.txt"), + "The supplier agreement may be terminated with sixty days written notice.", + ) + .unwrap(); + std::fs::write( + folder.join("garden.txt"), + "Vegetable beds and rainwater collection in the garden.", + ) + .unwrap(); + std::fs::write( + folder.join("retry.rs"), + "pub fn retry_backoff(attempt: u32) -> u64 {\n 100 * 2_u64.pow(attempt)\n}\n", + ) + .unwrap(); + let db = fixture.path().join("index.sqlite"); + searchi_core::SearchIndex::open(&db) + .unwrap() + .index_folder(&folder) + .unwrap(); + let indexer = searchi_service::HybridSearch::open(&db, &models, &worker, 1).unwrap(); + assert_eq!(indexer.worker_info().unwrap().role, "indexing"); + drop(indexer); + let embedded = Command::new(env!("CARGO_BIN_EXE_searchi")) + .args(["embed", "--db"]) + .arg(&db) + .arg("--model-dir") + .arg(&models) + .arg("--embedding-worker") + .arg(&worker) + .arg("--text-only") + .env("SEARCHI_FORCE_CPU", "1") + .output() + .unwrap(); + assert!( + embedded.status.success(), + "{}", + String::from_utf8_lossy(&embedded.stderr) + ); + let output = Command::new(env!("CARGO_BIN_EXE_searchi")) + .args(["search", "--db"]) + .arg(&db) + .args(["cancel the supplier partnership", "--json"]) + .arg("--model-dir") + .arg(&models) + .arg("--embedding-worker") + .arg(&worker) + .env("SEARCHI_FORCE_CPU", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let cli: Value = serde_json::from_slice(&output.stdout).unwrap(); + assert_eq!(cli[0]["filename"], "supplier.txt"); + let mut server = Server::start_with_assets(&db, Some((&models, &worker))); + server.request(1,"initialize",json!({"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"query-role-regression","version":"1"}})); + server.send(json!({"jsonrpc":"2.0","method":"notifications/initialized"})); + let stats = server.request(2, "tools/call", json!({"name":"stats","arguments":{}})); + let status = &stats["result"]["structuredContent"]; + assert_eq!(status["worker"]["role"], "query"); + assert_eq!(status["worker"]["context_tokens"], 512); + assert_eq!(status["worker"]["device"], "cpu"); + assert_eq!(status["worker"]["media_supported"], true); + let found = server.request( + 3, + "tools/call", + json!({"name":"search","arguments":{"query":"cancel the supplier partnership","limit":2}}), + ); + let first = &found["result"]["structuredContent"]["results"][0]; + assert_eq!(first["filename"], "supplier.txt"); + let handle = first["source_handle"].as_str().unwrap(); + let source = server.request( + 4, + "tools/call", + json!({"name":"read_source","arguments":{"handle":handle}}), + ); + assert!(source["result"]["structuredContent"]["text"] + .as_str() + .unwrap() + .contains("sixty days")); + let code=server.request(5,"tools/call",json!({"name":"search","arguments":{"query":"retry failed requests with exponential backoff","code_only":true}})); + let first = &code["result"]["structuredContent"]["results"][0]; + assert_eq!(first["filename"], "retry.rs"); + assert_eq!(first["source"]["kind"], "line"); + let handle = first["source_handle"].as_str().unwrap(); + let source = server.request( + 6, + "tools/call", + json!({"name":"read_source","arguments":{"handle":handle}}), + ); + assert!(source["result"]["structuredContent"]["text"] + .as_str() + .unwrap() + .contains("2_u64.pow")); +} diff --git a/crates/searchi-semantic/native/worker.cpp b/crates/searchi-semantic/native/worker.cpp index 0b53a4a..3f294f0 100644 --- a/crates/searchi-semantic/native/worker.cpp +++ b/crates/searchi-semantic/native/worker.cpp @@ -35,6 +35,7 @@ constexpr size_t max_passage_bytes = 1024 * 1024; constexpr size_t max_text_bytes = 65536; constexpr size_t max_windows = 256; constexpr int context_tokens = 2048; +constexpr int query_context_tokens = 512; constexpr int dimensions = 768; static void require(bool condition, const char * message) { @@ -53,6 +54,8 @@ class Encoder { const llama_vocab * vocab = nullptr; mtmd_context * media = nullptr; std::string model_path; + bool query_only; + int allocated_context; void release() { if (media) mtmd_free(media); if (batch) llama_batch_ext_free(batch); @@ -69,7 +72,7 @@ class Encoder { model = llama_model_load_from_file(model_path.c_str(),mp); if (!model) return false; auto cp = llama_context_default_params(); - cp.n_ctx = context_tokens; cp.n_batch = context_tokens; cp.n_ubatch = context_tokens; + cp.n_ctx = allocated_context; cp.n_batch = allocated_context; cp.n_ubatch = allocated_context; cp.n_seq_max = 1; cp.kv_unified = true; cp.embeddings = true; cp.offload_kqv = gpu; cp.op_offload = gpu; cp.pooling_type = LLAMA_POOLING_TYPE_MEAN; @@ -88,7 +91,7 @@ public: std::string cpu_backend_module; bool vulkan_module_loaded = false; int threads; - Encoder(const std::string & path, int cpu_threads, bool want_gpu) : model_path(path),threads(cpu_threads) { + Encoder(const std::string & path, int cpu_threads, bool want_gpu, bool query_role) : model_path(path),query_only(query_role),allocated_context(query_role ? query_context_tokens : context_tokens),threads(cpu_threads) { llama_log_set([](ggml_log_level level, const char * text, void *) { if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr); },nullptr); @@ -331,6 +334,7 @@ public: } json handle(const json & request) { auto op = request.at("op").get(); + require(!query_only || op == "query" || op == "code_query", "Query worker only supports query and code_query operations"); if(op=="image" || op=="audio") { auto encoded=request.at("data").get(); require(encoded.size()<=12000000 && encoded.size()%4==0,"Invalid media payload length"); @@ -363,7 +367,7 @@ public: bool debug = request.value("debug_performance",false); json timings; auto prefix = op == "code_query" ? "task: code retrieval | query: " : "task: search result | query: "; - auto vector = embed(prefix+request.at("text").get(),512,debug ? &timings : nullptr); + auto vector = embed(prefix+request.at("text").get(),query_context_tokens,debug ? &timings : nullptr); json response = {{"vectors",json::array({vector})}}; if (debug) response["performance"] = timings; return response; @@ -382,17 +386,19 @@ public: int main(int argc, char ** argv) { try { - std::string path, device = "auto"; int threads = 1; + std::string path, device = "auto", role = "indexing"; int threads = 1; for (int i = 1; i < argc; ++i) { require(i+1 < argc,"Worker options require values"); std::string option = argv[i++], value = argv[i]; if (option == "--model") path = value; else if (option == "--threads") threads = std::stoi(value); else if (option == "--device") device = value; + else if (option == "--role") role = value; else throw std::runtime_error("Unknown worker option"); } require(!path.empty() && threads >= 1 && threads <= 8,"Model path required; threads must be 1–8"); require(device == "cpu" || device == "auto","Device must be cpu or auto"); + require(role == "query" || role == "indexing", "Role must be query or indexing"); bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1"; // Disable backend registration before llama initializes its registry. // A GPU-capable distribution can therefore be tested on CPU without @@ -406,9 +412,9 @@ int main(int argc, char ** argv) { require(setenv("GGML_DISABLE_VULKAN", "1", 1) == 0, "Cannot disable Vulkan for CPU startup"); #endif } - Encoder encoder(path,threads,device == "auto" && !force_cpu); + Encoder encoder(path,threads,device == "auto" && !force_cpu,role == "query"); std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions}, - {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true},{"code_query_supported",true}, + {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true},{"code_query_supported",true},{"role",role},{"context_tokens",role == "query" ? query_context_tokens : context_tokens}, {"cpu_backend_module",encoder.cpu_backend_module},{"vulkan_module_loaded",encoder.vulkan_module_loaded}}).dump() << '\n' << std::flush; for (;;) { std::string line; char c; diff --git a/crates/searchi-semantic/src/gemma.rs b/crates/searchi-semantic/src/gemma.rs index 7b710f6..5f321ef 100644 --- a/crates/searchi-semantic/src/gemma.rs +++ b/crates/searchi-semantic/src/gemma.rs @@ -29,6 +29,7 @@ pub struct ImageEmbedding { const MODEL_FILE: &str = "embeddinggemma-2-Q8_0.gguf"; const MODEL_BYTES: u64 = 309855456; const MODEL_SHA256: &str = "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135"; +const QUERY_CONTEXT_TOKENS: usize = 512; const MAX_REQUEST: usize = 16 * 1024 * 1024; const MAX_RESPONSE: usize = 4 * 1024 * 1024; @@ -50,6 +51,14 @@ pub struct WorkerInfo { pub media_supported: bool, #[serde(default)] pub code_query_supported: bool, + #[serde(default)] + pub role: String, + #[serde(default)] + pub context_tokens: usize, +} + +fn supports_query_role(info: &WorkerInfo) -> bool { + info.role == "query" && info.context_tokens == QUERY_CONTEXT_TOKENS } fn query_operation(info: &WorkerInfo, code: bool) -> Result<&'static str> { @@ -68,6 +77,7 @@ pub struct GemmaEmbedder { pub info: WorkerInfo, healthy: bool, force_cpu: bool, + query_only: bool, model_dir: PathBuf, worker_path: PathBuf, projector_verified: bool, @@ -138,11 +148,40 @@ impl GemmaEmbedder { Self::load_with_device(model_dir, worker, threads, Device::Auto) } + /// A foreground-only session with a 512-token native context. The same + /// hardware-detecting worker still provides 2048-token indexing sessions. + pub fn load_query( + model_dir: impl AsRef, + worker: impl AsRef, + threads: usize, + ) -> Result { + Self::load_query_with_device(model_dir, worker, threads, Device::Auto) + } + + pub fn load_query_with_device( + model_dir: impl AsRef, + worker: impl AsRef, + threads: usize, + device: Device, + ) -> Result { + Self::load_role(model_dir, worker, threads, device, true) + } + pub fn load_with_device( model_dir: impl AsRef, worker: impl AsRef, threads: usize, device: Device, + ) -> Result { + Self::load_role(model_dir, worker, threads, device, false) + } + + fn load_role( + model_dir: impl AsRef, + worker: impl AsRef, + threads: usize, + device: Device, + query_only: bool, ) -> Result { ensure!( (1..=8).contains(&threads), @@ -162,13 +201,26 @@ impl GemmaEmbedder { let force_cpu = std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1"); let cpu = force_cpu || matches!(device, Device::Cpu); start_with_cpu_fallback(cpu, |cpu| { - Self::start_verified(&model_dir, &worker, threads, cpu) + Self::start_verified(&model_dir, &worker, threads, cpu, query_only) + }) + .with_context(|| { + if query_only { + "Cannot initialize foreground query role; repair or rebuild the native runtime" + } else { + "Cannot initialize indexing runtime" + } }) } // Both attempts share the already validated model/worker paths. CPU fallback // does not repeat the expensive hash check or weaken the handshake contract. - fn start_verified(model_dir: &Path, worker: &Path, threads: usize, cpu: bool) -> Result { + fn start_verified( + model_dir: &Path, + worker: &Path, + threads: usize, + cpu: bool, + query_only: bool, + ) -> Result { let model = model_dir.join(MODEL_FILE); let mut command = Command::new(worker); command @@ -183,6 +235,9 @@ impl GemmaEmbedder { .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()); + if query_only { + command.args(["--role", "query"]); + } // Background helpers must not create a visible console on Windows. #[cfg(windows)] { @@ -235,7 +290,8 @@ impl GemmaEmbedder { return Err(error).context("Invalid Gemma worker handshake"); } }; - if !info.ready + if (query_only && (!supports_query_role(&info))) + || !info.ready || info.model != "google/embeddinggemma-2" || info.dimensions != GEMMA_DIMENSIONS || info.runtime != "b11461-no-logits-v1" @@ -246,7 +302,7 @@ impl GemmaEmbedder { let _ = child.kill(); let _ = child.wait(); let _ = reader_thread.join(); - anyhow::bail!("Gemma worker handshake violates configured model/device contract"); + anyhow::bail!("Gemma worker handshake violates configured model/device contract or requested role; repair or rebuild the native runtime"); } let (requests, writes, writer_thread) = spawn_writer(stdin); Ok(Self { @@ -259,6 +315,7 @@ impl GemmaEmbedder { info, healthy: true, force_cpu: cpu, + query_only, model_dir: model_dir.to_path_buf(), worker_path: worker.to_path_buf(), projector_verified: false, @@ -288,11 +345,12 @@ impl GemmaEmbedder { // Restart once as CPU and retry this same bounded request. let _ = self.child.wait(); let restart_started = self.query_profile.as_ref().map(|_| Instant::now()); - let replacement = Self::load_with_device( + let replacement = Self::load_role( &self.model_dir, &self.worker_path, self.info.threads, Device::Cpu, + self.query_only, ) .with_context(|| { format!("CPU worker restart failed after GPU transport error: {error}") @@ -406,6 +464,10 @@ impl GemmaEmbedder { } pub fn split_windows(&mut self, text: &str, title: &str) -> Result> { + ensure!( + !self.query_only, + "Foreground query worker cannot index documents or media; use an indexing session" + ); ensure!( !text.trim().is_empty() && text.len() <= 1024 * 1024, "Passage must be nonempty and at most 1 MiB" @@ -432,6 +494,10 @@ impl GemmaEmbedder { } pub fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result>> { + ensure!( + !self.query_only, + "Foreground query worker cannot index documents or media; use an indexing session" + ); ensure!( !texts.is_empty() && texts.len() <= 8, "Document batch size must be between 1 and 8" @@ -545,6 +611,10 @@ impl GemmaEmbedder { paths: &[&Path], budget: u32, ) -> Result> { + ensure!( + !self.query_only, + "Foreground query worker cannot index documents or media; use an indexing session" + ); ensure!( !paths.is_empty() && paths.len() <= 8, "Image batch must contain 1–8 paths" @@ -567,6 +637,10 @@ impl GemmaEmbedder { /// One mono 16 kHz normalized PCM window, at most 30 seconds. Callers /// segment entire recordings and preserve real timestamps; no truncation. pub fn embed_audio_pcm(&mut self, samples: &[f32]) -> Result> { + ensure!( + !self.query_only, + "Foreground query worker cannot index documents or media; use an indexing session" + ); ensure!( !samples.is_empty() && samples.len() <= 480000, "Audio window must contain 1–480000 mono 16 kHz samples" @@ -851,9 +925,12 @@ mod tests { runtime: "fixture".into(), media_supported: false, code_query_supported: false, + role: "indexing".into(), + context_tokens: 2048, }, healthy: true, force_cpu: true, + query_only: false, model_dir: PathBuf::new(), worker_path: PathBuf::new(), projector_verified: false, @@ -861,6 +938,31 @@ mod tests { } } + #[test] + fn query_role_capability_is_explicit_and_independent_of_media_support() { + let mut info: WorkerInfo = serde_json::from_value(json!({ + "ready":true,"model":"google/embeddinggemma-2","dimensions":768, + "device":"cpu","threads":1,"runtime":"b11461-no-logits-v1", + "media_supported":true,"code_query_supported":true + })) + .unwrap(); + assert!( + !supports_query_role(&info), + "Legacy worker cannot claim reduced query allocation" + ); + info.role = "query".into(); + info.context_tokens = 2048; + assert!(!supports_query_role(&info)); + info.context_tokens = 512; + assert!(supports_query_role(&info)); + assert!( + info.media_supported, + "Binary capability remains available to full sessions" + ); + info.role = "indexing".into(); + assert!(!supports_query_role(&info)); + } + #[test] fn disconnected_writer_queue_is_transport_failure_and_reaps_helper() { let mut fixture = stalled_protocol_helper(false); @@ -962,7 +1064,7 @@ mod tests { let info: WorkerInfo = serde_json::from_value(handshake)?; ensure!( info.model == "google/embeddinggemma-2", - "Gemma worker handshake violates configured model/device contract" + "Gemma worker handshake violates configured model/device contract or requested role; repair or rebuild the native runtime" ); Ok(()) }); diff --git a/crates/searchi-semantic/tests/query_context_cpu.rs b/crates/searchi-semantic/tests/query_context_cpu.rs new file mode 100644 index 0000000..5dccaf2 --- /dev/null +++ b/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(()) +} diff --git a/crates/searchi-service/src/lib.rs b/crates/searchi-service/src/lib.rs index 869d202..b5a7e35 100644 --- a/crates/searchi-service/src/lib.rs +++ b/crates/searchi-service/src/lib.rs @@ -314,6 +314,17 @@ impl HybridSearch { Self::with_provider(db, Box::new(embedder)) } + /// Foreground-only Google session. Indexing must use `open` instead. + pub fn open_query( + db: impl AsRef, + model_dir: impl AsRef, + runtime_library: impl AsRef, + threads: usize, + ) -> Result { + let embedder = GemmaEmbedder::load_query(model_dir, runtime_library, threads)?; + Self::with_provider(db, Box::new(embedder)) + } + /// Explicit provider injection for future model generations, never fallback. pub fn with_provider( db: impl AsRef, diff --git a/docs/QUERY-CONTEXT.md b/docs/QUERY-CONTEXT.md new file mode 100644 index 0000000..2590c7f --- /dev/null +++ b/docs/QUERY-CONTEXT.md @@ -0,0 +1,92 @@ +# Foreground query context + +The normal hardware-detecting Google worker now accepts `--role query`. +That role allocates 512 tokens for `n_ctx`, `n_batch` and `n_ubatch`, matching +searchi's existing query token limit. The default `indexing` role retains all +2048 tokens and document/image/audio operations. This is one binary, not a +separate CPU distribution. Hardware selection and CPU fallback are unchanged. + +Desktop foreground sessions, CLI semantic searches and the retained MCP search +session use `HybridSearch::open_query`. Background document/media indexing +continues to use `open`. Role is retained when a failed GPU process is restarted +as CPU. A query-role worker rejects indexing operations before Rust verifies a +projector or decodes images, and the native protocol independently rejects them. +Invalid oversized queries leave the session usable. + +The handshake reports `role` and `context_tokens`. A legacy worker cannot +silently claim the smaller allocation: update/rebuild the worker with the app. +`media_supported` still describes the installed binary's compiled capability; +it remains true for query sessions so desktop status does not incorrectly hide +media indexing. It does not mean the foreground session loads the projector. +Query vectors and model/index fingerprints remain unchanged. + +## Actual Windows CPU evidence, 2026-10-13 + +Normal Windows x64 CPU/Vulkan worker, pinned llama b11461-no-logits-v1, +verified Google Q8_0 backbone, CPU forced and no Vulkan module loaded. Windows +11 build 26300 VM, 8 logical processors, one inference thread. Six fresh worker +processes, three per role, alternating order; cached filesystem/model bytes. +These are small paired allocation measurements, not cold-disk, UI or p95 proof. + +| Median measurement | Full 2048 | Query 512 | +| --- | ---: | ---: | +| Startup to ready |716.69 ms|710.10 ms| +| Ready private commitment |421,359,616 bytes|238,362,624 bytes| +| Ready resident working set |118,394,880 bytes|118,276,096 bytes| +| After first query private commitment |422,735,872 bytes|239,734,784 bytes| +| After first query resident working set |274,632,704 bytes|273,989,632 bytes| +| First query IPC/inference |106.57 ms|103.20 ms| + +Private commitment decreased by about 183 MB (43.4%). Resident memory and startup +were effectively unchanged in this sample. Do not advertise an 183 MB RSS saving +or a search-speed improvement from these results. Model verification took 195.86 ms +once before the process trials; startup excludes Rust verification and desktop +background preparation. Working sets include mapped shared model pages. + +The measured worker SHA256 was +`3bd88fdb26f5d62d400ed101277f637af028c757c619547f6333b2912137c60b`. +Raw samples are retained locally at `.searchi/qa/query-context-windows.json`. + +The actual Google integration `query_context_cpu` passed in 27.39 s: +English/German remembered-document ranking unchanged across both contexts; +CodeRetrieval vector parity for English/German/Arabic/Bengali; actual tokenizer +count exactly 512 with vector parity, rejection at 513 followed by a successful +query; foreground indexing/media guards and full-context long-document windows. +Three bilingual ranking vectors had cosine 0.99999999 or better. This tiny +regression is correctness proof, not independent retrieval-quality evaluation. +The benchmark additionally sends a direct native document operation to the +query role, verifies rejection is nonfatal, and then performs CodeRetrieval. + +Final worker rebuild after consolidating the query-size constant has SHA256 +`1272797af94c21db1ccabaa3ad996ab57d64f1bf9d33280c61d65c51360deab7`. +Six actual CPU integration tests passed with this final binary: two code +fixtures (16.52 s), two text/debug/reference fixtures (59.18 s), image/audio +reference parity (52.20 s), and the query-context test again (26.05 s). +Image reference cosines were 0.99994150, 0.99977568 and 0.99917617; audio +reference cosine was 0.99975052. The full indexing/media role remains usable. +Native startup/memory values above belong to the explicitly identified earlier +binary, before this equivalent constant cleanup; they are not installer checks. + +An actual CLI/MCP boundary regression also passed (36.36 s): the CLI indexed +three real document/code sources with the full role, CLI semantic search returned +the expected supplier file, and the retained MCP session reported query role, +512-token context and CPU. Both remembered-document and code-only searches +returned source handles whose guarded excerpts matched original text/line data. +Compiled media support remained visible to clients. This is a tiny integration +fixture, not a representative agent-library quality or performance benchmark. + +Reproduce with verified assets and the current normal worker: + +```powershell +$env:SEARCHI_FORCE_CPU = '1' +$env:SEARCHI_GEMMA_MODEL_DIR = 'PATH/TO/VERIFIED/MODELS' +$env:SEARCHI_GEMMA_WORKER = 'PATH/TO/searchi-gemma-worker.exe' +cargo test -p searchi-semantic --test query_context_cpu --locked -- --ignored --nocapture +cargo test -p searchi-cli --test mcp actual_google_cli_and_mcp_use_query_role_with_cited_results --locked -- --ignored --nocapture +python tools/embeddinggemma2/query_context_benchmark.py --worker "$env:SEARCHI_GEMMA_WORKER" --models "$env:SEARCHI_GEMMA_MODEL_DIR" --output .searchi/qa/query-context.json --repeats 3 +``` + +The benchmark reports Windows private commitment/working set or Linux +PSS/private resident pages separately. Native Linux/macOS and actual GPU context +allocation/parity, pressure behavior, long code queries, representative query +latency and updated installed desktop packaging remain release checks. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 73d7fa5..88d6dd4 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -178,6 +178,13 @@ Audit and measured priorities: [EMBEDDINGGEMMA2-AUDIT.md](EMBEDDINGGEMMA2-AUDIT. - [ ] Measure representative mixed-media contention, single-thread latency, extraction/decoder/ANN work and native UI p95 before closing this gate. - [ ] Attribute startup cost; measure smaller query context and selective encoders. + - [x] Query-only 512-token context in the normal worker, desktop/CLI/MCP + integration and actual Windows CPU parity/boundary/reference tests. + Paired samples reduced private commitment by 183 MB; resident memory and + cached startup were essentially unchanged. See [QUERY-CONTEXT.md](QUERY-CONTEXT.md). + - [ ] Verify native Linux/macOS, GPU parity and updated installed packages. + - [ ] Evaluate text readiness before the larger media projector finishes + downloading; preserve verified assets and honest media setup status. - [ ] Buffered bounded media IPC and linear UTF-8 source-window offset processing. - [ ] Evaluate 256d MRL and content-specific vision budgets on independent data. - [ ] Validate GPU precision, synchronization-aware timings, parity and fallback diff --git a/tools/embeddinggemma2/query_context_benchmark.py b/tools/embeddinggemma2/query_context_benchmark.py new file mode 100644 index 0000000..6767d35 --- /dev/null +++ b/tools/embeddinggemma2/query_context_benchmark.py @@ -0,0 +1,98 @@ +"""CPU-only paired real-worker context startup/memory measurements. + +No downloads or GPU execution. Windows reports private commitment and resident +working set separately; Linux reports PSS/private resident pages. Neither should +be summed as model allocations across workers. Requires an idle reference host. +""" +import argparse +import ctypes +import hashlib +import json +import os +from pathlib import Path +import platform +import statistics +import time +from benchmark import ROOT, Session, vectors + + +def memory(pid): + if os.name == "nt": + from ctypes import wintypes + class Counters(ctypes.Structure): + _fields_ = [("cb", wintypes.DWORD), ("PageFaultCount", wintypes.DWORD)] + [(name, ctypes.c_size_t) for name in ["PeakWorkingSetSize", "WorkingSetSize", "QuotaPeakPagedPoolUsage", "QuotaPagedPoolUsage", "QuotaPeakNonPagedPoolUsage", "QuotaNonPagedPoolUsage", "PagefileUsage", "PeakPagefileUsage", "PrivateUsage"]] + kernel = ctypes.WinDLL("kernel32", use_last_error=True) + kernel.OpenProcess.restype = wintypes.HANDLE + kernel.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + psapi = ctypes.WinDLL("psapi", use_last_error=True) + psapi.GetProcessMemoryInfo.argtypes = [wintypes.HANDLE, ctypes.POINTER(Counters), wintypes.DWORD] + handle = kernel.OpenProcess(0x1000 | 0x10, False, pid) + if not handle: + raise ctypes.WinError(ctypes.get_last_error()) + try: + info = Counters() + info.cb = ctypes.sizeof(info) + if not psapi.GetProcessMemoryInfo(handle, ctypes.byref(info), info.cb): + raise ctypes.WinError(ctypes.get_last_error()) + return {"private_commit_bytes": info.PrivateUsage, "working_set_bytes": info.WorkingSetSize, "page_fault_count": info.PageFaultCount} + finally: + kernel.CloseHandle(handle) + values = {} + for line in Path(f"/proc/{pid}/smaps_rollup").read_text().splitlines(): + parts = line.split() + if parts[0] in {"Pss:", "Private_Clean:", "Private_Dirty:", "Rss:"}: + values[parts[0][:-1] + "_bytes"] = int(parts[1]) * 1024 + return values + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--worker", type=Path, required=True) + parser.add_argument("--models", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + args = parser.parse_args() + if not 1 <= args.repeats <= 20: + parser.error("repeats must be 1..20") + manifest = json.loads((ROOT / "tools/embeddinggemma2/model-manifest.json").read_text()) + entry = manifest["files"][0] + model = args.models / entry["path"] + started = time.perf_counter() + assert model.stat().st_size == entry["bytes"] + with model.open("rb") as stream: + assert hashlib.file_digest(stream, "sha256").hexdigest() == entry["sha256"] + verification_ms = (time.perf_counter() - started) * 1000 + samples = [] + for repeat in range(args.repeats): + # Alternate ordering to expose warm filesystem/cache order effects. + for role in (["indexing", "query"] if repeat % 2 == 0 else ["query", "indexing"]): + session = Session([str(args.worker.resolve()), "--model", str(model.resolve()), "--threads", "1", "--device", "cpu", "--role", role]) + try: + hello = session.next() + startup_ms = (time.perf_counter() - session.started) * 1000 + assert hello["device"] == "cpu" and not hello["vulkan_module_loaded"] + assert hello["role"] == role and hello["context_tokens"] == (512 if role == "query" else 2048) + sample = {"repeat": repeat, "role": role, "hello": hello, "startup_ms": startup_ms, "ready_memory": memory(session.process.pid)} + response, first_ms = session.worker({"op": "query", "text": "cancel the supplier partnership", "debug_performance": True}) + vectors(response, 1) + sample.update(first_query_ms=first_ms, worker=response["performance"], after_query_memory=memory(session.process.pid)) + if role == "query": + session.send({"op": "documents", "texts": ["test"], "title": "none"}) + rejected = session.next() + assert not rejected["ok"] and not rejected["fatal"] + assert "only supports" in rejected["error"] + vectors(session.worker({"op": "code_query", "text": "retry failed requests"})[0], 1) + samples.append(sample) + finally: + session.close() + report = {"platform": platform.platform(), "cpu_count": os.cpu_count(), "verification_ms": verification_ms, "cache_note": "Cached model bytes; startup excludes Rust verification and desktop preparation. No GPU executed. Repeated fresh processes, not cold disk.", "worker_sha256": hashlib.sha256(args.worker.read_bytes()).hexdigest(), "samples": samples} + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + for role in ["indexing", "query"]: + rows = [s for s in samples if s["role"] == role] + print(role, "startup median ms", statistics.median(s["startup_ms"] for s in rows), "ready memory", rows[0]["ready_memory"], "after query", rows[0]["after_query_memory"]) + + +if __name__ == "__main__": + main()