rig_core/embeddings/
builder.rs

1//! The module defines the [EmbeddingsBuilder] struct which accumulates objects to be embedded
2//! and batch generates the embeddings for each object when built.
3//! Only types that implement the [Embed] trait can be added to the [EmbeddingsBuilder].
4
5use std::{cmp::max, ops::Range};
6
7use futures::{StreamExt, stream};
8
9use crate::{
10    completion::Usage,
11    embeddings::{
12        Embed, EmbedError, Embedding, EmbeddingError, EmbeddingModel, EmbeddingResponse,
13        embed::TextEmbedder,
14    },
15};
16
17/// Builder for creating embeddings from one or more documents of type `T`.
18/// Note: `T` can be any type that implements the [Embed] trait.
19///
20/// Using the builder is preferred over using [EmbeddingModel::embed_text] directly as
21/// it will batch the documents in a single request to the model provider.
22///
23/// # Example
24/// ```no_run
25/// use rig_core::{
26///     client::{EmbeddingsClient, ProviderClient},
27///     embeddings::EmbeddingsBuilder,
28///     providers::openai,
29/// };
30///
31/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
32/// // Create OpenAI client
33/// let openai_client = openai::Client::from_env()?;
34///
35/// let model = openai_client.embedding_model(openai::TEXT_EMBEDDING_3_SMALL);
36///
37/// let embeddings = EmbeddingsBuilder::new(model.clone())
38///     .documents(vec![
39///         "1. *flurbo* (noun): A green alien that lives on cold planets.".to_string(),
40///         "2. *flurbo* (noun): A fictional digital currency.".to_string(),
41///         "1. *glarb-glarb* (noun): An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
42///         "2. *glarb-glarb* (noun): A fictional creature from marshlands.".to_string(),
43///         "1. *linlingdong* (noun): A term used by inhabitants of the sombrero galaxy to describe humans.".to_string(),
44///         "2. *linlingdong* (noun): A rare instrument.".to_string(),
45///     ])?
46///     .build()
47///     .await?;
48/// # Ok(())
49/// # }
50/// ```
51pub struct EmbeddingsBuilder<M, T>
52where
53    M: EmbeddingModel,
54    T: Embed,
55{
56    model: M,
57    documents: Vec<(T, Vec<String>)>,
58}
59
60impl<M, T> EmbeddingsBuilder<M, T>
61where
62    M: EmbeddingModel,
63    T: Embed,
64{
65    /// Create a new embedding builder with the given embedding model
66    pub fn new(model: M) -> Self {
67        Self {
68            model,
69            documents: vec![],
70        }
71    }
72
73    /// Add a document to be embedded to the builder. `document` must implement the [Embed] trait.
74    pub fn document(mut self, document: T) -> Result<Self, EmbedError> {
75        let mut embedder = TextEmbedder::default();
76        document.embed(&mut embedder)?;
77
78        self.documents.push((document, embedder.texts));
79
80        Ok(self)
81    }
82
83    /// Add multiple documents to be embedded to the builder. `documents` must be iterable
84    /// with items that implement the [Embed] trait.
85    pub fn documents(self, documents: impl IntoIterator<Item = T>) -> Result<Self, EmbedError> {
86        let builder = documents
87            .into_iter()
88            .try_fold(self, |builder, doc| builder.document(doc))?;
89
90        Ok(builder)
91    }
92}
93
94impl<M, T> EmbeddingsBuilder<M, T>
95where
96    M: EmbeddingModel,
97    T: Embed + Send,
98{
99    /// Generate embeddings for all documents in the builder.
100    ///
101    /// Returns `(document, embeddings)` pairs. A document may produce one or many
102    /// embeddings depending on how its [`Embed`] implementation uses [`TextEmbedder`].
103    ///
104    /// # Order
105    ///
106    /// Both levels are ordered, and callers may rely on it:
107    ///
108    /// - pairs come back in the order the documents were added — positional
109    ///   callers depend on this, for example
110    ///   [`InMemoryVectorStore::add_documents`](crate::vector_store::in_memory_store::InMemoryVectorStore::add_documents),
111    ///   which derives its document ids from this sequence; and
112    /// - each document's embeddings come back in the order its [`Embed`] impl
113    ///   produced the texts.
114    ///
115    /// Neither depends on how the texts were batched or on which batch the
116    /// provider answered first. Both have been silently violated before
117    /// (rig#2344, rig#2345), so treat the guarantee as load-bearing rather than
118    /// incidental.
119    ///
120    /// The second bullet inherits one assumption this type cannot check:
121    /// providers pair a batch's embeddings to its texts positionally, so a
122    /// provider that reordered *within* a single response would still be
123    /// believed. That is the provider's contract, not this builder's.
124    ///
125    /// # Errors
126    ///
127    /// Alongside whatever the provider and the transport return, two cases
128    /// originate here:
129    ///
130    /// - **A document that produces no text** fails the whole build rather than
131    ///   coming back with an empty list. This is easy to hit by accident: an
132    ///   empty collection in an `#[embed]` field embeds nothing, because
133    ///   [`Embed`] is implemented for `Vec<T>` element-wise.
134    /// - **A provider returning fewer embeddings than the texts it was sent**
135    ///   fails rather than handing back a short list, since a short list cannot
136    ///   be told apart from a document that legitimately has fewer texts.
137    ///
138    /// Both name the offending document.
139    pub async fn build(self) -> Result<Vec<(T, Vec<Embedding>)>, EmbeddingError> {
140        let (result, _usage) = self.build_with_usage().await?;
141        Ok(result)
142    }
143
144    /// Generate embeddings for all documents in the builder and return accumulated token usage.
145    ///
146    /// Returns `(document, embeddings)` pairs and the total token usage across all
147    /// batches. A document may produce one or many embeddings depending on how its
148    /// [`Embed`] implementation uses [`TextEmbedder`].
149    ///
150    /// Ordering is guaranteed at both levels, and the same two errors originate
151    /// here; both are described on [`Self::build`].
152    pub async fn build_with_usage(
153        self,
154    ) -> Result<(Vec<(T, Vec<Embedding>)>, Usage), EmbeddingError> {
155        use stream::TryStreamExt;
156
157        // Flatten every document's texts into one slot-indexed list, recording
158        // the contiguous slot range each document owns.
159        //
160        // The slot index is what makes ordering independent of completion
161        // order at *both* levels. Keying by document alone was not enough
162        // (rig#2345): `chunks` splits on a flat text count, so one document's
163        // texts can straddle a batch boundary, `buffer_unordered` yields
164        // batches as they finish, and appending to a per-document list then
165        // recorded completion order — a straddling document got its own
166        // embeddings back shuffled. A batch now writes each embedding into its
167        // own slot, so when a batch finishes cannot affect where anything
168        // lands.
169        let mut docs: Vec<T> = Vec::with_capacity(self.documents.len());
170        let mut spans: Vec<Range<usize>> = Vec::with_capacity(self.documents.len());
171        let mut texts: Vec<String> = Vec::new();
172
173        for (doc, doc_texts) in self.documents {
174            let start = texts.len();
175            texts.extend(doc_texts);
176            spans.push(start..texts.len());
177            docs.push(doc);
178        }
179
180        let total_texts = texts.len();
181
182        // Compute the embeddings.
183        let (slots, usage) = stream::iter(texts.into_iter().enumerate())
184            // Chunk them into batches. Each batch size is at most the embedding API limit per request.
185            .chunks(M::MAX_DOCUMENTS)
186            // Generate the embeddings for each batch with usage tracking.
187            .map(|chunk| async {
188                let (slots, batch): (Vec<usize>, Vec<String>) = chunk.into_iter().unzip();
189
190                let response: EmbeddingResponse = self.model.embed_texts_with_usage(batch).await?;
191                Ok::<_, EmbeddingError>((
192                    slots
193                        .into_iter()
194                        .zip(response.embeddings)
195                        .collect::<Vec<_>>(),
196                    response.usage,
197                ))
198            })
199            // Parallelize the embeddings generation over 10 concurrent requests
200            .buffer_unordered(max(1, 1024 / M::MAX_DOCUMENTS))
201            // Write each embedding into the slot its text came from, and
202            // accumulate usage.
203            .try_fold(
204                (
205                    (0..total_texts)
206                        .map(|_| None)
207                        .collect::<Vec<Option<Embedding>>>(),
208                    Usage::default(),
209                ),
210                |(mut slots, mut usage_acc), (chunk_embeddings, chunk_usage)| async move {
211                    for (slot, embedding) in chunk_embeddings {
212                        // Every slot came from this function's own `enumerate`
213                        // and the `zip` above truncates to the shorter side, so
214                        // this index is in range by construction — including
215                        // when a provider answers with more embeddings than it
216                        // was sent. `get_mut` rather than `slots[slot]` only
217                        // because `clippy::indexing_slicing` is denied here.
218                        if let Some(place) = slots.get_mut(slot) {
219                            *place = Some(embedding);
220                        }
221                    }
222                    usage_acc += chunk_usage;
223                    Ok((slots, usage_acc))
224                },
225            )
226            .await?;
227
228        // Hand each document the contiguous run of slots its texts occupied,
229        // in text order.
230        let mut slots = slots.into_iter();
231        let mut result = Vec::with_capacity(docs.len());
232
233        for (index, (doc, span)) in docs.into_iter().zip(spans).enumerate() {
234            // A document that embedded no text has no embeddings to return;
235            // this has always been an error rather than an empty list.
236            if span.is_empty() {
237                return Err(crate::embeddings::EmbeddingError::ResponseError(format!(
238                    "document {index} produced no text to embed, so it has no \
239                     embeddings to return; an empty collection in an `#[embed]` \
240                     field embeds nothing"
241                )));
242            }
243
244            // An empty slot means the provider returned fewer embeddings than
245            // the texts sent in some batch. Previously `zip` dropped the
246            // surplus texts and the document came back with a short list;
247            // naming the slot turns silent loss into a located error.
248            let embeddings = slots
249                .by_ref()
250                .take(span.len())
251                .collect::<Option<Vec<Embedding>>>()
252                .ok_or_else(|| {
253                    crate::embeddings::EmbeddingError::ResponseError(format!(
254                        "provider returned fewer embeddings than texts sent: \
255                         document {index} is missing at least one of its {} texts \
256                         (slots {}..{} of {total_texts})",
257                        span.len(),
258                        span.start,
259                        span.end
260                    ))
261                })?;
262
263            result.push((doc, embeddings));
264        }
265
266        Ok((result, usage))
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use crate::embeddings::embed::{EmbedError, TextEmbedder};
273    use crate::embeddings::{Embed, Embedding, EmbeddingError, EmbeddingModel};
274    use crate::test_utils::{MockEmbeddingModel, MockMultiTextDocument, MockTextDocument};
275
276    use super::EmbeddingsBuilder;
277
278    fn definitions_multiple_text() -> Vec<MockMultiTextDocument> {
279        vec![
280            MockMultiTextDocument::new(
281                "doc0",
282                [
283                    "A green alien that lives on cold planets.",
284                    "A fictional digital currency that originated in the animated series Rick and Morty.",
285                ],
286            ),
287            MockMultiTextDocument::new(
288                "doc1",
289                [
290                    "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.",
291                    "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.",
292                ],
293            ),
294        ]
295    }
296
297    fn definitions_multiple_text_2() -> Vec<MockMultiTextDocument> {
298        vec![
299            MockMultiTextDocument::new("doc2", ["Another fake definitions"]),
300            MockMultiTextDocument::new("doc3", ["Some fake definition"]),
301        ]
302    }
303
304    fn definitions_single_text() -> Vec<MockTextDocument> {
305        vec![
306            MockTextDocument::new("doc0", "A green alien that lives on cold planets."),
307            MockTextDocument::new(
308                "doc1",
309                "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.",
310            ),
311        ]
312    }
313
314    #[tokio::test]
315    async fn test_build_multiple_text() {
316        let fake_definitions = definitions_multiple_text();
317
318        let fake_model = MockEmbeddingModel;
319        let result = EmbeddingsBuilder::new(fake_model)
320            .documents(fake_definitions)
321            .unwrap()
322            .build()
323            .await
324            .unwrap();
325
326        assert_eq!(result.len(), 2);
327
328        let first_definition = &result[0];
329        assert_eq!(first_definition.0.id, "doc0");
330        assert_eq!(first_definition.1.len(), 2);
331        assert_eq!(
332            first_definition.1.first().map(|e| e.document.as_str()),
333            Some("A green alien that lives on cold planets.")
334        );
335
336        let second_definition = &result[1];
337        assert_eq!(second_definition.0.id, "doc1");
338        assert_eq!(second_definition.1.len(), 2);
339        assert_eq!(
340            second_definition.1.get(1).map(|e| e.document.as_str()),
341            Some(
342                "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy."
343            )
344        )
345    }
346
347    #[tokio::test]
348    async fn test_build_single_text() {
349        let fake_definitions = definitions_single_text();
350
351        let fake_model = MockEmbeddingModel;
352        let result = EmbeddingsBuilder::new(fake_model)
353            .documents(fake_definitions)
354            .unwrap()
355            .build()
356            .await
357            .unwrap();
358
359        assert_eq!(result.len(), 2);
360
361        let first_definition = &result[0];
362        assert_eq!(first_definition.0.id, "doc0");
363        assert_eq!(first_definition.1.len(), 1);
364        assert_eq!(
365            first_definition.1.first().map(|e| e.document.as_str()),
366            Some("A green alien that lives on cold planets.")
367        );
368
369        let second_definition = &result[1];
370        assert_eq!(second_definition.0.id, "doc1");
371        assert_eq!(second_definition.1.len(), 1);
372        assert_eq!(
373            second_definition.1.first().map(|e| e.document.as_str()),
374            Some(
375                "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land."
376            )
377        )
378    }
379
380    #[tokio::test]
381    async fn test_build_multiple_and_single_text() {
382        let fake_definitions = definitions_multiple_text();
383        let fake_definitions_single = definitions_multiple_text_2();
384
385        let fake_model = MockEmbeddingModel;
386        let result = EmbeddingsBuilder::new(fake_model)
387            .documents(fake_definitions)
388            .unwrap()
389            .documents(fake_definitions_single)
390            .unwrap()
391            .build()
392            .await
393            .unwrap();
394
395        assert_eq!(result.len(), 4);
396
397        let second_definition = &result[1];
398        assert_eq!(second_definition.0.id, "doc1");
399        assert_eq!(second_definition.1.len(), 2);
400        assert_eq!(
401            second_definition.1.first().map(|e| e.document.as_str()),
402            Some(
403                "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land."
404            )
405        );
406
407        let third_definition = &result[2];
408        assert_eq!(third_definition.0.id, "doc2");
409        assert_eq!(third_definition.1.len(), 1);
410        assert_eq!(
411            third_definition.1.first().map(|e| e.document.as_str()),
412            Some("Another fake definitions")
413        )
414    }
415
416    #[tokio::test]
417    async fn test_build_string() {
418        let bindings = definitions_multiple_text();
419        let fake_definitions = bindings.iter().map(|def| def.texts.clone());
420
421        let fake_model = MockEmbeddingModel;
422        let result = EmbeddingsBuilder::new(fake_model)
423            .documents(fake_definitions)
424            .unwrap()
425            .build()
426            .await
427            .unwrap();
428
429        assert_eq!(result.len(), 2);
430
431        let first_definition = &result[0];
432        assert_eq!(first_definition.1.len(), 2);
433        assert_eq!(
434            first_definition.1.first().map(|e| e.document.as_str()),
435            Some("A green alien that lives on cold planets.")
436        );
437
438        let second_definition = &result[1];
439        assert_eq!(second_definition.1.len(), 2);
440        assert_eq!(
441            second_definition.1.get(1).map(|e| e.document.as_str()),
442            Some(
443                "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy."
444            )
445        )
446    }
447
448    #[tokio::test]
449    async fn test_build_preserves_input_order_across_batches() {
450        // More documents than MockEmbeddingModel::MAX_DOCUMENTS (5) to exercise
451        // the chunked, buffered batch path, and assert that the returned
452        // sequence matches the input order exactly.
453        let texts: Vec<String> = (0..12).map(|i| format!("text-{i:02}")).collect();
454
455        let fake_model = MockEmbeddingModel;
456        let result = EmbeddingsBuilder::new(fake_model)
457            .documents(texts.clone())
458            .unwrap()
459            .build()
460            .await
461            .unwrap();
462
463        assert_eq!(result.len(), texts.len());
464        for (i, (doc, embeddings)) in result.into_iter().enumerate() {
465            assert_eq!(doc, texts[i]);
466            assert_eq!(embeddings.len(), 1);
467            assert_eq!(embeddings[0].document, texts[i]);
468        }
469    }
470
471    /// A model whose *first* batch is slow, so later batches finish first.
472    ///
473    /// `buffer_unordered` yields batches as they complete, which is the only
474    /// way to observe rig#2345 deterministically: without a delay the batches
475    /// happen to finish in submission order and the defect hides.
476    #[derive(Clone)]
477    struct SlowFirstBatchModel {
478        calls: std::sync::Arc<std::sync::atomic::AtomicUsize>,
479    }
480
481    impl SlowFirstBatchModel {
482        fn new() -> Self {
483            Self {
484                calls: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
485            }
486        }
487    }
488
489    impl EmbeddingModel for SlowFirstBatchModel {
490        const MAX_DOCUMENTS: usize = 5;
491
492        type Client = crate::client::Nothing;
493
494        fn make(_: &Self::Client, _: impl Into<String>, _: Option<usize>) -> Self {
495            Self::new()
496        }
497
498        fn ndims(&self) -> usize {
499            10
500        }
501
502        async fn embed_texts(
503            &self,
504            documents: impl IntoIterator<Item = String> + crate::wasm_compat::WasmCompatSend,
505        ) -> Result<Vec<Embedding>, EmbeddingError> {
506            let nth = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
507            if nth == 0 {
508                tokio::time::sleep(std::time::Duration::from_millis(150)).await;
509            }
510            Ok(documents
511                .into_iter()
512                .map(|document| Embedding {
513                    document,
514                    vec: vec![0.0; 10],
515                })
516                .collect())
517        }
518    }
519
520    /// A model whose batches finish in reverse submission order: batch `n`
521    /// sleeps longer the earlier it was submitted.
522    ///
523    /// `SlowFirstBatchModel` only inverts the *first* boundary, so a document
524    /// straddling a later one still comes back correct even unfixed. This
525    /// inverts every boundary.
526    #[derive(Clone)]
527    struct DescendingLatencyModel {
528        batches: std::sync::Arc<std::sync::atomic::AtomicUsize>,
529    }
530
531    impl DescendingLatencyModel {
532        fn new() -> Self {
533            Self {
534                batches: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
535            }
536        }
537    }
538
539    impl EmbeddingModel for DescendingLatencyModel {
540        const MAX_DOCUMENTS: usize = 5;
541
542        type Client = crate::client::Nothing;
543
544        fn make(_: &Self::Client, _: impl Into<String>, _: Option<usize>) -> Self {
545            Self::new()
546        }
547
548        fn ndims(&self) -> usize {
549            10
550        }
551
552        async fn embed_texts(
553            &self,
554            documents: impl IntoIterator<Item = String> + crate::wasm_compat::WasmCompatSend,
555        ) -> Result<Vec<Embedding>, EmbeddingError> {
556            let nth = self
557                .batches
558                .fetch_add(1, std::sync::atomic::Ordering::SeqCst) as u64;
559            tokio::time::sleep(std::time::Duration::from_millis(
560                120u64.saturating_sub(nth * 40),
561            ))
562            .await;
563            Ok(documents
564                .into_iter()
565                .map(|document| Embedding {
566                    document,
567                    vec: vec![0.0; 10],
568                })
569                .collect())
570        }
571    }
572
573    /// A document contributing `n` texts, each naming **both** its owner and
574    /// its position: `d{doc}t0 .. d{doc}t{n-1}`.
575    ///
576    /// The owner half is load-bearing. Documents of the same length would
577    /// otherwise produce byte-identical text lists, and a test asserting
578    /// `["t0", "t1", "t2"]` per document could not tell a document's own run
579    /// from a neighbour's — it would pass even if every document were handed
580    /// the next one's embeddings wholesale.
581    #[derive(Debug)]
582    struct NTexts {
583        doc: usize,
584        n: usize,
585    }
586
587    impl NTexts {
588        fn new(doc: usize, n: usize) -> Self {
589            Self { doc, n }
590        }
591
592        /// The texts this document is expected to get back, in order.
593        fn expected(&self) -> Vec<String> {
594            (0..self.n).map(|i| format!("d{}t{i}", self.doc)).collect()
595        }
596    }
597
598    impl Embed for NTexts {
599        fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
600            for i in 0..self.n {
601                embedder.embed(format!("d{}t{i}", self.doc));
602            }
603            Ok(())
604        }
605    }
606
607    /// The texts a document actually got back, in order.
608    fn returned(embeddings: &[Embedding]) -> Vec<String> {
609        embeddings
610            .iter()
611            .map(|embedding| embedding.document.clone())
612            .collect()
613    }
614
615    /// rig#2345 — a document whose texts straddle a `MAX_DOCUMENTS` boundary
616    /// must get its embeddings back in text order.
617    ///
618    /// Six texts against a limit of 5 splits into `[d0t0..d0t4]` and `[d0t5]`;
619    /// the delayed first batch makes the trailing one finish first. Before the
620    /// slot index this returned the sixth text's embedding at index 0.
621    #[tokio::test]
622    async fn test_build_preserves_text_order_within_a_straddling_document() {
623        let doc = NTexts::new(0, 6);
624        let expected = doc.expected();
625
626        let result = EmbeddingsBuilder::new(SlowFirstBatchModel::new())
627            .document(doc)
628            .unwrap()
629            .build()
630            .await
631            .unwrap();
632
633        assert_eq!(result.len(), 1);
634        assert_eq!(returned(&result[0].1), expected);
635    }
636
637    /// The same guarantee with **more than one** straddle inverted at once:
638    /// every document's texts land in its own list, in order, none borrowed
639    /// from a neighbour. Texts carry their owner, so "borrowed from a
640    /// neighbour" is something this can actually observe.
641    ///
642    /// 4 documents x 3 texts = 12 slots over a limit of 5 gives batches
643    /// `[0,5) [5,10) [10,12)`, so doc1 (slots 3..6) and doc3 (slots 9..12) each
644    /// straddle. `DescendingLatencyModel` inverts both boundaries — with a
645    /// model that only delays the first batch, doc3's two batches still arrive
646    /// in submission order and it comes back correct even unfixed.
647    #[tokio::test]
648    async fn test_build_preserves_text_order_across_many_straddling_documents() {
649        let docs: Vec<NTexts> = (0..4).map(|doc| NTexts::new(doc, 3)).collect();
650        let expected: Vec<Vec<String>> = docs.iter().map(NTexts::expected).collect();
651
652        let result = EmbeddingsBuilder::new(DescendingLatencyModel::new())
653            .documents(docs)
654            .unwrap()
655            .build()
656            .await
657            .unwrap();
658
659        assert_eq!(result.len(), 4);
660        for (index, (_, embeddings)) in result.iter().enumerate() {
661            assert_eq!(
662                returned(embeddings),
663                expected[index],
664                "document {index} did not get its own texts, in order"
665            );
666        }
667    }
668
669    /// A document that embeds no text has no embeddings to return. This has
670    /// always been an error rather than an empty list, and the slot rewrite
671    /// keeps it that way — that behavior is what this pins, and it holds on
672    /// both sides of the fix.
673    ///
674    /// The wording changed: the message now names the document and says what
675    /// caused it, where before it was the unlocated `"missing embedding for
676    /// document after batch merge"`. Only the second assertion below is new
677    /// behavior.
678    #[tokio::test]
679    async fn test_build_rejects_a_document_that_embeds_no_text() {
680        let error = EmbeddingsBuilder::new(MockEmbeddingModel)
681            .document(NTexts::new(0, 0))
682            .unwrap()
683            .build()
684            .await
685            .expect_err("a document with no texts has no embeddings");
686
687        assert!(
688            matches!(error, EmbeddingError::ResponseError(_)),
689            "unexpected error variant: {error:?}"
690        );
691        assert!(
692            error.to_string().contains("document 0 produced no text"),
693            "error should name the offending document: {error}"
694        );
695    }
696
697    /// The same, for a document that is not the first — the index in the
698    /// message has to be the document's own, not a constant.
699    #[tokio::test]
700    async fn test_build_names_the_document_that_embeds_no_text() {
701        let error = EmbeddingsBuilder::new(MockEmbeddingModel)
702            .documents(vec![
703                NTexts::new(0, 2),
704                NTexts::new(1, 2),
705                NTexts::new(2, 0),
706            ])
707            .unwrap()
708            .build()
709            .await
710            .expect_err("a document with no texts has no embeddings");
711
712        assert!(
713            error.to_string().contains("document 2 produced no text"),
714            "error should name document 2: {error}"
715        );
716    }
717
718    /// A model that batches one text at a time, so *every* multi-text document
719    /// straddles, and answers later texts faster than earlier ones.
720    #[derive(Clone, Default)]
721    struct OneAtATimeReversedLatency;
722
723    impl EmbeddingModel for OneAtATimeReversedLatency {
724        const MAX_DOCUMENTS: usize = 1;
725
726        type Client = crate::client::Nothing;
727
728        fn make(_: &Self::Client, _: impl Into<String>, _: Option<usize>) -> Self {
729            Self
730        }
731
732        fn ndims(&self) -> usize {
733            10
734        }
735
736        async fn embed_texts(
737            &self,
738            documents: impl IntoIterator<Item = String> + crate::wasm_compat::WasmCompatSend,
739        ) -> Result<Vec<Embedding>, EmbeddingError> {
740            let documents: Vec<String> = documents.into_iter().collect();
741            // Earlier texts wait longer, so completion order is close to the
742            // reverse of submission order. Texts are named `d{doc}t{i}`, so the
743            // position is what follows the last `t`; if that ever stops
744            // parsing every batch waits 0ms, the completion order stops being
745            // inverted, and this test quietly stops proving anything — hence
746            // the assert rather than `unwrap_or(0)`.
747            let position = documents
748                .first()
749                .and_then(|text| text.rsplit_once('t'))
750                .and_then(|(_, n)| n.parse::<u64>().ok());
751            assert!(
752                position.is_some(),
753                "could not read a text position out of {documents:?}; \
754                 this mock cannot invert completion order without it"
755            );
756            let delay = position.map_or(0, |n| 60u64.saturating_sub(n * 10));
757            tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
758            Ok(documents
759                .into_iter()
760                .map(|document| Embedding {
761                    document,
762                    vec: vec![0.0; 10],
763                })
764                .collect())
765        }
766    }
767
768    /// Worst case for the span arithmetic: `MAX_DOCUMENTS = 1` means every text
769    /// is its own batch, all of them run concurrently, and they finish in
770    /// roughly reverse order. Nothing about the result may depend on that.
771    #[tokio::test]
772    async fn test_build_order_survives_one_text_per_batch_finishing_backwards() {
773        let doc = NTexts::new(0, 6);
774        let expected = doc.expected();
775
776        let result = EmbeddingsBuilder::new(OneAtATimeReversedLatency)
777            .document(doc)
778            .unwrap()
779            .build()
780            .await
781            .unwrap();
782
783        assert_eq!(result.len(), 1);
784        assert_eq!(returned(&result[0].1), expected);
785    }
786
787    /// Documents that tile the batch size exactly, so every document boundary
788    /// is also a batch boundary — the case where an off-by-one in the span
789    /// arithmetic would hand a document its neighbour's run.
790    #[tokio::test]
791    async fn test_build_order_when_documents_tile_the_batch_size_exactly() {
792        // 3 documents x 5 texts, MAX_DOCUMENTS = 5: batches align exactly with
793        // document boundaries.
794        let docs: Vec<NTexts> = (0..3).map(|doc| NTexts::new(doc, 5)).collect();
795        let expected: Vec<Vec<String>> = docs.iter().map(NTexts::expected).collect();
796
797        let result = EmbeddingsBuilder::new(SlowFirstBatchModel::new())
798            .documents(docs)
799            .unwrap()
800            .build()
801            .await
802            .unwrap();
803
804        assert_eq!(result.len(), 3);
805        for (index, (_, embeddings)) in result.iter().enumerate() {
806            assert_eq!(
807                returned(embeddings),
808                expected[index],
809                "document {index} did not get its own run"
810            );
811        }
812    }
813}
rig_core/embeddings/builder.rs

rig_core/embeddings/
builder.rs
