rig-core 0.35.0

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>
where
    D: Serialize,
{
    /// Embeddings of the documents.
    embeddings: HashMap<String, (D, OneOrMany<Embedding>)>,

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

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

impl<D> InMemoryVectorStoreBuilder<D>
where
    D: Serialize + Eq,
{
    /// 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)
    }
}