feat: add multimodal search and granular performance diagnostics

crates/searchi-semantic/src/gemma.rs

@@ -1,24 +1,35 @@
 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,
+    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 MAX_REQUEST: usize = 2 * 1024 * 1024;
+const MAX_REQUEST: usize = 16 * 1024 * 1024;
 const MAX_RESPONSE: usize = 4 * 1024 * 1024;
 
 #[derive(Debug, Clone, Copy)]
@@ -35,6 +46,8 @@ pub struct WorkerInfo {
     pub device: String,
     pub threads: usize,
     pub runtime: String,
+    #[serde(default)]
+    pub media_supported: bool,
 }
 
 /// Persistent isolated native worker. No HTTP or remote model fallback.
@@ -47,6 +60,8 @@ pub struct GemmaEmbedder {
     force_cpu: bool,
     model_dir: PathBuf,
     worker_path: PathBuf,
+    projector_verified: bool,
+    query_profile: Option<BTreeMap<String, u64>>,
 }
 
 impl Drop for GemmaEmbedder {
@@ -161,6 +176,8 @@ impl GemmaEmbedder {
             force_cpu: cpu,
             model_dir,
             worker_path: worker,
+            projector_verified: false,
+            query_profile: None,
         })
     }
 
@@ -169,10 +186,12 @@ impl GemmaEmbedder {
             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 2 MiB protocol budget"
+            "Gemma request exceeds 16 MiB protocol budget"
         );
         let response = match self.exchange(&encoded) {
             Ok(response) => response,
@@ -180,6 +199,7 @@ impl GemmaEmbedder {
                 // 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_with_device(
                     &self.model_dir,
                     &self.worker_path,
@@ -189,7 +209,13 @@ impl GemmaEmbedder {
                 .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),
@@ -201,6 +227,11 @@ impl GemmaEmbedder {
             );
             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;
@@ -215,18 +246,21 @@ impl GemmaEmbedder {
     }
 
     fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
+        let write_started = self.query_profile.as_ref().map(|_| Instant::now());
         let sent = (|| {
             self.stdin.write_all(encoded)?;
             self.stdin.write_all(b"\n")?;
             self.stdin.flush()
         })()
         .context("Cannot write local Gemma worker request");
+        self.record_query_duration("request_write_us", write_started);
         if let Err(error) = sent {
             self.healthy = false;
             self.info.ready = false;
             let _ = self.child.kill();
             return Err(error);
         }
+        let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
         let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
             Ok(Ok(value)) => value,
             other => {
@@ -236,6 +270,7 @@ impl GemmaEmbedder {
                 anyhow::bail!("Gemma worker request failed: {other:?}");
             }
         };
+        self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
         Ok(response)
     }
 
@@ -281,15 +316,265 @@ impl GemmaEmbedder {
     }
 
     pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_profile = None;
+        self.embed_query_inner(query, 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)
+    }
+
+    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) -> Result<Vec<f32>> {
         ensure!(
             !query.trim().is_empty() && query.len() <= 4096,
             "Query must be nonempty and at most 4096 UTF-8 bytes"
         );
-        let response = self.request(json!({"op":"query","text":query}))?;
+        let response =
+            self.request(json!({"op":"query","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!(
+            !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!(
+            !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
 }
@@ -331,6 +616,10 @@ fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
 }
 
 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",
@@ -338,7 +627,7 @@ fn verify_model(path: &Path) -> Result<()> {
         )
     })?;
     ensure!(
-        file.metadata()?.len() == MODEL_BYTES,
+        file.metadata()?.len() == bytes,
         "EmbeddingGemma 2 model size mismatch"
     );
     let mut reader = BufReader::new(file);
@@ -352,7 +641,7 @@ fn verify_model(path: &Path) -> Result<()> {
         hash.update(&buffer[..count]);
     }
     ensure!(
-        format!("{:x}", hash.finalize()) == MODEL_SHA256,
+        format!("{:x}", hash.finalize()) == sha,
         "EmbeddingGemma 2 model SHA256 mismatch"
     );
     Ok(())
@@ -362,6 +651,101 @@ fn verify_model(path: &Path) -> Result<()> {
 mod tests {
     use super::*;
     #[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);
+        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());
     }