lattice-embed 0.12.0

//! Native, pure-Rust embedding service backed by `lattice-inference`.
//!
//! Model loading is lazy and cancellation-safe; BERT-family and Qwen models take different
//! loading and batching paths. See `docs/service.md` for lifecycle and persistence details.

use super::{
    EmbeddingRole, EmbeddingService, EmbeddingsWithReport, MAX_TEXT_BYTES, TokenCount,
    ValidatedTextBatch,
};
use crate::error::{EmbedError, Result};
use crate::model::{EmbeddingModel, ModelConfig};
use async_trait::async_trait;
use lattice_inference::model::embeddinggemma2::EmbeddingGemma2Model;
use lattice_inference::{BertModel, EmbeddingTokenCounts, QwenModel};
use std::sync::{Arc, OnceLock};
use tracing::{info, warn};

/// Loaded model: BERT-family (encoder), Qwen (decoder) or EmbeddingGemma 2 (bidirectional Gemma encoder).
enum LoadedModel {
    Bert(Arc<BertModel>),
    Qwen(Arc<QwenModel>),
    EmbeddingGemma2 {
        model: Arc<EmbeddingGemma2Model>,
        output_dim: Option<usize>,
    },
}

// Wrapped model types provide the required thread-safety — see docs/service.md.

impl LoadedModel {
    fn encode_batch(&self, texts: &[&str]) -> std::result::Result<Vec<Vec<f32>>, String> {
        match self {
            LoadedModel::Bert(m) => m.encode_batch(texts).map_err(|e| e.to_string()),
            // For Qwen, use per-item encode() which checks the cache.
            LoadedModel::Qwen(m) => {
                let mut results = Vec::with_capacity(texts.len());
                for text in texts {
                    results.push(m.encode(text).map_err(|e| e.to_string())?);
                }
                Ok(results)
            }
            LoadedModel::EmbeddingGemma2 { model, output_dim } => texts
                .iter()
                .map(|text| model.encode(text, *output_dim).map_err(|e| e.to_string()))
                .collect(),
        }
    }

    /// Token counts per text, from the tokenizer's own pre-truncation length.
    fn count_batch(&self, texts: &[&str]) -> std::result::Result<Vec<TokenCount>, String> {
        match self {
            LoadedModel::Bert(m) => Ok(texts
                .iter()
                .map(|text| token_count(m.embedding_token_counts(text)))
                .collect()),
            LoadedModel::Qwen(m) => Ok(texts
                .iter()
                .map(|text| token_count(m.embedding_token_counts(text)))
                .collect()),
            LoadedModel::EmbeddingGemma2 { model, .. } => texts
                .iter()
                .map(|text| {
                    model
                        .tokenize_with_counts(text)
                        .map(|(_, counts)| token_count(counts))
                        .map_err(|e| e.to_string())
                })
                .collect(),
        }
    }

    /// Embeddings from the same computation as [`LoadedModel::encode_batch`], plus counts.
    ///
    /// EmbeddingGemma 2 tokenizes once and embeds those ids, which is exactly what its
    /// `encode` does. BERT and Qwen tokenize for the counts and again inside their
    /// encode entry points, which take text rather than ids.
    fn encode_batch_with_report(
        &self,
        texts: &[&str],
    ) -> std::result::Result<(Vec<Vec<f32>>, Vec<TokenCount>), String> {
        match self {
            LoadedModel::Bert(_) | LoadedModel::Qwen(_) => {
                let counts = self.count_batch(texts)?;
                Ok((self.encode_batch(texts)?, counts))
            }
            LoadedModel::EmbeddingGemma2 { model, output_dim } => {
                let mut embeddings = Vec::with_capacity(texts.len());
                let mut counts = Vec::with_capacity(texts.len());
                for text in texts {
                    let (ids, text_counts) = model
                        .tokenize_with_counts(text)
                        .map_err(|e| e.to_string())?;
                    embeddings.push(
                        model
                            .encode_ids(&ids, *output_dim)
                            .map_err(|e| e.to_string())?,
                    );
                    counts.push(token_count(text_counts));
                }
                Ok((embeddings, counts))
            }
        }
    }

