Skip to main content

lattice_embed/
lib.rs

1//! **Stability tier**: Unstable
2//!
3//! The SIMD dispatch layer (`simd/`) and model-loading API are actively evolving.
4//! Platform consumers should use this crate via `lattice-engine`, not directly.
5//! The 21 `unsafe` blocks are gated SIMD intrinsic calls (AVX512/AVX2/NEON); the
6//! 1 `dead_code_allows` retains a superseded dot-product fallback for reference.
7//! See `foundation/STABILITY.md` for the full policy.
8//!
9//! # lattice-embed
10
11#![warn(missing_docs)]
12#![allow(clippy::clone_on_copy)]
13//!
14//! Vector embedding generation with SIMD-accelerated operations for the lattice-runtime substrate.
15//!
16//! This crate provides embedding generation services that convert text into
17//! high-dimensional vector representations suitable for semantic search and
18//! similarity matching.
19//!
20//! ## Features
21//!
22//! - **Native Embeddings**: Generate embeddings locally using pure Rust inference (default)
23//! - **BGE Models**: Support for BGE family of models (small/base/large)
24//! - **Async API**: Full async/await support with tokio
25//! - **SIMD Acceleration**: AVX2/NEON optimized vector operations
26//!
27//! ## Quick Start
28//!
29//! ```rust,no_run
30//! use lattice_embed::{EmbeddingService, EmbeddingModel, NativeEmbeddingService};
31//!
32//! #[tokio::main]
33//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
34//!     let service = NativeEmbeddingService::default();
35//!
36//!     let embedding = service.embed_one(
37//!         "The quick brown fox jumps over the lazy dog",
38//!         EmbeddingModel::default(),
39//!     ).await?;
40//!
41//!     println!("Embedding dimension: {}", embedding.len());
42//!     // Output: Embedding dimension: 384
43//!
44//!     Ok(())
45//! }
46//! ```
47//!
48//! ## Batch Processing
49//!
50//! ```rust,no_run
51//! use lattice_embed::{EmbeddingService, EmbeddingModel, NativeEmbeddingService};
52//!
53//! #[tokio::main]
54//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
55//!     let service = NativeEmbeddingService::default();
56//!
57//!     let texts = vec![
58//!         "First document".to_string(),
59//!         "Second document".to_string(),
60//!         "Third document".to_string(),
61//!     ];
62//!
63//!     let embeddings = service.embed(&texts, EmbeddingModel::BgeSmallEnV15).await?;
64//!     assert_eq!(embeddings.len(), 3);
65//!
66//!     Ok(())
67//! }
68//! ```
69//!
70//! ## Available Models
71//!
72//! | Model | Dimensions | Use Case |
73//! |-------|------------|----------|
74//! | `BgeSmallEnV15` | 384 | Fast, general purpose (default) |
75//! | `BgeBaseEnV15` | 768 | Balanced quality/speed |
76//! | `BgeLargeEnV15` | 1024 | Highest quality |
77
78pub mod backfill;
79mod cache;
80mod error;
81pub mod migration;
82mod model;
83mod service;
84pub mod simd;
85pub mod types;
86
87pub use cache::{CacheStats, DEFAULT_CACHE_CAPACITY, EmbeddingCache, ShardStats};
88pub use error::{EmbedError, Result};
89pub use model::{EmbeddingModel, MIN_MRL_OUTPUT_DIM, ModelConfig, ModelProvenance};
90pub use service::{DEFAULT_MAX_BATCH_SIZE, EmbeddingService, MAX_TEXT_CHARS};
91pub use simd::{SimdConfig, simd_config};
92
93#[cfg(feature = "native")]
94pub use service::{CachedEmbeddingService, NativeEmbeddingService};
95
96/// Utility functions for vector operations.
97///
98/// All functions in this module are SIMD-accelerated when available (AVX2 on x86_64, NEON on aarch64).
99/// Runtime feature detection ensures automatic fallback to scalar implementations
100/// on systems without SIMD support.
101pub mod utils {
102    use crate::simd;
103
104    /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
105    ///
106    /// Compute cosine similarity between two vectors.
107    ///
108    /// Uses SIMD acceleration (AVX2/NEON) when available, with automatic scalar fallback.
109    ///
110    /// Returns a value between -1.0 and 1.0, where 1.0 indicates
111    /// identical direction and -1.0 indicates opposite direction.
112    ///
113    /// # Performance
114    ///
115    /// | Dimension | Scalar | SIMD |
116    /// |-----------|--------|------|
117    /// | 384 | ~650ns | ~90ns |
118    /// | 768 | ~1300ns | ~180ns |
119    /// | 1024 | ~1700ns | ~240ns |
120    ///
121    /// # Example
122    ///
123    /// ```rust
124    /// use lattice_embed::utils::cosine_similarity;
125    ///
126    /// let a = vec![1.0, 0.0, 0.0];
127    /// let b = vec![1.0, 0.0, 0.0];
128    /// assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.0001);
129    ///
130    /// let c = vec![1.0, 0.0, 0.0];
131    /// let d = vec![0.0, 1.0, 0.0];
132    /// assert!((cosine_similarity(&c, &d) - 0.0).abs() < 0.0001);
133    /// ```
134    #[inline]
135    pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
136        simd::cosine_similarity(a, b)
137    }
138
139    /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
140    ///
141    /// Compute dot product between two vectors.
142    ///
143    /// For normalized vectors (like embeddings), this equals cosine similarity.
144    /// Using `dot_product` directly on pre-normalized vectors is faster than
145    /// `cosine_similarity` since it skips norm computation.
146    ///
147    /// # Example
148    ///
149    /// ```rust
150    /// use lattice_embed::utils::dot_product;
151    ///
152    /// let a = vec![1.0, 2.0, 3.0];
153    /// let b = vec![4.0, 5.0, 6.0];
154    /// assert!((dot_product(&a, &b) - 32.0).abs() < 0.0001);
155    /// ```
156    #[inline]
157    pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
158        simd::dot_product(a, b)
159    }
160
161    /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
162    ///
163    /// Normalize a vector to unit length (L2 normalization).
164    ///
165    /// Modifies the vector in place. After normalization, the vector
166    /// will have magnitude 1.0.
167    ///
168    /// # Example
169    ///
170    /// ```rust
171    /// use lattice_embed::utils::normalize;
172    ///
173    /// let mut v = vec![3.0, 4.0];
174    /// normalize(&mut v);
175    ///
176    /// let magnitude: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
177    /// assert!((magnitude - 1.0).abs() < 0.0001);
178    /// ```
179    #[inline]
180    pub fn normalize(vector: &mut [f32]) {
181        simd::normalize(vector)
182    }
183
184    /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
185    ///
186    /// Compute Euclidean distance between two vectors.
187    ///
188    /// # Example
189    ///
190    /// ```rust
191    /// use lattice_embed::utils::euclidean_distance;
192    ///
193    /// let a = vec![0.0, 0.0];
194    /// let b = vec![3.0, 4.0];
195    /// assert!((euclidean_distance(&a, &b) - 5.0).abs() < 0.0001);
196    /// ```
197    #[inline]
198    pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
199        simd::euclidean_distance(a, b)
200    }
201
202    /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
203    ///
204    /// Compute cosine similarities for many vector pairs (batch operation).
205    ///
206    /// More efficient than calling `cosine_similarity` in a loop.
207    #[inline]
208    pub fn batch_cosine_similarity(pairs: &[(&[f32], &[f32])]) -> Vec<f32> {
209        simd::batch_cosine_similarity(pairs)
210    }
211
212    /// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
213    ///
214    /// Compute dot products for many vector pairs (batch operation).
215    #[inline]
216    pub fn batch_dot_product(pairs: &[(&[f32], &[f32])]) -> Vec<f32> {
217        simd::batch_dot_product(pairs)
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    #[test]
226    fn test_cosine_similarity_identical() {
227        let a = vec![1.0, 2.0, 3.0];
228        let b = vec![1.0, 2.0, 3.0];
229        let sim = utils::cosine_similarity(&a, &b);
230        assert!((sim - 1.0).abs() < 0.0001);
231    }
232
233    #[test]
234    fn test_cosine_similarity_orthogonal() {
235        let a = vec![1.0, 0.0];
236        let b = vec![0.0, 1.0];
237        let sim = utils::cosine_similarity(&a, &b);
238        assert!(sim.abs() < 0.0001);
239    }
240
241    #[test]
242    fn test_cosine_similarity_opposite() {
243        let a = vec![1.0, 0.0];
244        let b = vec![-1.0, 0.0];
245        let sim = utils::cosine_similarity(&a, &b);
246        assert!((sim + 1.0).abs() < 0.0001);
247    }
248
249    #[test]
250    fn test_normalize() {
251        let mut v = vec![3.0, 4.0];
252        utils::normalize(&mut v);
253        let magnitude: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
254        assert!((magnitude - 1.0).abs() < 0.0001);
255    }
256
257    #[test]
258    fn test_euclidean_distance() {
259        let a = vec![0.0, 0.0, 0.0];
260        let b = vec![1.0, 0.0, 0.0];
261        let dist = utils::euclidean_distance(&a, &b);
262        assert!((dist - 1.0).abs() < 0.0001);
263    }
264
265    #[test]
266    fn test_model_default() {
267        let model = EmbeddingModel::default();
268        assert_eq!(model.dimensions(), 384);
269    }
270}