feat: integrate Google semantic search and branded desktop launcher

crates/searchi-semantic/src/gemma.rs

@@ -0,0 +1,398 @@
+use crate::TextWindow;
+use anyhow::{ensure, Context, Result};
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+use sha2::{Digest, Sha256};
+use std::{
+    fs::File,
+    io::{BufRead, BufReader, Read, Write},
+    path::{Path, PathBuf},
+    process::{Child, ChildStdin, Command, Stdio},
+    sync::mpsc::{self, Receiver},
+    thread,
+    time::Duration,
+};
+
+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";
+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_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,
+}
+
+/// Persistent isolated native worker. No HTTP or remote model fallback.
+pub struct GemmaEmbedder {
+    child: Child,
+    stdin: ChildStdin,
+    responses: Receiver<std::result::Result<Value, String>>,
+    pub info: WorkerInfo,
+    healthy: bool,
+    force_cpu: bool,
+    model_dir: PathBuf,
+    worker_path: PathBuf,
+}
+
+impl Drop for GemmaEmbedder {
+    fn drop(&mut self) {
+        let _ = self.child.kill();
+        let _ = self.child.wait();
+    }
+}
+
+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)
+    }
+
+    pub fn load_with_device(
+        model_dir: impl AsRef<Path>,
+        worker: impl AsRef<Path>,
+        threads: usize,
+        device: Device,
+    ) -> 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);
+        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());
+        // 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().context("Cannot start local Gemma worker")?;
+        let stdin = child.stdin.take().context("Worker stdin missing")?;
+        let stdout = child.stdout.take().context("Worker stdout missing")?;
+        let (sender, responses) = mpsc::sync_channel(1);
+        thread::spawn(move || {
+            let mut reader = BufReader::new(stdout);
+            loop {
+                let response = read_response(&mut reader);
+                let failed = response.is_err();
+                if sender.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();
+                anyhow::bail!("Gemma worker did not initialize: {other:?}");
+            }
+        };
+        let info: WorkerInfo = match serde_json::from_value(ready) {
+            Ok(info) => info,
+            Err(error) => {
+                let _ = child.kill();
+                let _ = child.wait();
+                return Err(error).context("Invalid Gemma worker handshake");
+            }
+        };
+        if !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();
+            anyhow::bail!("Gemma worker handshake violates configured model/device contract");
+        }
+        Ok(Self {
+            child,
+            stdin,
+            responses,
+            info,
+            healthy: true,
+            force_cpu: cpu,
+            model_dir,
+            worker_path: worker,
+        })
+    }
+
+    fn request(&mut self, request: Value) -> Result<Value> {
+        ensure!(
+            self.healthy,
+            "Gemma worker session failed; reopen it explicitly"
+        );
+        let encoded = serde_json::to_vec(&request)?;
+        ensure!(
+            encoded.len() < MAX_REQUEST,
+            "Gemma request exceeds 2 MiB protocol budget"
+        );
+        let response = match self.exchange(&encoded) {
+            Ok(response) => response,
+            Err(error) if can_retry_cpu(&self.info.device, self.force_cpu) => {
+                // 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 replacement = Self::load_with_device(
+                    &self.model_dir,
+                    &self.worker_path,
+                    self.info.threads,
+                    Device::Cpu,
+                )
+                .with_context(|| {
+                    format!("CPU worker restart failed after GPU transport error: {error}")
+                })?;
+                *self = replacement;
+                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["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 exchange(&mut self, encoded: &[u8]) -> Result<Value> {
+        let sent = (|| {
+            self.stdin.write_all(encoded)?;
+            self.stdin.write_all(b"\n")?;
+            self.stdin.flush()
+        })()
+        .context("Cannot write local Gemma worker request");
+        if let Err(error) = sent {
+            self.healthy = false;
+            self.info.ready = false;
+            let _ = self.child.kill();
+            return Err(error);
+        }
+        let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
+            Ok(Ok(value)) => value,
+            other => {
+                self.healthy = false;
+                self.info.ready = false;
+                let _ = self.child.kill();
+                anyhow::bail!("Gemma worker request failed: {other:?}");
+            }
+        };
+        Ok(response)
+    }
+
+    pub fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
+        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!(
+            !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>> {
+        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}))?;
+        Ok(extract_vectors(response, 1)?.remove(0))
+    }
+}
+
+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<()> {
+    let file = File::open(path).with_context(|| {
+        format!(
+            "Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
+            path.display()
+        )
+    })?;
+    ensure!(
+        file.metadata()?.len() == MODEL_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()) == MODEL_SHA256,
+        "EmbeddingGemma 2 model SHA256 mismatch"
+    );
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[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());
+    }
+}