rig_core/client/
embeddings.rs

1use crate::Embed;
2use crate::embeddings::{EmbeddingModel, EmbeddingsBuilder};
3
4/// A provider client with embedding capabilities.
5/// Clone is required for conversions between client types.
6pub trait EmbeddingsClient {
7    /// The type of EmbeddingModel used by the Client
8    type EmbeddingModel: EmbeddingModel;
9
10    /// Create an embedding model with the given model.
11    ///
12    /// # Example
13    /// ```no_run
14    /// use rig_core::prelude::*;
15    /// use rig_core::providers::openai::{Client, self};
16    ///
17    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
18    /// // Initialize the OpenAI client
19    /// let openai = Client::new("your-open-ai-api-key")?;
20    ///
21    /// let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_LARGE);
22    /// # Ok(())
23    /// # }
24    /// ```
25    fn embedding_model(&self, model: impl Into<String>) -> Self::EmbeddingModel;
26
27    /// Create an embedding model with the given model identifier string and the number of dimensions in the embedding generated by the model.
28    /// Use this when the provider supports a model whose dimensionality is not
29    /// known by Rig.
30    ///
31    /// # Example with OpenAI
32    /// ```no_run
33    /// use rig_core::prelude::*;
34    /// use rig_core::providers::openai::{Client, self};
35    ///
36    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
37    /// // Initialize the OpenAI client
38    /// let openai = Client::new("your-open-ai-api-key")?;
39    ///
40    /// let embedding_model = openai.embedding_model_with_ndims("model-unknown-to-rig", 3072);
41    /// # Ok(())
42    /// # }
43    /// ```
44    fn embedding_model_with_ndims(
45        &self,
46        model: impl Into<String>,
47        ndims: usize,
48    ) -> Self::EmbeddingModel;
49
50    /// Create an embedding builder with the given embedding model.
51    ///
52    /// # Example with OpenAI
53    /// ```no_run
54    /// use rig_core::prelude::*;
55    /// use rig_core::providers::openai::{Client, self};
56    ///
57    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
58    /// // Initialize the OpenAI client
59    /// let openai = Client::new("your-open-ai-api-key")?;
60    ///
61    /// let embeddings = openai.embeddings(openai::TEXT_EMBEDDING_3_LARGE)
62    ///     .documents(vec!["Hello, world!", "Goodbye, world!"])?
63    ///     .build()
64    ///     .await?;
65    /// # Ok(())
66    /// # }
67    /// ```
68    fn embeddings<D: Embed>(
69        &self,
70        model: impl Into<String>,
71    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
72        EmbeddingsBuilder::new(self.embedding_model(model))
73    }
74
75    /// Create an embedding builder with the given name and the number of dimensions in the embedding generated by the model.
76    ///
77    /// # Example with OpenAI
78    /// ```no_run
79    /// use rig_core::prelude::*;
80    /// use rig_core::providers::openai::{Client, self};
81    ///
82    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
83    /// // Initialize the OpenAI client
84    /// let openai = Client::new("your-open-ai-api-key")?;
85    ///
86    /// let embeddings = openai.embeddings_with_ndims(openai::TEXT_EMBEDDING_3_LARGE, 3072)
87    ///     .documents(vec!["Hello, world!", "Goodbye, world!"])?
88    ///     .build()
89    ///     .await?;
90    /// # Ok(())
91    /// # }
92    /// ```
93    fn embeddings_with_ndims<D: Embed>(
94        &self,
95        model: &str,
96        ndims: usize,
97    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
98        EmbeddingsBuilder::new(self.embedding_model_with_ndims(model, ndims))
99    }
100}
rig_core/client/embeddings.rs

rig_core/client/
embeddings.rs
