lattice_embed/
types.rs

1//! Vector-space identity and storage-format descriptors for `lattice-embed`.
2//!
3//! These ML-domain types identify compatible vector spaces and define their canonical
4//! byte representation for deterministic hashing.
5//!
6//! See docs/model.md for the embedding-space format and cache-key relationship.
7
8use serde::{Deserialize, Serialize};
9
10// ============================================================================
11// DistanceMetric
12// ============================================================================
13
14/// Distance metric used for vector similarity search.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
16#[serde(rename_all = "snake_case")]
17#[non_exhaustive]
18#[repr(u8)]
19pub enum DistanceMetric {
20    /// Cosine similarity (1 - cosine distance).
21    #[default]
22    Cosine = 1,
23    /// Dot product (inner product).
24    Dot = 2,
25    /// Euclidean (L2) distance.
26    L2 = 3,
27}
28
29impl DistanceMetric {
30    /// Return the wire byte for this variant.
31    #[inline]
32    pub const fn as_byte(self) -> u8 {
33        self as u8
34    }
35
36    /// Reconstruct from a wire byte. Returns `None` for unknown values.
37    #[inline]
38    pub const fn from_byte(b: u8) -> Option<Self> {
39        match b {
40            1 => Some(Self::Cosine),
41            2 => Some(Self::Dot),
42            3 => Some(Self::L2),
43            _ => None,
44        }
45    }
46}
47
48// ============================================================================
49// VectorDType
50// ============================================================================
51
52/// Element data type for stored vectors.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
54#[serde(rename_all = "snake_case")]
55#[non_exhaustive]
56#[repr(u8)]
57pub enum VectorDType {
58    /// 32-bit float.
59    #[default]
60    F32 = 1,
61    /// 16-bit float (half precision).
62    F16 = 2,
63    /// 8-bit signed integer (quantized).
64    I8 = 3,
65}
66
67impl VectorDType {
68    /// Return the wire byte for this variant.
69    #[inline]
70    pub const fn as_byte(self) -> u8 {
71        self as u8
72    }
73
74    /// Size in bytes per element.
75    #[inline]
76    pub const fn size_bytes(self) -> usize {
77        match self {
78            Self::F32 => 4,
79            Self::F16 => 2,
80            Self::I8 => 1,
81        }
82    }
83}
84
85// ============================================================================
86// VectorNorm
87// ============================================================================
88
89/// Normalization state of stored vectors.
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
91#[serde(rename_all = "snake_case")]
92#[non_exhaustive]
93#[repr(u8)]
94pub enum VectorNorm {
95    /// No normalization applied.
96    #[default]
97    None = 0,
98    /// Normalized to unit length (L2 norm = 1).
99    Unit = 1,
100}
101
102impl VectorNorm {
103    /// Return the wire byte for this variant.
104    #[inline]
105    pub const fn as_byte(self) -> u8 {
106        self as u8
107    }
108}
109
110// ============================================================================
111// EmbeddingKey
112// ============================================================================
113
114/// Identifies an embedding space (model + revision + dims + metric + dtype + norm).
115///
116/// Used for selecting vector store collections, caching, and embedding migration routing.
117/// `canonical_bytes()` produces a stable hash for deduplication.
118#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
119pub struct EmbeddingKey {
120    /// Provider/model name (e.g., "bge-small-en-v1.5").
121    pub model: Box<str>,
122    /// Provider-specific revision (semver, date tag, or commit hash).
123    pub revision: Box<str>,
124    /// Vector dimensionality.
125    pub dims: u32,
126    /// Distance metric for similarity.
127    pub metric: DistanceMetric,
128    /// Element data type.
129    pub dtype: VectorDType,
130    /// Normalization state.
131    pub norm: VectorNorm,
132}
133
134impl EmbeddingKey {
135    /// Create a new `EmbeddingKey`.
136    pub fn new(
137        model: impl Into<Box<str>>,
138        revision: impl Into<Box<str>>,
139        dims: u32,
140        metric: DistanceMetric,
141        dtype: VectorDType,
142        norm: VectorNorm,
143    ) -> Self {
144        Self {
145            model: model.into(),
146            revision: revision.into(),
147            dims,
148            metric,
149            dtype,
150            norm,
151        }
152    }
153
154    /// Returns deterministic bytes that identify this exact embedding space.
155    /// See [`docs/design.md`](../docs/design.md#vector-space-identity-wire-format) for the byte layout.
156    pub fn canonical_bytes(&self) -> Vec<u8> {
157        let mut buf = Vec::new();
158
159        let model_bytes = self.model.as_bytes();
160        buf.extend_from_slice(&(model_bytes.len() as u32).to_be_bytes());
161        buf.extend_from_slice(model_bytes);
162
163        let rev_bytes = self.revision.as_bytes();
164        buf.extend_from_slice(&(rev_bytes.len() as u32).to_be_bytes());
165        buf.extend_from_slice(rev_bytes);
166
167        buf.extend_from_slice(&self.dims.to_be_bytes());
168        buf.push(self.metric.as_byte());
169        buf.push(self.dtype.as_byte());
170        buf.push(self.norm.as_byte());
171
172        buf
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn test_distance_metric_byte_roundtrip() {
182        for (metric, byte) in [
183            (DistanceMetric::Cosine, 1u8),
184            (DistanceMetric::Dot, 2u8),
185            (DistanceMetric::L2, 3u8),
186        ] {
187            assert_eq!(metric.as_byte(), byte);
188            assert_eq!(DistanceMetric::from_byte(byte), Some(metric));
189        }
190        assert_eq!(DistanceMetric::from_byte(99), None);
191    }
192
193    #[test]
194    fn test_vector_dtype_size_bytes() {
195        assert_eq!(VectorDType::F32.size_bytes(), 4);
196        assert_eq!(VectorDType::F16.size_bytes(), 2);
197        assert_eq!(VectorDType::I8.size_bytes(), 1);
198    }
199
200    #[test]
201    fn test_vector_norm_defaults() {
202        assert_eq!(VectorNorm::default(), VectorNorm::None);
203    }
204
205    #[test]
206    fn test_embedding_key_canonical_bytes_deterministic() {
207        let k1 = EmbeddingKey::new(
208            "bge-small-en-v1.5",
209            "v1.5",
210            384,
211            DistanceMetric::Cosine,
212            VectorDType::F32,
213            VectorNorm::Unit,
214        );
215        let k2 = EmbeddingKey::new(
216            "bge-small-en-v1.5",
217            "v1.5",
218            384,
219            DistanceMetric::Cosine,
220            VectorDType::F32,
221            VectorNorm::Unit,
222        );
223        assert_eq!(k1.canonical_bytes(), k2.canonical_bytes());
224    }
225
226    #[test]
227    fn test_embedding_key_canonical_bytes_differs_by_field() {
228        let k1 = EmbeddingKey::new(
229            "model-a",
230            "v1",
231            384,
232            DistanceMetric::Cosine,
233            VectorDType::F32,
234            VectorNorm::Unit,
235        );
236        let k2 = EmbeddingKey::new(
237            "model-b",
238            "v1",
239            384,
240            DistanceMetric::Cosine,
241            VectorDType::F32,
242            VectorNorm::Unit,
243        );
244        assert_ne!(k1.canonical_bytes(), k2.canonical_bytes());
245
246        let k3 = EmbeddingKey::new(
247            "model-a",
248            "v1",
249            768,
250            DistanceMetric::Cosine,
251            VectorDType::F32,
252            VectorNorm::Unit,
253        );
254        assert_ne!(k1.canonical_bytes(), k3.canonical_bytes());
255    }
256}
lattice_embed/types.rs

lattice_embed/
types.rs
