Skip to main content

lattice_embed/service/
mod.rs

1//! Embedding service trait and implementations.
2
3#[cfg(feature = "native")]
4mod cached;
5#[cfg(feature = "native")]
6mod native;
7
8#[cfg(test)]
9mod tests;
10
11use crate::error::{EmbedError, Result};
12use crate::model::{EmbeddingModel, ModelConfig};
13use async_trait::async_trait;
14
15// Re-exports
16#[cfg(feature = "native")]
17pub use cached::CachedEmbeddingService;
18#[cfg(feature = "native")]
19pub use native::NativeEmbeddingService;
20
21/// **Stable**: default maximum batch size to prevent OOM.
22///
23/// This limit prevents accidentally passing huge batches that could exhaust memory.
24/// Can be overridden by using chunked calls if larger batches are needed.
25pub const DEFAULT_MAX_BATCH_SIZE: usize = 1000;
26
27/// **Stable**: maximum allowed text length in characters.
28///
29/// This limit prevents OOM attacks via extremely large input texts.
30/// 32KB is sufficient for most embedding use cases while preventing abuse.
31pub const MAX_TEXT_CHARS: usize = 32768;
32
33/// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
34///
35/// Trait for embedding generation services.
36///
37/// This trait defines the interface for services that can convert text
38/// into vector embeddings. Implementations may use local models (native Rust)
39/// or remote APIs.
40///
41/// # Example
42///
43/// ```rust,no_run
44/// use lattice_embed::{EmbeddingService, EmbeddingModel, NativeEmbeddingService};
45///
46/// #[tokio::main]
47/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
48///     let service = NativeEmbeddingService::default();
49///     let embedding = service.embed_one("Hello, world!", EmbeddingModel::default()).await?;
50///     assert_eq!(embedding.len(), 384);
51///     Ok(())
52/// }
53/// ```
54#[async_trait]
55pub trait EmbeddingService: Send + Sync {
56    /// **Stable**: generate embeddings for multiple texts.
57    ///
58    /// Returns a vector of embeddings, one for each input text, in the same order.
59    async fn embed(&self, texts: &[String], model: EmbeddingModel) -> Result<Vec<Vec<f32>>>;
60
61    /// **Stable**: generate an embedding for a single text.
62    ///
63    /// This is a convenience method that calls `embed` with a single-element slice.
64    async fn embed_one(&self, text: &str, model: EmbeddingModel) -> Result<Vec<f32>> {
65        let texts = vec![text.to_string()];
66        let mut embeddings = self.embed(&texts, model).await?;
67        embeddings
68            .pop()
69            .ok_or_else(|| EmbedError::Internal("no embedding generated".into()))
70    }
71
72    /// **Unstable**: returns the effective `ModelConfig` for a given model on this service.
73    ///
74    /// The default returns a config with no MRL truncation. `NativeEmbeddingService`
75    /// overrides this to expose the configured output dimension so `CachedEmbeddingService`
76    /// can include the actual dimension in cache keys.
77    fn model_config(&self, model: EmbeddingModel) -> ModelConfig {
78        ModelConfig::new(model)
79    }
80
81    /// **Stable**: check if the service supports a given model.
82    fn supports_model(&self, model: EmbeddingModel) -> bool;
83
84    /// **Stable**: get the name/identifier of this service.
85    fn name(&self) -> &'static str;
86}