searchi/crates/searchi-service/src/lib.rs

//! Shared offline hybrid search using real Google embeddings and persisted,
//! filtered HNSW retrieval. Production-scale validation remains a release gate.
mod store;

use anyhow::{ensure, Result};
use searchi_core::SearchIndex;
pub use searchi_core::{SearchOptions, SearchResult};
pub use searchi_semantic::WorkerInfo;
use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
use serde::Serialize;
use std::{collections::HashMap, path::Path};
use store::{Store, MAX_SCAN_CHUNKS};

const MAX_BATCH: usize = 8;
const BYTE_SLICE: usize = 64 * 1024;
const BYTE_OVERLAP: usize = 512;
const MAX_PASSAGE_WINDOWS: usize = 2048;
const CHUNK_POLICY: &str = "byte64k-overlap512-title512-v1";
fn generation_fingerprint(model: &str) -> String {
    format!("{model}:{CHUNK_POLICY}")
}

/// Model-specific query/document prompts and actual tokenizer windows belong
/// to the provider. Storage and ranking never assume a particular dimension.
pub trait EmbeddingProvider: Send {
    fn fingerprint(&self) -> &str;
    fn dimensions(&self) -> usize;
    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
    fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>>;
    fn embed_query(&mut self, query: &str) -> Result<Vec<f32>>;
    fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
        None
    }
}

impl EmbeddingProvider for GemmaEmbedder {
    fn fingerprint(&self) -> &str {
        GEMMA_FINGERPRINT
    }
    fn dimensions(&self) -> usize {
        GEMMA_DIMENSIONS
    }
    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
        self.split_windows(text, title)
    }
    fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
        self.embed_documents(texts, title)
    }
    fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
        self.embed_query(query)
    }
    fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
        Some(self.info.clone())
    }
}

#[derive(Debug, Serialize)]
pub struct EmbeddingReport {
    pub passages: usize,
    pub chunks: usize,
    pub generation: i64,
    pub model_fingerprint: String,
}

#[derive(Debug, Serialize)]
pub struct EmbeddingStatus {
    pub generation: Option<i64>,
    pub chunks: usize,
    pub model_fingerprint: Option<String>,
    pub stale: bool,
    /// Rebuild is generation-based: all current passages are pending when stale.
    pub pending: usize,
}

/// Shared local hybrid service. Google EmbeddingGemma 2 is the default and
/// only shipped provider; no alternative model is silently substituted.
pub struct HybridSearch {
    lexical: SearchIndex,
    embedder: Box<dyn EmbeddingProvider>,
    store: Store,
    generation_fingerprint: String,
}

