ic-rig 0.2.0

//! [`EmbeddingModel`] trait, [`Embedding`] struct, and [`EmbeddingError`].

use serde::{Deserialize, Serialize};
use thiserror::Error;

// ── Embedding ─────────────────────────────────────────────────────────────────

/// A single document and its vector representation.
///
/// The `document` field preserves the original text so embeddings can be
/// matched back to their source after batching or storage.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Embedding {
    /// The text that was embedded.
    pub document: String,
    /// The embedding vector returned by the model.
    pub vec: Vec<f64>,
}

impl PartialEq for Embedding {
    fn eq(&self, other: &Self) -> bool {
        self.document == other.document
    }
}

impl Eq for Embedding {}

// ── EmbeddingError ────────────────────────────────────────────────────────────

#[derive(Debug, Error)]
pub enum EmbeddingError {
    #[error("HTTP error: {0}")]
    Http(String),

    #[error("JSON error: {0}")]
    Json(#[from] serde_json::Error),

    #[error("Provider error ({status}): {message}")]
    Provider { status: u16, message: String },

    #[error("Response error: {0}")]
    Response(String),
}

// ── EmbeddingModel ────────────────────────────────────────────────────────────

/// Trait for models that can generate vector embeddings from text.
///
/// Implement this for each provider. The [`EmbeddingsBuilder`](super::EmbeddingsBuilder)
/// uses `MAX_DOCUMENTS` to chunk requests so you never exceed the provider's
/// per-request batch limit.
///
/// # Example
///
/// ```rust,ignore
/// use irig::embeddings::{EmbeddingModel, Embedding, EmbeddingError};
///
/// pub struct MyEmbedder<H> { client: H, model: String }
///
/// impl<H: HttpClient> EmbeddingModel for MyEmbedder<H> {
///     const MAX_DOCUMENTS: usize = 100;
///     type Error = EmbeddingError;
///
///     fn ndims(&self) -> usize { 1536 }
///
///     async fn embed_texts(&self, texts: Vec<String>) -> Result<Vec<Embedding>, EmbeddingError> {
///         // call API, return one Embedding per input text (same order)
///         todo!()
///     }
/// }
/// ```
pub trait EmbeddingModel {
    /// Maximum number of texts the provider accepts in a single request.
    const MAX_DOCUMENTS: usize;

    type Error: std::error::Error + 'static;

    /// Dimensionality of the output vectors.
    fn ndims(&self) -> usize;

    /// Embed a batch of texts. Must return exactly one [`Embedding`] per input
    /// text, in the same order.
    fn embed_texts(
        &self,
        texts: Vec<String>,
    ) -> impl std::future::Future<Output = Result<Vec<Embedding>, Self::Error>>;

    /// Embed a single text. Provided as a default for convenience.
    fn embed_text(
        &self,
        text: &str,
    ) -> impl std::future::Future<Output = Result<Embedding, Self::Error>> {
        async move {
            self.embed_texts(vec![text.to_owned()])
                .await?
                .into_iter()
                .next()
                // embed_texts guarantees one result per input, so this can't fail.
                .ok_or_else(|| unreachable!("embed_texts returned empty vec for one input"))
        }
    }
}