lattice-embed 0.9.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, MAX_TEXT_BYTES, ValidatedTextBatch};
use crate::error::{EmbedError, Result};
use crate::model::{EmbeddingModel, ModelConfig};
use async_trait::async_trait;
use lattice_inference::{BertModel, QwenModel};
use std::sync::{Arc, OnceLock};
use tracing::{info, warn};

/// Loaded model — either BERT-family (encoder) or Qwen (decoder).
enum LoadedModel {
    Bert(Arc<BertModel>),
    Qwen(Arc<QwenModel>),
}

// 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)
            }
        }
    }

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

/// **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.
    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)
    }

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

/// 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
    }

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