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