rig/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    /// ```
14    /// use rig::prelude::*;
15    /// use rig::providers::openai::{Client, self};
16    ///
17    /// // Initialize the OpenAI client
18    /// let openai = Client::new("your-open-ai-api-key");
19    ///
20    /// let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_LARGE);
21    /// ```
22    fn embedding_model(&self, model: impl Into<String>) -> Self::EmbeddingModel;
23
24    /// Create an embedding model with the given model identifier string and the number of dimensions in the embedding generated by the model.
25    /// This is the suggested method if you need to use a model not included in
26    /// `Self::EmbeddingModel::Models`
27    ///
28    /// # Example with OpenAI
29    /// ```
30    /// use rig::prelude::*;
31    /// use rig::providers::openai::{Client, self};
32    ///
33    /// // Initialize the OpenAI client
34    /// let openai = Client::new("your-open-ai-api-key");
35    ///
36    /// let embedding_model = openai.embedding_model("model-unknown-to-rig", 3072);
37    /// ```
38    fn embedding_model_with_ndims(
39        &self,
40        model: impl Into<String>,
41        ndims: usize,
42    ) -> Self::EmbeddingModel;
43
44    /// Create an embedding builder with the given embedding model.
45    ///
46    /// # Example with OpenAI
47    /// ```
48    /// use rig::prelude::*;
49    /// use rig::providers::openai::{Client, self};
50    ///
51    /// // Initialize the OpenAI client
52    /// let openai = Client::new("your-open-ai-api-key");
53    ///
54    /// let embeddings = openai.embeddings(openai::TEXT_EMBEDDING_3_LARGE)
55    ///     .simple_document("doc0", "Hello, world!")
56    ///     .simple_document("doc1", "Goodbye, world!")
57    ///     .build()
58    ///     .await
59    ///     .expect("Failed to embed documents");
60    /// ```
61    fn embeddings<D: Embed>(
62        &self,
63        model: impl Into<String>,
64    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
65        EmbeddingsBuilder::new(self.embedding_model(model))
66    }
67
68    /// Create an embedding builder with the given name and the number of dimensions in the embedding generated by the model.
69    ///
70    /// # Example with OpenAI
71    /// ```
72    /// use rig::prelude::*;
73    /// use rig::providers::openai::{Client, self};
74    ///
75    /// // Initialize the OpenAI client
76    /// let openai = Client::new("your-open-ai-api-key");
77    ///
78    /// let embeddings = openai.embeddings_with_ndims(openai::TEXT_EMBEDDING_3_LARGE, 3072)
79    ///     .simple_document("doc0", "Hello, world!")
80    ///     .simple_document("doc1", "Goodbye, world!")
81    ///     .build()
82    ///     .await
83    ///     .expect("Failed to embed documents");
84    /// ```
85    fn embeddings_with_ndims<D: Embed>(
86        &self,
87        model: &str,
88        ndims: usize,
89    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
90        EmbeddingsBuilder::new(self.embedding_model_with_ndims(model, ndims))
91    }
92}
rig/client/embeddings.rs

rig/client/
embeddings.rs
