searchi/crates/searchi-semantic/src/gemma.rs

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());
    }
}