rig_core/vector_store/
builder.rs

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

rig_core/vector_store/
builder.rs
