rig-core 0.44.0

//! Configuration builder for the in-memory vector store.
//!
//! ```
//! use rig_core::vector_store::builder::InMemoryVectorStoreBuilder;
//!
//! let store = InMemoryVectorStoreBuilder::<String>::new().build();
//! # let _ = store;
//! ```

use serde::Serialize;
use std::collections::HashMap;

use crate::embeddings::Embedding;

use super::{IndexStrategy, in_memory_store::InMemoryVectorStore};

/// Builder for creating an [`InMemoryVectorStore`] with custom configuration.
pub struct InMemoryVectorStoreBuilder<D> {
    /// Embeddings of the documents.
    embeddings: HashMap<String, (D, Vec<Embedding>)>,

    /// Index strategy for the vector store.
    index_strategy: IndexStrategy,
}

impl<D> Default for InMemoryVectorStoreBuilder<D>
where
    D: Serialize,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<D> InMemoryVectorStoreBuilder<D>
where
    D: Serialize,
{
    /// Create a new builder with default settings.
    ///
    /// The default index strategy is [`IndexStrategy::BruteForce`].
    pub fn new() -> Self {
        Self {
            embeddings: HashMap::new(),
            index_strategy: IndexStrategy::default(),
        }
    }

    /// Set the index strategy for the vector store.
    ///
    /// ```
    /// use rig_core::vector_store::{builder::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 have the form `"doc{n}"`, starting at the current document count and
    /// skipping occupied IDs so existing documents are never overwritten.
    pub fn documents(mut self, documents: impl IntoIterator<Item = (D, Vec<Embedding>)>) -> Self {
        let mut index = self.embeddings.len();
        for (doc, embeddings) in documents {
            let mut id = format!("doc{index}");
            while self.embeddings.contains_key(&id) {
                index += 1;
                id = format!("doc{index}");
            }
            self.embeddings.insert(id, (doc, embeddings));
            index += 1;
        }
        self
    }

    /// Adds documents with explicit IDs, replacing existing entries with matching IDs.
    pub fn documents_with_ids(
        mut self,
        documents: impl IntoIterator<Item = (impl ToString, D, Vec<Embedding>)>,
    ) -> Self {
        documents.into_iter().for_each(|(id, doc, embeddings)| {
            self.embeddings.insert(id.to_string(), (doc, embeddings));
        });
        self
    }

    /// Build the [`InMemoryVectorStore`] with the configured settings.
    pub fn build(self) -> InMemoryVectorStore<D> {
        InMemoryVectorStore::from_builder(self.embeddings, self.index_strategy)
    }
}