ic_rig/embeddings/
embedding.rs

1//! [`EmbeddingModel`] trait, [`Embedding`] struct, and [`EmbeddingError`].
2
3use serde::{Deserialize, Serialize};
4use thiserror::Error;
5
6// ── Embedding ─────────────────────────────────────────────────────────────────
7
8/// A single document and its vector representation.
9///
10/// The `document` field preserves the original text so embeddings can be
11/// matched back to their source after batching or storage.
12#[derive(Clone, Debug, Default, Serialize, Deserialize)]
13pub struct Embedding {
14    /// The text that was embedded.
15    pub document: String,
16    /// The embedding vector returned by the model.
17    pub vec: Vec<f64>,
18}
19
20impl PartialEq for Embedding {
21    fn eq(&self, other: &Self) -> bool {
22        self.document == other.document
23    }
24}
25
26impl Eq for Embedding {}
27
28// ── EmbeddingError ────────────────────────────────────────────────────────────
29
30#[derive(Debug, Error)]
31pub enum EmbeddingError {
32    #[error("HTTP error: {0}")]
33    Http(String),
34
35    #[error("JSON error: {0}")]
36    Json(#[from] serde_json::Error),
37
38    #[error("Provider error ({status}): {message}")]
39    Provider { status: u16, message: String },
40
41    #[error("Response error: {0}")]
42    Response(String),
43}
44
45// ── EmbeddingModel ────────────────────────────────────────────────────────────
46
47/// Trait for models that can generate vector embeddings from text.
48///
49/// Implement this for each provider. The [`EmbeddingsBuilder`](super::EmbeddingsBuilder)
50/// uses `MAX_DOCUMENTS` to chunk requests so you never exceed the provider's
51/// per-request batch limit.
52///
53/// # Example
54///
55/// ```rust,ignore
56/// use irig::embeddings::{EmbeddingModel, Embedding, EmbeddingError};
57///
58/// pub struct MyEmbedder<H> { client: H, model: String }
59///
60/// impl<H: HttpClient> EmbeddingModel for MyEmbedder<H> {
61///     const MAX_DOCUMENTS: usize = 100;
62///     type Error = EmbeddingError;
63///
64///     fn ndims(&self) -> usize { 1536 }
65///
66///     async fn embed_texts(&self, texts: Vec<String>) -> Result<Vec<Embedding>, EmbeddingError> {
67///         // call API, return one Embedding per input text (same order)
68///         todo!()
69///     }
70/// }
71/// ```
72pub trait EmbeddingModel {
73    /// Maximum number of texts the provider accepts in a single request.
74    const MAX_DOCUMENTS: usize;
75
76    type Error: std::error::Error + 'static;
77
78    /// Dimensionality of the output vectors.
79    fn ndims(&self) -> usize;
80
81    /// Embed a batch of texts. Must return exactly one [`Embedding`] per input
82    /// text, in the same order.
83    fn embed_texts(
84        &self,
85        texts: Vec<String>,
86    ) -> impl std::future::Future<Output = Result<Vec<Embedding>, Self::Error>>;
87
88    /// Embed a single text. Provided as a default for convenience.
89    fn embed_text(
90        &self,
91        text: &str,
92    ) -> impl std::future::Future<Output = Result<Embedding, Self::Error>> {
93        async move {
94            self.embed_texts(vec![text.to_owned()])
95                .await?
96                .into_iter()
97                .next()
98                // embed_texts guarantees one result per input, so this can't fail.
99                .ok_or_else(|| unreachable!("embed_texts returned empty vec for one input"))
100        }
101    }
102}
ic_rig/embeddings/embedding.rs

ic_rig/embeddings/
embedding.rs
