rig_core/vector_store/
builder.rs

1use serde::Serialize;
2use std::collections::HashMap;
3
4use crate::{OneOrMany, 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, OneOrMany<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(
65        mut self,
66        documents: impl IntoIterator<Item = (D, OneOrMany<Embedding>)>,
67    ) -> Self {
68        let current_index = self.embeddings.len();
69        documents
70            .into_iter()
71            .enumerate()
72            .for_each(|(i, (doc, embeddings))| {
73                self.embeddings
74                    .insert(format!("doc{}", i + current_index), (doc, embeddings));
75            });
76        self
77    }
78
79    /// Add documents with explicit IDs.
80    pub fn documents_with_ids(
81        mut self,
82        documents: impl IntoIterator<Item = (impl ToString, D, OneOrMany<Embedding>)>,
83    ) -> Self {
84        documents.into_iter().for_each(|(id, doc, embeddings)| {
85            self.embeddings.insert(id.to_string(), (doc, embeddings));
86        });
87        self
88    }
89
90    /// Add documents with IDs generated by a function.
91    pub fn documents_with_id_f(
92        mut self,
93        documents: impl IntoIterator<Item = (D, OneOrMany<Embedding>)>,
94        f: fn(&D) -> String,
95    ) -> Self {
96        documents.into_iter().for_each(|(doc, embeddings)| {
97            let id = f(&doc);
98            self.embeddings.insert(id, (doc, embeddings));
99        });
100        self
101    }
102
103    /// Build the [`InMemoryVectorStore`] with the configured settings.
104    pub fn build(self) -> InMemoryVectorStore<D> {
105        InMemoryVectorStore::from_builder(self.embeddings, self.index_strategy)
106    }
107}
rig_core/vector_store/builder.rs

rig_core/vector_store/
builder.rs
