gemini_genai_rs/client/
mod.rs

1//! Unified Gemini API client — wraps both Live (WebSocket) and REST API access.
2//!
3//! The [`Client`] struct provides a single entry point for all Gemini APIs.
4//! REST API modules are feature-gated behind their respective features
5//! (e.g., `generate`, `embed`, `models`) so that live-only users pay zero cost.
6
7#[cfg(feature = "http")]
8pub mod http;
9
10use std::sync::Arc;
11
12use crate::protocol::types::{ApiEndpoint, ModelId, SessionConfig};
13use crate::transport::ConnectBuilder;
14use crate::transport::auth::{
15    AuthProvider, GoogleAIAuth, GoogleAITokenAuth, RestAuth, ServiceEndpoint, VertexAIAuth,
16};
17
18/// Unified Gemini API client.
19///
20/// Mirrors the `GoogleGenAI` class from `@google/genai` (js-genai).
21/// Provides access to both Live (WebSocket) and REST APIs through a single
22/// authenticated entry point.
23///
24/// # Construction
25///
26/// ```no_run
27/// use gemini_genai_rs::Client;
28///
29/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
30/// // From API key (Google AI)
31/// let client = Client::from_api_key("your-api-key");
32///
33/// // From Vertex AI credentials
34/// let client = Client::from_vertex("project-id", "us-central1", "access-token");
35///
36/// // Live WebSocket session on the platform's default Live model
37/// let session = client.live(None).connect().await?;
38/// # Ok(())
39/// # }
40/// ```
41pub struct Client {
42    endpoint: ApiEndpoint,
43    model: ModelId,
44    auth: Arc<dyn RestAuth>,
45    /// Scheme, host and port that replace the platform's in REST URLs.
46    base_url: Option<String>,
47    #[cfg(feature = "http")]
48    http: http::HttpClient,
49}
50
51impl Client {
52    /// Create a client with Google AI API key authentication.
53    pub fn from_api_key(api_key: impl Into<String>) -> Self {
54        let key: String = api_key.into();
55        let endpoint = ApiEndpoint::google_ai(key.clone());
56        let auth: Arc<dyn RestAuth> = Arc::new(GoogleAIAuth::new(key));
57        Self {
58            endpoint,
59            model: ModelId::FLASH_LATEST,
60            auth,
61            base_url: None,
62            #[cfg(feature = "http")]
63            http: http::HttpClient::new(http::HttpConfig::default()),
64        }
65    }
66
67    /// Create a client with Google AI OAuth2 token authentication.
68    pub fn from_access_token(access_token: impl Into<String>) -> Self {
69        let token: String = access_token.into();
70        let endpoint = ApiEndpoint::google_ai_token(token.clone());
71        let auth: Arc<dyn RestAuth> = Arc::new(GoogleAITokenAuth::new(token));
72        Self {
73            endpoint,
74            model: ModelId::FLASH_LATEST,
75            auth,
76            base_url: None,
77            #[cfg(feature = "http")]
78            http: http::HttpClient::new(http::HttpConfig::default()),
79        }
80    }
81
82    /// Create a client with Vertex AI authentication.
83    pub fn from_vertex(
84        project: impl Into<String>,
85        location: impl Into<String>,
86        access_token: impl Into<String>,
87    ) -> Self {
88        let proj: String = project.into();
89        let loc: String = location.into();
90        let tok: String = access_token.into();
91        let endpoint = ApiEndpoint::vertex(proj.clone(), loc.clone(), tok.clone());
92        let auth: Arc<dyn RestAuth> = Arc::new(VertexAIAuth::new(proj, loc, tok));
93        Self {
94            endpoint,
95            model: ModelId::FLASH_LATEST,
96            auth,
97            base_url: None,
98            #[cfg(feature = "http")]
99            http: http::HttpClient::new(http::HttpConfig::default()),
100        }
101    }
102
103    /// Create a client with Vertex AI authentication and dynamic token refresh.
104    ///
105    /// The `refresher` closure is called on every REST API request and on
106    /// every Live connection attempt (including reconnects) to obtain a
107    /// fresh Bearer token. It should cache internally (see
108    /// `GcloudTokenProvider` in gemini-adk-rs for an example).
109    ///
110    /// This is the right constructor for anything that outlives a token's
111    /// ~1 h lifetime.
112    pub fn from_vertex_refreshable(
113        project: impl Into<String>,
114        location: impl Into<String>,
115        refresher: impl Fn() -> String + Send + Sync + 'static,
116    ) -> Self {
117        let proj: String = project.into();
118        let loc: String = location.into();
119        // One source feeds both sides: REST requests and every Live
120        // (re)connection attempt see a fresh token.
121        let refresher = Arc::new(refresher);
122        let live_refresher = refresher.clone();
123        let endpoint =
124            ApiEndpoint::vertex_refreshing(proj.clone(), loc.clone(), move || live_refresher());
125        let auth: Arc<dyn RestAuth> =
126            Arc::new(VertexAIAuth::with_token_refresher(proj, loc, move || {
127                refresher()
128            }));
129        Self {
130            endpoint,
131            model: ModelId::FLASH_LATEST,
132            auth,
133            base_url: None,
134            #[cfg(feature = "http")]
135            http: http::HttpClient::new(http::HttpConfig::default()),
136        }
137    }
138
139    /// Set the default model for all API calls.
140    pub fn model(mut self, model: impl Into<ModelId>) -> Self {
141        self.model = model.into();
142        self
143    }
144
145    /// Send REST requests to `base_url` instead of the platform's host.
146    ///
147    /// The path, API version and credentials stay as they are; only the
148    /// scheme, host and port change. Use it for a gateway, a proxy or a local
149    /// stand-in, as google-genai's `http_options.base_url` does. Live
150    /// sessions are not affected.
151    ///
152    /// ```
153    /// use gemini_genai_rs::{Client, prelude::ServiceEndpoint};
154    ///
155    /// let client = Client::from_api_key("key").base_url("http://127.0.0.1:8080");
156    /// let url = client.rest_url(ServiceEndpoint::GenerateContent);
157    /// assert!(url.starts_with("http://127.0.0.1:8080/v1beta/models/"));
158    /// ```
159    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
160        self.base_url = Some(base_url.into().trim_end_matches('/').to_string());
161        self
162    }
163
164    /// Configure the HTTP client (timeouts, retries, etc.).
165    #[cfg(feature = "http")]
166    pub fn http_config(mut self, config: http::HttpConfig) -> Self {
167        self.http = http::HttpClient::new(config);
168        self
169    }
170
171    /// Get a reference to the underlying auth provider.
172    pub fn auth(&self) -> &dyn AuthProvider {
173        &*self.auth
174    }
175
176    /// Get the default model.
177    pub fn default_model(&self) -> &ModelId {
178        &self.model
179    }
180
181    /// Build the REST URL for a given service endpoint, using the default model.
182    pub fn rest_url(&self, endpoint: ServiceEndpoint) -> String {
183        self.rebase(self.auth.rest_url(endpoint, Some(&self.model)))
184    }
185
186    /// Build the REST URL for a given service endpoint with a specific model.
187    pub fn rest_url_for(&self, endpoint: ServiceEndpoint, model: &ModelId) -> String {
188        self.rebase(self.auth.rest_url(endpoint, Some(model)))
189    }
190
191    /// `url` on [`Client::base_url`]'s host, when one is set.
192    fn rebase(&self, url: String) -> String {
193        let Some(base) = &self.base_url else {
194            return url;
195        };
196        // Everything after `scheme://host[:port]` is kept.
197        let path = url
198            .split_once("://")
199            .and_then(|(_, rest)| rest.find('/').map(|i| &rest[i..]))
200            .unwrap_or("");
201        format!("{base}{path}")
202    }
203
204    /// Get auth headers for REST API calls.
205    pub async fn auth_headers(&self) -> Result<Vec<(String, String)>, crate::session::AuthError> {
206        self.auth.auth_headers().await
207    }
208
209    /// A Live session on this client's credentials.
210    ///
211    /// `None` connects to the platform's default Live model (the REST default
212    /// model is a text model and would not do). Tune the session with
213    /// [`ConnectBuilder::configure`], then `.connect().await`.
214    pub fn live(&self, model: Option<ModelId>) -> ConnectBuilder {
215        let mut config = SessionConfig::from_endpoint(self.endpoint.clone());
216        config.model = model;
217        ConnectBuilder::new(config)
218    }
219
220    /// Get a reference to the HTTP client for making REST API calls.
221    #[cfg(feature = "http")]
222    pub fn http_client(&self) -> &http::HttpClient {
223        &self.http
224    }
225
226    /// Make a raw REST API request (low-level).
227    ///
228    /// Higher-level module methods (e.g., `generate_content()`) should be preferred.
229    #[cfg(feature = "http")]
230    pub async fn rest_request(
231        &self,
232        endpoint: ServiceEndpoint,
233        body: &impl serde::Serialize,
234    ) -> Result<serde_json::Value, http::HttpError> {
235        let url = self.rest_url(endpoint);
236        let headers = self.auth.auth_headers().await?;
237        self.http.post_json(&url, headers, body).await
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn client_from_api_key() {
247        let client = Client::from_api_key("test-key");
248        // The REST client's default is a text model: `generateContent` on a
249        // Live-only native-audio model 404s.
250        assert_eq!(client.default_model(), &ModelId::FLASH_LATEST);
251    }
252
253    #[test]
254    fn client_from_vertex() {
255        let client = Client::from_vertex("proj", "us-central1", "tok");
256        let url = client.auth().ws_url(&ModelId::FLASH_LATEST);
257        assert!(url.contains("us-central1-aiplatform.googleapis.com"));
258    }
259
260    #[test]
261    fn client_model_override() {
262        let client = Client::from_api_key("key").model("models/gemini-2.0-flash-live-001");
263        assert_eq!(
264            client.default_model(),
265            &ModelId::from_static("models/gemini-2.0-flash-live-001")
266        );
267    }
268
269    #[test]
270    fn client_rest_url_generate() {
271        let client = Client::from_api_key("my-key")
272            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
273        let url = client.rest_url(ServiceEndpoint::GenerateContent);
274        assert!(url.contains(":generateContent"));
275        assert!(!url.contains("my-key"), "the key rides in a header: {url}");
276    }
277
278    #[test]
279    fn client_rest_url_vertex() {
280        let client = Client::from_vertex("proj", "us-east1", "tok")
281            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
282        let url = client.rest_url(ServiceEndpoint::GenerateContent);
283        assert!(url.contains("us-east1-aiplatform.googleapis.com"));
284        assert!(url.contains(":generateContent"));
285    }
286
287    #[test]
288    fn base_url_replaces_only_the_host() {
289        let model = ModelId::from_static("gemini-3.5-flash-lite");
290        let google = Client::from_api_key("my-key").base_url("http://127.0.0.1:8766/");
291        assert_eq!(
292            google.rest_url_for(ServiceEndpoint::GenerateContent, &model),
293            "http://127.0.0.1:8766/v1beta/models/gemini-3.5-flash-lite:generateContent"
294        );
295        let vertex = Client::from_vertex("proj", "us-east1", "tok").base_url("https://gw.example");
296        let url = vertex.rest_url_for(ServiceEndpoint::GenerateContent, &model);
297        assert!(
298            url.starts_with("https://gw.example/v1beta1/projects/proj/locations/us-east1/"),
299            "{url}"
300        );
301        // Unset, the platform's host stays.
302        let plain =
303            Client::from_api_key("k").rest_url_for(ServiceEndpoint::GenerateContent, &model);
304        assert!(plain.starts_with("https://generativelanguage.googleapis.com/v1beta/"));
305    }
306
307    #[test]
308    fn live_session_builder_created() {
309        let client = Client::from_api_key("key");
310        let _builder = client.live(Some(ModelId::from_static(
311            "models/gemini-2.0-flash-live-001",
312        )));
313    }
314
315    #[tokio::test]
316    async fn client_from_vertex_refreshable() {
317        use std::sync::atomic::{AtomicU32, Ordering};
318        let call_count = Arc::new(AtomicU32::new(0));
319        let cc = call_count.clone();
320        let client = Client::from_vertex_refreshable("proj", "us-central1", move || {
321            cc.fetch_add(1, Ordering::SeqCst);
322            "refreshed-token".to_string()
323        });
324        // Nothing is fetched eagerly: a token minted at construction would be
325        // the stale one by the time a reconnect needs it.
326        assert_eq!(call_count.load(Ordering::SeqCst), 0);
327        // Every REST request consults the source …
328        let headers = client.auth_headers().await.unwrap();
329        assert_eq!(headers[0].1, "Bearer refreshed-token");
330        assert_eq!(call_count.load(Ordering::SeqCst), 1);
331        // … and so does every Live connection attempt, through the same source.
332        let live_config = SessionConfig::from_endpoint(client.endpoint.clone());
333        assert_eq!(
334            live_config.bearer_token().as_deref(),
335            Some("refreshed-token")
336        );
337        assert_eq!(call_count.load(Ordering::SeqCst), 2);
338    }
339}
gemini_genai_rs/client/mod.rs

gemini_genai_rs/client/
mod.rs
