1214 lines
46 KiB
Rust
1214 lines
46 KiB
Rust
use crate::TextWindow;
|
||
use anyhow::{ensure, Context, Result};
|
||
use base64::Engine;
|
||
use serde::{Deserialize, Serialize};
|
||
use serde_json::{json, Value};
|
||
use sha2::{Digest, Sha256};
|
||
use std::{
|
||
collections::BTreeMap,
|
||
fs::File,
|
||
io::{BufRead, BufReader, Read, Write},
|
||
path::{Path, PathBuf},
|
||
process::{Child, ChildStdin, Command, Stdio},
|
||
sync::mpsc::{self, Receiver},
|
||
thread,
|
||
time::{Duration, Instant},
|
||
};
|
||
|
||
pub const GEMMA_DIMENSIONS: usize = 768;
|
||
pub const GEMMA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mean:l2:query512:doc2048:windows-overlap128:prompts-v1";
|
||
pub const GEMMA_MEDIA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mtmd:mean:l2:rgb-catmullrom48-exif-largejpeg-scaled-v2:pcm16k30s-v1";
|
||
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ImageEmbedding {
|
||
pub vector: Vec<f32>,
|
||
pub encoded_width: u32,
|
||
pub encoded_height: u32,
|
||
pub token_budget: u32,
|
||
}
|
||
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;
|
||
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub enum Device {
|
||
Auto,
|
||
Cpu,
|
||
}
|
||
|
||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||
pub struct WorkerInfo {
|
||
pub ready: bool,
|
||
pub model: String,
|
||
pub dimensions: usize,
|
||
pub device: String,
|
||
pub threads: usize,
|
||
pub runtime: String,
|
||
#[serde(default)]
|
||
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> {
|
||
ensure!(!code || info.code_query_supported,"Installed embedding worker does not support code retrieval; repair/rebuild the native runtime");
|
||
Ok(if code { "code_query" } else { "query" })
|
||
}
|
||
|
||
/// Persistent isolated native worker. No HTTP or remote model fallback.
|
||
pub struct GemmaEmbedder {
|
||
child: Child,
|
||
requests: Option<mpsc::SyncSender<Vec<u8>>>,
|
||
writes: Receiver<std::result::Result<(), String>>,
|
||
writer_thread: Option<thread::JoinHandle<()>>,
|
||
reader_thread: Option<thread::JoinHandle<()>>,
|
||
responses: Receiver<std::result::Result<Value, String>>,
|
||
pub info: WorkerInfo,
|
||
healthy: bool,
|
||
force_cpu: bool,
|
||
query_only: bool,
|
||
model_dir: PathBuf,
|
||
worker_path: PathBuf,
|
||
projector_verified: bool,
|
||
query_profile: Option<BTreeMap<String, u64>>,
|
||
}
|
||
|
||
impl Drop for GemmaEmbedder {
|
||
fn drop(&mut self) {
|
||
self.stop_worker();
|
||
}
|
||
}
|
||
|
||
type WriterTransport = (
|
||
mpsc::SyncSender<Vec<u8>>,
|
||
Receiver<std::result::Result<(), String>>,
|
||
thread::JoinHandle<()>,
|
||
);
|
||
|
||
fn spawn_writer(mut stdin: ChildStdin) -> WriterTransport {
|
||
let (requests, pending) = mpsc::sync_channel::<Vec<u8>>(1);
|
||
let (completed, writes) = mpsc::sync_channel(1);
|
||
let writer = thread::spawn(move || {
|
||
while let Ok(encoded) = pending.recv() {
|
||
let result = stdin
|
||
.write_all(&encoded)
|
||
.and_then(|_| stdin.write_all(b"\n"))
|
||
.and_then(|_| stdin.flush())
|
||
.map_err(|error| format!("Cannot write local Gemma worker request: {error}"));
|
||
let failed = result.is_err();
|
||
if completed.try_send(result).is_err() || failed {
|
||
break;
|
||
}
|
||
}
|
||
});
|
||
(requests, writes, writer)
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
struct WorkerTransportFailure(String);
|
||
impl std::fmt::Display for WorkerTransportFailure {
|
||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
formatter.write_str(&self.0)
|
||
}
|
||
}
|
||
impl std::error::Error for WorkerTransportFailure {}
|
||
|
||
/// Retry only process/transport failure, never an incompatible protocol or model.
|
||
fn start_with_cpu_fallback<T>(cpu: bool, mut start: impl FnMut(bool) -> Result<T>) -> Result<T> {
|
||
match start(cpu) {
|
||
Err(error) if !cpu && error.downcast_ref::<WorkerTransportFailure>().is_some() => {
|
||
start(true).with_context(|| {
|
||
format!("CPU initialization failed after automatic worker startup failure: {error}")
|
||
})
|
||
}
|
||
result => result,
|
||
}
|
||
}
|
||
|
||
impl GemmaEmbedder {
|
||
/// Load the verified Google model and explicitly supplied local worker.
|
||
/// Production selects available acceleration with CPU fallback. Tests set
|
||
/// `SEARCHI_FORCE_CPU=1` or call `load_with_device(..., Device::Cpu)`.
|
||
pub fn load(
|
||
model_dir: impl AsRef<Path>,
|
||
worker: impl AsRef<Path>,
|
||
threads: usize,
|
||
) -> Result<Self> {
|
||
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),
|
||
"CPU threads must be between 1 and 8"
|
||
);
|
||
let model_dir = model_dir.as_ref().to_path_buf();
|
||
let model = model_dir.join(MODEL_FILE);
|
||
verify_model(&model)?;
|
||
let worker = worker
|
||
.as_ref()
|
||
.canonicalize()
|
||
.context("Cannot resolve local Gemma worker executable")?;
|
||
ensure!(
|
||
worker.is_file(),
|
||
"Gemma worker must be a local executable file"
|
||
);
|
||
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, 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,
|
||
query_only: bool,
|
||
) -> Result<Self> {
|
||
let model = model_dir.join(MODEL_FILE);
|
||
let mut command = Command::new(worker);
|
||
command
|
||
.args(["--model"])
|
||
.arg(model)
|
||
.args([
|
||
"--threads",
|
||
&threads.to_string(),
|
||
"--device",
|
||
if cpu { "cpu" } else { "auto" },
|
||
])
|
||
.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)]
|
||
{
|
||
use std::os::windows::process::CommandExt;
|
||
command.creation_flags(0x08000000); // CREATE_NO_WINDOW
|
||
}
|
||
let mut child = command.spawn().map_err(|error| {
|
||
WorkerTransportFailure(format!("Cannot start local Gemma worker: {error}"))
|
||
})?;
|
||
let (Some(stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else {
|
||
let _ = child.kill();
|
||
let _ = child.wait();
|
||
anyhow::bail!("Gemma worker protocol pipes missing");
|
||
};
|
||
let (sender, responses) = mpsc::sync_channel(1);
|
||
let reader_thread = thread::spawn(move || {
|
||
let mut reader = BufReader::new(stdout);
|
||
loop {
|
||
let response = read_response(&mut reader);
|
||
let failed = response.is_err();
|
||
if sender.try_send(response).is_err() || failed {
|
||
break;
|
||
}
|
||
}
|
||
});
|
||
let ready = match responses.recv_timeout(Duration::from_secs(120)) {
|
||
Ok(Ok(value)) => value,
|
||
other => {
|
||
let _ = child.kill();
|
||
let _ = child.wait();
|
||
let _ = reader_thread.join();
|
||
// Malformed/oversized JSON is a protocol violation, not evidence
|
||
// that hardware acceleration failed. Do not mask it with retry.
|
||
if matches!(&other, Ok(Err(error)) if error.starts_with("Invalid Gemma worker JSON:") || error == "Gemma worker response exceeds protocol byte budget")
|
||
{
|
||
anyhow::bail!("Invalid Gemma worker startup protocol: {other:?}");
|
||
}
|
||
return Err(WorkerTransportFailure(format!(
|
||
"Gemma worker did not initialize: {other:?}"
|
||
))
|
||
.into());
|
||
}
|
||
};
|
||
let info: WorkerInfo = match serde_json::from_value(ready) {
|
||
Ok(info) => info,
|
||
Err(error) => {
|
||
let _ = child.kill();
|
||
let _ = child.wait();
|
||
let _ = reader_thread.join();
|
||
return Err(error).context("Invalid Gemma worker handshake");
|
||
}
|
||
};
|
||
if (query_only && (!supports_query_role(&info)))
|
||
|| !info.ready
|
||
|| info.model != "google/embeddinggemma-2"
|
||
|| info.dimensions != GEMMA_DIMENSIONS
|
||
|| info.runtime != "b11461-no-logits-v1"
|
||
|| info.threads != threads
|
||
|| !["cpu", "gpu"].contains(&info.device.as_str())
|
||
|| (cpu && info.device != "cpu")
|
||
{
|
||
let _ = child.kill();
|
||
let _ = child.wait();
|
||
let _ = reader_thread.join();
|
||
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 {
|
||
child,
|
||
requests: Some(requests),
|
||
writes,
|
||
writer_thread: Some(writer_thread),
|
||
reader_thread: Some(reader_thread),
|
||
responses,
|
||
info,
|
||
healthy: true,
|
||
force_cpu: cpu,
|
||
query_only,
|
||
model_dir: model_dir.to_path_buf(),
|
||
worker_path: worker.to_path_buf(),
|
||
projector_verified: false,
|
||
query_profile: None,
|
||
})
|
||
}
|
||
|
||
fn request(&mut self, request: Value) -> Result<Value> {
|
||
ensure!(
|
||
self.healthy,
|
||
"Gemma worker session failed; reopen it explicitly"
|
||
);
|
||
let encode_started = self.query_profile.as_ref().map(|_| Instant::now());
|
||
let encoded = serde_json::to_vec(&request)?;
|
||
self.record_query_duration("request_encode_us", encode_started);
|
||
ensure!(
|
||
encoded.len() < MAX_REQUEST,
|
||
"Gemma request exceeds 16 MiB protocol budget"
|
||
);
|
||
let response = match self.exchange(&encoded) {
|
||
Ok(response) => response,
|
||
Err(error)
|
||
if can_retry_cpu(&self.info.device, self.force_cpu)
|
||
&& error.downcast_ref::<WorkerTransportFailure>().is_some() =>
|
||
{
|
||
// GPU runtimes can abort rather than return an encoder error.
|
||
// 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_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}")
|
||
})?;
|
||
let diagnostics = self.query_profile.take();
|
||
*self = replacement;
|
||
self.query_profile = diagnostics;
|
||
self.record_query_duration("cpu_restart_us", restart_started);
|
||
if let Some(profile) = &mut self.query_profile {
|
||
profile.insert("cpu_restart_count".into(), 1);
|
||
}
|
||
self.exchange(&encoded)?
|
||
}
|
||
Err(error) => return Err(error),
|
||
};
|
||
if let Some(device) = response["device"].as_str() {
|
||
ensure!(
|
||
["cpu", "gpu"].contains(&device) && (!self.force_cpu || device == "cpu"),
|
||
"Worker response violates configured device policy"
|
||
);
|
||
self.info.device = device.into();
|
||
}
|
||
if response["fallback"].as_str() == Some("gpu_inference_failure") {
|
||
if let Some(profile) = &mut self.query_profile {
|
||
profile.insert("worker_cpu_retry_count".into(), 1);
|
||
}
|
||
}
|
||
if response["fatal"] == true {
|
||
self.healthy = false;
|
||
self.info.ready = false;
|
||
let _ = self.child.kill();
|
||
}
|
||
ensure!(
|
||
response["ok"] == true,
|
||
"Gemma worker error: {}",
|
||
response["error"].as_str().unwrap_or("unknown error")
|
||
);
|
||
Ok(response)
|
||
}
|
||
|
||
fn stop_worker(&mut self) {
|
||
self.requests.take();
|
||
// The shipped native worker does not spawn descendants. Killing/reaping
|
||
// it closes its pipe ends and releases any blocked writer/read thread.
|
||
let _ = self.child.kill();
|
||
let _ = self.child.wait();
|
||
if let Some(writer) = self.writer_thread.take() {
|
||
let _ = writer.join();
|
||
}
|
||
if let Some(reader) = self.reader_thread.take() {
|
||
let _ = reader.join();
|
||
}
|
||
}
|
||
|
||
fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
|
||
self.exchange_with_timeout(encoded, Duration::from_secs(180))
|
||
}
|
||
|
||
fn exchange_with_timeout(&mut self, encoded: &[u8], timeout: Duration) -> Result<Value> {
|
||
let deadline = Instant::now() + timeout;
|
||
let write_started = self.query_profile.as_ref().map(|_| Instant::now());
|
||
let result = (|| {
|
||
self.requests
|
||
.as_ref()
|
||
.ok_or_else(|| WorkerTransportFailure("Gemma worker writer stopped".into()))?
|
||
.try_send(encoded.to_vec())
|
||
.map_err(|_| {
|
||
WorkerTransportFailure("Gemma worker request queue unavailable".into())
|
||
})?;
|
||
match self
|
||
.writes
|
||
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
|
||
{
|
||
Ok(Ok(())) => (),
|
||
other => {
|
||
return Err(WorkerTransportFailure(format!(
|
||
"Gemma worker request write failed or timed out: {other:?}"
|
||
))
|
||
.into())
|
||
}
|
||
}
|
||
self.record_query_duration("request_write_us", write_started);
|
||
let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
|
||
let response = match self
|
||
.responses
|
||
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
|
||
{
|
||
Ok(Ok(value)) => value,
|
||
Ok(Err(error))
|
||
if error.starts_with("Invalid Gemma worker JSON:")
|
||
|| error == "Gemma worker response exceeds protocol byte budget" =>
|
||
{
|
||
anyhow::bail!("Invalid Gemma worker response protocol: {error}")
|
||
}
|
||
other => {
|
||
return Err(WorkerTransportFailure(format!(
|
||
"Gemma worker request failed or timed out: {other:?}"
|
||
))
|
||
.into())
|
||
}
|
||
};
|
||
self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
|
||
Ok(response)
|
||
})();
|
||
if result.is_err() {
|
||
self.healthy = false;
|
||
self.info.ready = false;
|
||
self.stop_worker();
|
||
}
|
||
result
|
||
}
|
||
|
||
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"
|
||
);
|
||
let response = self.request(json!({"op":"windows","text":text,"title":title}))?;
|
||
let windows: Vec<TextWindow> = serde_json::from_value(response["windows"].clone())?;
|
||
ensure!(
|
||
!windows.is_empty() && windows.len() <= 256,
|
||
"Invalid Gemma window count"
|
||
);
|
||
for window in &windows {
|
||
ensure!(
|
||
text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
|
||
"Gemma window does not match original source range"
|
||
);
|
||
}
|
||
ensure!(
|
||
windows
|
||
.last()
|
||
.is_some_and(|window| window.end_byte == text.len()),
|
||
"Gemma windows did not cover the passage tail"
|
||
);
|
||
Ok(windows)
|
||
}
|
||
|
||
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"
|
||
);
|
||
ensure!(
|
||
texts
|
||
.iter()
|
||
.all(|text| !text.trim().is_empty() && text.len() <= 65536),
|
||
"Document window must be nonempty and at most 64 KiB"
|
||
);
|
||
let response = self.request(json!({"op":"documents","texts":texts,"title":title}))?;
|
||
extract_vectors(response, texts.len())
|
||
}
|
||
|
||
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
|
||
self.query_profile = None;
|
||
self.embed_query_inner(query, false, false)
|
||
}
|
||
|
||
/// Opt-in diagnostics contain only durations and counts, never query text.
|
||
pub fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
|
||
self.query_profile = Some(BTreeMap::new());
|
||
self.embed_query_inner(query, true, false)
|
||
}
|
||
|
||
/// Google code retrieval task; unavailable workers fail explicitly.
|
||
pub fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
|
||
self.query_profile = None;
|
||
self.embed_query_inner(query, false, true)
|
||
}
|
||
pub fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
|
||
self.query_profile = Some(BTreeMap::new());
|
||
self.embed_query_inner(query, true, true)
|
||
}
|
||
pub fn query_diagnostics(&self) -> BTreeMap<String, u64> {
|
||
self.query_profile.clone().unwrap_or_default()
|
||
}
|
||
|
||
fn record_query_duration(&mut self, name: &str, started: Option<Instant>) {
|
||
if let (Some(profile), Some(started)) = (&mut self.query_profile, started) {
|
||
let micros = started.elapsed().as_micros().min(u64::MAX as u128) as u64;
|
||
*profile.entry(name.into()).or_default() += micros;
|
||
}
|
||
}
|
||
|
||
fn embed_query_inner(&mut self, query: &str, debug: bool, code: bool) -> Result<Vec<f32>> {
|
||
ensure!(
|
||
!query.trim().is_empty() && query.len() <= 4096,
|
||
"Query must be nonempty and at most 4096 UTF-8 bytes"
|
||
);
|
||
let op = query_operation(&self.info, code)?;
|
||
let response = self.request(json!({"op":op,"text":query,"debug_performance":debug}))?;
|
||
if debug {
|
||
if let Some(profile) = &mut self.query_profile {
|
||
for name in [
|
||
"tokenization_us",
|
||
"batch_prepare_us",
|
||
"inference_us",
|
||
"pool_normalize_us",
|
||
"token_count",
|
||
] {
|
||
if let Some(value) = response["performance"][name].as_u64() {
|
||
profile.insert(format!("worker_{name}"), value);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
let validation_started = self.query_profile.as_ref().map(|_| Instant::now());
|
||
let result = extract_vectors(response, 1)?.remove(0);
|
||
self.record_query_duration("response_validate_us", validation_started);
|
||
Ok(result)
|
||
}
|
||
|
||
fn projector(&mut self) -> Result<PathBuf> {
|
||
ensure!(
|
||
self.info.media_supported,
|
||
"Installed embedding worker does not support media; repair/rebuild the native runtime"
|
||
);
|
||
let path = self.model_dir.join("mmproj-embeddinggemma-2-Q8_0.gguf");
|
||
if !self.projector_verified {
|
||
verify_artifact(
|
||
&path,
|
||
554821024,
|
||
"c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f",
|
||
)?;
|
||
self.projector_verified = true;
|
||
}
|
||
Ok(path)
|
||
}
|
||
|
||
/// Genuine Google pixel embeddings, serial bounded image operations.
|
||
pub fn embed_images(&mut self, paths: &[&Path]) -> Result<Vec<Vec<f32>>> {
|
||
self.embed_images_with_budget(paths, 280)
|
||
}
|
||
|
||
/// 70 tokens for video frames, 140 for photos, 280 or 1120 for screenshots.
|
||
pub fn embed_images_with_budget(
|
||
&mut self,
|
||
paths: &[&Path],
|
||
budget: u32,
|
||
) -> Result<Vec<Vec<f32>>> {
|
||
Ok(self
|
||
.embed_images_with_metadata(paths, budget)?
|
||
.into_iter()
|
||
.map(|image| image.vector)
|
||
.collect())
|
||
}
|
||
|
||
pub fn embed_images_with_metadata(
|
||
&mut self,
|
||
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"
|
||
);
|
||
let projector = self.projector()?;
|
||
let mut vectors = Vec::with_capacity(paths.len());
|
||
for path in paths {
|
||
let pixels = decode_image(path, budget)?;
|
||
let response=self.request(json!({"op":"image","projector":projector,"width":pixels.width(),"height":pixels.height(),"data":base64::engine::general_purpose::STANDARD.encode(pixels.as_raw())}))?;
|
||
vectors.push(ImageEmbedding {
|
||
vector: extract_vectors(response, 1)?.remove(0),
|
||
encoded_width: pixels.width(),
|
||
encoded_height: pixels.height(),
|
||
token_budget: budget,
|
||
});
|
||
}
|
||
Ok(vectors)
|
||
}
|
||
|
||
/// 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"
|
||
);
|
||
ensure!(
|
||
samples.iter().all(|v| v.is_finite() && v.abs() <= 1.0),
|
||
"Invalid normalized audio PCM"
|
||
);
|
||
ensure!(
|
||
cfg!(target_endian = "little"),
|
||
"Native PCM worker requires little-endian host"
|
||
);
|
||
let projector = self.projector()?;
|
||
let bytes: Vec<u8> = samples.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||
let response=self.request(json!({"op":"audio","projector":projector,"data":base64::engine::general_purpose::STANDARD.encode(bytes)}))?;
|
||
Ok(extract_vectors(response, 1)?.remove(0))
|
||
}
|
||
}
|
||
|
||
fn decode_image(path: &Path, budget: u32) -> Result<image::RgbImage> {
|
||
ensure!(
|
||
(70..=1120).contains(&budget),
|
||
"Image token budget must be 70–1120"
|
||
);
|
||
ensure!(
|
||
std::fs::metadata(path)?.len() <= 64 * 1024 * 1024,
|
||
"Image file exceeds 64 MiB"
|
||
);
|
||
let mut reader = image::ImageReader::open(path)?.with_guessed_format()?;
|
||
let is_jpeg = reader.format() == Some(image::ImageFormat::Jpeg);
|
||
let mut limits = image::Limits::default();
|
||
limits.max_alloc = Some(256 * 1024 * 1024);
|
||
limits.max_image_width = Some(16384);
|
||
limits.max_image_height = Some(16384);
|
||
reader.limits(limits);
|
||
use image::ImageDecoder;
|
||
let mut decoder = reader
|
||
.into_decoder()
|
||
.context("Cannot initialize bounded image decoder")?;
|
||
let (source_width, source_height) = decoder.dimensions();
|
||
ensure!(
|
||
u64::from(source_width) * u64::from(source_height) <= 64 * 1024 * 1024
|
||
&& decoder.total_bytes() <= 256 * 1024 * 1024,
|
||
"Decoded image exceeds memory/pixel budget"
|
||
);
|
||
let orientation = decoder
|
||
.orientation()
|
||
.context("Cannot read image EXIF orientation")?;
|
||
if is_jpeg && u64::from(source_width) * u64::from(source_height) > 16 * 1024 * 1024 {
|
||
drop(decoder);
|
||
return decode_large_jpeg(path, budget, source_width, source_height, orientation);
|
||
}
|
||
ensure!(u64::from(source_width)*u64::from(source_height)<=16*1024*1024,"Non-JPEG image exceeds16megapixels; bounded scaled decoding is currently available for JPEG");
|
||
let mut decoded = image::DynamicImage::from_decoder(decoder)
|
||
.context("Cannot decode supported image pixels")?;
|
||
decoded.apply_orientation(orientation);
|
||
ensure!(
|
||
u64::from(decoded.width()) * u64::from(decoded.height()) <= 16 * 1024 * 1024,
|
||
"Image exceeds 16 megapixels"
|
||
);
|
||
let (w, h) = image_target(decoded.width(), decoded.height(), budget)?;
|
||
Ok(image::imageops::resize(
|
||
&decoded.to_rgb8(),
|
||
w,
|
||
h,
|
||
image::imageops::FilterType::CatmullRom,
|
||
))
|
||
}
|
||
|
||
fn decode_large_jpeg(
|
||
path: &Path,
|
||
budget: u32,
|
||
width: u32,
|
||
height: u32,
|
||
orientation: image::metadata::Orientation,
|
||
) -> Result<image::RgbImage> {
|
||
let (target_width, target_height) = image_target(width, height, budget)?;
|
||
let mut decoder = jpeg_decoder::Decoder::new(BufReader::new(File::open(path)?));
|
||
decoder.set_max_decoding_buffer_size(256 * 1024 * 1024);
|
||
decoder
|
||
.read_info()
|
||
.context("Cannot read scaled JPEG header")?;
|
||
let source = decoder.info().context("Missing scaled JPEG header")?;
|
||
ensure!(
|
||
u32::from(source.width) == width && u32::from(source.height) == height,
|
||
"JPEG changed during decoding"
|
||
);
|
||
ensure!(
|
||
matches!(
|
||
source.pixel_format,
|
||
jpeg_decoder::PixelFormat::RGB24 | jpeg_decoder::PixelFormat::L8
|
||
),
|
||
"Large JPEG requires supported RGB24 or grayscale scaled decoding"
|
||
);
|
||
let (w, h) = decoder.scale(u16::try_from(target_width)?, u16::try_from(target_height)?)?;
|
||
ensure!(
|
||
u64::from(w) * u64::from(h) * 4 <= 256 * 1024 * 1024,
|
||
"Scaled JPEG output exceeds memory budget"
|
||
);
|
||
let pixels = decoder
|
||
.decode()
|
||
.context("Cannot decode large JPEG with scaled IDCT")?;
|
||
let mut decoded = match source.pixel_format {
|
||
jpeg_decoder::PixelFormat::RGB24 => image::DynamicImage::ImageRgb8(
|
||
image::RgbImage::from_raw(u32::from(w), u32::from(h), pixels)
|
||
.context("Invalid scaled JPEG RGB length")?,
|
||
),
|
||
jpeg_decoder::PixelFormat::L8 => image::DynamicImage::ImageLuma8(
|
||
image::GrayImage::from_raw(u32::from(w), u32::from(h), pixels)
|
||
.context("Invalid scaled JPEG grayscale length")?,
|
||
),
|
||
other => anyhow::bail!(
|
||
"Large JPEG pixel format {other:?} is not supported by bounded scaled decoder"
|
||
),
|
||
};
|
||
decoded.apply_orientation(orientation);
|
||
let (w, h) = image_target(decoded.width(), decoded.height(), budget)?;
|
||
Ok(image::imageops::resize(
|
||
&decoded.into_rgb8(),
|
||
w,
|
||
h,
|
||
image::imageops::FilterType::CatmullRom,
|
||
))
|
||
}
|
||
|
||
fn image_target(width: u32, height: u32, budget: u32) -> Result<(u32, u32)> {
|
||
ensure!(
|
||
width > 0 && height > 0 && (70..=1120).contains(&budget),
|
||
"Invalid image dimensions/budget"
|
||
);
|
||
let scale = (f64::from(budget) * 2304.0 / (f64::from(width) * f64::from(height))).sqrt();
|
||
let mut w = ((f64::from(width) * scale / 48.0).floor() as u32) * 48;
|
||
let mut h = ((f64::from(height) * scale / 48.0).floor() as u32) * 48;
|
||
if w == 0 {
|
||
w = 48;
|
||
h = (height / width * 48).min(budget * 48);
|
||
}
|
||
if h == 0 {
|
||
h = 48;
|
||
w = (width / height * 48).min(budget * 48);
|
||
}
|
||
ensure!(
|
||
w >= 48 && h >= 48 && u64::from(w) * u64::from(h) <= u64::from(budget) * 2304,
|
||
"Image target outside token budget"
|
||
);
|
||
Ok((w, h))
|
||
}
|
||
|
||
fn can_retry_cpu(actual_device: &str, force_cpu: bool) -> bool {
|
||
actual_device == "gpu" && !force_cpu
|
||
}
|
||
|
||
fn read_response(reader: &mut impl BufRead) -> std::result::Result<Value, String> {
|
||
let mut line = Vec::new();
|
||
match reader
|
||
.take((MAX_RESPONSE + 1) as u64)
|
||
.read_until(b'\n', &mut line)
|
||
{
|
||
Ok(0) => Err("Gemma worker closed its output".into()),
|
||
Ok(_) if line.len() > MAX_RESPONSE || line.last() != Some(&b'\n') => {
|
||
Err("Gemma worker response exceeds protocol byte budget".into())
|
||
}
|
||
Ok(_) => serde_json::from_slice(&line)
|
||
.map_err(|error| format!("Invalid Gemma worker JSON: {error}")),
|
||
Err(error) => Err(format!("Gemma worker read failed: {error}")),
|
||
}
|
||
}
|
||
|
||
fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
|
||
let vectors: Vec<Vec<f32>> = serde_json::from_value(response["vectors"].clone())?;
|
||
ensure!(
|
||
vectors.len() == count,
|
||
"Gemma returned incorrect embedding batch length"
|
||
);
|
||
for vector in &vectors {
|
||
ensure!(
|
||
vector.len() == GEMMA_DIMENSIONS && vector.iter().all(|value| value.is_finite()),
|
||
"Invalid Gemma embedding shape or values"
|
||
);
|
||
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
|
||
ensure!(
|
||
(norm - 1.0).abs() < 1e-4,
|
||
"Gemma embedding is not L2 normalized"
|
||
);
|
||
}
|
||
Ok(vectors)
|
||
}
|
||
|
||
fn verify_model(path: &Path) -> Result<()> {
|
||
verify_artifact(path, MODEL_BYTES, MODEL_SHA256)
|
||
}
|
||
|
||
fn verify_artifact(path: &Path, bytes: u64, sha: &str) -> Result<()> {
|
||
let file = File::open(path).with_context(|| {
|
||
format!(
|
||
"Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
|
||
path.display()
|
||
)
|
||
})?;
|
||
ensure!(
|
||
file.metadata()?.len() == bytes,
|
||
"EmbeddingGemma 2 model size mismatch"
|
||
);
|
||
let mut reader = BufReader::new(file);
|
||
let mut hash = Sha256::new();
|
||
let mut buffer = [0; 65536];
|
||
loop {
|
||
let count = reader.read(&mut buffer)?;
|
||
if count == 0 {
|
||
break;
|
||
}
|
||
hash.update(&buffer[..count]);
|
||
}
|
||
ensure!(
|
||
format!("{:x}", hash.finalize()) == sha,
|
||
"EmbeddingGemma 2 model SHA256 mismatch"
|
||
);
|
||
Ok(())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
// This helper is an inert pipe-protocol fixture, never an embedding provider.
|
||
fn stalled_protocol_helper(read_request: bool) -> GemmaEmbedder {
|
||
#[cfg(windows)]
|
||
let mut command = {
|
||
use std::os::windows::process::CommandExt;
|
||
let mut command = Command::new("powershell.exe");
|
||
command.args([
|
||
"-NoProfile",
|
||
"-NonInteractive",
|
||
"-Command",
|
||
if read_request {
|
||
"[Console]::In.ReadLine() | Out-Null; Start-Sleep -Seconds 20"
|
||
} else {
|
||
"Start-Sleep -Seconds 20"
|
||
},
|
||
]);
|
||
command.creation_flags(0x08000000);
|
||
command
|
||
};
|
||
#[cfg(not(windows))]
|
||
let mut command = {
|
||
let mut command = Command::new("sh");
|
||
command.args([
|
||
"-c",
|
||
if read_request {
|
||
"read line; exec sleep 20"
|
||
} else {
|
||
"exec sleep 20"
|
||
},
|
||
]);
|
||
command
|
||
};
|
||
let mut child = command
|
||
.stdin(Stdio::piped())
|
||
.stdout(Stdio::piped())
|
||
.stderr(Stdio::null())
|
||
.spawn()
|
||
.unwrap();
|
||
let (requests, writes, writer_thread) = spawn_writer(child.stdin.take().unwrap());
|
||
let stdout = child.stdout.take().unwrap();
|
||
let (sender, responses) = mpsc::sync_channel(1);
|
||
let reader_thread = thread::spawn(move || {
|
||
let response = read_response(&mut BufReader::new(stdout));
|
||
let _ = sender.try_send(response);
|
||
});
|
||
GemmaEmbedder {
|
||
child,
|
||
requests: Some(requests),
|
||
writes,
|
||
writer_thread: Some(writer_thread),
|
||
reader_thread: Some(reader_thread),
|
||
responses,
|
||
info: WorkerInfo {
|
||
ready: true,
|
||
model: "pipe protocol fixture".into(),
|
||
dimensions: 0,
|
||
device: "cpu".into(),
|
||
threads: 1,
|
||
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,
|
||
query_profile: None,
|
||
}
|
||
}
|
||
|
||
#[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);
|
||
let (sender, receiver) = mpsc::sync_channel(1);
|
||
drop(receiver);
|
||
fixture.requests = Some(sender);
|
||
let error = fixture
|
||
.exchange_with_timeout(b"{}", Duration::from_millis(300))
|
||
.unwrap_err();
|
||
assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
|
||
assert!(fixture.child.try_wait().unwrap().is_some());
|
||
assert!(!fixture.healthy && fixture.requests.is_none());
|
||
assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn entire_request_deadline_kills_and_joins_worker_that_never_reads_stdin() {
|
||
let mut fixture = stalled_protocol_helper(false);
|
||
let started = Instant::now();
|
||
let error = fixture
|
||
.exchange_with_timeout(&vec![b'x'; MAX_REQUEST - 1], Duration::from_millis(300))
|
||
.unwrap_err();
|
||
assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
|
||
assert!(error.to_string().contains("write failed or timed out"));
|
||
assert!(started.elapsed() < Duration::from_secs(5));
|
||
assert!(fixture.child.try_wait().unwrap().is_some());
|
||
assert!(!fixture.healthy && !fixture.info.ready);
|
||
assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn response_wait_uses_same_deadline_and_reaps_silent_reader() {
|
||
let mut fixture = stalled_protocol_helper(true);
|
||
let started = Instant::now();
|
||
let error = fixture
|
||
.exchange_with_timeout(b"{}", Duration::from_millis(300))
|
||
.unwrap_err();
|
||
assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
|
||
assert!(error.to_string().contains("request failed or timed out"));
|
||
assert!(started.elapsed() < Duration::from_secs(5));
|
||
assert!(fixture.child.try_wait().unwrap().is_some());
|
||
assert!(fixture.requests.is_none());
|
||
assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn automatic_startup_transport_failure_retries_cpu_once() {
|
||
let mut attempts = Vec::new();
|
||
let result = start_with_cpu_fallback(false, |cpu| {
|
||
attempts.push(cpu);
|
||
if cpu {
|
||
Ok("cpu ready")
|
||
} else {
|
||
Err(WorkerTransportFailure("worker closed its output before ready".into()).into())
|
||
}
|
||
})
|
||
.unwrap();
|
||
assert_eq!(result, "cpu ready");
|
||
assert_eq!(attempts, [false, true]);
|
||
}
|
||
|
||
#[test]
|
||
fn startup_failure_never_loops_or_retries_explicit_cpu() {
|
||
for cpu in [false, true] {
|
||
let mut attempts = Vec::new();
|
||
let result: Result<()> = start_with_cpu_fallback(cpu, |selected_cpu| {
|
||
attempts.push(selected_cpu);
|
||
Err(WorkerTransportFailure("worker exited".into()).into())
|
||
});
|
||
let error = format!("{:#}", result.unwrap_err());
|
||
assert_eq!(attempts, if cpu { vec![true] } else { vec![false, true] });
|
||
assert!(error.contains("worker exited"));
|
||
if !cpu {
|
||
assert!(error.contains("CPU initialization failed"));
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn malformed_startup_protocol_is_not_retried() {
|
||
let mut attempts = 0;
|
||
let result: Result<()> = start_with_cpu_fallback(false, |_| {
|
||
attempts += 1;
|
||
let mut response = BufReader::new(&b"not JSON\n"[..]);
|
||
let error = read_response(&mut response).unwrap_err();
|
||
assert!(error.starts_with("Invalid Gemma worker JSON:"));
|
||
anyhow::bail!("Invalid Gemma worker startup protocol: {error}");
|
||
});
|
||
assert!(result.unwrap_err().to_string().contains("startup protocol"));
|
||
assert_eq!(attempts, 1);
|
||
}
|
||
|
||
#[test]
|
||
fn incompatible_startup_contract_is_not_retried() {
|
||
let mut attempts = 0;
|
||
let result: Result<()> = start_with_cpu_fallback(false, |_| {
|
||
attempts += 1;
|
||
let handshake = json!({"ready":true,"model":"other-model","dimensions":768,"device":"cpu","threads":1,"runtime":"other-runtime"});
|
||
let info: WorkerInfo = serde_json::from_value(handshake)?;
|
||
ensure!(
|
||
info.model == "google/embeddinggemma-2",
|
||
"Gemma worker handshake violates configured model/device contract or requested role; repair or rebuild the native runtime"
|
||
);
|
||
Ok(())
|
||
});
|
||
assert!(result
|
||
.unwrap_err()
|
||
.to_string()
|
||
.contains("model/device contract"));
|
||
assert_eq!(attempts, 1);
|
||
}
|
||
|
||
#[test]
|
||
fn old_text_worker_handshake_does_not_claim_media_support() {
|
||
let mut handshake = json!({"ready":true,"model":"google/embeddinggemma-2","dimensions":768,"device":"cpu","threads":1,"runtime":"b11461-no-logits-v1"});
|
||
let old: WorkerInfo = serde_json::from_value(handshake.clone()).unwrap();
|
||
assert!(!old.media_supported);
|
||
assert!(!old.code_query_supported);
|
||
assert!(query_operation(&old, true).is_err());
|
||
assert_eq!(query_operation(&old, false).unwrap(), "query");
|
||
let mut new = old.clone();
|
||
new.code_query_supported = true;
|
||
assert_eq!(query_operation(&new, true).unwrap(), "code_query");
|
||
handshake["media_supported"] = json!(true);
|
||
assert!(
|
||
serde_json::from_value::<WorkerInfo>(handshake)
|
||
.unwrap()
|
||
.media_supported
|
||
);
|
||
}
|
||
#[test]
|
||
fn media_image_grid_is_bounded_for_extreme_aspects_and_full_budget() {
|
||
for budget in [70, 140, 280, 1120] {
|
||
for (width, height) in [(1, 16384), (16384, 1), (1920, 1080), (500, 500)] {
|
||
let (w, h) = image_target(width, height, budget).unwrap();
|
||
assert_eq!(w % 48, 0);
|
||
assert_eq!(h % 48, 0);
|
||
assert!(u64::from(w) * u64::from(h) <= u64::from(budget) * 2304);
|
||
}
|
||
}
|
||
assert!(image_target(0, 100, 280).is_err());
|
||
assert!(image_target(100, 100, 1121).is_err());
|
||
}
|
||
#[test]
|
||
fn artifact_verification_rejects_projector_size_and_checksum_mismatch() {
|
||
let path = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/LICENSE-DigUp.txt");
|
||
let bytes = std::fs::read(&path).unwrap();
|
||
let sha = format!("{:x}", Sha256::digest(&bytes));
|
||
assert!(verify_artifact(&path, bytes.len() as u64, &sha).is_ok());
|
||
assert!(verify_artifact(&path, bytes.len() as u64 + 1, &sha).is_err());
|
||
assert!(verify_artifact(&path, bytes.len() as u64, &"0".repeat(64)).is_err());
|
||
}
|
||
#[test]
|
||
#[ignore = "Explicit48MP JPEG resource test; no model needed, run release CPU"]
|
||
fn large_phone_jpeg_scaled_decode_preserves_exif_and_bounds() {
|
||
use image::ImageEncoder;
|
||
let image = image::RgbImage::from_fn(8000, 6000, |x, _| {
|
||
if x < 4000 {
|
||
image::Rgb([240, 10, 10])
|
||
} else {
|
||
image::Rgb([10, 10, 240])
|
||
}
|
||
});
|
||
let mut encoded = Vec::new();
|
||
let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut encoded, 80);
|
||
// Little-endian TIFF orientation tag6: rotate90clockwise.
|
||
encoder
|
||
.set_exif_metadata(vec![
|
||
b'I', b'I', 42, 0, 8, 0, 0, 0, 1, 0, 0x12, 1, 3, 0, 1, 0, 0, 0, 6, 0, 0, 0, 0, 0,
|
||
0, 0,
|
||
])
|
||
.unwrap();
|
||
encoder.encode_image(&image).unwrap();
|
||
drop(image);
|
||
let path =
|
||
std::env::temp_dir().join(format!("searchi-phone48mp-{}.jpg", std::process::id()));
|
||
std::fs::write(&path, &encoded).unwrap();
|
||
drop(encoded);
|
||
let started = std::time::Instant::now();
|
||
let result = decode_image(&path, 280);
|
||
let _ = std::fs::remove_file(&path);
|
||
let pixels = result.unwrap();
|
||
println!(
|
||
"48MPactualJPEG scaleddecode+EXIF+resize {:?}; output{}x{} RGB{}bytes",
|
||
started.elapsed(),
|
||
pixels.width(),
|
||
pixels.height(),
|
||
pixels.as_raw().len()
|
||
);
|
||
assert!(pixels.height() > pixels.width());
|
||
assert!(pixels.get_pixel(pixels.width() / 2, pixels.height() / 4)[0] > 200);
|
||
assert!(pixels.get_pixel(pixels.width() / 2, pixels.height() * 3 / 4)[2] > 200);
|
||
assert!(u64::from(pixels.width()) * u64::from(pixels.height()) <= 280 * 2304);
|
||
}
|
||
#[test]
|
||
fn oversized_jpeg_header_rejected_before_pixel_decode() {
|
||
let mut encoded = Vec::new();
|
||
image::codecs::jpeg::JpegEncoder::new(&mut encoded)
|
||
.encode_image(&image::RgbImage::new(16, 16))
|
||
.unwrap();
|
||
let sof = encoded.windows(2).position(|p| p == [0xff, 0xc0]).unwrap();
|
||
encoded[sof + 5..sof + 7].copy_from_slice(&9000u16.to_be_bytes());
|
||
encoded[sof + 7..sof + 9].copy_from_slice(&9000u16.to_be_bytes());
|
||
let path =
|
||
std::env::temp_dir().join(format!("searchi-too-large-{}.jpg", std::process::id()));
|
||
std::fs::write(&path, encoded).unwrap();
|
||
let result = decode_image(&path, 280);
|
||
let _ = std::fs::remove_file(&path);
|
||
assert!(result
|
||
.unwrap_err()
|
||
.to_string()
|
||
.contains("memory/pixel budget"));
|
||
}
|
||
#[test]
|
||
fn missing_google_model_fails_without_starting_worker() {
|
||
assert!(GemmaEmbedder::load("missing-google-model", "missing-worker", 1).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn retry_policy_never_enables_gpu_or_retries_cpu_as_gpu() {
|
||
assert!(can_retry_cpu("gpu", false));
|
||
assert!(!can_retry_cpu("gpu", true));
|
||
assert!(!can_retry_cpu("cpu", false));
|
||
assert!(!can_retry_cpu("cpu", true));
|
||
}
|
||
|
||
#[test]
|
||
fn protocol_rejects_oversized_truncated_and_malformed_frames() {
|
||
assert!(read_response(&mut std::io::Cursor::new(vec![b'x'; MAX_RESPONSE + 1])).is_err());
|
||
assert!(read_response(&mut std::io::Cursor::new(b"{\"ok\":true}")).is_err());
|
||
assert!(read_response(&mut std::io::Cursor::new(b"not-json\n")).is_err());
|
||
assert_eq!(
|
||
read_response(&mut std::io::Cursor::new(b"{\"ok\":true}\n")).unwrap()["ok"],
|
||
true
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn protocol_rejects_wrong_embedding_shape_count_and_normalization() {
|
||
assert!(extract_vectors(json!({"vectors":[]}), 1).is_err());
|
||
assert!(extract_vectors(json!({"vectors":[[1.0,0.0]]}), 1).is_err());
|
||
assert!(extract_vectors(json!({"vectors":[vec![0.0;GEMMA_DIMENSIONS]]}), 1).is_err());
|
||
let mut invalid = vec![0.0; GEMMA_DIMENSIONS];
|
||
invalid[0] = 1.0;
|
||
invalid[1] = 1.0;
|
||
assert!(extract_vectors(json!({"vectors":[invalid]}), 1).is_err());
|
||
}
|
||
}
|