1use 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#[derive(Debug, Clone, Serialize, PartialEq)]
15#[non_exhaustive]
16#[serde(untagged)]
17pub enum PayloadValue {
18 String(String),
20 Integer(i64),
22 Float(f64),
27 Bool(bool),
29}
30
31impl PayloadValue {
32 #[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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
166pub struct PointPayload {
167 pub fields: HashMap<String, PayloadValue>,
169}
170
171impl PointPayload {
172 #[must_use]
174 pub fn new() -> Self {
175 Self {
176 fields: HashMap::new(),
177 }
178 }
179
180 #[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 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#[derive(Debug, Clone)]
240pub struct SearchResult {
241 pub id: String,
243 pub score: f32,
245 pub payload: PointPayload,
247}
248
249#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
251#[non_exhaustive]
252#[serde(rename_all = "lowercase")]
253pub enum SimilarityMetric {
254 #[default]
256 Cosine,
257 Dot,
259 L2,
261}
262
263#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
265pub struct FilterCondition {
266 pub field: String,
268 pub equals: PayloadValue,
270}
271
272#[derive(Debug, Clone, Default, Serialize, Deserialize)]
274pub struct SearchFilter {
275 #[serde(default)]
277 pub must: Vec<FilterCondition>,
278}
279
280impl SearchFilter {
281 #[must_use]
283 pub fn new() -> Self {
284 Self::default()
285 }
286
287 #[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 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#[async_trait]
344pub trait VectorStore: Send + Sync {
345 async fn ensure_collection(&self, collection: &str, dimensions: usize) -> AppResult<()>;
347
348 async fn upsert(
350 &self,
351 collection: &str,
352 id: &str,
353 vector: Vec<f32>,
354 payload: PointPayload,
355 ) -> AppResult<()>;
356
357 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 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}