rig_core/vector_store/
builder.rs

1//! Configuration builder for the in-memory vector store.
2//!
3//! ```
4//! use rig_core::vector_store::builder::InMemoryVectorStoreBuilder;
5//!
6//! let store = InMemoryVectorStoreBuilder::<String>::new().build();
7//! # let _ = store;
8//! ```
9
10use serde::Serialize;
11use std::collections::HashMap;
12
13use crate::embeddings::Embedding;
14
15use super::{IndexStrategy, in_memory_store::InMemoryVectorStore};
16
17/// Builder for creating an [`InMemoryVectorStore`] with custom configuration.
18pub struct InMemoryVectorStoreBuilder<D> {
19    /// Embeddings of the documents.
20    embeddings: HashMap<String, (D, Vec<Embedding>)>,
21
22    /// Index strategy for the vector store.
23    index_strategy: IndexStrategy,
24}
25
26impl<D> Default for InMemoryVectorStoreBuilder<D>
27where
28    D: Serialize + Eq,
29{
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl<D> InMemoryVectorStoreBuilder<D>
36where
37    D: Serialize + Eq,
38{
39    /// Create a new builder with default settings.
40    ///
41    /// The default index strategy is [`IndexStrategy::BruteForce`].
42    pub fn new() -> Self {
43        Self {
44            embeddings: HashMap::new(),
45            index_strategy: IndexStrategy::default(),
46        }
47    }
48
49    /// Set the index strategy for the vector store.
50    ///
51    /// ```
52    /// use rig_core::vector_store::{builder::InMemoryVectorStoreBuilder, IndexStrategy};
53    ///
54    /// let store = InMemoryVectorStoreBuilder::<String>::new()
55    ///     .index_strategy(IndexStrategy::LSH {
56    ///         num_tables: 5,
57    ///         num_hyperplanes: 10,
58    ///     })
59    ///     .build();
60    /// ```
61    pub fn index_strategy(mut self, index_strategy: IndexStrategy) -> Self {
62        self.index_strategy = index_strategy;
63        self
64    }
65
66    /// Add documents with auto-generated IDs.
67    /// IDs have the form `"doc{n}"`, starting at the current document count and
68    /// skipping occupied IDs so existing documents are never overwritten.
69    pub fn documents(mut self, documents: impl IntoIterator<Item = (D, Vec<Embedding>)>) -> Self {
70        let mut index = self.embeddings.len();
71        for (doc, embeddings) in documents {
72            let mut id = format!("doc{index}");
73            while self.embeddings.contains_key(&id) {
74                index += 1;
75                id = format!("doc{index}");
76            }
77            self.embeddings.insert(id, (doc, embeddings));
78            index += 1;
79        }
80        self
81    }
82
83    /// Adds documents with explicit IDs, replacing existing entries with matching IDs.
84    pub fn documents_with_ids(
85        mut self,
86        documents: impl IntoIterator<Item = (impl ToString, D, Vec<Embedding>)>,
87    ) -> Self {
88        documents.into_iter().for_each(|(id, doc, embeddings)| {
89            self.embeddings.insert(id.to_string(), (doc, embeddings));
90        });
91        self
92    }
93
94    /// Build the [`InMemoryVectorStore`] with the configured settings.
95    pub fn build(self) -> InMemoryVectorStore<D> {
96        InMemoryVectorStore::from_builder(self.embeddings, self.index_strategy)
97    }
98}
rig_core/vector_store/builder.rs

rig_core/vector_store/
builder.rs