impl HybridSearch {
    /// Actual model/runtime/device reported by the persistent native worker.
    pub fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
        self.embedder.worker_info()
    }
    /// Cheap local status check. No model, runtime worker, or network is loaded.
    /// Corpus revision triggers detect file inserts, updates, and deletions.
    pub fn embedding_status(db: impl AsRef<Path>) -> Result<EmbeddingStatus> {
        let connection =
            rusqlite::Connection::open_with_flags(db, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
        connection.busy_timeout(std::time::Duration::from_secs(5))?;
        let exists: bool = connection.query_row(
            "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name='semantic_meta')",
            [],
            |row| row.get(0),
        )?;
        let pending: usize =
            connection.query_row("SELECT count(*) FROM passages", [], |row| row.get(0))?;
        if !exists {
            return Ok(EmbeddingStatus {
                generation: None,
                chunks: 0,
                model_fingerprint: None,
                stale: true,
                pending,
            });
        }
        let version: i64 = connection.query_row(
            "SELECT schema_version FROM semantic_meta WHERE singleton=1",
            [],
            |row| row.get(0),
        )?;
        if version < 5 {
            return Ok(EmbeddingStatus {
                generation: None,
                chunks: 0,
                model_fingerprint: None,
                stale: true,
                pending,
            });
        }
        ensure!(
            version == 5,
            "Unsupported semantic schema version {version}"
        );
        let (generation,fingerprint,revision,current):(Option<i64>,Option<String>,Option<i64>,i64)=connection.query_row(
            "SELECT g.id,g.fingerprint,CASE WHEN g.ann_blob IS NOT NULL AND g.dimensions=768 THEN g.corpus_revision END,m.corpus_revision FROM semantic_meta m LEFT JOIN semantic_generations g ON g.id=m.active_generation WHERE m.singleton=1",[],
            |row| Ok((row.get(0)?,row.get(1)?,row.get(2)?,row.get(3)?)))?;
        let chunks: usize = connection.query_row(
            "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
            [generation],
            |row| row.get(0),
        )?;
        let stale = revision != Some(current)
            || fingerprint.as_deref() != Some(generation_fingerprint(GEMMA_FINGERPRINT).as_str());
        Ok(EmbeddingStatus {
            generation,
            chunks,
            model_fingerprint: fingerprint,
            stale,
            pending: if stale { pending } else { 0 },
        })
    }

    pub fn open(
        db: impl AsRef<Path>,
        model_dir: impl AsRef<Path>,
        runtime_library: impl AsRef<Path>,
        threads: usize,
    ) -> Result<Self> {
        let embedder = GemmaEmbedder::load(model_dir, runtime_library, threads)?;
        Self::with_provider(db, Box::new(embedder))
    }

    /// Explicit provider injection for future model generations, never fallback.
    pub fn with_provider(
        db: impl AsRef<Path>,
        embedder: Box<dyn EmbeddingProvider>,
    ) -> Result<Self> {
        let lexical = SearchIndex::open(&db)?;
        let store = Store::open(db.as_ref(), embedder.dimensions())?;
        let generation_fingerprint = generation_fingerprint(embedder.fingerprint());
        Ok(Self {
            lexical,
            embedder,
            store,
            generation_fingerprint,
        })
    }

    /// Execute real query inference and load the active ANN generation before
    /// foreground use. Call on a background worker after assets are verified.
    pub fn warmup(&mut self) -> Result<()> {
        self.lexical.prewarm()?;
        let vector = self.embedder.embed_query("warm up local file search")?;
        let _ = self.store.search(
            &vector,
            &self.generation_fingerprint,
            1,
            &SearchOptions::default(),
        )?;
        Ok(())
    }

    /// Build a replacement generation in bounded CPU batches. Inference or
    /// storage errors leave the previous generation active. A crashed build is
    /// reclaimed on the next explicit rebuild under an OS advisory lock.
    pub fn index_embeddings(&mut self) -> Result<EmbeddingReport> {
        let _lock = self.store.build_lock()?;
        let generation = self.store.begin(&self.generation_fingerprint)?;
        let built = self.build(generation);
        match built {
            Ok(report) => {
                self.store.activate(generation)?;
                Ok(report)
            }
            Err(error) => {
                self.store.abandon(generation)?;
                Err(error)
            }
        }
    }

    fn build(&mut self, generation: i64) -> Result<EmbeddingReport> {
        let mut report = EmbeddingReport {
            passages: 0,
            chunks: 0,
            generation,
            model_fingerprint: self.generation_fingerprint.clone(),
        };
        let mut after_id = 0;
        loop {
            let passages = self.lexical.passage_batch(after_id, 64)?;
            if passages.is_empty() {
                break;
            }
            for passage in passages {
                after_id = passage.id;
                if passage.text.trim().is_empty() {
                    continue;
                }
                let title = bounded_title(&passage.filename);
                let windows = passage_windows(&mut *self.embedder, &passage.text, title)?;
                ensure!(
                    report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
                    "Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
                );
                for (batch_index, batch) in windows.chunks(MAX_BATCH).enumerate() {
                    let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
                    let vectors = self.embedder.embed_documents(&texts, title)?;
                    self.store.insert(
                        generation,
                        &passage,
                        batch_index * MAX_BATCH,
                        batch,
                        &vectors,
                    )?;
                }
                report.chunks += windows.len();
                report.passages += 1;
            }
        }
        Ok(report)
    }

    /// Filter each retrieval lane before ranking, preserve filename priority,
    /// then fuse keyword and semantic ranks with reciprocal-rank fusion (k=60).
    pub fn search(
        &mut self,
        query: &str,
        limit: usize,
        options: &SearchOptions,
    ) -> Result<Vec<SearchResult>> {
        let candidates = limit.saturating_mul(4).clamp(20, 1000);
        // Core validates filters and query budgets even for empty requests.
        let mut lexical = self
            .lexical
            .search_with_options(query, candidates, options)?;
        if query.trim().is_empty() || limit == 0 {
            return Ok(Vec::new());
        }
        promote_identifiers(query.trim(), &mut lexical);
        let vector = self.embedder.embed_query(query.trim())?;
        let semantic =
            self.store
                .search(&vector, &self.generation_fingerprint, candidates, options)?;
        Ok(fuse(lexical, semantic, limit.min(1000)))
    }

    /// Semantic lane for reproducible relevance evaluation and diagnostics.
    pub fn search_semantic(
        &mut self,
        query: &str,
        limit: usize,
        options: &SearchOptions,
    ) -> Result<Vec<SearchResult>> {
        let _ = self.lexical.search_with_options(query, 0, options)?;
        if query.trim().is_empty() || limit == 0 {
            return Ok(Vec::new());
        }
        let vector = self.embedder.embed_query(query.trim())?;
        self.store.search(
            &vector,
            &self.generation_fingerprint,
            limit.min(1000),
            options,
        )
    }
}

