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    #[cfg(feature = "http")]
46    http: http::HttpClient,
47}
48
49impl Client {
50    /// Create a client with Google AI API key authentication.
51    pub fn from_api_key(api_key: impl Into<String>) -> Self {
52        let key: String = api_key.into();
53        let endpoint = ApiEndpoint::google_ai(key.clone());
54        let auth: Arc<dyn RestAuth> = Arc::new(GoogleAIAuth::new(key));
55        Self {
56            endpoint,
57            model: ModelId::FLASH_LATEST,
58            auth,
59            #[cfg(feature = "http")]
60            http: http::HttpClient::new(http::HttpConfig::default()),
61        }
62    }
63
64    /// Create a client with Google AI OAuth2 token authentication.
65    pub fn from_access_token(access_token: impl Into<String>) -> Self {
66        let token: String = access_token.into();
67        let endpoint = ApiEndpoint::google_ai_token(token.clone());
68        let auth: Arc<dyn RestAuth> = Arc::new(GoogleAITokenAuth::new(token));
69        Self {
70            endpoint,
71            model: ModelId::FLASH_LATEST,
72            auth,
73            #[cfg(feature = "http")]
74            http: http::HttpClient::new(http::HttpConfig::default()),
75        }
76    }
77
78    /// Create a client with Vertex AI authentication.
79    pub fn from_vertex(
80        project: impl Into<String>,
81        location: impl Into<String>,
82        access_token: impl Into<String>,
83    ) -> Self {
84        let proj: String = project.into();
85        let loc: String = location.into();
86        let tok: String = access_token.into();
87        let endpoint = ApiEndpoint::vertex(proj.clone(), loc.clone(), tok.clone());
88        let auth: Arc<dyn RestAuth> = Arc::new(VertexAIAuth::new(proj, loc, tok));
89        Self {
90            endpoint,
91            model: ModelId::FLASH_LATEST,
92            auth,
93            #[cfg(feature = "http")]
94            http: http::HttpClient::new(http::HttpConfig::default()),
95        }
96    }
97
98    /// Create a client with Vertex AI authentication and dynamic token refresh.
99    ///
100    /// The `refresher` closure is called on every REST API request and on
101    /// every Live connection attempt (including reconnects) to obtain a
102    /// fresh Bearer token. It should cache internally (see
103    /// `GcloudTokenProvider` in gemini-adk-rs for an example).
104    ///
105    /// This is the right constructor for anything that outlives a token's
106    /// ~1 h lifetime.
107    pub fn from_vertex_refreshable(
108        project: impl Into<String>,
109        location: impl Into<String>,
110        refresher: impl Fn() -> String + Send + Sync + 'static,
111    ) -> Self {
112        let proj: String = project.into();
113        let loc: String = location.into();
114        // One source feeds both sides: REST requests and every Live
115        // (re)connection attempt see a fresh token.
116        let refresher = Arc::new(refresher);
117        let live_refresher = refresher.clone();
118        let endpoint =
119            ApiEndpoint::vertex_refreshing(proj.clone(), loc.clone(), move || live_refresher());
120        let auth: Arc<dyn RestAuth> =
121            Arc::new(VertexAIAuth::with_token_refresher(proj, loc, move || {
122                refresher()
123            }));
124        Self {
125            endpoint,
126            model: ModelId::FLASH_LATEST,
127            auth,
128            #[cfg(feature = "http")]
129            http: http::HttpClient::new(http::HttpConfig::default()),
130        }
131    }
132
133    /// Set the default model for all API calls.
134    pub fn model(mut self, model: impl Into<ModelId>) -> Self {
135        self.model = model.into();
136        self
137    }
138
139    /// Configure the HTTP client (timeouts, retries, etc.).
140    #[cfg(feature = "http")]
141    pub fn http_config(mut self, config: http::HttpConfig) -> Self {
142        self.http = http::HttpClient::new(config);
143        self
144    }
145
146    /// Get a reference to the underlying auth provider.
147    pub fn auth(&self) -> &dyn AuthProvider {
148        &*self.auth
149    }
150
151    /// Get the default model.
152    pub fn default_model(&self) -> &ModelId {
153        &self.model
154    }
155
156    /// Build the REST URL for a given service endpoint, using the default model.
157    pub fn rest_url(&self, endpoint: ServiceEndpoint) -> String {
158        self.auth.rest_url(endpoint, Some(&self.model))
159    }
160
161    /// Build the REST URL for a given service endpoint with a specific model.
162    pub fn rest_url_for(&self, endpoint: ServiceEndpoint, model: &ModelId) -> String {
163        self.auth.rest_url(endpoint, Some(model))
164    }
165
166    /// Get auth headers for REST API calls.
167    pub async fn auth_headers(&self) -> Result<Vec<(String, String)>, crate::session::AuthError> {
168        self.auth.auth_headers().await
169    }
170
171    /// A Live session on this client's credentials.
172    ///
173    /// `None` connects to the platform's default Live model (the REST default
174    /// model is a text model and would not do). Tune the session with
175    /// [`ConnectBuilder::configure`], then `.connect().await`.
176    pub fn live(&self, model: Option<ModelId>) -> ConnectBuilder {
177        let mut config = SessionConfig::from_endpoint(self.endpoint.clone());
178        config.model = model;
179        ConnectBuilder::new(config)
180    }
181
182    /// Get a reference to the HTTP client for making REST API calls.
183    #[cfg(feature = "http")]
184    pub fn http_client(&self) -> &http::HttpClient {
185        &self.http
186    }
187
188    /// Make a raw REST API request (low-level).
189    ///
190    /// Higher-level module methods (e.g., `generate_content()`) should be preferred.
191    #[cfg(feature = "http")]
192    pub async fn rest_request(
193        &self,
194        endpoint: ServiceEndpoint,
195        body: &impl serde::Serialize,
196    ) -> Result<serde_json::Value, http::HttpError> {
197        let url = self.rest_url(endpoint);
198        let headers = self.auth.auth_headers().await?;
199        self.http.post_json(&url, headers, body).await
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn client_from_api_key() {
209        let client = Client::from_api_key("test-key");
210        // The REST client's default is a text model: `generateContent` on a
211        // Live-only native-audio model 404s.
212        assert_eq!(client.default_model(), &ModelId::FLASH_LATEST);
213    }
214
215    #[test]
216    fn client_from_vertex() {
217        let client = Client::from_vertex("proj", "us-central1", "tok");
218        let url = client.auth().ws_url(&ModelId::FLASH_LATEST);
219        assert!(url.contains("us-central1-aiplatform.googleapis.com"));
220    }
221
222    #[test]
223    fn client_model_override() {
224        let client = Client::from_api_key("key").model("models/gemini-2.0-flash-live-001");
225        assert_eq!(
226            client.default_model(),
227            &ModelId::from_static("models/gemini-2.0-flash-live-001")
228        );
229    }
230
231    #[test]
232    fn client_rest_url_generate() {
233        let client = Client::from_api_key("my-key")
234            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
235        let url = client.rest_url(ServiceEndpoint::GenerateContent);
236        assert!(url.contains(":generateContent"));
237        assert!(!url.contains("my-key"), "the key rides in a header: {url}");
238    }
239
240    #[test]
241    fn client_rest_url_vertex() {
242        let client = Client::from_vertex("proj", "us-east1", "tok")
243            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
244        let url = client.rest_url(ServiceEndpoint::GenerateContent);
245        assert!(url.contains("us-east1-aiplatform.googleapis.com"));
246        assert!(url.contains(":generateContent"));
247    }
248
249    #[test]
250    fn live_session_builder_created() {
251        let client = Client::from_api_key("key");
252        let _builder = client.live(Some(ModelId::from_static(
253            "models/gemini-2.0-flash-live-001",
254        )));
255    }
256
257    #[tokio::test]
258    async fn client_from_vertex_refreshable() {
259        use std::sync::atomic::{AtomicU32, Ordering};
260        let call_count = Arc::new(AtomicU32::new(0));
261        let cc = call_count.clone();
262        let client = Client::from_vertex_refreshable("proj", "us-central1", move || {
263            cc.fetch_add(1, Ordering::SeqCst);
264            "refreshed-token".to_string()
265        });
266        // Nothing is fetched eagerly: a token minted at construction would be
267        // the stale one by the time a reconnect needs it.
268        assert_eq!(call_count.load(Ordering::SeqCst), 0);
269        // Every REST request consults the source …
270        let headers = client.auth_headers().await.unwrap();
271        assert_eq!(headers[0].1, "Bearer refreshed-token");
272        assert_eq!(call_count.load(Ordering::SeqCst), 1);
273        // … and so does every Live connection attempt, through the same source.
274        let live_config = SessionConfig::from_endpoint(client.endpoint.clone());
275        assert_eq!(
276            live_config.bearer_token().as_deref(),
277            Some("refreshed-token")
278        );
279        assert_eq!(call_count.load(Ordering::SeqCst), 2);
280    }
281}
gemini_genai_rs/client/mod.rs

gemini_genai_rs/client/
mod.rs
