rig-core 0.37.0

An opinionated library for building LLM powered applications.
Documentation
use crate::Embed;
use crate::embeddings::{EmbeddingModel, EmbeddingsBuilder};

/// A provider client with embedding capabilities.
/// Clone is required for conversions between client types.
pub trait EmbeddingsClient {
    /// The type of EmbeddingModel used by the Client
    type EmbeddingModel: EmbeddingModel;

    /// Create an embedding model with the given model.
    ///
    /// # Example
    /// ```no_run
    /// use rig_core::prelude::*;
    /// use rig_core::providers::openai::{Client, self};
    ///
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key")?;
    ///
    /// let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_LARGE);
    /// # Ok(())
    /// # }
    /// ```
    fn embedding_model(&self, model: impl Into<String>) -> Self::EmbeddingModel;

    /// Create an embedding model with the given model identifier string and the number of dimensions in the embedding generated by the model.
    /// Use this when the provider supports a model whose dimensionality is not
    /// known by Rig.
    ///
    /// # Example with OpenAI
    /// ```no_run
    /// use rig_core::prelude::*;
    /// use rig_core::providers::openai::{Client, self};
    ///
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key")?;
    ///
    /// let embedding_model = openai.embedding_model_with_ndims("model-unknown-to-rig", 3072);
    /// # Ok(())
    /// # }
    /// ```
    fn embedding_model_with_ndims(
        &self,
        model: impl Into<String>,
        ndims: usize,
    ) -> Self::EmbeddingModel;

    /// Create an embedding builder with the given embedding model.
    ///
    /// # Example with OpenAI
    /// ```no_run
    /// use rig_core::prelude::*;
    /// use rig_core::providers::openai::{Client, self};
    ///
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key")?;
    ///
    /// let embeddings = openai.embeddings(openai::TEXT_EMBEDDING_3_LARGE)
    ///     .documents(vec!["Hello, world!", "Goodbye, world!"])?
    ///     .build()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    fn embeddings<D: Embed>(
        &self,
        model: impl Into<String>,
    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
        EmbeddingsBuilder::new(self.embedding_model(model))
    }

    /// Create an embedding builder with the given name and the number of dimensions in the embedding generated by the model.
    ///
    /// # Example with OpenAI
    /// ```no_run
    /// use rig_core::prelude::*;
    /// use rig_core::providers::openai::{Client, self};
    ///
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key")?;
    ///
    /// let embeddings = openai.embeddings_with_ndims(openai::TEXT_EMBEDDING_3_LARGE, 3072)
    ///     .documents(vec!["Hello, world!", "Goodbye, world!"])?
    ///     .build()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    fn embeddings_with_ndims<D: Embed>(
        &self,
        model: &str,
        ndims: usize,
    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
        EmbeddingsBuilder::new(self.embedding_model_with_ndims(model, ndims))
    }
}