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}