Skip to main content

use_vector_store/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../README.md")]
3
4use core::fmt;
5use std::error::Error;
6
7macro_rules! string_newtype {
8    ($(#[$meta:meta])* $name:ident) => {
9        $(#[$meta])*
10        #[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
11        pub struct $name(String);
12
13        impl $name {
14            /// Creates a new string-backed primitive.
15            pub fn new(value: impl Into<String>) -> Self {
16                Self(value.into())
17            }
18
19            /// Returns the stored string value.
20            pub fn as_str(&self) -> &str {
21                &self.0
22            }
23        }
24
25        impl AsRef<str> for $name {
26            fn as_ref(&self) -> &str {
27                self.as_str()
28            }
29        }
30
31        impl From<String> for $name {
32            fn from(value: String) -> Self {
33                Self::new(value)
34            }
35        }
36
37        impl From<&str> for $name {
38            fn from(value: &str) -> Self {
39                Self::new(value)
40            }
41        }
42
43        impl fmt::Display for $name {
44            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
45                formatter.write_str(self.as_str())
46            }
47        }
48    };
49}
50
51string_newtype! {
52    /// A vector record identifier.
53    VectorId
54}
55string_newtype! {
56    /// A vector collection name.
57    VectorCollectionName
58}
59
60/// A vector dimension count.
61#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
62pub struct VectorDimension(usize);
63
64impl VectorDimension {
65    /// Creates a vector dimension count.
66    pub const fn new(value: usize) -> Self {
67        Self(value)
68    }
69
70    /// Returns the dimension count.
71    pub const fn value(self) -> usize {
72        self.0
73    }
74}
75
76impl fmt::Display for VectorDimension {
77    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
78        write!(formatter, "{}", self.0)
79    }
80}
81
82/// A vector embedding payload.
83#[derive(Clone, Debug, Default, PartialEq)]
84pub struct Embedding(Vec<f32>);
85
86impl Embedding {
87    /// Creates an embedding from components.
88    pub fn new(values: Vec<f32>) -> Self {
89        Self(values)
90    }
91
92    /// Returns embedding components.
93    pub fn values(&self) -> &[f32] {
94        &self.0
95    }
96
97    /// Returns the embedding dimension.
98    pub fn dimension(&self) -> VectorDimension {
99        VectorDimension::new(self.0.len())
100    }
101
102    /// Returns whether the embedding has no components.
103    pub fn is_empty(&self) -> bool {
104        self.0.is_empty()
105    }
106}
107
108/// Error returned when a vector dimension does not match its embedding.
109#[derive(Clone, Copy, Debug, Eq, PartialEq)]
110pub struct InvalidDimensionError {
111    expected: VectorDimension,
112    actual: VectorDimension,
113}
114
115impl InvalidDimensionError {
116    /// Creates a dimension mismatch error.
117    pub const fn new(expected: VectorDimension, actual: VectorDimension) -> Self {
118        Self { expected, actual }
119    }
120
121    /// Returns the expected dimension.
122    pub const fn expected(self) -> VectorDimension {
123        self.expected
124    }
125
126    /// Returns the actual dimension.
127    pub const fn actual(self) -> VectorDimension {
128        self.actual
129    }
130}
131
132impl fmt::Display for InvalidDimensionError {
133    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
134        write!(
135            formatter,
136            "vector dimension mismatch: expected {}, got {}",
137            self.expected, self.actual
138        )
139    }
140}
141
142impl Error for InvalidDimensionError {}
143
144/// Similarity metric labels used by vector stores.
145#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
146pub enum SimilarityMetric {
147    Cosine,
148    DotProduct,
149    Euclidean,
150    Manhattan,
151    Hamming,
152    #[default]
153    Unknown,
154}
155
156impl SimilarityMetric {
157    /// Returns a stable lowercase label.
158    pub const fn as_str(self) -> &'static str {
159        match self {
160            Self::Cosine => "cosine",
161            Self::DotProduct => "dot-product",
162            Self::Euclidean => "euclidean",
163            Self::Manhattan => "manhattan",
164            Self::Hamming => "hamming",
165            Self::Unknown => "unknown",
166        }
167    }
168}
169
170impl fmt::Display for SimilarityMetric {
171    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
172        formatter.write_str(self.as_str())
173    }
174}
175
176/// String metadata attached to a vector record.
177#[derive(Clone, Debug, Default, Eq, PartialEq)]
178pub struct VectorMetadata {
179    entries: Vec<(String, String)>,
180}
181
182impl VectorMetadata {
183    /// Creates empty metadata.
184    pub const fn new() -> Self {
185        Self {
186            entries: Vec::new(),
187        }
188    }
189
190    /// Adds a metadata entry.
191    pub fn with_entry(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
192        self.entries.push((key.into(), value.into()));
193        self
194    }
195
196    /// Returns metadata entries.
197    pub fn entries(&self) -> &[(String, String)] {
198        &self.entries
199    }
200}
201
202/// A vector store record.
203#[derive(Clone, Debug, PartialEq)]
204pub struct VectorRecord {
205    id: VectorId,
206    embedding: Embedding,
207    dimension: Option<VectorDimension>,
208    similarity_metric: Option<SimilarityMetric>,
209    metadata: VectorMetadata,
210}
211
212impl VectorRecord {
213    /// Creates a vector record.
214    pub fn new(id: VectorId, embedding: Embedding) -> Self {
215        Self {
216            id,
217            embedding,
218            dimension: None,
219            similarity_metric: None,
220            metadata: VectorMetadata::new(),
221        }
222    }
223
224    /// Sets and validates the expected dimension.
225    pub fn with_dimension(
226        mut self,
227        dimension: VectorDimension,
228    ) -> Result<Self, InvalidDimensionError> {
229        let actual = self.embedding.dimension();
230        if dimension != actual {
231            return Err(InvalidDimensionError::new(dimension, actual));
232        }
233        self.dimension = Some(dimension);
234        Ok(self)
235    }
236
237    /// Sets the similarity metric.
238    pub const fn with_similarity_metric(mut self, similarity_metric: SimilarityMetric) -> Self {
239        self.similarity_metric = Some(similarity_metric);
240        self
241    }
242
243    /// Sets vector metadata.
244    pub fn with_metadata(mut self, metadata: VectorMetadata) -> Self {
245        self.metadata = metadata;
246        self
247    }
248
249    /// Returns the record identifier.
250    pub const fn id(&self) -> &VectorId {
251        &self.id
252    }
253
254    /// Returns the embedding.
255    pub const fn embedding(&self) -> &Embedding {
256        &self.embedding
257    }
258
259    /// Returns the validated dimension, if present.
260    pub const fn dimension(&self) -> Option<VectorDimension> {
261        self.dimension
262    }
263
264    /// Returns the similarity metric, if present.
265    pub const fn similarity_metric(&self) -> Option<SimilarityMetric> {
266        self.similarity_metric
267    }
268
269    /// Returns vector metadata.
270    pub const fn metadata(&self) -> &VectorMetadata {
271        &self.metadata
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use super::{
278        Embedding, InvalidDimensionError, SimilarityMetric, VectorCollectionName, VectorDimension,
279        VectorId, VectorMetadata, VectorRecord,
280    };
281
282    #[test]
283    fn constructs_vector_labels_and_embedding() {
284        let id = VectorId::new("review_embedding");
285        let collection = VectorCollectionName::new("reviews");
286        let embedding = Embedding::new(vec![0.1, 0.2, 0.3]);
287
288        assert_eq!(id.to_string(), "review_embedding");
289        assert_eq!(collection.as_ref(), "reviews");
290        assert_eq!(embedding.dimension(), VectorDimension::new(3));
291    }
292
293    #[test]
294    fn validates_vector_dimensions() -> Result<(), InvalidDimensionError> {
295        let metadata = VectorMetadata::new().with_entry("source", "review");
296        let record = VectorRecord::new(VectorId::new("review_1"), Embedding::new(vec![1.0, 0.0]))
297            .with_dimension(VectorDimension::new(2))?
298            .with_similarity_metric(SimilarityMetric::Cosine)
299            .with_metadata(metadata);
300
301        assert_eq!(record.dimension(), Some(VectorDimension::new(2)));
302        assert_eq!(record.similarity_metric(), Some(SimilarityMetric::Cosine));
303        assert_eq!(record.metadata().entries().len(), 1);
304        assert_eq!(SimilarityMetric::DotProduct.to_string(), "dot-product");
305
306        Ok(())
307    }
308
309    #[test]
310    fn rejects_dimension_mismatches() {
311        let result = VectorRecord::new(VectorId::new("review_1"), Embedding::new(vec![1.0, 0.0]))
312            .with_dimension(VectorDimension::new(3));
313
314        assert_eq!(
315            result,
316            Err(InvalidDimensionError::new(
317                VectorDimension::new(3),
318                VectorDimension::new(2)
319            ))
320        );
321    }
322}