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}