rig_core/vector_store/
mod.rs

1//! Vector-store insertion and similarity-search interfaces.
2//! [`VectorStoreIndex`] implementations also implement [`Tool`] when their
3//! filter is deserializable and uses JSON values.
4//!
5//! ```
6//! use rig_core::vector_store::in_memory_store::InMemoryVectorStore;
7//!
8//! let store = InMemoryVectorStore::<String>::from_documents([]);
9//! assert!(store.is_empty());
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::{Tool, ToolContext},
22    vector_store::request::{FilterError, SearchFilter},
23    wasm_compat::{WasmCompatSend, WasmCompatSync},
24};
25
26pub mod in_memory_store;
27pub mod lsh;
28pub mod request;
29
30/// How the in-memory vector store picks the documents a search scores.
31#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
32pub enum IndexStrategy {
33    /// Scores every document: exact results.
34    #[default]
35    BruteForce,
36    /// Scores only the documents sharing an [`lsh::LSHIndex`] bucket with the
37    /// query, exactly. Faster on large stores, but a true neighbor in no
38    /// shared bucket is missed, so a search can return fewer or other results
39    /// than [`Self::BruteForce`].
40    LSH {
41        /// Hash tables; more find more true neighbors. `0` scores every
42        /// document, like [`Self::BruteForce`].
43        num_tables: usize,
44        /// Hyperplanes per table (at most 64); more make smaller buckets.
45        num_hyperplanes: usize,
46    },
47}
48
49/// Errors from vector store operations.
50#[derive(Debug, thiserror::Error)]
51pub enum VectorStoreError {
52    /// Embedding generation failed while preparing a vector query or insert.
53    #[error("Embedding error: {0}")]
54    EmbeddingError(#[from] ProviderError),
55
56    /// JSON serialization or deserialization failed.
57    #[error("Json error: {0}")]
58    JsonError(#[from] serde_json::Error),
59
60    #[cfg(not(target_family = "wasm"))]
61    /// Backend-specific datastore error.
62    #[error("Datastore error: {0}")]
63    DatastoreError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
64
65    /// Filter construction or translation failed.
66    #[error("Filter error: {0}")]
67    FilterError(#[from] FilterError),
68
69    #[cfg(target_family = "wasm")]
70    /// Backend-specific datastore error.
71    #[error("Datastore error: {0}")]
72    DatastoreError(#[from] Box<dyn std::error::Error + 'static>),
73
74    /// A document was missing an ID required by the backend.
75    #[error("Missing Id: {0}")]
76    MissingIdError(String),
77
78    /// HTTP request failed for an external vector store service.
79    ///
80    /// Non-success responses arrive as the status-bearing
81    /// [`crate::http_client::Error`] variants, so callers can still inspect
82    /// the status code; response-less transport failures arrive as
83    /// [`crate::http_client::Error::Instance`].
84    #[error("HTTP request error: {0}")]
85    Http(#[from] crate::http_client::Error),
86
87    /// External vector store service returned an error response.
88    #[error("External call to API returned an error. Error code: {0} Message: {1}")]
89    ExternalAPIError(StatusCode, String),
90
91    /// The requested sample count exceeds what the backend can express.
92    #[error("Requested {requested} samples, but this vector store returns at most {max}")]
93    SamplesOutOfRange {
94        /// Sample count from the search request.
95        requested: u64,
96        /// Largest sample count the backend accepts.
97        max: u64,
98    },
99}
100
101impl VectorStoreError {
102    /// Wraps a backend error as [`VectorStoreError::DatastoreError`].
103    ///
104    /// Handles the wasm/non-wasm trait-bound split in one place; use as
105    /// `.map_err(VectorStoreError::datastore)`.
106    #[cfg(not(target_family = "wasm"))]
107    pub fn datastore(e: impl std::error::Error + Send + Sync + 'static) -> Self {
108        Self::DatastoreError(Box::new(e))
109    }
110
111    /// Wraps a backend error as [`VectorStoreError::DatastoreError`].
112    #[cfg(target_family = "wasm")]
113    pub fn datastore(e: impl std::error::Error + 'static) -> Self {
114        Self::DatastoreError(Box::new(e))
115    }
116}
117
118/// Serializes each document to JSON once, then applies `f` to every
119/// `(document, embedding)` pair, flattening the results into a single vector.
120///
121/// This is the shared shape of most [`InsertDocuments`] implementations:
122/// build one backend record per embedding, carrying the owning document's
123/// serialized form.
124pub fn flatten_embedded<Doc: Serialize, R>(
125    documents: Vec<(Doc, Vec<Embedding>)>,
126    mut f: impl FnMut(&Value, Embedding) -> Result<R, VectorStoreError>,
127) -> Result<Vec<R>, VectorStoreError> {
128    let mut records = Vec::new();
129    for (document, embeddings) in documents {
130        let json_document = serde_json::to_value(&document)?;
131        for embedding in embeddings {
132            records.push(f(&json_document, embedding)?);
133        }
134    }
135    Ok(records)
136}
137
138/// Trait for inserting documents and embeddings into a vector store.
139pub trait InsertDocuments: WasmCompatSend + WasmCompatSync {
140    /// Insert precomputed embeddings for each document.
141    ///
142    /// Callers must supply at least one embedding per document. Empty-list
143    /// behavior is backend-dependent and is not uniformly validated.
144    fn insert_documents<Doc: Serialize + Embed + WasmCompatSend>(
145        &self,
146        documents: Vec<(Doc, Vec<Embedding>)>,
147    ) -> impl std::future::Future<Output = Result<(), VectorStoreError>> + WasmCompatSend;
148}
149
150/// Trait for querying a vector store by similarity.
151pub trait VectorStoreIndex: WasmCompatSend + WasmCompatSync {
152    /// The filter type for this backend.
153    type Filter: SearchFilter + WasmCompatSend + WasmCompatSync;
154
155    /// Returns the top N most similar documents, most similar first.
156    fn top_n<T: DeserializeOwned + WasmCompatSend>(
157        &self,
158        req: VectorSearchRequest<Self::Filter>,
159    ) -> impl std::future::Future<Output = Result<Vec<VectorSearchResult<T>>, VectorStoreError>>
160    + WasmCompatSend;
161
162    /// Returns the top N most similar document IDs, most similar first.
163    fn top_n_ids(
164        &self,
165        req: VectorSearchRequest<Self::Filter>,
166    ) -> impl std::future::Future<Output = Result<Vec<VectorSearchIdResult>, VectorStoreError>>
167    + WasmCompatSend;
168}
169
170/// One document returned by [`VectorStoreIndex::top_n`].
171#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
172pub struct VectorSearchResult<T> {
173    /// Backend similarity score. Its scale and direction follow the backend's metric.
174    pub score: f64,
175    /// Backend document ID.
176    pub id: String,
177    /// The stored document, deserialized as `T`.
178    pub document: T,
179}
180
181/// One document ID returned by [`VectorStoreIndex::top_n_ids`].
182#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
183pub struct VectorSearchIdResult {
184    /// Backend similarity score. Its scale and direction follow the backend's metric.
185    pub score: f64,
186    /// Backend document ID.
187    pub id: String,
188}
189
190impl<T, F> Tool for T
191where
192    F: SearchFilter<Value = serde_json::Value>
193        + WasmCompatSend
194        + WasmCompatSync
195        + serde::de::DeserializeOwned,
196    T: VectorStoreIndex<Filter = F>,
197{
198    const NAME: &'static str = "search_vector_store";
199    type Error = VectorStoreError;
200    type Args = VectorSearchRequest<F>;
201    type Output = Vec<VectorSearchResult<Value>>;
202
203    fn description(&self) -> String {
204        "Retrieves the most relevant documents from a vector store based on a query.".to_string()
205    }
206
207    fn parameters(&self) -> serde_json::Value {
208        json!({
209            "type": "object",
210            "properties": {
211                "query": {
212                    "type": "string",
213                    "description": "The query string to search for relevant documents in the vector store."
214                },
215                "samples": {
216                    "type": "integer",
217                    "description": "The maximum number of samples / documents to retrieve.",
218                    "default": 5,
219                    "minimum": 1
220                },
221                "threshold": {
222                    "type": "number",
223                    "description": "Similarity search threshold. If present, any result with a distance less than this may be omitted from the final result."
224                }
225            },
226            "required": ["query", "samples"]
227        })
228    }
229
230    async fn call(
231        &self,
232        _: &mut ToolContext,
233        args: Self::Args,
234    ) -> Result<Self::Output, Self::Error> {
235        self.top_n(args).await
236    }
237}
238
239#[cfg(test)]
240mod tests;
rig_core/vector_store/mod.rs

rig_core/vector_store/
mod.rs
