karl/searchi
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases 10 Wiki Activity

Use measured query-only context across desktop CLI and MCP

Browse Source
This commit is contained in:
karl committed 2026-10-13 07:30:40 +02:00
1 parent 95c72c5fd4
commit 88cd9c7f5a
12 files changed
+582 -20
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/searchi-semantic/src/gemma.rs
+108 -6
View File
@@ -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<Path>,
worker: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
Self::load_query_with_device(model_dir, worker, threads, Device::Auto)
}
pub fn load_query_with_device(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
device: Device,
) -> Result<Self> {
Self::load_role(model_dir, worker, threads, device, true)
}
pub fn load_with_device(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
device: Device,
) -> Result<Self> {
Self::load_role(model_dir, worker, threads, device, false)
}
fn load_role(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
device: Device,
query_only: bool,
) -> Result<Self> {
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<Self> {
fn start_verified(
model_dir: &Path,
worker: &Path,
threads: usize,
cpu: bool,
query_only: bool,
) -> Result<Self> {
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<Vec<TextWindow>> {
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<Vec<Vec<f32>>> {
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<Vec<ImageEmbedding>> {
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<Vec<f32>> {
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(())
});