fn bounded_title(title: &str) -> &str {
    let mut end = title.len().min(512);
    while !title.is_char_boundary(end) {
        end -= 1;
    }
    &title[..end]
}

fn byte_slices(text: &str) -> Result<Vec<(usize, usize)>> {
    ensure!(
        text.len() <= 2 * 1024 * 1024,
        "Passage exceeds 2 MiB service byte budget"
    );
    let mut slices = Vec::new();
    let mut start = 0;
    while start < text.len() {
        let mut end = (start + BYTE_SLICE).min(text.len());
        while !text.is_char_boundary(end) {
            end -= 1;
        }
        ensure!(end > start, "UTF-8 slice made no progress");
        slices.push((start, end));
        if end == text.len() {
            break;
        }
        start = end.saturating_sub(BYTE_OVERLAP);
        while !text.is_char_boundary(start) {
            start -= 1;
        }
    }
    Ok(slices)
}

fn passage_windows(
    provider: &mut dyn EmbeddingProvider,
    text: &str,
    title: &str,
) -> Result<Vec<TextWindow>> {
    let mut windows = Vec::new();
    for (start, end) in byte_slices(text)? {
        let fragment = &text[start..end];
        if fragment.trim().is_empty() {
            continue;
        }
        let mut slice = provider.split_windows(fragment, title)?;
        for window in &mut slice {
            window.start_byte += start;
            window.end_byte += start;
        }
        windows.extend(slice);
        ensure!(
            windows.len() <= MAX_PASSAGE_WINDOWS,
            "Passage exceeds {MAX_PASSAGE_WINDOWS} semantic window budget"
        );
    }
    Ok(windows)
}

fn promote_identifiers(query: &str, results: &mut [SearchResult]) {
    if query.chars().any(char::is_whitespace)
        || !query.chars().any(|c| c.is_ascii_digit() || c == '_')
    {
        return;
    }
    for result in results {
        if result.match_kind == "content" && literal_identifier_match(&result.snippet, query) {
            result.match_kind = "exact_identifier".into();
        }
    }
}

fn literal_identifier_match(text: &str, query: &str) -> bool {
    text.match_indices(query).any(|(start, _)| {
        let end = start + query.len();
        let is_identifier = |c: char| c.is_alphanumeric() || c == '_' || c == '-';
        !text[..start].chars().next_back().is_some_and(is_identifier)
            && !text[end..].chars().next().is_some_and(is_identifier)
    })
}

fn fuse(
    lexical: Vec<SearchResult>,
    semantic: Vec<SearchResult>,
    limit: usize,
) -> Vec<SearchResult> {
    let mut results = HashMap::<String, SearchResult>::new();
    for (rank, mut result) in lexical.into_iter().enumerate() {
        let weight = if result.match_kind == "filename" {
            1.15
        } else {
            1.0
        };
        result.score = weight / (61.0 + rank as f64);
        results.insert(result.path.clone(), result);
    }
    for (rank, mut result) in semantic.into_iter().enumerate() {
        let contribution = 1.0 / (61.0 + rank as f64);
        if let Some(existing) = results.get_mut(&result.path) {
            existing.score += contribution;
            if existing.match_kind == "content" {
                existing.match_kind = "hybrid".into();
            }
        } else {
            result.score = contribution;
            results.insert(result.path.clone(), result);
        }
    }
    let tier = |result: &SearchResult| match result.match_kind.as_str() {
        "exact_filename" => 0,
        "exact_identifier" => 1,
        _ => 2,
    };
    let mut results: Vec<_> = results.into_values().collect();
    results.sort_by(|a, b| {
        tier(a)
            .cmp(&tier(b))
            .then_with(|| b.score.total_cmp(&a.score))
            .then_with(|| a.path.cmp(&b.path))
    });
    results.truncate(limit);
    results
}

