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