feat: LSH for in memory vector store

rig-core/Cargo.toml

@@ -43,6 +43,7 @@ tokio = { workspace = true, features = ["rt", "sync"] }
 http = "1.3.1"
 tracing-futures = { version = "0.2.5", features = ["futures-03"] }
 serenity = { version = "0.12.4", optional = true }
+fastrand = "2.3.0"
 eventsource-stream = "0.2.3"
 pin-project-lite = "0.2.16"
 futures-timer = "3.0.3"
@@ -165,4 +166,4 @@ name = "request_hook"
 
 [[example]]
 name = "discord_bot"
-required-features = ["discord-bot"]
+required-features = ["discord-bot"]

rig-core/examples/rag.rs

@@ -63,7 +63,6 @@ async fn main() -> Result<(), anyhow::Error> {
 
     // Create vector store with the embeddings
     let vector_store = InMemoryVectorStore::from_documents(embeddings);
-
     // Create vector store index
     let index = vector_store.index(embedding_model);
     let rag_agent = openai_client.agent("gpt-4")

rig-core/src/vector_store/builder.rs

@@ -0,0 +1,94 @@
+use serde::Serialize;
+use std::collections::HashMap;
+
+use crate::{OneOrMany, embeddings::Embedding};
+
+use super::{IndexStrategy, in_memory_store::InMemoryVectorStore};
+
+/// Builder for creating an [InMemoryVectorStore] with custom configuration.
+pub struct InMemoryVectorStoreBuilder<D: Serialize> {
+    embeddings: HashMap<String, (D, OneOrMany<Embedding>)>,
+    index_strategy: IndexStrategy,
+}
+
+impl<D: Serialize + Eq> Default for InMemoryVectorStoreBuilder<D> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<D: Serialize + Eq> InMemoryVectorStoreBuilder<D> {
+    /// Create a new builder with default settings.
+    /// Default index strategy is BruteForce.
+    pub fn new() -> Self {
+        Self {
+            embeddings: HashMap::new(),
+            index_strategy: IndexStrategy::default(),
+        }
+    }
+
+    /// Set the index strategy for the vector store.
+    ///
+    /// # Examples
+    ///
+    /// ```ignore
+    /// use rig::vector_store::{InMemoryVectorStoreBuilder, IndexStrategy};
+    ///
+    /// let store = InMemoryVectorStoreBuilder::<String>::new()
+    ///     .index_strategy(IndexStrategy::LSH {
+    ///         num_tables: 5,
+    ///         num_hyperplanes: 10,
+    ///     })
+    ///     .build();
+    /// ```
+    pub fn index_strategy(mut self, index_strategy: IndexStrategy) -> Self {
+        self.index_strategy = index_strategy;
+        self
+    }
+
+    /// Add documents with auto-generated IDs.
+    /// IDs will have the form `"doc{n}"` where `n` is the index.
+    pub fn documents(
+        mut self,
+        documents: impl IntoIterator<Item = (D, OneOrMany<Embedding>)>,
+    ) -> Self {
+        let current_index = self.embeddings.len();
+        documents
+            .into_iter()
+            .enumerate()
+            .for_each(|(i, (doc, embeddings))| {
+                self.embeddings
+                    .insert(format!("doc{}", i + current_index), (doc, embeddings));
+            });
+        self
+    }
+
+    /// Add documents with explicit IDs.
+    pub fn documents_with_ids(
+        mut self,
+        documents: impl IntoIterator<Item = (impl ToString, D, OneOrMany<Embedding>)>,
+    ) -> Self {
+        documents.into_iter().for_each(|(id, doc, embeddings)| {
+            self.embeddings.insert(id.to_string(), (doc, embeddings));
+        });
+        self
+    }
+
+    /// Add documents with IDs generated by a function.
+    pub fn documents_with_id_f(
+        mut self,
+        documents: impl IntoIterator<Item = (D, OneOrMany<Embedding>)>,
+        f: fn(&D) -> String,
+    ) -> Self {
+        documents.into_iter().for_each(|(doc, embeddings)| {
+            let id = f(&doc);
+            self.embeddings.insert(id, (doc, embeddings));
+        });
+        self
+    }
+
+    /// Build the [InMemoryVectorStore] with the configured settings.
+    pub fn build(self) -> InMemoryVectorStore<D> {
+        InMemoryVectorStore::from_builder(self.embeddings, self.index_strategy)
+    }
+}