rig_core/test_utils/
embeddings.rs

1//! Embedding helpers for deterministic tests.
2
3use crate::driver::{Exchange, Local, Model, Opened, Opening, Step, Transport};
4use crate::embeddings::{Embedding, EmbeddingResponse};
5use crate::error::ProviderError;
6use crate::wire::Capabilities;
7
8/// The mock embedding runtime: every text embeds to one fixed vector, five
9/// texts to a request at ten dimensions. It is the transport of a [`Local`]
10/// embedding wire ([`Self::model`]).
11#[derive(Clone, Copy, Debug, Default, PartialEq)]
12pub struct MockEmbeddings;
13
14/// A deterministic embedding model that returns a fixed vector for each input document.
15pub type MockEmbeddingModel = Model<Local<crate::operation::Embedding>, MockEmbeddings>;
16
17impl MockEmbeddings {
18    /// The mock embedding model: its local wire over this runtime.
19    pub fn model() -> MockEmbeddingModel {
20        Model::new(
21            Local::new(super::MOCK_PROVIDER).with_capabilities(Capabilities::embedding(5, 10)),
22            Self,
23        )
24    }
25}
26
27impl Transport<Local<crate::operation::Embedding>> for MockEmbeddings {
28    fn send(
29        &self,
30        texts: Vec<String>,
31        _exchange: Exchange,
32    ) -> Opening<Step<crate::operation::Embedding>> {
33        let response = Self::embed(texts);
34        Opening::ready(Opened::new(futures::stream::iter([
35            Ok::<_, ProviderError>(Step::End(response)),
36        ])))
37    }
38}
39
40impl MockEmbeddings {
41    /// The reply this runtime gives for `texts`: one fixed ten-dimension
42    /// vector per text, in order.
43    pub fn embed(texts: Vec<String>) -> EmbeddingResponse {
44        EmbeddingResponse {
45            provider: super::MOCK_PROVIDER.to_owned(),
46            ..EmbeddingResponse::new(
47                texts
48                    .into_iter()
49                    .map(|document| Embedding {
50                        document,
51                        vec: vec![0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
52                    })
53                    .collect(),
54            )
55        }
56    }
57}
rig_core/test_utils/embeddings.rs

rig_core/test_utils/
embeddings.rs
