rig_core/vector_store/
request.rs

1//! Types for constructing vector search queries.
2//!
3//! - [`VectorSearchRequest`]: Query parameters (text, result count, threshold, filters).
4//! - [`SearchFilter`]: Trait for backend-agnostic filter expressions.
5//! - [`Filter`]: Canonical, serializable filter representation.
6
7use serde::{Deserialize, Serialize};
8
9use super::VectorStoreError;
10use crate::markers::{Missing, Provided};
11
12/// A vector search request for querying a [`super::VectorStoreIndex`].
13///
14/// The type parameter `F` specifies the filter type (defaults to [`Filter<serde_json::Value>`]).
15/// Use [`VectorSearchRequest::builder()`] to construct instances.
16#[derive(Clone, Serialize, Deserialize, Debug)]
17pub struct VectorSearchRequest<F = Filter<serde_json::Value>> {
18    /// The query text to embed and search with.
19    query: String,
20    /// Maximum number of results to return.
21    samples: u64,
22    /// Minimum similarity score for results.
23    threshold: Option<f64>,
24    /// Backend-specific parameters as a JSON object.
25    additional_params: Option<serde_json::Value>,
26    /// Filter expression to narrow results by metadata.
27    filter: Option<F>,
28}
29
30impl<Filter> VectorSearchRequest<Filter> {
31    /// Creates a [`VectorSearchRequestBuilder`] which you can use to instantiate this struct.
32    pub fn builder() -> VectorSearchRequestBuilder<Filter> {
33        VectorSearchRequestBuilder::<Filter>::default()
34    }
35
36    /// The query to be embedded and used in similarity search.
37    pub fn query(&self) -> &str {
38        &self.query
39    }
40
41    /// Returns the maximum number of results to return.
42    pub fn samples(&self) -> u64 {
43        self.samples
44    }
45
46    /// Returns the optional similarity threshold.
47    pub fn threshold(&self) -> Option<f64> {
48        self.threshold
49    }
50
51    /// Returns a reference to the optional filter expression.
52    pub fn filter(&self) -> &Option<Filter> {
53        &self.filter
54    }
55
56    /// Transforms the filter type using the provided function.
57    ///
58    /// This is useful for converting between filter representations, such as
59    /// translating the canonical [`super::request::Filter`] to a backend-specific filter type.
60    pub fn map_filter<T, F>(self, f: F) -> VectorSearchRequest<T>
61    where
62        F: Fn(Filter) -> T,
63    {
64        VectorSearchRequest {
65            query: self.query,
66            samples: self.samples,
67            threshold: self.threshold,
68            additional_params: self.additional_params,
69            filter: self.filter.map(f),
70        }
71    }
72
73    /// Transforms the filter type using a provided function which can additionally return a result.
74    ///
75    /// Useful for converting between filter representations where the conversion can potentially fail (eg, unrepresentable or invalid values).
76    pub fn try_map_filter<T, F>(self, f: F) -> Result<VectorSearchRequest<T>, FilterError>
77    where
78        F: Fn(Filter) -> Result<T, FilterError>,
79    {
80        let filter = self.filter.map(f).transpose()?;
81
82        Ok(VectorSearchRequest {
83            query: self.query,
84            samples: self.samples,
85            threshold: self.threshold,
86            additional_params: self.additional_params,
87            filter,
88        })
89    }
90}
91
92/// Errors from constructing or converting filter expressions.
93#[derive(Debug, Clone, thiserror::Error)]
94pub enum FilterError {
95    #[error("Expected: {expected}, got: {got}")]
96    Expected { expected: String, got: String },
97
98    #[error("Cannot compile '{0}' to the backend's filter type")]
99    TypeError(String),
100
101    #[error("Missing field '{0}'")]
102    MissingField(String),
103
104    #[error("'{0}' must {1}")]
105    Must(String, String),
106
107    // NOTE: Uses String because `serde_json::Error` is not `Clone`.
108    #[error("Filter serialization failed: {0}")]
109    Serialization(String),
110}
111
112/// Trait for constructing filter expressions in vector search queries.
113///
114/// Uses [tagless final](https://nrinaudo.github.io/articles/tagless_final.html) encoding
115/// for backend-agnostic filters. Use `SearchFilter::eq(...)` etc. directly and let
116/// type inference resolve the concrete filter type.
117pub trait SearchFilter {
118    type Value;
119
120    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self;
121    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self;
122    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self;
123    fn and(self, rhs: Self) -> Self;
124    fn or(self, rhs: Self) -> Self;
125}
126
127/// A rendered SQL-style condition together with its positional bind parameters.
128///
129/// Shared by the SQL-flavoured vector stores, whose filter algebra differs only
130/// in the parameter type `P` and in the placeholder token their driver expects
131/// (`$` for Postgres, `?` for CQL). The placeholder is therefore supplied by the
132/// caller on every leaf constructor rather than baked into this type.
133///
134/// Parameters are collected left to right in the order their placeholders appear
135/// in [`SqlCondition::condition`], which is the order drivers bind them in.
136#[derive(Clone, Debug, Serialize, Deserialize)]
137pub struct SqlCondition<P> {
138    condition: String,
139    params: Vec<P>,
140}
141
142/// Hand-written so that `P` needs no [`Default`] of its own: a parameterless
143/// empty condition is meaningful for every parameter type.
144impl<P> Default for SqlCondition<P> {
145    fn default() -> Self {
146        Self {
147            condition: String::new(),
148            params: Vec::new(),
149        }
150    }
151}
152
153impl<P> SqlCondition<P> {
154    /// Renders `<key> <op> <placeholder>` bound to a single parameter, e.g.
155    /// `price >= $`.
156    pub fn binary(key: impl AsRef<str>, op: &str, placeholder: &str, value: P) -> Self {
157        Self {
158            condition: format!("{} {op} {placeholder}", key.as_ref()),
159            params: vec![value],
160        }
161    }
162
163    /// Renders `<key> <op> (<placeholder>, ...)` with one placeholder per value,
164    /// e.g. `id IN (?, ?)`.
165    pub fn list(key: impl AsRef<str>, op: &str, placeholder: &str, values: Vec<P>) -> Self {
166        let placeholders = vec![placeholder; values.len()].join(", ");
167
168        Self {
169            condition: format!("{} {op} ({placeholders})", key.as_ref()),
170            params: values,
171        }
172    }
173
174    /// Wraps an already-rendered, parameterless condition such as `id is null`.
175    pub fn raw(condition: impl Into<String>) -> Self {
176        Self {
177            condition: condition.into(),
178            params: Vec::new(),
179        }
180    }
181
182    /// Conjoins two conditions as `(lhs) AND (rhs)`, concatenating their parameters.
183    pub fn and(self, rhs: Self) -> Self {
184        self.combine("AND", rhs)
185    }
186
187    /// Disjoins two conditions as `(lhs) OR (rhs)`, concatenating their parameters.
188    pub fn or(self, rhs: Self) -> Self {
189        self.combine("OR", rhs)
190    }
191
192    /// Negates the condition as `NOT (condition)`, keeping its parameters.
193    #[allow(clippy::should_implement_trait)]
194    pub fn not(self) -> Self {
195        Self {
196            condition: format!("NOT ({})", self.condition),
197            ..self
198        }
199    }
200
201    fn combine(self, joiner: &str, rhs: Self) -> Self {
202        Self {
203            condition: format!("({}) {joiner} ({})", self.condition, rhs.condition),
204            params: self.params.into_iter().chain(rhs.params).collect(),
205        }
206    }
207
208    /// The rendered condition, with placeholders as the caller supplied them.
209    pub fn condition(&self) -> &str {
210        &self.condition
211    }
212
213    /// The bind parameters, in placeholder order.
214    pub fn params(&self) -> &[P] {
215        &self.params
216    }
217
218    /// Consumes the condition, returning the rendered text and its parameters.
219    pub fn into_parts(self) -> (String, Vec<P>) {
220        (self.condition, self.params)
221    }
222}
223
224/// Converts the canonical JSON-valued [`Filter`] used by type-erased vector
225/// searches into a backend's native filter representation.
226///
227/// JSON-valued [`SearchFilter`] implementations receive this automatically.
228/// Backends with native value types implement the conversion once here rather
229/// than hand-writing both [`VectorStoreIndexDyn`](super::VectorStoreIndexDyn)
230/// methods.
231pub trait DynamicSearchFilter: SearchFilter + Sized {
232    /// Converts a canonical dynamic filter into this backend's filter type.
233    fn from_dynamic_filter(filter: Filter<serde_json::Value>) -> Result<Self, FilterError>;
234
235    /// Normalizes a document returned through the type-erased search surface.
236    ///
237    /// Native backend filters retain their documents exactly. JSON-valued
238    /// filters override this to preserve the dynamic surface's historical
239    /// payload pruning.
240    fn normalize_dynamic_document(document: serde_json::Value) -> serde_json::Value {
241        document
242    }
243}
244
245impl<F> DynamicSearchFilter for F
246where
247    F: SearchFilter<Value = serde_json::Value>,
248{
249    fn from_dynamic_filter(filter: Filter<serde_json::Value>) -> Result<Self, FilterError> {
250        Ok(filter.interpret())
251    }
252
253    fn normalize_dynamic_document(document: serde_json::Value) -> serde_json::Value {
254        prune_document(document).unwrap_or_default()
255    }
256}
257
258fn prune_document(document: serde_json::Value) -> Option<serde_json::Value> {
259    match document {
260        serde_json::Value::Object(mut map) => {
261            let new_map = map
262                .iter_mut()
263                .filter_map(|(key, value)| {
264                    prune_document(value.take()).map(|value| (key.clone(), value))
265                })
266                .collect::<serde_json::Map<_, _>>();
267
268            Some(serde_json::Value::Object(new_map))
269        }
270        serde_json::Value::Array(vec) if vec.len() > 400 => None,
271        serde_json::Value::Array(vec) => Some(serde_json::Value::Array(
272            vec.into_iter().filter_map(prune_document).collect(),
273        )),
274        value => Some(value),
275    }
276}
277
278/// Canonical, serializable filter representation.
279///
280/// Use for serialization, runtime inspection, or translating between backends via
281/// [`Filter::interpret`]. Prefer [`SearchFilter`] trait methods for writing queries.
282#[derive(Debug, Clone, Serialize, Deserialize)]
283#[serde(rename_all = "lowercase")]
284pub enum Filter<V>
285where
286    V: std::fmt::Debug + Clone,
287{
288    Eq(String, V),
289    Gt(String, V),
290    Lt(String, V),
291    And(Box<Self>, Box<Self>),
292    Or(Box<Self>, Box<Self>),
293}
294
295impl<V> SearchFilter for Filter<V>
296where
297    V: std::fmt::Debug + Clone + Serialize + for<'de> Deserialize<'de>,
298{
299    type Value = V;
300
301    /// Select values where the entry at `key` is equal to `value`
302    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
303        Self::Eq(key.as_ref().to_owned(), value)
304    }
305
306    /// Select values where the entry at `key` is greater than `value`
307    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
308        Self::Gt(key.as_ref().to_owned(), value)
309    }
310
311    /// Select values where the entry at `key` is less than `value`
312    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
313        Self::Lt(key.as_ref().to_owned(), value)
314    }
315
316    /// Select values where the entry satisfies `self` *and* `rhs`
317    fn and(self, rhs: Self) -> Self {
318        Self::And(self.into(), rhs.into())
319    }
320
321    /// Select values where the entry satisfies `self` *or* `rhs`
322    fn or(self, rhs: Self) -> Self {
323        Self::Or(self.into(), rhs.into())
324    }
325}
326
327impl<V> Filter<V>
328where
329    V: std::fmt::Debug + Clone,
330{
331    /// Converts this filter into a backend-specific filter type.
332    pub fn interpret<F>(self) -> F
333    where
334        F: SearchFilter<Value = V>,
335    {
336        self.interpret_with(|v| v)
337    }
338
339    /// Converts this filter into a backend-specific filter type, converting
340    /// each leaf value with `conv`.
341    pub fn interpret_with<F, W>(self, conv: impl Fn(V) -> W + Copy) -> F
342    where
343        F: SearchFilter<Value = W>,
344    {
345        match self.try_interpret(|v| Ok::<W, std::convert::Infallible>(conv(v))) {
346            Ok(filter) => filter,
347            Err(never) => match never {},
348        }
349    }
350
351    /// Converts this filter into a backend-specific filter type, converting
352    /// each leaf value with `conv`. Fails on the first value that fails to convert.
353    pub fn try_interpret<F, W, E>(self, conv: impl Fn(V) -> Result<W, E> + Copy) -> Result<F, E>
354    where
355        F: SearchFilter<Value = W>,
356    {
357        Ok(match self {
358            Self::Eq(key, val) => F::eq(key, conv(val)?),
359            Self::Gt(key, val) => F::gt(key, conv(val)?),
360            Self::Lt(key, val) => F::lt(key, conv(val)?),
361            Self::And(lhs, rhs) => F::and(lhs.try_interpret(conv)?, rhs.try_interpret(conv)?),
362            Self::Or(lhs, rhs) => F::or(lhs.try_interpret(conv)?, rhs.try_interpret(conv)?),
363        })
364    }
365}
366
367impl Filter<serde_json::Value> {
368    /// Tests whether a JSON document satisfies this filter.
369    ///
370    /// Leaf filters (`Eq`/`Gt`/`Lt`) look their key up in `value` (expected to be
371    /// a JSON object) and compare the resulting field against the filter operand.
372    /// A missing field, or an operand that is not order-comparable with the field,
373    /// never satisfies the leaf. `And`/`Or` combine leaf results.
374    pub fn satisfies(&self, value: &serde_json::Value) -> bool {
375        use Filter::*;
376        use serde_json::{Value, Value::*};
377        use std::cmp::Ordering;
378
379        fn compare_pair(l: &Value, r: &Value) -> Option<Ordering> {
380            match (l, r) {
381                // Compare integers exactly; fall back to f64 only for floats or
382                // mixed int/float operands. Trying `as_f64` first (as the old
383                // code did) would lose precision for integers beyond 2^53.
384                (Number(l), Number(r)) => {
385                    if let (Some(l), Some(r)) = (l.as_i64(), r.as_i64()) {
386                        Some(l.cmp(&r))
387                    } else if let (Some(l), Some(r)) = (l.as_u64(), r.as_u64()) {
388                        Some(l.cmp(&r))
389                    } else {
390                        l.as_f64()
391                            .zip(r.as_f64())
392                            .and_then(|(l, r)| l.partial_cmp(&r))
393                    }
394                }
395                (String(l), String(r)) => Some(l.cmp(r)),
396                (Null, Null) => Some(Ordering::Equal),
397                (Bool(l), Bool(r)) => Some(l.cmp(r)),
398                _ => None,
399            }
400        }
401
402        match self {
403            // Numbers compare numerically so `5` matches `5.0`, consistent with
404            // `Gt`/`Lt`; other JSON types fall back to structural equality so
405            // strings/bools/arrays/objects still match exactly.
406            Eq(k, v) => value
407                .get(k)
408                .is_some_and(|field| compare_pair(field, v) == Some(Ordering::Equal) || field == v),
409            Gt(k, v) => value
410                .get(k)
411                .and_then(|field| compare_pair(field, v))
412                .is_some_and(|ord| ord == Ordering::Greater),
413            Lt(k, v) => value
414                .get(k)
415                .and_then(|field| compare_pair(field, v))
416                .is_some_and(|ord| ord == Ordering::Less),
417            And(l, r) => l.satisfies(value) && r.satisfies(value),
418            Or(l, r) => l.satisfies(value) || r.satisfies(value),
419        }
420    }
421}
422
423/// Builder for [`VectorSearchRequest`]. Requires `query` and `samples`.
424#[derive(Clone, Serialize, Deserialize, Debug)]
425pub struct VectorSearchRequestBuilder<F = Filter<serde_json::Value>, Q = Missing, S = Missing> {
426    query: Q,
427    samples: S,
428    threshold: Option<f64>,
429    additional_params: Option<serde_json::Value>,
430    filter: Option<F>,
431}
432
433impl<F> Default for VectorSearchRequestBuilder<F, Missing, Missing> {
434    fn default() -> Self {
435        Self {
436            query: Missing,
437            samples: Missing,
438            threshold: None,
439            additional_params: None,
440            filter: None,
441        }
442    }
443}
444
445impl<F, Q, S> VectorSearchRequestBuilder<F, Q, S>
446where
447    F: SearchFilter,
448{
449    /// Sets the query text. Required.
450    pub fn query<T>(self, query: T) -> VectorSearchRequestBuilder<F, Provided<String>, S>
451    where
452        T: Into<String>,
453    {
454        VectorSearchRequestBuilder {
455            query: Provided(query.into()),
456            samples: self.samples,
457            threshold: self.threshold,
458            additional_params: self.additional_params,
459            filter: self.filter,
460        }
461    }
462
463    /// Sets the maximum number of results. Required.
464    pub fn samples(self, samples: u64) -> VectorSearchRequestBuilder<F, Q, Provided<u64>> {
465        VectorSearchRequestBuilder {
466            query: self.query,
467            samples: Provided(samples),
468            threshold: self.threshold,
469            additional_params: self.additional_params,
470            filter: self.filter,
471        }
472    }
473
474    /// Sets the minimum similarity threshold.
475    pub fn threshold(mut self, threshold: f64) -> Self {
476        self.threshold = Some(threshold);
477        self
478    }
479
480    /// Sets backend-specific parameters.
481    pub fn additional_params(
482        mut self,
483        params: serde_json::Value,
484    ) -> Result<Self, VectorStoreError> {
485        self.additional_params = Some(params);
486        Ok(self)
487    }
488
489    /// Sets a filter expression.
490    pub fn filter(mut self, filter: F) -> Self {
491        self.filter = Some(filter);
492        self
493    }
494}
495
496/// Only implement `build()` when both `query` and `samples` have been provided.
497impl<F> VectorSearchRequestBuilder<F, Provided<String>, Provided<u64>> {
498    /// Builds the request
499    pub fn build(self) -> VectorSearchRequest<F> {
500        VectorSearchRequest {
501            query: self.query.0,
502            samples: self.samples.0,
503            threshold: self.threshold,
504            additional_params: self.additional_params,
505            filter: self.filter,
506        }
507    }
508}
509
510#[cfg(test)]
511mod tests {
512    use super::{Filter, SearchFilter};
513    use serde_json::json;
514
515    type F = Filter<serde_json::Value>;
516
517    #[test]
518    fn eq_matches_field_within_multi_field_document() {
519        let doc = json!({ "category": "fruit", "text": "banana" });
520        assert!(F::eq("category", json!("fruit")).satisfies(&doc));
521        assert!(!F::eq("category", json!("veg")).satisfies(&doc));
522        // A field that does not exist never matches.
523        assert!(!F::eq("missing", json!("fruit")).satisfies(&doc));
524    }
525
526    #[test]
527    fn gt_and_lt_compare_the_named_field() {
528        let doc = json!({ "price": 10, "text": "banana" });
529        assert!(F::gt("price", json!(5)).satisfies(&doc));
530        assert!(!F::gt("price", json!(10)).satisfies(&doc));
531        assert!(F::lt("price", json!(20)).satisfies(&doc));
532        assert!(!F::lt("price", json!(10)).satisfies(&doc));
533        // Missing / non-comparable fields never satisfy an ordering filter.
534        assert!(!F::gt("missing", json!(1)).satisfies(&doc));
535        assert!(!F::gt("text", json!(1)).satisfies(&doc));
536    }
537
538    #[test]
539    fn eq_matches_integer_and_float_representations() {
540        // A field stored as a float still matches an integer operand and vice
541        // versa, consistent with Gt/Lt numeric coercion.
542        assert!(F::eq("score", json!(5)).satisfies(&json!({ "score": 5.0 })));
543        assert!(F::eq("score", json!(5.0)).satisfies(&json!({ "score": 5 })));
544        assert!(!F::eq("score", json!(6)).satisfies(&json!({ "score": 5.0 })));
545        // Non-numeric fields still use structural equality.
546        assert!(F::eq("tag", json!("a")).satisfies(&json!({ "tag": "a" })));
547        assert!(F::eq("tags", json!(["a", "b"])).satisfies(&json!({ "tags": ["a", "b"] })));
548        assert!(!F::eq("tags", json!(["a"])).satisfies(&json!({ "tags": ["a", "b"] })));
549    }
550
551    #[test]
552    fn ordering_compares_large_integers_exactly() {
553        // Integers beyond 2^53 must not collapse to the same f64.
554        let doc = json!({ "id": 9007199254740993_u64 }); // 2^53 + 1
555        assert!(F::gt("id", json!(9007199254740992_u64)).satisfies(&doc)); // > 2^53
556        assert!(!F::gt("id", json!(9007199254740993_u64)).satisfies(&doc));
557        assert!(F::lt("id", json!(9007199254740994_u64)).satisfies(&doc));
558    }
559
560    #[test]
561    fn and_or_combine_leaf_filters() {
562        let doc = json!({ "category": "fruit", "price": 10 });
563        let both = F::eq("category", json!("fruit")).and(F::gt("price", json!(5)));
564        assert!(both.satisfies(&doc));
565
566        let missing_branch = F::eq("category", json!("fruit")).and(F::gt("price", json!(50)));
567        assert!(!missing_branch.satisfies(&doc));
568
569        let either = F::eq("category", json!("veg")).or(F::lt("price", json!(50)));
570        assert!(either.satisfies(&doc));
571    }
572
573    #[test]
574    fn try_interpret_converts_nested_leaf_values() {
575        let f: Filter<i64> =
576            Filter::Eq("a".into(), 1).and(Filter::Gt("b".into(), 2).or(Filter::Lt("c".into(), 3)));
577        let out: Filter<String> = f
578            .try_interpret(|v| Ok::<_, std::convert::Infallible>(v.to_string()))
579            .unwrap();
580        match out {
581            Filter::And(lhs, rhs) => {
582                assert!(matches!(*lhs, Filter::Eq(ref k, ref v) if k == "a" && v == "1"));
583                match *rhs {
584                    Filter::Or(l, r) => {
585                        assert!(matches!(*l, Filter::Gt(ref k, ref v) if k == "b" && v == "2"));
586                        assert!(matches!(*r, Filter::Lt(ref k, ref v) if k == "c" && v == "3"));
587                    }
588                    other => panic!("expected Or, got {other:?}"),
589                }
590            }
591            other => panic!("expected And, got {other:?}"),
592        }
593    }
594
595    #[test]
596    fn try_interpret_propagates_conversion_errors() {
597        let f: Filter<i64> = Filter::Eq("a".into(), 1).and(Filter::Gt("b".into(), -2));
598        let out: Result<Filter<u64>, String> =
599            f.try_interpret(|v| u64::try_from(v).map_err(|e| e.to_string()));
600        assert!(out.is_err());
601    }
602}
rig_core/vector_store/request.rs

rig_core/vector_store/
request.rs
