Skip to main content

rskit_vectorstore/
store.rs

1//! Vector store trait definition.
2
3use std::collections::HashMap;
4use std::fmt;
5
6use async_trait::async_trait;
7use rskit_errors::{AppError, AppResult, ErrorCode};
8use serde::de::Visitor;
9use serde::{Deserialize, Serialize};
10
11use crate::VectorStoreLimits;
12
13/// Typed scalar payload value stored alongside vector points.
14#[derive(Debug, Clone, Serialize, PartialEq)]
15#[non_exhaustive]
16#[serde(untagged)]
17pub enum PayloadValue {
18    /// UTF-8 string value.
19    String(String),
20    /// Signed integer value.
21    Integer(i64),
22    /// Floating-point value.
23    ///
24    /// Values are validated as finite by store/adaptor request validation before
25    /// they are accepted for storage or filtering.
26    Float(f64),
27    /// Boolean value.
28    Bool(bool),
29}
30
31impl PayloadValue {
32    /// Return the value as a string when it is string-typed.
33    #[must_use]
34    pub fn as_str(&self) -> Option<&str> {
35        match self {
36            Self::String(value) => Some(value),
37            Self::Integer(_) | Self::Float(_) | Self::Bool(_) => None,
38        }
39    }
40
41    pub(crate) fn encoded_len(&self) -> usize {
42        match self {
43            Self::String(value) => value.len(),
44            Self::Integer(_) => std::mem::size_of::<i64>(),
45            Self::Float(_) => std::mem::size_of::<f64>(),
46            Self::Bool(_) => std::mem::size_of::<bool>(),
47        }
48    }
49}
50
51impl<'de> Deserialize<'de> for PayloadValue {
52    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
53    where
54        D: serde::Deserializer<'de>,
55    {
56        deserializer.deserialize_any(PayloadValueVisitor)
57    }
58}
59
60struct PayloadValueVisitor;
61
62impl Visitor<'_> for PayloadValueVisitor {
63    type Value = PayloadValue;
64
65    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
66        formatter.write_str("a string, signed integer, finite float, or boolean payload value")
67    }
68
69    fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E>
70    where
71        E: serde::de::Error,
72    {
73        Ok(PayloadValue::Bool(value))
74    }
75
76    fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E>
77    where
78        E: serde::de::Error,
79    {
80        Ok(PayloadValue::Integer(value))
81    }
82
83    fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
84    where
85        E: serde::de::Error,
86    {
87        i64::try_from(value)
88            .map(PayloadValue::Integer)
89            .map_err(|_| {
90                <E as serde::de::Error>::custom("unsigned payload integers must fit in i64")
91            })
92    }
93
94    fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
95    where
96        E: serde::de::Error,
97    {
98        Ok(PayloadValue::Float(value))
99    }
100
101    fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
102    where
103        E: serde::de::Error,
104    {
105        Ok(PayloadValue::String(value.to_owned()))
106    }
107
108    fn visit_string<E>(self, value: String) -> Result<Self::Value, E>
109    where
110        E: serde::de::Error,
111    {
112        Ok(PayloadValue::String(value))
113    }
114}
115
116impl From<String> for PayloadValue {
117    fn from(value: String) -> Self {
118        Self::String(value)
119    }
120}
121
122impl From<&str> for PayloadValue {
123    fn from(value: &str) -> Self {
124        Self::String(value.to_owned())
125    }
126}
127
128impl From<i64> for PayloadValue {
129    fn from(value: i64) -> Self {
130        Self::Integer(value)
131    }
132}
133
134impl From<i32> for PayloadValue {
135    fn from(value: i32) -> Self {
136        Self::Integer(i64::from(value))
137    }
138}
139
140impl From<u32> for PayloadValue {
141    fn from(value: u32) -> Self {
142        Self::Integer(i64::from(value))
143    }
144}
145
146impl From<f64> for PayloadValue {
147    fn from(value: f64) -> Self {
148        Self::Float(value)
149    }
150}
151
152impl From<f32> for PayloadValue {
153    fn from(value: f32) -> Self {
154        Self::Float(f64::from(value))
155    }
156}
157
158impl From<bool> for PayloadValue {
159    fn from(value: bool) -> Self {
160        Self::Bool(value)
161    }
162}
163
164/// Payload stored alongside each vector point.
165#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
166pub struct PointPayload {
167    /// Metadata fields carried with the point.
168    pub fields: HashMap<String, PayloadValue>,
169}
170
171impl PointPayload {
172    /// Create an empty payload.
173    #[must_use]
174    pub fn new() -> Self {
175        Self {
176            fields: HashMap::new(),
177        }
178    }
179
180    /// Add a payload field.
181    #[must_use]
182    pub fn with_field(mut self, key: impl Into<String>, value: impl Into<PayloadValue>) -> Self {
183        self.fields.insert(key.into(), value.into());
184        self
185    }
186
187    /// Validate field count and approximate scalar payload bytes against limits.
188    pub fn validate_limits(&self, limits: &VectorStoreLimits) -> AppResult<()> {
189        if self.fields.len() > limits.max_payload_fields {
190            return Err(AppError::new(
191                ErrorCode::InvalidInput,
192                format!(
193                    "vector payload has {} fields, exceeding max_payload_fields {}",
194                    self.fields.len(),
195                    limits.max_payload_fields
196                ),
197            ));
198        }
199        let mut total_bytes = 0usize;
200        for (key, payload_value) in &self.fields {
201            if let PayloadValue::Float(float_value) = payload_value
202                && !float_value.is_finite()
203            {
204                return Err(AppError::new(
205                    ErrorCode::InvalidInput,
206                    "vector payload float values must be finite",
207                ));
208            }
209            total_bytes = total_bytes
210                .checked_add(key.len())
211                .and_then(|bytes| bytes.checked_add(payload_value.encoded_len()))
212                .ok_or_else(|| {
213                    AppError::new(
214                        ErrorCode::InvalidInput,
215                        "vector payload byte size overflowed validation bounds",
216                    )
217                })?;
218        }
219        if total_bytes > limits.max_payload_bytes {
220            return Err(AppError::new(
221                ErrorCode::InvalidInput,
222                format!(
223                    "vector payload is {total_bytes} bytes, exceeding max_payload_bytes {}",
224                    limits.max_payload_bytes
225                ),
226            ));
227        }
228        Ok(())
229    }
230}
231
232impl Default for PointPayload {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238/// A single search result from the vector store.
239#[derive(Debug, Clone)]
240pub struct SearchResult {
241    /// Point identifier.
242    pub id: String,
243    /// Backend-specific similarity score.
244    pub score: f32,
245    /// Payload attached to the point.
246    pub payload: PointPayload,
247}
248
249/// Canonical vector distance/similarity metrics.
250#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
251#[non_exhaustive]
252#[serde(rename_all = "lowercase")]
253pub enum SimilarityMetric {
254    /// Cosine similarity.
255    #[default]
256    Cosine,
257    /// Dot product.
258    Dot,
259    /// Euclidean L2 distance.
260    L2,
261}
262
263/// Exact-match metadata filter condition.
264#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
265pub struct FilterCondition {
266    /// Payload field path.
267    pub field: String,
268    /// Exact value to match.
269    pub equals: PayloadValue,
270}
271
272/// Optional filters for search queries.
273#[derive(Debug, Clone, Default, Serialize, Deserialize)]
274pub struct SearchFilter {
275    /// Filter by exact field match (e.g., platform = "youtube").
276    #[serde(default)]
277    pub must: Vec<FilterCondition>,
278}
279
280impl SearchFilter {
281    /// Create an empty filter.
282    #[must_use]
283    pub fn new() -> Self {
284        Self::default()
285    }
286
287    /// Add an exact-match condition to the `must` list.
288    #[must_use]
289    pub fn must_match(mut self, field: impl Into<String>, value: impl Into<PayloadValue>) -> Self {
290        self.must.push(FilterCondition {
291            field: field.into(),
292            equals: value.into(),
293        });
294        self
295    }
296
297    /// Validate filter condition count and approximate scalar bytes against limits.
298    pub fn validate_limits(&self, limits: &VectorStoreLimits) -> AppResult<()> {
299        if self.must.len() > limits.max_filter_conditions {
300            return Err(AppError::new(
301                ErrorCode::InvalidInput,
302                format!(
303                    "vector search filter has {} conditions, exceeding max_filter_conditions {}",
304                    self.must.len(),
305                    limits.max_filter_conditions
306                ),
307            ));
308        }
309        let mut total_bytes = 0usize;
310        for condition in &self.must {
311            if let PayloadValue::Float(float_value) = &condition.equals
312                && !float_value.is_finite()
313            {
314                return Err(AppError::new(
315                    ErrorCode::InvalidInput,
316                    "vector filter float values must be finite",
317                ));
318            }
319            total_bytes = total_bytes
320                .checked_add(condition.field.len())
321                .and_then(|bytes| bytes.checked_add(condition.equals.encoded_len()))
322                .ok_or_else(|| {
323                    AppError::new(
324                        ErrorCode::InvalidInput,
325                        "vector filter byte size overflowed validation bounds",
326                    )
327                })?;
328        }
329        if total_bytes > limits.max_payload_bytes {
330            return Err(AppError::new(
331                ErrorCode::InvalidInput,
332                format!(
333                    "vector search filter is {total_bytes} bytes, exceeding max_payload_bytes {}",
334                    limits.max_payload_bytes
335                ),
336            ));
337        }
338        Ok(())
339    }
340}
341
342/// Trait for vector similarity search stores.
343#[async_trait]
344pub trait VectorStore: Send + Sync {
345    /// Ensure a collection exists, creating it if necessary.
346    async fn ensure_collection(&self, collection: &str, dimensions: usize) -> AppResult<()>;
347
348    /// Insert or update a vector point.
349    async fn upsert(
350        &self,
351        collection: &str,
352        id: &str,
353        vector: Vec<f32>,
354        payload: PointPayload,
355    ) -> AppResult<()>;
356
357    /// Search for similar vectors.
358    async fn search(
359        &self,
360        collection: &str,
361        vector: Vec<f32>,
362        limit: usize,
363        filter: Option<SearchFilter>,
364    ) -> AppResult<Vec<SearchResult>>;
365
366    /// Delete a point by ID.
367    async fn delete(&self, collection: &str, id: &str) -> AppResult<()>;
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373
374    #[test]
375    fn payload_value_serializes_as_json_scalar() {
376        assert_eq!(
377            serde_json::to_value(PayloadValue::String("doc".to_owned())).unwrap(),
378            serde_json::json!("doc")
379        );
380        assert_eq!(
381            serde_json::to_value(PayloadValue::Integer(42)).unwrap(),
382            serde_json::json!(42)
383        );
384        assert_eq!(
385            serde_json::to_value(PayloadValue::Float(1.5)).unwrap(),
386            serde_json::json!(1.5)
387        );
388        assert_eq!(
389            serde_json::to_value(PayloadValue::Bool(true)).unwrap(),
390            serde_json::json!(true)
391        );
392    }
393
394    #[test]
395    fn payload_value_deserializes_from_json_scalar() {
396        assert_eq!(
397            serde_json::from_value::<PayloadValue>(serde_json::json!("doc")).unwrap(),
398            PayloadValue::String("doc".to_owned())
399        );
400        assert_eq!(
401            serde_json::from_value::<PayloadValue>(serde_json::json!(42)).unwrap(),
402            PayloadValue::Integer(42)
403        );
404        assert_eq!(
405            serde_json::from_value::<PayloadValue>(serde_json::json!(1.5)).unwrap(),
406            PayloadValue::Float(1.5)
407        );
408        assert_eq!(
409            serde_json::from_value::<PayloadValue>(serde_json::json!(true)).unwrap(),
410            PayloadValue::Bool(true)
411        );
412    }
413
414    #[test]
415    fn payload_value_rejects_unsigned_integer_over_i64_max() {
416        let value = serde_json::Value::Number(serde_json::Number::from(u64::MAX));
417
418        assert!(serde_json::from_value::<PayloadValue>(value).is_err());
419    }
420}