lattice_embed/service/
native.rs

1//! Native, pure-Rust embedding service backed by `lattice-inference`.
2//!
3//! Model loading is lazy and cancellation-safe; BERT-family and Qwen models take different
4//! loading and batching paths. See `docs/service.md` for lifecycle and persistence details.
5
6use super::{EmbeddingRole, EmbeddingService, MAX_TEXT_BYTES, ValidatedTextBatch};
7use crate::error::{EmbedError, Result};
8use crate::model::{EmbeddingModel, ModelConfig};
9use async_trait::async_trait;
10use lattice_inference::{BertModel, QwenModel};
11use std::sync::{Arc, OnceLock};
12use tracing::{info, warn};
13
14/// Loaded model — either BERT-family (encoder) or Qwen (decoder).
15enum LoadedModel {
16    Bert(Arc<BertModel>),
17    Qwen(Arc<QwenModel>),
18}
19
20// Wrapped model types provide the required thread-safety — see docs/service.md.
21
22impl LoadedModel {
23    fn encode_batch(&self, texts: &[&str]) -> std::result::Result<Vec<Vec<f32>>, String> {
24        match self {
25            LoadedModel::Bert(m) => m.encode_batch(texts).map_err(|e| e.to_string()),
26            // For Qwen, use per-item encode() which checks the cache.
27            LoadedModel::Qwen(m) => {
28                let mut results = Vec::with_capacity(texts.len());
29                for text in texts {
30                    results.push(m.encode(text).map_err(|e| e.to_string())?);
31                }
32                Ok(results)
33            }
34        }
35    }
36
37    fn cache_size(&self) -> usize {
38        match self {
39            LoadedModel::Qwen(m) => m.cache_size(),
40            _ => 0,
41        }
42    }
43}
44
45/// **Unstable**: model-loading API still evolving; signature may change as lattice-inference matures.
46///
47/// Pure-Rust local embedding service for one BERT-family or Qwen model configuration.
48///
49/// It memoizes the load result independently of cancellation.
50/// See [`docs/service.md`](../../docs/service.md#nativeembeddingservice-implementation-notes) for loading and architecture details.
51pub struct NativeEmbeddingService {
52    model: Arc<OnceLock<std::result::Result<LoadedModel, String>>>,
53    model_config: ModelConfig,
54}
55
56impl Default for NativeEmbeddingService {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62const LATTICE_EMBED_DIM: &str = "LATTICE_EMBED_DIM";
63
64fn model_config_from_env(model: EmbeddingModel) -> Result<ModelConfig> {
65    let output_dim = match std::env::var(LATTICE_EMBED_DIM) {
66        Ok(raw) if raw.trim().is_empty() => None,
67        Ok(raw) => {
68            let dim = raw.trim().parse::<usize>().map_err(|e| {
69                EmbedError::InvalidInput(format!("invalid {LATTICE_EMBED_DIM}={raw:?}: {e}"))
70            })?;
71            Some(dim)
72        }
73        Err(std::env::VarError::NotPresent) => None,
74        Err(e) => {
75            return Err(EmbedError::InvalidInput(format!(
76                "invalid {LATTICE_EMBED_DIM}: {e}"
77            )));
78        }
79    };
80    ModelConfig::try_new(model, output_dim)
81}
82
83impl NativeEmbeddingService {
84    /// **Unstable**: constructor signature may change; use `EmbeddingService` trait for stable API.
85    pub fn new() -> Self {
86        Self {
87            model: Arc::new(OnceLock::new()),
88            model_config: ModelConfig::new(EmbeddingModel::default()),
89        }
90    }
91
92    /// **Unstable**: constructor signature may change; use `EmbeddingService` trait for stable API.
93    pub fn with_model(model_type: EmbeddingModel) -> Self {
94        Self {
95            model: Arc::new(OnceLock::new()),
96            model_config: ModelConfig::new(model_type),
97        }
98    }
99
100    /// **Unstable**: create with explicit model config (model + optional MRL truncation dim).
101    pub fn with_model_config(model_config: ModelConfig) -> Result<Self> {
102        model_config.validate()?;
103        Ok(Self {
104            model: Arc::new(OnceLock::new()),
105            model_config,
106        })
107    }
108
109    /// **Unstable**: create with model config read from `LATTICE_EMBED_DIM` env var.
110    pub fn with_model_from_env(model_type: EmbeddingModel) -> Result<Self> {
111        let config = model_config_from_env(model_type)?;
112        Ok(Self {
113            model: Arc::new(OnceLock::new()),
114            model_config: config,
115        })
116    }
117
118    /// **Unstable**: persistence API may be moved to a separate manager type.
119    pub fn save_cache(&self) -> Result<usize> {
120        let Some(Ok(model)) = self.model.get() else {
121            return Ok(0);
122        };
123        match model {
124            LoadedModel::Qwen(m) => {
125                let model_name = self.model_config.model.to_string();
126                let path = embedding_cache_path(&model_name, m.dimensions());
127                m.cache_save(&path)
128                    .map_err(|e| EmbedError::InferenceFailed(e.to_string()))
129            }
130            _ => Ok(0),
131        }
132    }
133
134    /// **Unstable**: internal diagnostic; may be removed or moved to metrics.
135    pub fn cache_size(&self) -> usize {
136        self.model
137            .get()
138            .and_then(|r| r.as_ref().ok())
139            .map(LoadedModel::cache_size)
140            .unwrap_or(0)
141    }
142
143    /// **Unstable**: preload the model without performing an encode pass.
144    ///
145    /// Returns the same loading errors as the first `embed` call.
146    /// See [`docs/service.md`](../../docs/service.md#nativeembeddingservice-implementation-notes) for preload behavior.
147    pub async fn ensure_loaded(&self) -> Result<()> {
148        self.ensure_model().await.map(|_| ())
149    }
150
151    /// Ensures the shared model load has completed.
152    ///
153    /// See [`docs/service.md`](../../docs/service.md#nativeembeddingservice-implementation-notes) for the cancellation invariant.
154    async fn ensure_model(&self) -> Result<&LoadedModel> {
155        // Fast path: already loaded.
156        if let Some(result) = self.model.get() {
157            return result
158                .as_ref()
159                .map_err(|e| EmbedError::ModelInitialization(e.clone()));
160        }
161
162        // The shared lock outlives a cancelled caller — see docs/service.md.
163        let model_lock = self.model.clone();
164        let model_config = self.model_config;
165
166        tokio::task::spawn_blocking(move || {
167            // Serialize initialization on the blocking thread pool.
168            model_lock.get_or_init(|| load_model_sync(model_config));
169        })
170        .await
171        .map_err(|e| EmbedError::ModelInitialization(e.to_string()))?;
172
173        self.model
174            .get()
175            .expect("set by spawn_blocking")
176            .as_ref()
177            .map_err(|e| EmbedError::ModelInitialization(e.clone()))
178    }
179
180    /// Encodes text that has already been prepared, meaning any retrieval
181    /// instruction is applied and the caller-facing contract is checked.
182    ///
183    /// The length bound here is a memory backstop, not the published cap: it
184    /// admits the caller's allowance plus whatever instruction this model
185    /// prepends, which is zero bytes for symmetric models.
186    async fn encode_prepared(
187        &self,
188        texts: &[&str],
189        model: EmbeddingModel,
190    ) -> Result<Vec<Vec<f32>>> {
191        if model != self.model_config.model {
192            return Err(EmbedError::InvalidInput(format!(
193                "requested model {:?} but this service is loaded with {:?}",
194                model, self.model_config.model
195            )));
196        }
197        super::validate_texts_bounded(
198            texts,
199            MAX_TEXT_BYTES.saturating_add(model.max_instruction_bytes()),
200        )?;
201
202        let loaded = self.ensure_model().await?;
203        loaded
204            .encode_batch(texts)
205            .map_err(EmbedError::InferenceFailed)
206    }
207
208    async fn encode_prevalidated_with_role(
209        &self,
210        texts: ValidatedTextBatch<'_>,
211        model: EmbeddingModel,
212        role: EmbeddingRole,
213    ) -> Result<Vec<Vec<f32>>> {
214        let prefix = role.instruction(model);
215        if prefix.is_none()
216            && let Some(borrowed) = texts.borrowed()
217        {
218            return self.encode_prepared(borrowed, model).await;
219        }
220
221        if prefix.is_none() {
222            let borrowed = (0..texts.len())
223                .map(|index| texts.get(index))
224                .collect::<Vec<_>>();
225            return self.encode_prepared(&borrowed, model).await;
226        }
227
228        let prepared = texts.to_owned_with_prefix(prefix);
229        let borrowed = prepared.iter().map(String::as_str).collect::<Vec<_>>();
230        self.encode_prepared(&borrowed, model).await
231    }
232}
233
234/// Synchronous model loading (runs on blocking thread pool).
235fn load_model_sync(model_config: ModelConfig) -> std::result::Result<LoadedModel, String> {
236    match model_config.model {
237        EmbeddingModel::BgeSmallEnV15
238        | EmbeddingModel::BgeBaseEnV15
239        | EmbeddingModel::BgeLargeEnV15
240        | EmbeddingModel::MultilingualE5Small
241        | EmbeddingModel::MultilingualE5Base
242        | EmbeddingModel::AllMiniLmL6V2
243        | EmbeddingModel::ParaphraseMultilingualMiniLmL12V2 => {
244            let model_name = match model_config.model {
245                EmbeddingModel::BgeSmallEnV15 => "bge-small-en-v1.5",
246                EmbeddingModel::BgeBaseEnV15 => "bge-base-en-v1.5",
247                EmbeddingModel::BgeLargeEnV15 => "bge-large-en-v1.5",
248                EmbeddingModel::MultilingualE5Small => "multilingual-e5-small",
249                EmbeddingModel::MultilingualE5Base => "multilingual-e5-base",
250                EmbeddingModel::AllMiniLmL6V2 => "all-minilm-l6-v2",
251                EmbeddingModel::ParaphraseMultilingualMiniLmL12V2 => {
252                    "paraphrase-multilingual-minilm-l12-v2"
253                }
254                _ => unreachable!(),
255            };
256            info!(model = model_name, "loading native BERT embedding model");
257            let mut bert = BertModel::from_pretrained(model_name).map_err(|e| e.to_string())?;
258            // Route each model family through its correct pooling strategy.
259            // BGE uses CLS pooling; E5 and MiniLM use mean pooling.
260            if let Some(pooling) = model_config.model.bert_pooling() {
261                bert.set_pooling(pooling);
262            }
263            Ok(LoadedModel::Bert(Arc::new(bert)))
264        }
265        EmbeddingModel::Qwen3Embedding0_6B | EmbeddingModel::Qwen3Embedding4B => {
266            load_qwen_model(model_config)
267        }
268        other => Err(format!("unsupported model: {other:?}")),
269    }
270}
271
272fn load_qwen_model(model_config: ModelConfig) -> std::result::Result<LoadedModel, String> {
273    model_config.validate().map_err(|e| e.to_string())?;
274    let model_type = model_config.model;
275    let model_name = model_type.to_string();
276    info!(
277        model = %model_name,
278        output_dim = ?model_config.output_dim,
279        "loading Qwen embedding model"
280    );
281    let model_dir = qwen_model_dir(model_type).map_err(|e| e.to_string())?;
282    let mut model = QwenModel::from_directory(&model_dir).map_err(|e| e.to_string())?;
283    model.set_output_dim(model_config.output_dim);
284    let cache_path = embedding_cache_path(&model_name, model.dimensions());
285    match model.cache_load(&cache_path) {
286        Ok(n) if n > 0 => {
287            info!(entries = n, path = %cache_path.display(), "loaded embedding cache")
288        }
289        Ok(_) => {}
290        Err(e) => {
291            warn!(
292                path = %cache_path.display(),
293                error = %e,
294                "embedding cache failed integrity check, ignoring (will regenerate on next save)"
295            )
296        }
297    }
298    Ok(LoadedModel::Qwen(Arc::new(model)))
299}
300
301/// Path for persistent embedding cache: ~/.lattice/cache/embed_{model}_{dim}d.bin
302fn embedding_cache_path(model: &str, dim: usize) -> std::path::PathBuf {
303    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into());
304    std::path::PathBuf::from(home)
305        .join(".lattice")
306        .join("cache")
307        .join(format!("embed_{model}_{dim}d.bin"))
308}
309
310/// Locate Qwen3-Embedding model directory for the given model variant.
311fn qwen_model_dir(model_type: EmbeddingModel) -> Result<std::path::PathBuf> {
312    // Check env override first — applies to whichever Qwen model is loaded.
313    if let Ok(dir) = std::env::var("LATTICE_QWEN_MODEL_DIR") {
314        return Ok(std::path::PathBuf::from(dir));
315    }
316
317    let slug = match model_type {
318        EmbeddingModel::Qwen3Embedding0_6B => "qwen3-embedding-0.6b",
319        EmbeddingModel::Qwen3Embedding4B => "qwen3-embedding-4b",
320        other => {
321            return Err(EmbedError::ModelInitialization(format!(
322                "not a Qwen model: {other}"
323            )));
324        }
325    };
326
327    let home = std::env::var("HOME")
328        .map_err(|_| EmbedError::ModelInitialization("HOME not set".into()))?;
329    let dir = std::path::PathBuf::from(home)
330        .join(".lattice")
331        .join("models")
332        .join(slug);
333
334    if dir.join("model.safetensors").exists() || dir.join("model.safetensors.index.json").exists() {
335        Ok(dir)
336    } else {
337        Err(EmbedError::ModelInitialization(format!(
338            "Qwen3 model not found at {dir}. Download it with:\n  huggingface-cli download {repo} --local-dir {dir}",
339            dir = dir.display(),
340            repo = model_type.model_id()
341        )))
342    }
343}
344
345#[async_trait]
346impl EmbeddingService for NativeEmbeddingService {
347    async fn embed(&self, texts: &[String], model: EmbeddingModel) -> Result<Vec<Vec<f32>>> {
348        let texts = ValidatedTextBatch::new(texts)?;
349        self.encode_prevalidated_with_role(texts, model, EmbeddingRole::Generic)
350            .await
351    }
352
353    async fn embed_with_role(
354        &self,
355        texts: &[String],
356        model: EmbeddingModel,
357        role: EmbeddingRole,
358    ) -> Result<Vec<Vec<f32>>> {
359        let texts = ValidatedTextBatch::new(texts)?;
360        self.encode_prevalidated_with_role(texts, model, role).await
361    }
362
363    async fn embed_with_role_prevalidated(
364        &self,
365        texts: ValidatedTextBatch<'_>,
366        model: EmbeddingModel,
367        role: EmbeddingRole,
368    ) -> Result<Vec<Vec<f32>>> {
369        self.encode_prevalidated_with_role(texts, model, role).await
370    }
371
372    fn model_config(&self, model: EmbeddingModel) -> ModelConfig {
373        if model == self.model_config.model {
374            self.model_config
375        } else {
376            ModelConfig::new(model)
377        }
378    }
379
380    fn supports_model(&self, model: EmbeddingModel) -> bool {
381        model == self.model_config.model
382    }
383
384    fn name(&self) -> &'static str {
385        "native-bert"
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    #[test]
394    fn test_cache_path_contains_dim_in_filename() {
395        let path = embedding_cache_path("qwen3-embedding-4b", 1024);
396        let filename = path.file_name().unwrap().to_str().unwrap();
397        assert_eq!(filename, "embed_qwen3-embedding-4b_1024d.bin");
398    }
399
400    #[test]
401    fn test_cache_path_different_dims_produce_different_paths() {
402        let path_1024 = embedding_cache_path("qwen3-embedding-4b", 1024);
403        let path_2560 = embedding_cache_path("qwen3-embedding-4b", 2560);
404        assert_ne!(path_1024, path_2560);
405        assert!(path_1024.to_string_lossy().contains("1024d"));
406        assert!(path_2560.to_string_lossy().contains("2560d"));
407    }
408
409    #[test]
410    fn test_cache_path_model_slug_differentiates_variants() {
411        let path_4b = embedding_cache_path("qwen3-embedding-4b", 2560);
412        let path_06b = embedding_cache_path("qwen3-embedding-0.6b", 1024);
413        assert_ne!(path_4b, path_06b);
414        assert!(path_4b.to_string_lossy().contains("qwen3-embedding-4b"));
415        assert!(path_06b.to_string_lossy().contains("qwen3-embedding-0.6b"));
416    }
417
418    #[test]
419    fn test_cache_path_same_model_same_dim_same_path() {
420        let p1 = embedding_cache_path("qwen3-embedding-4b", 1024);
421        let p2 = embedding_cache_path("qwen3-embedding-4b", 1024);
422        assert_eq!(p1, p2);
423    }
424}
lattice_embed/service/native.rs

lattice_embed/service/
native.rs
