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    /// Any additional parameters that are required by the vector store.
13    additional_params: Option<serde_json::Value>,
14}
15
16impl VectorSearchRequest {
17    /// Creates a [`VectorSearchRequestBuilder`] which you can use to instantiate this struct.
18    pub fn builder() -> VectorSearchRequestBuilder {
19        VectorSearchRequestBuilder::default()
20    }
21
22    /// The query to be embedded and used in similarity search.
23    pub fn query(&self) -> &str {
24        &self.query
25    }
26
27    /// 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.
28    pub fn samples(&self) -> u64 {
29        self.samples
30    }
31}
32
33/// The builder struct to instantiate [`VectorSearchRequest`].
34#[derive(Clone, Serialize, Deserialize, Debug, Default)]
35pub struct VectorSearchRequestBuilder {
36    query: Option<String>,
37    samples: Option<u64>,
38    additional_params: Option<serde_json::Value>,
39}
40
41impl VectorSearchRequestBuilder {
42    /// Set the query (that will then be embedded )
43    pub fn query<T>(mut self, query: T) -> Self
44    where
45        T: Into<String>,
46    {
47        self.query = Some(query.into());
48        self
49    }
50
51    pub fn samples(mut self, samples: u64) -> Self {
52        self.samples = Some(samples);
53        self
54    }
55
56    pub fn additional_params(
57        mut self,
58        params: serde_json::Value,
59    ) -> Result<Self, VectorStoreError> {
60        self.additional_params = Some(params);
61        Ok(self)
62    }
63
64    pub fn build(self) -> Result<VectorSearchRequest, VectorStoreError> {
65        let Some(query) = self.query else {
66            return Err(VectorStoreError::BuilderError(
67                "`query` is a required variable for building a vector search request".into(),
68            ));
69        };
70
71        let Some(samples) = self.samples else {
72            return Err(VectorStoreError::BuilderError(
73                "`samples` is a required variable for building a vector search request".into(),
74            ));
75        };
76
77        let additional_params = if let Some(params) = self.additional_params {
78            if !params.is_object() {
79                return Err(VectorStoreError::BuilderError(
80                    "Expected JSON object for additional params, got something else".into(),
81                ));
82            }
83            Some(params)
84        } else {
85            None
86        };
87
88        Ok(VectorSearchRequest {
89            query,
90            samples,
91            additional_params,
92        })
93    }
94}
rig/vector_store/request.rs

rig/vector_store/
request.rs
