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