rig/providers/gemini/
client.rs

1use super::{
2    completion::CompletionModel, embedding::EmbeddingModel, transcription::TranscriptionModel,
3};
4use crate::client::{
5    ClientBuilderError, CompletionClient, EmbeddingsClient, ProviderClient, TranscriptionClient,
6    impl_conversion_traits,
7};
8use crate::{
9    Embed,
10    embeddings::{self},
11};
12use serde::Deserialize;
13
14// ================================================================
15// Google Gemini Client
16// ================================================================
17const GEMINI_API_BASE_URL: &str = "https://generativelanguage.googleapis.com";
18
19pub struct ClientBuilder<'a> {
20    api_key: &'a str,
21    base_url: &'a str,
22    http_client: Option<reqwest::Client>,
23}
24
25impl<'a> ClientBuilder<'a> {
26    pub fn new(api_key: &'a str) -> Self {
27        Self {
28            api_key,
29            base_url: GEMINI_API_BASE_URL,
30            http_client: None,
31        }
32    }
33
34    pub fn base_url(mut self, base_url: &'a str) -> Self {
35        self.base_url = base_url;
36        self
37    }
38
39    pub fn custom_client(mut self, client: reqwest::Client) -> Self {
40        self.http_client = Some(client);
41        self
42    }
43
44    pub fn build(self) -> Result<Client, ClientBuilderError> {
45        let mut default_headers = reqwest::header::HeaderMap::new();
46        default_headers.insert(
47            reqwest::header::CONTENT_TYPE,
48            "application/json".parse().unwrap(),
49        );
50        let http_client = if let Some(http_client) = self.http_client {
51            http_client
52        } else {
53            reqwest::Client::builder().build()?
54        };
55
56        Ok(Client {
57            base_url: self.base_url.to_string(),
58            api_key: self.api_key.to_string(),
59            default_headers,
60            http_client,
61        })
62    }
63}
64#[derive(Clone)]
65pub struct Client {
66    base_url: String,
67    api_key: String,
68    default_headers: reqwest::header::HeaderMap,
69    http_client: reqwest::Client,
70}
71
72impl std::fmt::Debug for Client {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("Client")
75            .field("base_url", &self.base_url)
76            .field("http_client", &self.http_client)
77            .field("default_headers", &self.default_headers)
78            .field("api_key", &"<REDACTED>")
79            .finish()
80    }
81}
82
83impl Client {
84    /// Create a new Google Gemini client builder.
85    ///
86    /// # Example
87    /// ```
88    /// use rig::providers::gemini::{ClientBuilder, self};
89    ///
90    /// // Initialize the Google Gemini client
91    /// let gemini_client = Client::builder("your-google-gemini-api-key")
92    ///    .build()
93    /// ```
94    pub fn builder(api_key: &str) -> ClientBuilder<'_> {
95        ClientBuilder::new(api_key)
96    }
97
98    /// Create a new Google Gemini client. For more control, use the `builder` method.
99    ///
100    /// # Panics
101    /// - If the reqwest client cannot be built (if the TLS backend cannot be initialized).
102    pub fn new(api_key: &str) -> Self {
103        Self::builder(api_key)
104            .build()
105            .expect("Gemini client should build")
106    }
107
108    pub(crate) fn post(&self, path: &str) -> reqwest::RequestBuilder {
109        // API key gets inserted as query param - no need to add bearer auth or headers
110        let url = format!("{}/{}?key={}", self.base_url, path, self.api_key).replace("//", "/");
111
112        tracing::debug!("POST {}/{}?key={}", self.base_url, path, "****");
113        self.http_client
114            .post(url)
115            .headers(self.default_headers.clone())
116    }
117
118    pub(crate) fn post_sse(&self, path: &str) -> reqwest::RequestBuilder {
119        let url =
120            format!("{}/{}?alt=sse&key={}", self.base_url, path, self.api_key).replace("//", "/");
121
122        tracing::debug!("POST {}/{}?alt=sse&key={}", self.base_url, path, "****");
123        self.http_client
124            .post(url)
125            .headers(self.default_headers.clone())
126    }
127}
128
129impl ProviderClient for Client {
130    /// Create a new Google Gemini client from the `GEMINI_API_KEY` environment variable.
131    /// Panics if the environment variable is not set.
132    fn from_env() -> Self {
133        let api_key = std::env::var("GEMINI_API_KEY").expect("GEMINI_API_KEY not set");
134        Self::new(&api_key)
135    }
136
137    fn from_val(input: crate::client::ProviderValue) -> Self {
138        let crate::client::ProviderValue::Simple(api_key) = input else {
139            panic!("Incorrect provider value type")
140        };
141        Self::new(&api_key)
142    }
143}
144
145impl CompletionClient for Client {
146    type CompletionModel = CompletionModel;
147
148    /// Create a completion model with the given name.
149    /// Gemini-specific parameters can be set using the [GenerationConfig](crate::providers::gemini::completion::gemini_api_types::GenerationConfig) struct.
150    /// [Gemini API Reference](https://ai.google.dev/api/generate-content#generationconfig)
151    fn completion_model(&self, model: &str) -> CompletionModel {
152        CompletionModel::new(self.clone(), model)
153    }
154}
155
156impl EmbeddingsClient for Client {
157    type EmbeddingModel = EmbeddingModel;
158
159    /// Create an embedding model with the given name.
160    /// Note: default embedding dimension of 0 will be used if model is not known.
161    /// If this is the case, it's better to use function `embedding_model_with_ndims`
162    ///
163    /// # Example
164    /// ```
165    /// use rig::providers::gemini::{Client, self};
166    ///
167    /// // Initialize the Google Gemini client
168    /// let gemini = Client::new("your-google-gemini-api-key");
169    ///
170    /// let embedding_model = gemini.embedding_model(gemini::embedding::EMBEDDING_GECKO_001);
171    /// ```
172    fn embedding_model(&self, model: &str) -> EmbeddingModel {
173        EmbeddingModel::new(self.clone(), model, None)
174    }
175
176    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
177    ///
178    /// # Example
179    /// ```
180    /// use rig::providers::gemini::{Client, self};
181    ///
182    /// // Initialize the Google Gemini client
183    /// let gemini = Client::new("your-google-gemini-api-key");
184    ///
185    /// let embedding_model = gemini.embedding_model_with_ndims("model-unknown-to-rig", 1024);
186    /// ```
187    fn embedding_model_with_ndims(&self, model: &str, ndims: usize) -> EmbeddingModel {
188        EmbeddingModel::new(self.clone(), model, Some(ndims))
189    }
190
191    /// Create an embedding builder with the given embedding model.
192    ///
193    /// # Example
194    /// ```
195    /// use rig::providers::gemini::{Client, self};
196    ///
197    /// // Initialize the Google Gemini client
198    /// let gemini = Client::new("your-google-gemini-api-key");
199    ///
200    /// let embeddings = gemini.embeddings(gemini::embedding::EMBEDDING_GECKO_001)
201    ///     .simple_document("doc0", "Hello, world!")
202    ///     .simple_document("doc1", "Goodbye, world!")
203    ///     .build()
204    ///     .await
205    ///     .expect("Failed to embed documents");
206    /// ```
207    fn embeddings<D: Embed>(
208        &self,
209        model: &str,
210    ) -> embeddings::EmbeddingsBuilder<EmbeddingModel, D> {
211        embeddings::EmbeddingsBuilder::new(self.embedding_model(model))
212    }
213}
214
215impl TranscriptionClient for Client {
216    type TranscriptionModel = TranscriptionModel;
217
218    /// Create a transcription model with the given name.
219    /// Gemini-specific parameters can be set using the [GenerationConfig](crate::providers::gemini::completion::gemini_api_types::GenerationConfig) struct.
220    /// [Gemini API Reference](https://ai.google.dev/api/generate-content#generationconfig)
221    fn transcription_model(&self, model: &str) -> TranscriptionModel {
222        TranscriptionModel::new(self.clone(), model)
223    }
224}
225
226impl_conversion_traits!(
227    AsImageGeneration,
228    AsAudioGeneration for Client
229);
230
231#[derive(Debug, Deserialize)]
232pub struct ApiErrorResponse {
233    pub message: String,
234}
235
236#[derive(Debug, Deserialize)]
237#[serde(untagged)]
238pub enum ApiResponse<T> {
239    Ok(T),
240    Err(ApiErrorResponse),
241}
rig/providers/gemini/client.rs

rig/providers/gemini/
client.rs
