rig/vector_store/
request.rs

1use serde::{Deserialize, Serialize};
2
3use super::VectorStoreError;
4
5/// A vector search request - used in the [`super::VectorStoreIndex`] trait.
6#[derive(Clone, Serialize, Deserialize, Debug)]
7pub struct VectorSearchRequest {
8    /// The query to be embedded and used in similarity search.
9    query: String,
10    /// The maximum number of samples that may be returned. If adding a similarity search threshold, you may receive less than the inputted number if there aren't enough results that satisfy the threshold.
11    samples: u64,
12    /// Similarity search threshold. If present, any result with a distance less than this may be omitted from the final result.
13    threshold: Option<f64>,
14    /// Any additional parameters that are required by the vector store.
15    additional_params: Option<serde_json::Value>,
16}
17
18impl VectorSearchRequest {
19    /// Creates a [`VectorSearchRequestBuilder`] which you can use to instantiate this struct.
20    pub fn builder() -> VectorSearchRequestBuilder {
21        VectorSearchRequestBuilder::default()
22    }
23
24    /// The query to be embedded and used in similarity search.
25    pub fn query(&self) -> &str {
26        &self.query
27    }
28
29    /// The maximum number of samples that may be returned. If adding a similarity search threshold, you may receive less than the inputted number if there aren't enough results that satisfy the threshold.
30    pub fn samples(&self) -> u64 {
31        self.samples
32    }
33
34    pub fn threshold(&self) -> Option<f64> {
35        self.threshold
36    }
37}
38
39/// The builder struct to instantiate [`VectorSearchRequest`].
40#[derive(Clone, Serialize, Deserialize, Debug, Default)]
41pub struct VectorSearchRequestBuilder {
42    query: Option<String>,
43    samples: Option<u64>,
44    threshold: Option<f64>,
45    additional_params: Option<serde_json::Value>,
46}
47
48impl VectorSearchRequestBuilder {
49    /// Set the query (that will then be embedded )
50    pub fn query<T>(mut self, query: T) -> Self
51    where
52        T: Into<String>,
53    {
54        self.query = Some(query.into());
55        self
56    }
57
58    pub fn samples(mut self, samples: u64) -> Self {
59        self.samples = Some(samples);
60        self
61    }
62
63    pub fn threshold(mut self, threshold: f64) -> Self {
64        self.threshold = Some(threshold);
65        self
66    }
67
68    pub fn additional_params(
69        mut self,
70        params: serde_json::Value,
71    ) -> Result<Self, VectorStoreError> {
72        self.additional_params = Some(params);
73        Ok(self)
74    }
75
76    pub fn build(self) -> Result<VectorSearchRequest, VectorStoreError> {
77        let Some(query) = self.query else {
78            return Err(VectorStoreError::BuilderError(
79                "`query` is a required variable for building a vector search request".into(),
80            ));
81        };
82
83        let Some(samples) = self.samples else {
84            return Err(VectorStoreError::BuilderError(
85                "`samples` is a required variable for building a vector search request".into(),
86            ));
87        };
88
89        let additional_params = if let Some(params) = self.additional_params {
90            if !params.is_object() {
91                return Err(VectorStoreError::BuilderError(
92                    "Expected JSON object for additional params, got something else".into(),
93                ));
94            }
95            Some(params)
96        } else {
97            None
98        };
99
100        Ok(VectorSearchRequest {
101            query,
102            samples,
103            threshold: self.threshold,
104            additional_params,
105        })
106    }
107}
rig/vector_store/request.rs

rig/vector_store/
request.rs