    fn cache_size(&self) -> usize {
        match self {
            LoadedModel::Qwen(m) => m.cache_size(),
            _ => 0,
        }
    }
}

fn token_count(counts: EmbeddingTokenCounts) -> TokenCount {
    TokenCount::new(counts.before_truncation, counts.embedded)
}

/// Borrowed views of the prepared texts: the instruction-prefixed copies when a role added
/// one, otherwise the caller's own text.
fn prepared_views<'a>(
    texts: &'a ValidatedTextBatch<'_>,
    prefixed: Option<&'a [String]>,
) -> Vec<&'a str> {
    match prefixed {
        Some(prefixed) => prefixed.iter().map(String::as_str).collect(),
        None => (0..texts.len()).map(|index| texts.get(index)).collect(),
    }
}

/// **Unstable**: model-loading API still evolving; signature may change as lattice-inference matures.
///
/// Pure-Rust local embedding service for one BERT-family or Qwen model configuration.
///
/// It memoizes the load result independently of cancellation.
/// See [`docs/service.md`](../../docs/service.md#nativeembeddingservice-implementation-notes) for loading and architecture details.
pub struct NativeEmbeddingService {
    model: Arc<OnceLock<std::result::Result<LoadedModel, String>>>,
    model_config: ModelConfig,
}

impl Default for NativeEmbeddingService {
    fn default() -> Self {
        Self::new()
    }
}

const LATTICE_EMBED_DIM: &str = "LATTICE_EMBED_DIM";

fn model_config_from_env(model: EmbeddingModel) -> Result<ModelConfig> {
    let output_dim = match std::env::var(LATTICE_EMBED_DIM) {
        Ok(raw) if raw.trim().is_empty() => None,
        Ok(raw) => {
            let dim = raw.trim().parse::<usize>().map_err(|e| {
                EmbedError::InvalidInput(format!("invalid {LATTICE_EMBED_DIM}={raw:?}: {e}"))
            })?;
            Some(dim)
        }
        Err(std::env::VarError::NotPresent) => None,
        Err(e) => {
            return Err(EmbedError::InvalidInput(format!(
                "invalid {LATTICE_EMBED_DIM}: {e}"
            )));
        }
    };
    ModelConfig::try_new(model, output_dim)
}

impl NativeEmbeddingService {
    /// **Unstable**: constructor signature may change; use `EmbeddingService` trait for stable API.
    pub fn new() -> Self {
        Self {
            model: Arc::new(OnceLock::new()),
            model_config: ModelConfig::new(EmbeddingModel::default()),
        }
    }

    /// **Unstable**: constructor signature may change; use `EmbeddingService` trait for stable API.
    pub fn with_model(model_type: EmbeddingModel) -> Self {
        Self {
            model: Arc::new(OnceLock::new()),
            model_config: ModelConfig::new(model_type),
        }
    }

    /// **Unstable**: create with explicit model config (model + optional MRL truncation dim).
    pub fn with_model_config(model_config: ModelConfig) -> Result<Self> {
        model_config.validate()?;
        Ok(Self {
            model: Arc::new(OnceLock::new()),
            model_config,
        })
    }

    /// **Unstable**: create with model config read from `LATTICE_EMBED_DIM` env var.
    pub fn with_model_from_env(model_type: EmbeddingModel) -> Result<Self> {
        let config = model_config_from_env(model_type)?;
        Ok(Self {
            model: Arc::new(OnceLock::new()),
            model_config: config,
        })
    }

