rig_core/vector_store/
mod.rs

1//! Vector store abstractions for semantic search and retrieval.
2//!
3//! # Core Traits
4//!
5//! - [`VectorStoreIndex`]: Query a vector store for similar documents.
6//! - [`InsertDocuments`]: Insert documents and their embeddings.
7//! - [`VectorStoreIndexDyn`]: Type-erased version for dynamic contexts.
8//!
9//! Use [`VectorSearchRequest`] to build queries. See [`request`] for filtering.
10//!
11//! Types implementing [`VectorStoreIndex`] automatically implement [`Tool`].
12
13pub use request::VectorSearchRequest;
14use reqwest::StatusCode;
15use serde::{Deserialize, Serialize};
16use serde_json::{Value, json};
17
18use crate::{
19    Embed, OneOrMany,
20    embeddings::{Embedding, EmbeddingError},
21    tool::Tool,
22    vector_store::request::{Filter, FilterError, SearchFilter},
23    wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync},
24};
25
26pub mod builder;
27pub mod in_memory_store;
28pub mod lsh;
29pub mod request;
30
31/// Errors from vector store operations.
32#[derive(Debug, thiserror::Error)]
33pub enum VectorStoreError {
34    /// Embedding generation failed while preparing a vector query or insert.
35    #[error("Embedding error: {0}")]
36    EmbeddingError(#[from] EmbeddingError),
37
38    /// JSON serialization or deserialization failed.
39    #[error("Json error: {0}")]
40    JsonError(#[from] serde_json::Error),
41
42    #[cfg(not(target_family = "wasm"))]
43    /// Backend-specific datastore error.
44    #[error("Datastore error: {0}")]
45    DatastoreError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
46
47    /// Filter construction or translation failed.
48    #[error("Filter error: {0}")]
49    FilterError(#[from] FilterError),
50
51    #[cfg(target_family = "wasm")]
52    /// Backend-specific datastore error.
53    #[error("Datastore error: {0}")]
54    DatastoreError(#[from] Box<dyn std::error::Error + 'static>),
55
56    /// A document was missing an ID required by the backend.
57    #[error("Missing Id: {0}")]
58    MissingIdError(String),
59
60    /// HTTP request failed for an external vector store service.
61    #[error("HTTP request error: {0}")]
62    ReqwestError(#[from] reqwest::Error),
63
64    /// External vector store service returned an error response.
65    #[error("External call to API returned an error. Error code: {0} Message: {1}")]
66    ExternalAPIError(StatusCode, String),
67
68    /// A vector search request builder received invalid input.
69    #[error("Error while building VectorSearchRequest: {0}")]
70    BuilderError(String),
71}
72
73/// Trait for inserting documents and embeddings into a vector store.
74pub trait InsertDocuments: WasmCompatSend + WasmCompatSync {
75    /// Insert precomputed embeddings for each document.
76    fn insert_documents<Doc: Serialize + Embed + WasmCompatSend>(
77        &self,
78        documents: Vec<(Doc, OneOrMany<Embedding>)>,
79    ) -> impl std::future::Future<Output = Result<(), VectorStoreError>> + WasmCompatSend;
80}
81
82/// Trait for querying a vector store by similarity.
83pub trait VectorStoreIndex: WasmCompatSend + WasmCompatSync {
84    /// The filter type for this backend.
85    type Filter: SearchFilter + WasmCompatSend + WasmCompatSync;
86
87    /// Returns the top N most similar documents as `(score, id, document)` tuples.
88    fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
89        &self,
90        req: VectorSearchRequest<Self::Filter>,
91    ) -> impl std::future::Future<Output = Result<Vec<(f64, String, T)>, VectorStoreError>>
92    + WasmCompatSend;
93
94    /// Returns the top N most similar document IDs as `(score, id)` tuples.
95    fn top_n_ids(
96        &self,
97        req: VectorSearchRequest<Self::Filter>,
98    ) -> impl std::future::Future<Output = Result<Vec<(f64, String)>, VectorStoreError>> + WasmCompatSend;
99}
100
101/// Type-erased `top_n` result: `(score, id, document)` tuples as JSON values.
102pub type TopNResults = Result<Vec<(f64, String, Value)>, VectorStoreError>;
103
104/// Type-erased [`VectorStoreIndex`] for dynamic dispatch.
105pub trait VectorStoreIndexDyn: WasmCompatSend + WasmCompatSync {
106    /// Returns the top N documents for a JSON-serializable request.
107    fn top_n<'a>(
108        &'a self,
109        req: VectorSearchRequest<Filter<serde_json::Value>>,
110    ) -> WasmBoxedFuture<'a, TopNResults>;
111
112    /// Returns only the top N document IDs for a JSON-serializable request.
113    fn top_n_ids<'a>(
114        &'a self,
115        req: VectorSearchRequest<Filter<serde_json::Value>>,
116    ) -> WasmBoxedFuture<'a, Result<Vec<(f64, String)>, VectorStoreError>>;
117}
118
119impl<I: VectorStoreIndex<Filter = F>, F> VectorStoreIndexDyn for I
120where
121    F: std::fmt::Debug
122        + Clone
123        + SearchFilter<Value = serde_json::Value>
124        + WasmCompatSend
125        + WasmCompatSync
126        + Serialize
127        + for<'de> Deserialize<'de>
128        + 'static,
129{
130    fn top_n<'a>(
131        &'a self,
132        req: VectorSearchRequest<Filter<serde_json::Value>>,
133    ) -> WasmBoxedFuture<'a, TopNResults> {
134        let req = req.map_filter(Filter::interpret);
135
136        Box::pin(async move {
137            Ok(self
138                .top_n::<serde_json::Value>(req)
139                .await?
140                .into_iter()
141                .map(|(score, id, doc)| (score, id, prune_document(doc).unwrap_or_default()))
142                .collect::<Vec<_>>())
143        })
144    }
145
146    fn top_n_ids<'a>(
147        &'a self,
148        req: VectorSearchRequest<Filter<serde_json::Value>>,
149    ) -> WasmBoxedFuture<'a, Result<Vec<(f64, String)>, VectorStoreError>> {
150        let req = req.map_filter(Filter::interpret);
151
152        Box::pin(self.top_n_ids(req))
153    }
154}
155
156fn prune_document(document: serde_json::Value) -> Option<serde_json::Value> {
157    match document {
158        Value::Object(mut map) => {
159            let new_map = map
160                .iter_mut()
161                .filter_map(|(key, value)| {
162                    prune_document(value.take()).map(|value| (key.clone(), value))
163                })
164                .collect::<serde_json::Map<_, _>>();
165
166            Some(Value::Object(new_map))
167        }
168        Value::Array(vec) if vec.len() > 400 => None,
169        Value::Array(vec) => Some(Value::Array(
170            vec.into_iter().filter_map(prune_document).collect(),
171        )),
172        Value::Number(num) => Some(Value::Number(num)),
173        Value::String(s) => Some(Value::String(s)),
174        Value::Bool(b) => Some(Value::Bool(b)),
175        Value::Null => Some(Value::Null),
176    }
177}
178
179/// The output of vector store queries invoked via [`Tool`]
180#[derive(Serialize, Deserialize, Debug)]
181pub struct VectorStoreOutput {
182    /// Similarity score returned by the vector store.
183    pub score: f64,
184    /// Document ID returned by the vector store.
185    pub id: String,
186    /// Serialized document payload.
187    pub document: Value,
188}
189
190impl<T, F> Tool for T
191where
192    F: SearchFilter<Value = serde_json::Value>
193        + WasmCompatSend
194        + WasmCompatSync
195        + for<'de> Deserialize<'de>,
196    T: VectorStoreIndex<Filter = F>,
197{
198    const NAME: &'static str = "search_vector_store";
199
200    type Error = VectorStoreError;
201    type Args = VectorSearchRequest<F>;
202    type Output = Vec<VectorStoreOutput>;
203
204    fn description(&self) -> String {
205        "Retrieves the most relevant documents from a vector store based on a query.".to_string()
206    }
207
208    fn parameters(&self) -> serde_json::Value {
209        json!({
210            "type": "object",
211            "properties": {
212                "query": {
213                    "type": "string",
214                    "description": "The query string to search for relevant documents in the vector store."
215                },
216                "samples": {
217                    "type": "integer",
218                    "description": "The maximum number of samples / documents to retrieve.",
219                    "default": 5,
220                    "minimum": 1
221                },
222                "threshold": {
223                    "type": "number",
224                    "description": "Similarity search threshold. If present, any result with a distance less than this may be omitted from the final result."
225                }
226            },
227            "required": ["query", "samples"]
228        })
229    }
230
231    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
232        let results = self.top_n(args).await?;
233        Ok(results
234            .into_iter()
235            .map(|(score, id, document)| VectorStoreOutput {
236                score,
237                id,
238                document,
239            })
240            .collect())
241    }
242}
243
244/// Index strategy for the super::InMemoryVectorStore
245#[derive(Clone, Debug, Default)]
246pub enum IndexStrategy {
247    /// Checks all documents in the vector store to find the most relevant documents.
248    #[default]
249    BruteForce,
250
251    /// Uses LSH to find candidates then computes exact distances.
252    LSH {
253        /// Number of tables to use for LSH.
254        num_tables: usize,
255        /// Number of hyperplanes to use for LSH.
256        num_hyperplanes: usize,
257    },
258}
rig_core/vector_store/mod.rs

rig_core/vector_store/
mod.rs
