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 vector queries for runtime-defined retrieval policies.
8//!
9//! Use [`VectorSearchRequest`] to build queries. See [`request`] for filtering.
10//!
11//! Types implementing [`VectorStoreIndex`] automatically implement [`PortableTool`].
12
13pub use request::VectorSearchRequest;
14use reqwest::StatusCode;
15use serde::{Deserialize, Serialize};
16use serde_json::{Value, json};
17
18use crate::{
19    Embed,
20    embeddings::{Embedding, EmbeddingError},
21    tool::PortableTool,
22    vector_store::request::{DynamicSearchFilter, 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
73impl VectorStoreError {
74    /// Wraps a backend error as [`VectorStoreError::DatastoreError`].
75    ///
76    /// Handles the wasm/non-wasm trait-bound split in one place; use as
77    /// `.map_err(VectorStoreError::datastore)`.
78    #[cfg(not(target_family = "wasm"))]
79    pub fn datastore(e: impl std::error::Error + Send + Sync + 'static) -> Self {
80        Self::DatastoreError(Box::new(e))
81    }
82
83    /// Wraps a backend error as [`VectorStoreError::DatastoreError`].
84    #[cfg(target_family = "wasm")]
85    pub fn datastore(e: impl std::error::Error + 'static) -> Self {
86        Self::DatastoreError(Box::new(e))
87    }
88}
89
90/// Serializes each document to JSON once, then applies `f` to every
91/// `(document, embedding)` pair, flattening the results into a single vector.
92///
93/// This is the shared shape of most [`InsertDocuments`] implementations:
94/// build one backend record per embedding, carrying the owning document's
95/// serialized form.
96pub fn flatten_embedded<Doc: Serialize, R>(
97    documents: Vec<(Doc, Vec<Embedding>)>,
98    mut f: impl FnMut(&Value, Embedding) -> Result<R, VectorStoreError>,
99) -> Result<Vec<R>, VectorStoreError> {
100    let mut records = Vec::new();
101    for (document, embeddings) in documents {
102        let json_document = serde_json::to_value(&document)?;
103        for embedding in embeddings {
104            records.push(f(&json_document, embedding)?);
105        }
106    }
107    Ok(records)
108}
109
110/// Trait for inserting documents and embeddings into a vector store.
111pub trait InsertDocuments: WasmCompatSend + WasmCompatSync {
112    /// Insert precomputed embeddings for each document.
113    ///
114    /// **Every document must carry at least one embedding.** The embedding
115    /// list was non-empty by construction until it became a `Vec`; the
116    /// requirement did not go away, it moved to the caller. Implementors do
117    /// not guard it, and what an empty list does varies by store — some
118    /// silently insert nothing, some store a document no similarity search
119    /// can ever return, some surface a confusing driver error. Embeddings
120    /// produced by `EmbeddingsBuilder` always satisfy this; only hand-built
121    /// tuples can violate it.
122    fn insert_documents<Doc: Serialize + Embed + WasmCompatSend>(
123        &self,
124        documents: Vec<(Doc, Vec<Embedding>)>,
125    ) -> impl std::future::Future<Output = Result<(), VectorStoreError>> + WasmCompatSend;
126}
127
128/// Trait for querying a vector store by similarity.
129pub trait VectorStoreIndex: WasmCompatSend + WasmCompatSync {
130    /// The filter type for this backend.
131    type Filter: SearchFilter + WasmCompatSend + WasmCompatSync;
132
133    /// Returns the top N most similar documents as `(score, id, document)` tuples.
134    fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
135        &self,
136        req: VectorSearchRequest<Self::Filter>,
137    ) -> impl std::future::Future<Output = Result<Vec<(f64, String, T)>, VectorStoreError>>
138    + WasmCompatSend;
139
140    /// Returns the top N most similar document IDs as `(score, id)` tuples.
141    fn top_n_ids(
142        &self,
143        req: VectorSearchRequest<Self::Filter>,
144    ) -> impl std::future::Future<Output = Result<Vec<(f64, String)>, VectorStoreError>> + WasmCompatSend;
145}
146
147/// Type-erased `top_n` result: `(score, id, document)` tuples as JSON values.
148pub type TopNResults = Result<Vec<(f64, String, Value)>, VectorStoreError>;
149
150/// Type-erased [`VectorStoreIndex`] for dynamic dispatch.
151pub trait VectorStoreIndexDyn: WasmCompatSend + WasmCompatSync {
152    /// Returns the top N documents for a JSON-serializable request.
153    fn top_n<'a>(
154        &'a self,
155        req: VectorSearchRequest<Filter<serde_json::Value>>,
156    ) -> WasmBoxedFuture<'a, TopNResults>;
157
158    /// Returns only the top N document IDs for a JSON-serializable request.
159    fn top_n_ids<'a>(
160        &'a self,
161        req: VectorSearchRequest<Filter<serde_json::Value>>,
162    ) -> WasmBoxedFuture<'a, Result<Vec<(f64, String)>, VectorStoreError>>;
163}
164
165impl<I, F> VectorStoreIndexDyn for I
166where
167    I: VectorStoreIndex<Filter = F>,
168    F: DynamicSearchFilter + WasmCompatSend + WasmCompatSync + 'static,
169{
170    fn top_n<'a>(
171        &'a self,
172        req: VectorSearchRequest<Filter<serde_json::Value>>,
173    ) -> WasmBoxedFuture<'a, TopNResults> {
174        Box::pin(async move {
175            let req = req.try_map_filter(F::from_dynamic_filter)?;
176            Ok(self
177                .top_n::<serde_json::Value>(req)
178                .await?
179                .into_iter()
180                .map(|(score, id, doc)| (score, id, F::normalize_dynamic_document(doc)))
181                .collect::<Vec<_>>())
182        })
183    }
184
185    fn top_n_ids<'a>(
186        &'a self,
187        req: VectorSearchRequest<Filter<serde_json::Value>>,
188    ) -> WasmBoxedFuture<'a, Result<Vec<(f64, String)>, VectorStoreError>> {
189        Box::pin(async move {
190            let req = req.try_map_filter(F::from_dynamic_filter)?;
191            self.top_n_ids(req).await
192        })
193    }
194}
195
196/// The output of vector store queries invoked via [`PortableTool`]
197#[derive(Serialize, Deserialize, Debug)]
198pub struct VectorStoreOutput {
199    /// Similarity score returned by the vector store.
200    pub score: f64,
201    /// Document ID returned by the vector store.
202    pub id: String,
203    /// Serialized document payload.
204    pub document: Value,
205}
206
207impl<T, F> PortableTool for T
208where
209    F: SearchFilter<Value = serde_json::Value>
210        + WasmCompatSend
211        + WasmCompatSync
212        + for<'de> Deserialize<'de>,
213    T: VectorStoreIndex<Filter = F>,
214{
215    const NAME: &'static str = "search_vector_store";
216    type Error = VectorStoreError;
217    type Args = VectorSearchRequest<F>;
218    type Output = Vec<VectorStoreOutput>;
219
220    fn description(&self) -> String {
221        "Retrieves the most relevant documents from a vector store based on a query.".to_string()
222    }
223
224    fn parameters(&self) -> serde_json::Value {
225        json!({
226            "type": "object",
227            "properties": {
228                "query": {
229                    "type": "string",
230                    "description": "The query string to search for relevant documents in the vector store."
231                },
232                "samples": {
233                    "type": "integer",
234                    "description": "The maximum number of samples / documents to retrieve.",
235                    "default": 5,
236                    "minimum": 1
237                },
238                "threshold": {
239                    "type": "number",
240                    "description": "Similarity search threshold. If present, any result with a distance less than this may be omitted from the final result."
241                }
242            },
243            "required": ["query", "samples"]
244        })
245    }
246
247    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
248        let results = self.top_n(args).await?;
249        Ok(results
250            .into_iter()
251            .map(|(score, id, document)| VectorStoreOutput {
252                score,
253                id,
254                document,
255            })
256            .collect())
257    }
258}
259
260/// Index strategy for the super::InMemoryVectorStore
261#[derive(Clone, Debug, Default)]
262pub enum IndexStrategy {
263    /// Checks all documents in the vector store to find the most relevant documents.
264    #[default]
265    BruteForce,
266
267    /// Uses LSH to find candidates then computes exact distances.
268    LSH {
269        /// Number of tables to use for LSH.
270        num_tables: usize,
271        /// Number of hyperplanes to use for LSH.
272        num_hyperplanes: usize,
273    },
274}
275
276#[cfg(test)]
277mod tests {
278    use std::sync::{Arc, Mutex};
279
280    use super::*;
281    use crate::vector_store::request::Filter;
282
283    struct TestIndex {
284        queries: Arc<Mutex<Vec<String>>>,
285    }
286
287    #[derive(Clone)]
288    struct NativeFilter;
289
290    impl SearchFilter for NativeFilter {
291        type Value = String;
292
293        fn eq(_key: impl AsRef<str>, _value: Self::Value) -> Self {
294            Self
295        }
296
297        fn gt(_key: impl AsRef<str>, _value: Self::Value) -> Self {
298            Self
299        }
300
301        fn lt(_key: impl AsRef<str>, _value: Self::Value) -> Self {
302            Self
303        }
304
305        fn and(self, _rhs: Self) -> Self {
306            self
307        }
308
309        fn or(self, _rhs: Self) -> Self {
310            self
311        }
312    }
313
314    impl DynamicSearchFilter for NativeFilter {
315        fn from_dynamic_filter(filter: Filter<serde_json::Value>) -> Result<Self, FilterError> {
316            filter.try_interpret(|value| match value {
317                Value::String(value) => Ok(value),
318                other => Err(FilterError::Expected {
319                    expected: "string".to_owned(),
320                    got: other.to_string(),
321                }),
322            })
323        }
324    }
325
326    struct NativeIndex;
327
328    impl VectorStoreIndex for NativeIndex {
329        type Filter = NativeFilter;
330
331        async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
332            &self,
333            _req: VectorSearchRequest<Self::Filter>,
334        ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
335            let document = serde_json::from_value(Value::Array(vec![Value::Null; 401]))?;
336            Ok(vec![(0.9, "doc-1".to_owned(), document)])
337        }
338
339        async fn top_n_ids(
340            &self,
341            _req: VectorSearchRequest<Self::Filter>,
342        ) -> Result<Vec<(f64, String)>, VectorStoreError> {
343            Ok(vec![(0.9, "doc-1".to_owned())])
344        }
345    }
346
347    impl VectorStoreIndex for TestIndex {
348        type Filter = Filter<serde_json::Value>;
349
350        async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
351            &self,
352            req: VectorSearchRequest,
353        ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
354            self.queries
355                .lock()
356                .expect("query recorder lock")
357                .push(req.query().to_string());
358            let document = serde_json::from_value(json!({ "answer": 42 }))?;
359            Ok(vec![(0.9, "doc-1".to_string(), document)])
360        }
361
362        async fn top_n_ids(
363            &self,
364            _req: VectorSearchRequest,
365        ) -> Result<Vec<(f64, String)>, VectorStoreError> {
366            Ok(vec![(0.9, "doc-1".to_string())])
367        }
368    }
369
370    #[tokio::test]
371    async fn vector_store_index_remains_a_tool() {
372        let queries = Arc::new(Mutex::new(Vec::new()));
373        let index = TestIndex {
374            queries: queries.clone(),
375        };
376        let request = VectorSearchRequest::builder()
377            .query("answer")
378            .samples(1)
379            .build();
380        let output = <TestIndex as PortableTool>::call(&index, request)
381            .await
382            .expect("vector tool call should succeed");
383
384        assert_eq!(<TestIndex as PortableTool>::NAME, "search_vector_store");
385        assert_eq!(
386            *queries.lock().expect("query recorder lock"),
387            vec!["answer"]
388        );
389        assert_eq!(output.len(), 1);
390        let result = output.first().expect("one vector result");
391        assert_eq!(result.score, 0.9);
392        assert_eq!(result.id, "doc-1");
393        assert_eq!(result.document, json!({ "answer": 42 }));
394    }
395
396    #[tokio::test]
397    async fn dynamic_native_filter_preserves_backend_documents() {
398        let request = VectorSearchRequest::builder()
399            .query("answer")
400            .samples(1)
401            .filter(Filter::eq("tag", json!("example")))
402            .build();
403
404        let results = VectorStoreIndexDyn::top_n(&NativeIndex, request)
405            .await
406            .expect("dynamic vector search should succeed");
407
408        assert_eq!(results[0].2.as_array().map(Vec::len), Some(401));
409    }
410
411    #[test]
412    fn datastore_wraps_backend_errors() {
413        let err = VectorStoreError::datastore(std::io::Error::other("db down"));
414        assert!(matches!(err, VectorStoreError::DatastoreError(_)));
415        assert_eq!(err.to_string(), "Datastore error: db down");
416    }
417}
rig_core/vector_store/mod.rs

rig_core/vector_store/
mod.rs
