Use measured query-only context across desktop CLI and MCP
This commit is contained in:
1 parent
95c72c5fd4
commit
88cd9c7f5a
12 files changed
+582
-20
No files matched your search
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()?;
|
||||
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
@@ -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<std::string>();
|
||||
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<std::string>();
|
||||
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<std::string>(),512,debug ? &timings : nullptr);
|
||||
auto vector = embed(prefix+request.at("text").get<std::string>(),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;
|
||||
|
||||
@@ -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(())
|
||||
});
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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<Path>,
|
||||
model_dir: impl AsRef<Path>,
|
||||
runtime_library: impl AsRef<Path>,
|
||||
threads: usize,
|
||||
) -> Result<Self> {
|
||||
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<Path>,
|
||||
|
||||
@@ -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.
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user