ic-rig 0.2.0

//! [`EmbeddingsBuilder`] — batch embedding with automatic chunking.
//!
//! The builder collects documents, extracts their texts via [`Embed`], chunks
//! them to respect `MAX_DOCUMENTS`, calls the model in sequential batches,
//! then reassembles the results back to the originating document.
//!
//! # Why sequential batching?
//!
//! Rig uses `buffer_unordered(10)` from the `futures` crate to parallelise
//! batch requests. `irig` deliberately avoids the `futures` dependency and
//! runs batches sequentially instead. On ICP, HTTP outcalls are already
//! individually async — if you need parallelism at the application level,
//! issue multiple agent/embedding calls from your canister code.
//!
//! # Example
//!
//! ```rust,ignore
//! let results: Vec<(Article, Vec<Embedding>)> =
//!     EmbeddingsBuilder::new(model)
//!         .documents(articles)?
//!         .build()
//!         .await?;
//!
//! for (article, embeddings) in results {
//!     // embeddings[0] = title vector, embeddings[1] = body vector
//!     store.upsert(article.id, embeddings);
//! }
//! ```

use super::{
    Embed, EmbedError, Embedding, EmbeddingError, EmbeddingModel,
    embed::TextEmbedder,
};

// ── Builder ───────────────────────────────────────────────────────────────────

/// Accumulates documents, embeds them in efficient batches, and returns each
/// document paired with its embedding(s).
pub struct EmbeddingsBuilder<M, T> {
    model: M,
    /// Each entry is `(document, [texts_to_embed])`. One document may produce
    /// multiple texts (and therefore multiple embeddings).
    documents: Vec<(T, Vec<String>)>,
}

impl<M: EmbeddingModel, T: Embed> EmbeddingsBuilder<M, T> {
    pub fn new(model: M) -> Self {
        Self { model, documents: Vec::new() }
    }

    /// Add a single document.
    pub fn document(mut self, doc: T) -> Result<Self, EmbedError> {
        let mut embedder = TextEmbedder::default();
        doc.embed(&mut embedder)?;
        self.documents.push((doc, embedder.texts));
        Ok(self)
    }

    /// Add multiple documents.
    pub fn documents(self, docs: impl IntoIterator<Item = T>) -> Result<Self, EmbedError> {
        docs.into_iter().try_fold(self, |b, doc| b.document(doc))
    }

    /// Embed all queued documents.
    ///
    /// Returns `Vec<(T, Vec<Embedding>)>`. Each `Vec<Embedding>` has one entry
    /// per text the document pushed via [`Embed::embed`] — same order.
    pub async fn build(self) -> Result<Vec<(T, Vec<Embedding>)>, EmbeddingError> {
        // Flatten: (doc_index, text) pairs in insertion order.
        let mut flat: Vec<(usize, String)> = Vec::new();
        for (i, (_, texts)) in self.documents.iter().enumerate() {
            for text in texts {
                flat.push((i, text.clone()));
            }
        }

        // Embed in chunks, sequentially.
        // `embeddings_by_doc[i]` accumulates the Vec<Embedding> for document i.
        let mut embeddings_by_doc: Vec<Vec<Embedding>> =
            (0..self.documents.len()).map(|_| Vec::new()).collect();

        for chunk in flat.chunks(M::MAX_DOCUMENTS) {
            let (ids, texts): (Vec<usize>, Vec<String>) = chunk
                .iter()
                .cloned()
                .unzip();

            let batch = self
                .model
                .embed_texts(texts)
                .await
                .map_err(|e| EmbeddingError::Response(e.to_string()))?;

            if batch.len() != ids.len() {
                return Err(EmbeddingError::Response(format!(
                    "model returned {} embeddings for {} inputs",
                    batch.len(),
                    ids.len(),
                )));
            }

            for (doc_idx, embedding) in ids.into_iter().zip(batch) {
                embeddings_by_doc[doc_idx].push(embedding);
            }
        }

        // Pair each document with its embeddings.
        let result = self
            .documents
            .into_iter()
            .zip(embeddings_by_doc)
            .map(|((doc, _), embeddings)| (doc, embeddings))
            .collect();

        Ok(result)
    }
}