rig/client/
embeddings.rs

1use crate::Embed;
2#[allow(deprecated)]
3use crate::embeddings::embedding::EmbeddingModelDyn;
4use crate::embeddings::{EmbeddingModel, EmbeddingsBuilder};
5
6/// A provider client with embedding capabilities.
7/// Clone is required for conversions between client types.
8pub trait EmbeddingsClient {
9    /// The type of EmbeddingModel used by the Client
10    type EmbeddingModel: EmbeddingModel;
11
12    /// Create an embedding model with the given model.
13    ///
14    /// # Example
15    /// ```
16    /// use rig::prelude::*;
17    /// use rig::providers::openai::{Client, self};
18    ///
19    /// // Initialize the OpenAI client
20    /// let openai = Client::new("your-open-ai-api-key");
21    ///
22    /// let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_LARGE);
23    /// ```
24    fn embedding_model(&self, model: impl Into<String>) -> Self::EmbeddingModel;
25
26    /// Create an embedding model with the given model identifier string and the number of dimensions in the embedding generated by the model.
27    /// This is the suggested method if you need to use a model not included in
28    /// `Self::EmbeddingModel::Models`
29    ///
30    /// # Example with OpenAI
31    /// ```
32    /// use rig::prelude::*;
33    /// use rig::providers::openai::{Client, self};
34    ///
35    /// // Initialize the OpenAI client
36    /// let openai = Client::new("your-open-ai-api-key");
37    ///
38    /// let embedding_model = openai.embedding_model("model-unknown-to-rig", 3072);
39    /// ```
40    fn embedding_model_with_ndims(
41        &self,
42        model: impl Into<String>,
43        ndims: usize,
44    ) -> Self::EmbeddingModel;
45
46    /// Create an embedding builder with the given embedding model.
47    ///
48    /// # Example with OpenAI
49    /// ```
50    /// use rig::prelude::*;
51    /// use rig::providers::openai::{Client, self};
52    ///
53    /// // Initialize the OpenAI client
54    /// let openai = Client::new("your-open-ai-api-key");
55    ///
56    /// let embeddings = openai.embeddings(openai::TEXT_EMBEDDING_3_LARGE)
57    ///     .simple_document("doc0", "Hello, world!")
58    ///     .simple_document("doc1", "Goodbye, world!")
59    ///     .build()
60    ///     .await
61    ///     .expect("Failed to embed documents");
62    /// ```
63    fn embeddings<D: Embed>(
64        &self,
65        model: impl Into<String>,
66    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
67        EmbeddingsBuilder::new(self.embedding_model(model))
68    }
69
70    /// Create an embedding builder with the given name and the number of dimensions in the embedding generated by the model.
71    ///
72    /// # Example with OpenAI
73    /// ```
74    /// use rig::prelude::*;
75    /// use rig::providers::openai::{Client, self};
76    ///
77    /// // Initialize the OpenAI client
78    /// let openai = Client::new("your-open-ai-api-key");
79    ///
80    /// let embeddings = openai.embeddings_with_ndims(openai::TEXT_EMBEDDING_3_LARGE, 3072)
81    ///     .simple_document("doc0", "Hello, world!")
82    ///     .simple_document("doc1", "Goodbye, world!")
83    ///     .build()
84    ///     .await
85    ///     .expect("Failed to embed documents");
86    /// ```
87    fn embeddings_with_ndims<D: Embed>(
88        &self,
89        model: &str,
90        ndims: usize,
91    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
92        EmbeddingsBuilder::new(self.embedding_model_with_ndims(model, ndims))
93    }
94}
95
96#[allow(deprecated)]
97#[deprecated(
98    since = "0.25.0",
99    note = "`DynClientBuilder` and related features have been deprecated and will be removed in a future release. In this case, use `EmbeddingsClient` instead."
100)]
101pub trait EmbeddingsClientDyn {
102    /// Create an embedding model with the given name.
103    /// Note: default embedding dimension of 0 will be used if model is not known.
104    /// If this is the case, it's better to use function `embedding_model_with_ndims`
105    fn embedding_model<'a>(&self, model: &str) -> Box<dyn EmbeddingModelDyn + 'a>;
106
107    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
108    fn embedding_model_with_ndims<'a>(
109        &self,
110        model: &str,
111        ndims: usize,
112    ) -> Box<dyn EmbeddingModelDyn + 'a>;
113}
114
115#[allow(deprecated)]
116impl<M, T> EmbeddingsClientDyn for T
117where
118    T: EmbeddingsClient<EmbeddingModel = M>,
119    M: EmbeddingModel + 'static,
120{
121    fn embedding_model<'a>(&self, model: &str) -> Box<dyn EmbeddingModelDyn + 'a> {
122        Box::new(self.embedding_model(model))
123    }
124
125    fn embedding_model_with_ndims<'a>(
126        &self,
127        model: &str,
128        ndims: usize,
129    ) -> Box<dyn EmbeddingModelDyn + 'a> {
130        Box::new(self.embedding_model_with_ndims(model, ndims))
131    }
132}
rig/client/embeddings.rs

rig/client/
embeddings.rs
