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/// Canonical, serializable filter representation.
128///
129/// Use for serialization, runtime inspection, or translating between backends via
130/// [`Filter::interpret`]. Prefer [`SearchFilter`] trait methods for writing queries.
131#[derive(Debug, Clone, Serialize, Deserialize)]
132#[serde(rename_all = "lowercase")]
133pub enum Filter<V>
134where
135    V: std::fmt::Debug + Clone,
136{
137    Eq(String, V),
138    Gt(String, V),
139    Lt(String, V),
140    And(Box<Self>, Box<Self>),
141    Or(Box<Self>, Box<Self>),
142}
143
144impl<V> SearchFilter for Filter<V>
145where
146    V: std::fmt::Debug + Clone + Serialize + for<'de> Deserialize<'de>,
147{
148    type Value = V;
149
150    /// Select values where the entry at `key` is equal to `value`
151    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
152        Self::Eq(key.as_ref().to_owned(), value)
153    }
154
155    /// Select values where the entry at `key` is greater than `value`
156    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
157        Self::Gt(key.as_ref().to_owned(), value)
158    }
159
160    /// Select values where the entry at `key` is less than `value`
161    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
162        Self::Lt(key.as_ref().to_owned(), value)
163    }
164
165    /// Select values where the entry satisfies `self` *and* `rhs`
166    fn and(self, rhs: Self) -> Self {
167        Self::And(self.into(), rhs.into())
168    }
169
170    /// Select values where the entry satisfies `self` *or* `rhs`
171    fn or(self, rhs: Self) -> Self {
172        Self::Or(self.into(), rhs.into())
173    }
174}
175
176impl<V> Filter<V>
177where
178    V: std::fmt::Debug + Clone,
179{
180    /// Converts this filter into a backend-specific filter type.
181    pub fn interpret<F>(self) -> F
182    where
183        F: SearchFilter<Value = V>,
184    {
185        match self {
186            Self::Eq(key, val) => F::eq(key, val),
187            Self::Gt(key, val) => F::gt(key, val),
188            Self::Lt(key, val) => F::lt(key, val),
189            Self::And(lhs, rhs) => F::and(lhs.interpret(), rhs.interpret()),
190            Self::Or(lhs, rhs) => F::or(lhs.interpret(), rhs.interpret()),
191        }
192    }
193}
194
195impl Filter<serde_json::Value> {
196    /// Tests whether a JSON document satisfies this filter.
197    ///
198    /// Leaf filters (`Eq`/`Gt`/`Lt`) look their key up in `value` (expected to be
199    /// a JSON object) and compare the resulting field against the filter operand.
200    /// A missing field, or an operand that is not order-comparable with the field,
201    /// never satisfies the leaf. `And`/`Or` combine leaf results.
202    pub fn satisfies(&self, value: &serde_json::Value) -> bool {
203        use Filter::*;
204        use serde_json::{Value, Value::*};
205        use std::cmp::Ordering;
206
207        fn compare_pair(l: &Value, r: &Value) -> Option<Ordering> {
208            match (l, r) {
209                // Compare integers exactly; fall back to f64 only for floats or
210                // mixed int/float operands. Trying `as_f64` first (as the old
211                // code did) would lose precision for integers beyond 2^53.
212                (Number(l), Number(r)) => {
213                    if let (Some(l), Some(r)) = (l.as_i64(), r.as_i64()) {
214                        Some(l.cmp(&r))
215                    } else if let (Some(l), Some(r)) = (l.as_u64(), r.as_u64()) {
216                        Some(l.cmp(&r))
217                    } else {
218                        l.as_f64()
219                            .zip(r.as_f64())
220                            .and_then(|(l, r)| l.partial_cmp(&r))
221                    }
222                }
223                (String(l), String(r)) => Some(l.cmp(r)),
224                (Null, Null) => Some(Ordering::Equal),
225                (Bool(l), Bool(r)) => Some(l.cmp(r)),
226                _ => None,
227            }
228        }
229
230        match self {
231            // Numbers compare numerically so `5` matches `5.0`, consistent with
232            // `Gt`/`Lt`; other JSON types fall back to structural equality so
233            // strings/bools/arrays/objects still match exactly.
234            Eq(k, v) => value
235                .get(k)
236                .is_some_and(|field| compare_pair(field, v) == Some(Ordering::Equal) || field == v),
237            Gt(k, v) => value
238                .get(k)
239                .and_then(|field| compare_pair(field, v))
240                .is_some_and(|ord| ord == Ordering::Greater),
241            Lt(k, v) => value
242                .get(k)
243                .and_then(|field| compare_pair(field, v))
244                .is_some_and(|ord| ord == Ordering::Less),
245            And(l, r) => l.satisfies(value) && r.satisfies(value),
246            Or(l, r) => l.satisfies(value) || r.satisfies(value),
247        }
248    }
249}
250
251/// Builder for [`VectorSearchRequest`]. Requires `query` and `samples`.
252#[derive(Clone, Serialize, Deserialize, Debug)]
253pub struct VectorSearchRequestBuilder<F = Filter<serde_json::Value>, Q = Missing, S = Missing> {
254    query: Q,
255    samples: S,
256    threshold: Option<f64>,
257    additional_params: Option<serde_json::Value>,
258    filter: Option<F>,
259}
260
261impl<F> Default for VectorSearchRequestBuilder<F, Missing, Missing> {
262    fn default() -> Self {
263        Self {
264            query: Missing,
265            samples: Missing,
266            threshold: None,
267            additional_params: None,
268            filter: None,
269        }
270    }
271}
272
273impl<F, Q, S> VectorSearchRequestBuilder<F, Q, S>
274where
275    F: SearchFilter,
276{
277    /// Sets the query text. Required.
278    pub fn query<T>(self, query: T) -> VectorSearchRequestBuilder<F, Provided<String>, S>
279    where
280        T: Into<String>,
281    {
282        VectorSearchRequestBuilder {
283            query: Provided(query.into()),
284            samples: self.samples,
285            threshold: self.threshold,
286            additional_params: self.additional_params,
287            filter: self.filter,
288        }
289    }
290
291    /// Sets the maximum number of results. Required.
292    pub fn samples(self, samples: u64) -> VectorSearchRequestBuilder<F, Q, Provided<u64>> {
293        VectorSearchRequestBuilder {
294            query: self.query,
295            samples: Provided(samples),
296            threshold: self.threshold,
297            additional_params: self.additional_params,
298            filter: self.filter,
299        }
300    }
301
302    /// Sets the minimum similarity threshold.
303    pub fn threshold(mut self, threshold: f64) -> Self {
304        self.threshold = Some(threshold);
305        self
306    }
307
308    /// Sets backend-specific parameters.
309    pub fn additional_params(
310        mut self,
311        params: serde_json::Value,
312    ) -> Result<Self, VectorStoreError> {
313        self.additional_params = Some(params);
314        Ok(self)
315    }
316
317    /// Sets a filter expression.
318    pub fn filter(mut self, filter: F) -> Self {
319        self.filter = Some(filter);
320        self
321    }
322}
323
324/// Only implement `build()` when both `query` and `samples` have been provided.
325impl<F> VectorSearchRequestBuilder<F, Provided<String>, Provided<u64>> {
326    /// Builds the request
327    pub fn build(self) -> VectorSearchRequest<F> {
328        VectorSearchRequest {
329            query: self.query.0,
330            samples: self.samples.0,
331            threshold: self.threshold,
332            additional_params: self.additional_params,
333            filter: self.filter,
334        }
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use super::{Filter, SearchFilter};
341    use serde_json::json;
342
343    type F = Filter<serde_json::Value>;
344
345    #[test]
346    fn eq_matches_field_within_multi_field_document() {
347        let doc = json!({ "category": "fruit", "text": "banana" });
348        assert!(F::eq("category", json!("fruit")).satisfies(&doc));
349        assert!(!F::eq("category", json!("veg")).satisfies(&doc));
350        // A field that does not exist never matches.
351        assert!(!F::eq("missing", json!("fruit")).satisfies(&doc));
352    }
353
354    #[test]
355    fn gt_and_lt_compare_the_named_field() {
356        let doc = json!({ "price": 10, "text": "banana" });
357        assert!(F::gt("price", json!(5)).satisfies(&doc));
358        assert!(!F::gt("price", json!(10)).satisfies(&doc));
359        assert!(F::lt("price", json!(20)).satisfies(&doc));
360        assert!(!F::lt("price", json!(10)).satisfies(&doc));
361        // Missing / non-comparable fields never satisfy an ordering filter.
362        assert!(!F::gt("missing", json!(1)).satisfies(&doc));
363        assert!(!F::gt("text", json!(1)).satisfies(&doc));
364    }
365
366    #[test]
367    fn eq_matches_integer_and_float_representations() {
368        // A field stored as a float still matches an integer operand and vice
369        // versa, consistent with Gt/Lt numeric coercion.
370        assert!(F::eq("score", json!(5)).satisfies(&json!({ "score": 5.0 })));
371        assert!(F::eq("score", json!(5.0)).satisfies(&json!({ "score": 5 })));
372        assert!(!F::eq("score", json!(6)).satisfies(&json!({ "score": 5.0 })));
373        // Non-numeric fields still use structural equality.
374        assert!(F::eq("tag", json!("a")).satisfies(&json!({ "tag": "a" })));
375        assert!(F::eq("tags", json!(["a", "b"])).satisfies(&json!({ "tags": ["a", "b"] })));
376        assert!(!F::eq("tags", json!(["a"])).satisfies(&json!({ "tags": ["a", "b"] })));
377    }
378
379    #[test]
380    fn ordering_compares_large_integers_exactly() {
381        // Integers beyond 2^53 must not collapse to the same f64.
382        let doc = json!({ "id": 9007199254740993_u64 }); // 2^53 + 1
383        assert!(F::gt("id", json!(9007199254740992_u64)).satisfies(&doc)); // > 2^53
384        assert!(!F::gt("id", json!(9007199254740993_u64)).satisfies(&doc));
385        assert!(F::lt("id", json!(9007199254740994_u64)).satisfies(&doc));
386    }
387
388    #[test]
389    fn and_or_combine_leaf_filters() {
390        let doc = json!({ "category": "fruit", "price": 10 });
391        let both = F::eq("category", json!("fruit")).and(F::gt("price", json!(5)));
392        assert!(both.satisfies(&doc));
393
394        let missing_branch = F::eq("category", json!("fruit")).and(F::gt("price", json!(50)));
395        assert!(!missing_branch.satisfies(&doc));
396
397        let either = F::eq("category", json!("veg")).or(F::lt("price", json!(50)));
398        assert!(either.satisfies(&doc));
399    }
400}
rig_core/vector_store/request.rs

rig_core/vector_store/
request.rs
