rig/client/
embeddings.rs

1use crate::Embed;
2use crate::client::{AsEmbeddings, ProviderClient};
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: ProviderClient + Clone {
9    /// The type of EmbeddingModel used by the Client
10    type EmbeddingModel: EmbeddingModel;
11
12    /// Create an embedding model with the given name.
13    /// Note: default embedding dimension of 0 will be used if model is not known.
14    /// If this is the case, it's better to use function `embedding_model_with_ndims`
15    ///
16    /// # Example
17    /// ```
18    /// use rig::prelude::*;
19    /// use rig::providers::openai::{Client, self};
20    ///
21    /// // Initialize the OpenAI client
22    /// let openai = Client::new("your-open-ai-api-key");
23    ///
24    /// let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_LARGE);
25    /// ```
26    fn embedding_model(&self, model: &str) -> Self::EmbeddingModel;
27
28    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
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(&self, model: &str, ndims: usize) -> Self::EmbeddingModel;
41
42    /// Create an embedding builder with the given embedding model.
43    ///
44    /// # Example with OpenAI
45    /// ```
46    /// use rig::prelude::*;
47    /// use rig::providers::openai::{Client, self};
48    ///
49    /// // Initialize the OpenAI client
50    /// let openai = Client::new("your-open-ai-api-key");
51    ///
52    /// let embeddings = openai.embeddings(openai::TEXT_EMBEDDING_3_LARGE)
53    ///     .simple_document("doc0", "Hello, world!")
54    ///     .simple_document("doc1", "Goodbye, world!")
55    ///     .build()
56    ///     .await
57    ///     .expect("Failed to embed documents");
58    /// ```
59    fn embeddings<D: Embed>(&self, model: &str) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
60        EmbeddingsBuilder::new(self.embedding_model(model))
61    }
62
63    /// Create an embedding builder with the given name and the number of dimensions in the embedding generated by the model.
64    ///
65    /// # Example with OpenAI
66    /// ```
67    /// use rig::prelude::*;
68    /// use rig::providers::openai::{Client, self};
69    ///
70    /// // Initialize the OpenAI client
71    /// let openai = Client::new("your-open-ai-api-key");
72    ///
73    /// let embeddings = openai.embeddings_with_ndims(openai::TEXT_EMBEDDING_3_LARGE, 3072)
74    ///     .simple_document("doc0", "Hello, world!")
75    ///     .simple_document("doc1", "Goodbye, world!")
76    ///     .build()
77    ///     .await
78    ///     .expect("Failed to embed documents");
79    /// ```
80    fn embeddings_with_ndims<D: Embed>(
81        &self,
82        model: &str,
83        ndims: usize,
84    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
85        EmbeddingsBuilder::new(self.embedding_model_with_ndims(model, ndims))
86    }
87}
88
89pub trait EmbeddingsClientDyn: ProviderClient {
90    /// Create an embedding model with the given name.
91    /// Note: default embedding dimension of 0 will be used if model is not known.
92    /// If this is the case, it's better to use function `embedding_model_with_ndims`
93    fn embedding_model<'a>(&self, model: &str) -> Box<dyn EmbeddingModelDyn + 'a>;
94
95    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
96    fn embedding_model_with_ndims<'a>(
97        &self,
98        model: &str,
99        ndims: usize,
100    ) -> Box<dyn EmbeddingModelDyn + 'a>;
101}
102
103impl<T: EmbeddingsClient<EmbeddingModel = M>, M: EmbeddingModel + 'static> EmbeddingsClientDyn
104    for T
105{
106    fn embedding_model<'a>(&self, model: &str) -> Box<dyn EmbeddingModelDyn + 'a> {
107        Box::new(self.embedding_model(model))
108    }
109
110    fn embedding_model_with_ndims<'a>(
111        &self,
112        model: &str,
113        ndims: usize,
114    ) -> Box<dyn EmbeddingModelDyn + 'a> {
115        Box::new(self.embedding_model_with_ndims(model, ndims))
116    }
117}
118
119impl<T: EmbeddingsClientDyn + Clone + 'static> AsEmbeddings for T {
120    fn as_embeddings(&self) -> Option<Box<dyn EmbeddingsClientDyn>> {
121        Some(Box::new(self.clone()))
122    }
123}
rig/client/embeddings.rs

rig/client/
embeddings.rs
