ic_rig/embeddings/
builder.rs

1//! [`EmbeddingsBuilder`] — batch embedding with automatic chunking.
2//!
3//! The builder collects documents, extracts their texts via [`Embed`], chunks
4//! them to respect `MAX_DOCUMENTS`, calls the model in sequential batches,
5//! then reassembles the results back to the originating document.
6//!
7//! # Why sequential batching?
8//!
9//! Rig uses `buffer_unordered(10)` from the `futures` crate to parallelise
10//! batch requests. `irig` deliberately avoids the `futures` dependency and
11//! runs batches sequentially instead. On ICP, HTTP outcalls are already
12//! individually async — if you need parallelism at the application level,
13//! issue multiple agent/embedding calls from your canister code.
14//!
15//! # Example
16//!
17//! ```rust,ignore
18//! let results: Vec<(Article, Vec<Embedding>)> =
19//!     EmbeddingsBuilder::new(model)
20//!         .documents(articles)?
21//!         .build()
22//!         .await?;
23//!
24//! for (article, embeddings) in results {
25//!     // embeddings[0] = title vector, embeddings[1] = body vector
26//!     store.upsert(article.id, embeddings);
27//! }
28//! ```
29
30use super::{
31    Embed, EmbedError, Embedding, EmbeddingError, EmbeddingModel,
32    embed::TextEmbedder,
33};
34
35// ── Builder ───────────────────────────────────────────────────────────────────
36
37/// Accumulates documents, embeds them in efficient batches, and returns each
38/// document paired with its embedding(s).
39pub struct EmbeddingsBuilder<M, T> {
40    model: M,
41    /// Each entry is `(document, [texts_to_embed])`. One document may produce
42    /// multiple texts (and therefore multiple embeddings).
43    documents: Vec<(T, Vec<String>)>,
44}
45
46impl<M: EmbeddingModel, T: Embed> EmbeddingsBuilder<M, T> {
47    pub fn new(model: M) -> Self {
48        Self { model, documents: Vec::new() }
49    }
50
51    /// Add a single document.
52    pub fn document(mut self, doc: T) -> Result<Self, EmbedError> {
53        let mut embedder = TextEmbedder::default();
54        doc.embed(&mut embedder)?;
55        self.documents.push((doc, embedder.texts));
56        Ok(self)
57    }
58
59    /// Add multiple documents.
60    pub fn documents(self, docs: impl IntoIterator<Item = T>) -> Result<Self, EmbedError> {
61        docs.into_iter().try_fold(self, |b, doc| b.document(doc))
62    }
63
64    /// Embed all queued documents.
65    ///
66    /// Returns `Vec<(T, Vec<Embedding>)>`. Each `Vec<Embedding>` has one entry
67    /// per text the document pushed via [`Embed::embed`] — same order.
68    pub async fn build(self) -> Result<Vec<(T, Vec<Embedding>)>, EmbeddingError> {
69        // Flatten: (doc_index, text) pairs in insertion order.
70        let mut flat: Vec<(usize, String)> = Vec::new();
71        for (i, (_, texts)) in self.documents.iter().enumerate() {
72            for text in texts {
73                flat.push((i, text.clone()));
74            }
75        }
76
77        // Embed in chunks, sequentially.
78        // `embeddings_by_doc[i]` accumulates the Vec<Embedding> for document i.
79        let mut embeddings_by_doc: Vec<Vec<Embedding>> =
80            (0..self.documents.len()).map(|_| Vec::new()).collect();
81
82        for chunk in flat.chunks(M::MAX_DOCUMENTS) {
83            let (ids, texts): (Vec<usize>, Vec<String>) = chunk
84                .iter()
85                .cloned()
86                .unzip();
87
88            let batch = self
89                .model
90                .embed_texts(texts)
91                .await
92                .map_err(|e| EmbeddingError::Response(e.to_string()))?;
93
94            if batch.len() != ids.len() {
95                return Err(EmbeddingError::Response(format!(
96                    "model returned {} embeddings for {} inputs",
97                    batch.len(),
98                    ids.len(),
99                )));
100            }
101
102            for (doc_idx, embedding) in ids.into_iter().zip(batch) {
103                embeddings_by_doc[doc_idx].push(embedding);
104            }
105        }
106
107        // Pair each document with its embeddings.
108        let result = self
109            .documents
110            .into_iter()
111            .zip(embeddings_by_doc)
112            .map(|((doc, _), embeddings)| (doc, embeddings))
113            .collect();
114
115        Ok(result)
116    }
117}
ic_rig/embeddings/builder.rs

ic_rig/embeddings/
builder.rs
