rig_core/vector_store/
mod.rs

1//! Vector-store insertion and similarity-search interfaces.
2//! [`VectorStoreIndex`] implementations also implement [`PortableTool`] when their
3//! filter is deserializable and uses JSON values.
4//!
5//! ```
6//! use rig_core::vector_store::builder::InMemoryVectorStoreBuilder;
7//!
8//! let store = InMemoryVectorStoreBuilder::<String>::new().build();
9//! # let _ = store;
10//! ```
11
12use http::StatusCode;
13pub use request::VectorSearchRequest;
14use serde::{Deserialize, Serialize, de::DeserializeOwned};
15use serde_json::{Value, json};
16
17use crate::error::ProviderError;
18use crate::{
19    Embed,
20    embeddings::Embedding,
21    tool::PortableTool,
22    vector_store::request::{FilterError, SearchFilter},
23    wasm_compat::{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] ProviderError),
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    ///
62    /// Non-success responses arrive as the status-bearing
63    /// [`crate::http_client::Error`] variants, so callers can still inspect
64    /// the status code; response-less transport failures arrive as
65    /// [`crate::http_client::Error::Instance`].
66    #[error("HTTP request error: {0}")]
67    Http(#[from] crate::http_client::Error),
68
69    /// External vector store service returned an error response.
70    #[error("External call to API returned an error. Error code: {0} Message: {1}")]
71    ExternalAPIError(StatusCode, String),
72
73    /// The requested sample count exceeds what the backend can express.
74    #[error("Requested {requested} samples, but this vector store returns at most {max}")]
75    SamplesOutOfRange {
76        /// Sample count from the search request.
77        requested: u64,
78        /// Largest sample count the backend accepts.
79        max: u64,
80    },
81}
82
83impl VectorStoreError {
84    /// Wraps a backend error as [`VectorStoreError::DatastoreError`].
85    ///
86    /// Handles the wasm/non-wasm trait-bound split in one place; use as
87    /// `.map_err(VectorStoreError::datastore)`.
88    #[cfg(not(target_family = "wasm"))]
89    pub fn datastore(e: impl std::error::Error + Send + Sync + 'static) -> Self {
90        Self::DatastoreError(Box::new(e))
91    }
92
93    /// Wraps a backend error as [`VectorStoreError::DatastoreError`].
94    #[cfg(target_family = "wasm")]
95    pub fn datastore(e: impl std::error::Error + 'static) -> Self {
96        Self::DatastoreError(Box::new(e))
97    }
98}
99
100/// Serializes each document to JSON once, then applies `f` to every
101/// `(document, embedding)` pair, flattening the results into a single vector.
102///
103/// This is the shared shape of most [`InsertDocuments`] implementations:
104/// build one backend record per embedding, carrying the owning document's
105/// serialized form.
106pub fn flatten_embedded<Doc: Serialize, R>(
107    documents: Vec<(Doc, Vec<Embedding>)>,
108    mut f: impl FnMut(&Value, Embedding) -> Result<R, VectorStoreError>,
109) -> Result<Vec<R>, VectorStoreError> {
110    let mut records = Vec::new();
111    for (document, embeddings) in documents {
112        let json_document = serde_json::to_value(&document)?;
113        for embedding in embeddings {
114            records.push(f(&json_document, embedding)?);
115        }
116    }
117    Ok(records)
118}
119
120/// Trait for inserting documents and embeddings into a vector store.
121pub trait InsertDocuments: WasmCompatSend + WasmCompatSync {
122    /// Insert precomputed embeddings for each document.
123    ///
124    /// Callers must supply at least one embedding per document. Empty-list
125    /// behavior is backend-dependent and is not uniformly validated.
126    fn insert_documents<Doc: Serialize + Embed + WasmCompatSend>(
127        &self,
128        documents: Vec<(Doc, Vec<Embedding>)>,
129    ) -> impl std::future::Future<Output = Result<(), VectorStoreError>> + WasmCompatSend;
130}
131
132/// Trait for querying a vector store by similarity.
133pub trait VectorStoreIndex: WasmCompatSend + WasmCompatSync {
134    /// The filter type for this backend.
135    type Filter: SearchFilter + WasmCompatSend + WasmCompatSync;
136
137    /// Returns the top N most similar documents, most similar first.
138    fn top_n<T: DeserializeOwned + WasmCompatSend>(
139        &self,
140        req: VectorSearchRequest<Self::Filter>,
141    ) -> impl std::future::Future<Output = Result<Vec<VectorSearchResult<T>>, VectorStoreError>>
142    + WasmCompatSend;
143
144    /// Returns the top N most similar document IDs, most similar first.
145    fn top_n_ids(
146        &self,
147        req: VectorSearchRequest<Self::Filter>,
148    ) -> impl std::future::Future<Output = Result<Vec<VectorSearchIdResult>, VectorStoreError>>
149    + WasmCompatSend;
150}
151
152/// One document returned by [`VectorStoreIndex::top_n`].
153#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
154pub struct VectorSearchResult<T> {
155    /// Backend similarity score. Its scale and direction follow the backend's metric.
156    pub score: f64,
157    /// Backend document ID.
158    pub id: String,
159    /// The stored document, deserialized as `T`.
160    pub document: T,
161}
162
163/// One document ID returned by [`VectorStoreIndex::top_n_ids`].
164#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
165pub struct VectorSearchIdResult {
166    /// Backend similarity score. Its scale and direction follow the backend's metric.
167    pub score: f64,
168    /// Backend document ID.
169    pub id: String,
170}
171
172impl<T, F> PortableTool for T
173where
174    F: SearchFilter<Value = serde_json::Value>
175        + WasmCompatSend
176        + WasmCompatSync
177        + serde::de::DeserializeOwned,
178    T: VectorStoreIndex<Filter = F>,
179{
180    const NAME: &'static str = "search_vector_store";
181    type Error = VectorStoreError;
182    type Args = VectorSearchRequest<F>;
183    type Output = Vec<VectorSearchResult<Value>>;
184
185    fn description(&self) -> String {
186        "Retrieves the most relevant documents from a vector store based on a query.".to_string()
187    }
188
189    fn parameters(&self) -> serde_json::Value {
190        json!({
191            "type": "object",
192            "properties": {
193                "query": {
194                    "type": "string",
195                    "description": "The query string to search for relevant documents in the vector store."
196                },
197                "samples": {
198                    "type": "integer",
199                    "description": "The maximum number of samples / documents to retrieve.",
200                    "default": 5,
201                    "minimum": 1
202                },
203                "threshold": {
204                    "type": "number",
205                    "description": "Similarity search threshold. If present, any result with a distance less than this may be omitted from the final result."
206                }
207            },
208            "required": ["query", "samples"]
209        })
210    }
211
212    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
213        self.top_n(args).await
214    }
215}
216
217/// Candidate-selection strategy for the in-memory vector store.
218#[derive(Clone, Debug, Default)]
219pub enum IndexStrategy {
220    /// Checks all documents in the vector store to find the most relevant documents.
221    #[default]
222    BruteForce,
223
224    /// Uses LSH to find candidates then computes exact distances.
225    LSH {
226        /// Number of tables to use for LSH.
227        num_tables: usize,
228        /// Number of hyperplanes to use for LSH.
229        num_hyperplanes: usize,
230    },
231}
232
233#[cfg(test)]
234mod tests;
rig_core/vector_store/mod.rs

rig_core/vector_store/
mod.rs