#[cfg(test)]
mod tests {
    use super::*;
    fn result(path: &str, kind: &str) -> SearchResult {
        SearchResult {
            path: path.into(),
            filename: path.into(),
            snippet: "evidence".into(),
            score: 1.0,
            match_kind: kind.into(),
            source: None,
            extraction_error: None,
        }
    }
    #[test]
    fn exact_filename_wins_even_when_another_result_matches_both_lanes() {
        let results = fuse(
            vec![
                result("invoice.pdf", "exact_filename"),
                result("notes.txt", "content"),
            ],
            vec![
                result("notes.txt", "semantic"),
                result("another.txt", "semantic"),
            ],
            10,
        );
        assert_eq!(results[0].path, "invoice.pdf");
        assert_eq!(results[1].match_kind, "hybrid");
        assert_eq!(results.len(), 3);
    }

    #[test]
    fn identifiers_use_literal_boundaries_and_survive_semantic_competition() {
        let mut lexical = vec![result("support.txt", "content")];
        lexical[0].snippet = "Error E_SEARCH_071 means upgrade the app".into();
        promote_identifiers("E_SEARCH_071", &mut lexical);
        let results = fuse(lexical, vec![result("wrong.txt", "semantic")], 10);
        assert_eq!(results[0].match_kind, "exact_identifier");
        assert!(!literal_identifier_match(
            "E_SEARCH_071_EXTRA",
            "E_SEARCH_071"
        ));
    }

    #[test]
    fn keyword_semantic_consensus_beats_weak_partial_filename() {
        let results = fuse(
            vec![
                result("meeting-cancellation.txt", "filename"),
                result("supplier-agreement.txt", "content"),
            ],
            vec![result("supplier-agreement.txt", "semantic")],
            10,
        );
        assert_eq!(results[0].path, "supplier-agreement.txt");
        assert_eq!(results[0].match_kind, "hybrid");
        assert_eq!(results[1].match_kind, "filename");
    }

    #[test]
    fn byte_slices_preserve_utf8_overlap_offsets_and_complete_tail() {
        let text = format!("{}TAIL", "Ä🐈中文 ".repeat(10000));
        let slices = byte_slices(&text).unwrap();
        assert!(slices.len() > 1);
        assert_eq!(slices[0].0, 0);
        assert_eq!(slices.last().unwrap().1, text.len());
        for (start, end) in &slices {
            assert!(text.is_char_boundary(*start) && text.is_char_boundary(*end));
            assert!(end - start <= BYTE_SLICE);
        }
        for pair in slices.windows(2) {
            assert!(pair[1].0 < pair[0].1 && pair[1].1 > pair[0].1);
        }
        assert!(text[slices.last().unwrap().0..].ends_with("TAIL"));
        assert!(byte_slices(&"x".repeat(2 * 1024 * 1024 + 1)).is_err());
        let title = "中文".repeat(120);
        assert!(bounded_title(&title).len() <= 512 && title.starts_with(bounded_title(&title)));
    }

    #[test]
    #[ignore = "Requires explicitly provisioned Google model and CPU worker"]
    fn google_cpu_preserves_and_retrieves_content_after_64k_byte_boundary() -> Result<()> {
        let model =
            std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
        let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
        let mut provider =
            GemmaEmbedder::load_with_device(model, worker, 1, searchi_semantic::Device::Cpu)?;
        let text=format!("{}\nÜbernachtungskosten auf Dienstreisen werden gegen Vorlage der Hotelrechnung erstattet.","Routine system checks were completed successfully. ".repeat(1500));
        assert!(text.len() > BYTE_SLICE);
        let windows = passage_windows(&mut provider, &text, "routine-log.txt")?;
        let last = windows.last().unwrap();
        assert_eq!(last.end_byte, text.len());
        assert!(last.start_byte >= BYTE_SLICE - BYTE_OVERLAP);
        assert!(text.find("Übernachtungskosten").unwrap() > BYTE_SLICE);
        assert!(last.text.contains("Hotelrechnung"));
        assert_eq!(
            text.get(last.start_byte..last.end_byte),
            Some(last.text.as_str())
        );
        let vectors = provider.embed_documents(
            &[windows[0].text.as_str(), last.text.as_str()],
            "routine-log.txt",
        )?;
        let query = provider.embed_query(
            "Where does it say that hotel stays during work travel can be reimbursed?",
        )?;
        let cosine = |vector: &[f32]| vector.iter().zip(&query).map(|(a, b)| a * b).sum::<f32>();
        assert!(
            cosine(&vectors[1]) > cosine(&vectors[0]),
            "Real model did not retrieve the tail window"
        );
        Ok(())
    }
}