    /// **Unstable**: persistence API may be moved to a separate manager type.
    pub fn save_cache(&self) -> Result<usize> {
        let Some(Ok(model)) = self.model.get() else {
            return Ok(0);
        };
        match model {
            LoadedModel::Qwen(m) => {
                let model_name = self.model_config.model.to_string();
                let path = embedding_cache_path(&model_name, m.dimensions());
                m.cache_save(&path)
                    .map_err(|e| EmbedError::InferenceFailed(e.to_string()))
            }
            _ => Ok(0),
        }
    }

    /// **Unstable**: internal diagnostic; may be removed or moved to metrics.
    pub fn cache_size(&self) -> usize {
        self.model
            .get()
            .and_then(|r| r.as_ref().ok())
            .map(LoadedModel::cache_size)
            .unwrap_or(0)
    }

    /// **Unstable**: preload the model without performing an encode pass.
    ///
    /// Returns the same loading errors as the first `embed` call.
    /// See [`docs/service.md`](../../docs/service.md#nativeembeddingservice-implementation-notes) for preload behavior.
    pub async fn ensure_loaded(&self) -> Result<()> {
        self.ensure_model().await.map(|_| ())
    }

    /// Ensures the shared model load has completed.
    ///
    /// See [`docs/service.md`](../../docs/service.md#nativeembeddingservice-implementation-notes) for the cancellation invariant.
    #[expect(
        clippy::expect_used,
        reason = "the OnceCell is initialized by the spawn_blocking call that precedes this read"
    )]
    async fn ensure_model(&self) -> Result<&LoadedModel> {
        // Fast path: already loaded.
        if let Some(result) = self.model.get() {
            return result
                .as_ref()
                .map_err(|e| EmbedError::ModelInitialization(e.clone()));
        }

        // The shared lock outlives a cancelled caller — see docs/service.md.
        let model_lock = self.model.clone();
        let model_config = self.model_config;

        tokio::task::spawn_blocking(move || {
            // Serialize initialization on the blocking thread pool.
            model_lock.get_or_init(|| load_model_sync(model_config));
        })
        .await
        .map_err(|e| EmbedError::ModelInitialization(e.to_string()))?;

        self.model
            .get()
            .expect("set by spawn_blocking")
            .as_ref()
            .map_err(|e| EmbedError::ModelInitialization(e.clone()))
    }

    /// Encodes text that has already been prepared, meaning any retrieval
    /// instruction is applied and the caller-facing contract is checked.
    ///
    /// The length bound here is a memory backstop, not the published cap: it
    /// admits the caller's allowance plus whatever instruction this model
    /// prepends, which is zero bytes for symmetric models.
    async fn encode_prepared(
        &self,
        texts: &[&str],
        model: EmbeddingModel,
    ) -> Result<Vec<Vec<f32>>> {
        if model != self.model_config.model {
            return Err(EmbedError::InvalidInput(format!(
                "requested model {:?} but this service is loaded with {:?}",
                model, self.model_config.model
            )));
        }
        super::validate_texts_bounded(
            texts,
            MAX_TEXT_BYTES.saturating_add(model.max_instruction_bytes()),
        )?;

        let loaded = self.ensure_model().await?;
        loaded
            .encode_batch(texts)
            .map_err(EmbedError::InferenceFailed)
    }

    /// The checks `encode_prepared` makes before it touches the model, for the count and
    /// report paths.
    fn check_prepared(&self, texts: &[&str], model: EmbeddingModel) -> Result<()> {
        if model != self.model_config.model {
            return Err(EmbedError::InvalidInput(format!(
                "requested model {:?} but this service is loaded with {:?}",
                model, self.model_config.model
            )));
        }
        super::validate_texts_bounded(
            texts,
            MAX_TEXT_BYTES.saturating_add(model.max_instruction_bytes()),
        )
    }

    async fn count_prevalidated_with_role(
        &self,
        texts: ValidatedTextBatch<'_>,
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<TokenCount>> {
        let prefixed = role
            .instruction(model)
            .map(|prefix| texts.to_owned_with_prefix(Some(prefix)));
        let views = prepared_views(&texts, prefixed.as_deref());
        self.check_prepared(&views, model)?;
        let loaded = self.ensure_model().await?;
        loaded
            .count_batch(&views)
            .map_err(EmbedError::InferenceFailed)
    }

    async fn report_prevalidated_with_role(
        &self,
        texts: ValidatedTextBatch<'_>,
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<EmbeddingsWithReport> {
        let prefixed = role
            .instruction(model)
            .map(|prefix| texts.to_owned_with_prefix(Some(prefix)));
        let views = prepared_views(&texts, prefixed.as_deref());
        self.check_prepared(&views, model)?;
        let loaded = self.ensure_model().await?;
        let (embeddings, token_counts) = loaded
            .encode_batch_with_report(&views)
            .map_err(EmbedError::InferenceFailed)?;
        Ok(EmbeddingsWithReport::new(embeddings, token_counts))
    }

    async fn encode_prevalidated_with_role(
        &self,
        texts: ValidatedTextBatch<'_>,
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<Vec<f32>>> {
        let prefix = role.instruction(model);
        if prefix.is_none()
            && let Some(borrowed) = texts.borrowed()
        {
            return self.encode_prepared(borrowed, model).await;
        }

        if prefix.is_none() {
            let borrowed = (0..texts.len())
                .map(|index| texts.get(index))
                .collect::<Vec<_>>();
            return self.encode_prepared(&borrowed, model).await;
        }

        let prepared = texts.to_owned_with_prefix(prefix);
        let borrowed = prepared.iter().map(String::as_str).collect::<Vec<_>>();
        self.encode_prepared(&borrowed, model).await
    }
}

/// Synchronous model loading (runs on blocking thread pool).
fn load_model_sync(model_config: ModelConfig) -> std::result::Result<LoadedModel, String> {
    match model_config.model {
        EmbeddingModel::BgeSmallEnV15
        | EmbeddingModel::BgeBaseEnV15
        | EmbeddingModel::BgeLargeEnV15
        | EmbeddingModel::MultilingualE5Small
        | EmbeddingModel::MultilingualE5Base
        | EmbeddingModel::AllMiniLmL6V2
        | EmbeddingModel::ParaphraseMultilingualMiniLmL12V2 => {
            let model_name = match model_config.model {
                EmbeddingModel::BgeSmallEnV15 => "bge-small-en-v1.5",
                EmbeddingModel::BgeBaseEnV15 => "bge-base-en-v1.5",
                EmbeddingModel::BgeLargeEnV15 => "bge-large-en-v1.5",
                EmbeddingModel::MultilingualE5Small => "multilingual-e5-small",
                EmbeddingModel::MultilingualE5Base => "multilingual-e5-base",
                EmbeddingModel::AllMiniLmL6V2 => "all-minilm-l6-v2",
                EmbeddingModel::ParaphraseMultilingualMiniLmL12V2 => {
                    "paraphrase-multilingual-minilm-l12-v2"
                }
                _ => unreachable!(),
            };
            info!(model = model_name, "loading native BERT embedding model");
            let mut bert = BertModel::from_pretrained(model_name).map_err(|e| e.to_string())?;
            // Route each model family through its correct pooling strategy.
            // BGE uses CLS pooling; E5 and MiniLM use mean pooling.
            if let Some(pooling) = model_config.model.bert_pooling() {
                bert.set_pooling(pooling);
            }
            Ok(LoadedModel::Bert(Arc::new(bert)))
        }
        EmbeddingModel::Qwen3Embedding0_6B | EmbeddingModel::Qwen3Embedding4B => {
            load_qwen_model(model_config)
        }
        EmbeddingModel::EmbeddingGemma2 => load_embeddinggemma2_model(model_config),
        other => Err(format!("unsupported model: {other:?}")),
    }
}

fn load_qwen_model(model_config: ModelConfig) -> std::result::Result<LoadedModel, String> {
    model_config.validate().map_err(|e| e.to_string())?;
    let model_type = model_config.model;
    let model_name = model_type.to_string();
    info!(
        model = %model_name,
        output_dim = ?model_config.output_dim,
        "loading Qwen embedding model"
    );
    let model_dir = qwen_model_dir(model_type).map_err(|e| e.to_string())?;
    let mut model = QwenModel::from_directory(&model_dir).map_err(|e| e.to_string())?;
    model.set_output_dim(model_config.output_dim);
    let cache_path = embedding_cache_path(&model_name, model.dimensions());
    match model.cache_load(&cache_path) {
        Ok(n) if n > 0 => {
            info!(entries = n, path = %cache_path.display(), "loaded embedding cache")
        }
        Ok(_) => {}
        Err(e) => {
            warn!(
                path = %cache_path.display(),
                error = %e,
                "embedding cache failed integrity check, ignoring (will regenerate on next save)"
            )
        }
    }
    Ok(LoadedModel::Qwen(Arc::new(model)))
}

fn load_embeddinggemma2_model(
    model_config: ModelConfig,
) -> std::result::Result<LoadedModel, String> {
    model_config.validate().map_err(|e| e.to_string())?;
    info!(
        output_dim = ?model_config.output_dim,
        "loading EmbeddingGemma 2 embedding model"
    );
    let model_dir = embeddinggemma2_model_dir().map_err(|e| e.to_string())?;
    let model = EmbeddingGemma2Model::from_model_dir(&model_dir).map_err(|e| e.to_string())?;
    Ok(LoadedModel::EmbeddingGemma2 {
        model: Arc::new(model),
        output_dim: model_config.output_dim,
    })
}

/// Locate the EmbeddingGemma 2 model directory: `LATTICE_EMBEDDINGGEMMA2_MODEL_DIR`, else
/// `~/.lattice/models/embeddinggemma-2`.
fn embeddinggemma2_model_dir() -> Result<std::path::PathBuf> {
    if let Ok(dir) = std::env::var("LATTICE_EMBEDDINGGEMMA2_MODEL_DIR") {
        return Ok(std::path::PathBuf::from(dir));
    }
    let home = std::env::var("HOME")
        .map_err(|_| EmbedError::ModelInitialization("HOME not set".into()))?;
    let dir = std::path::PathBuf::from(home)
        .join(".lattice")
        .join("models")
        .join("embeddinggemma-2");
    if dir.join("model.safetensors").exists() {
        Ok(dir)
    } else {
        Err(EmbedError::ModelInitialization(format!(
            "EmbeddingGemma 2 model not found at {dir}. Download it with:\n  huggingface-cli download {repo} --local-dir {dir}",
            dir = dir.display(),
            repo = EmbeddingModel::EmbeddingGemma2.model_id()
        )))
    }
}

/// Path for persistent embedding cache: ~/.lattice/cache/embed_{model}_{dim}d.bin
fn embedding_cache_path(model: &str, dim: usize) -> std::path::PathBuf {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
    std::path::PathBuf::from(home)
        .join(".lattice")
        .join("cache")
        .join(format!("embed_{model}_{dim}d.bin"))
}

/// Locate Qwen3-Embedding model directory for the given model variant.
fn qwen_model_dir(model_type: EmbeddingModel) -> Result<std::path::PathBuf> {
    // Check env override first — applies to whichever Qwen model is loaded.
    if let Ok(dir) = std::env::var("LATTICE_QWEN_MODEL_DIR") {
        return Ok(std::path::PathBuf::from(dir));
    }

    let slug = match model_type {
        EmbeddingModel::Qwen3Embedding0_6B => "qwen3-embedding-0.6b",
        EmbeddingModel::Qwen3Embedding4B => "qwen3-embedding-4b",
        other => {
            return Err(EmbedError::ModelInitialization(format!(
                "not a Qwen model: {other}"
            )));
        }
    };

    let home = std::env::var("HOME")
        .map_err(|_| EmbedError::ModelInitialization("HOME not set".into()))?;
    let dir = std::path::PathBuf::from(home)
        .join(".lattice")
        .join("models")
        .join(slug);

    if dir.join("model.safetensors").exists() || dir.join("model.safetensors.index.json").exists() {
        Ok(dir)
    } else {
        Err(EmbedError::ModelInitialization(format!(
            "Qwen3 model not found at {dir}. Download it with:\n  huggingface-cli download {repo} --local-dir {dir}",
            dir = dir.display(),
            repo = model_type.model_id()
        )))
    }
}

#[async_trait]
impl EmbeddingService for NativeEmbeddingService {
    async fn embed(&self, texts: &[String], model: EmbeddingModel) -> Result<Vec<Vec<f32>>> {
        let texts = ValidatedTextBatch::new(texts)?;
        self.encode_prevalidated_with_role(texts, model, EmbeddingRole::Generic)
            .await
    }

    async fn embed_with_role(
        &self,
        texts: &[String],
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<Vec<f32>>> {
        let texts = ValidatedTextBatch::new(texts)?;
        self.encode_prevalidated_with_role(texts, model, role).await
    }

    async fn embed_with_role_prevalidated(
        &self,
        texts: ValidatedTextBatch<'_>,
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<Vec<f32>>> {
        self.encode_prevalidated_with_role(texts, model, role).await
    }

    async fn count_tokens(
        &self,
        texts: &[String],
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<TokenCount>> {
        let texts = ValidatedTextBatch::new(texts)?;
        self.count_prevalidated_with_role(texts, model, role).await
    }

    async fn embed_with_report(
        &self,
        texts: &[String],
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<EmbeddingsWithReport> {
        let texts = ValidatedTextBatch::new(texts)?;
        self.report_prevalidated_with_role(texts, model, role).await
    }

    fn model_config(&self, model: EmbeddingModel) -> ModelConfig {
        if model == self.model_config.model {
            self.model_config
        } else {
            ModelConfig::new(model)
        }
    }

    fn supports_model(&self, model: EmbeddingModel) -> bool {
        model == self.model_config.model
    }

    fn name(&self) -> &'static str {
        "native-bert"
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cache_path_contains_dim_in_filename() {
        let path = embedding_cache_path("qwen3-embedding-4b", 1024);
        let filename = path.file_name().unwrap().to_str().unwrap();
        assert_eq!(filename, "embed_qwen3-embedding-4b_1024d.bin");
    }

    #[test]
    fn test_cache_path_different_dims_produce_different_paths() {
        let path_1024 = embedding_cache_path("qwen3-embedding-4b", 1024);
        let path_2560 = embedding_cache_path("qwen3-embedding-4b", 2560);
        assert_ne!(path_1024, path_2560);
        assert!(path_1024.to_string_lossy().contains("1024d"));
        assert!(path_2560.to_string_lossy().contains("2560d"));
    }

    #[test]
    fn test_cache_path_model_slug_differentiates_variants() {
        let path_4b = embedding_cache_path("qwen3-embedding-4b", 2560);
        let path_06b = embedding_cache_path("qwen3-embedding-0.6b", 1024);
        assert_ne!(path_4b, path_06b);
        assert!(path_4b.to_string_lossy().contains("qwen3-embedding-4b"));
        assert!(path_06b.to_string_lossy().contains("qwen3-embedding-0.6b"));
    }

    #[test]
    fn test_cache_path_same_model_same_dim_same_path() {
        let p1 = embedding_cache_path("qwen3-embedding-4b", 1024);
        let p2 = embedding_cache_path("qwen3-embedding-4b", 1024);
        assert_eq!(p1, p2);
    }
}