adk_gemini/
client.rs

1use crate::{
2    backend,
3    batch::{BatchBuilder, BatchHandle},
4    cache::{CacheBuilder, CachedContentHandle},
5    embedding::{
6        BatchContentEmbeddingResponse, BatchEmbedContentsRequest, ContentEmbeddingResponse,
7        EmbedBuilder, EmbedContentRequest,
8    },
9    files::{
10        handle::FileHandle,
11        model::{File, ListFilesResponse},
12    },
13    generation::{ContentBuilder, GenerateContentRequest, GenerationConfig, GenerationResponse},
14};
15use eventsource_stream::EventStreamError;
16use futures::Stream;
17#[cfg(feature = "vertex")]
18use google_cloud_aiplatform_v1::client::PredictionService;
19#[cfg(feature = "vertex")]
20use google_cloud_auth::credentials::{self, Credentials};
21use mime::Mime;
22use reqwest::{ClientBuilder, header::InvalidHeaderValue};
23use serde::{Deserialize, Serialize};
24use snafu::{ResultExt, Snafu};
25use std::{
26    fmt::{self, Formatter},
27    sync::{Arc, LazyLock},
28};
29use tracing::{Level, Span, instrument};
30use url::Url;
31
32use crate::batch::model::*;
33use crate::cache::model::*;
34
35static DEFAULT_BASE_URL: LazyLock<Url> = LazyLock::new(|| {
36    Url::parse("https://generativelanguage.googleapis.com/v1beta/")
37        .expect("unreachable error: failed to parse default base URL")
38});
39static V1_BASE_URL: LazyLock<Url> = LazyLock::new(|| {
40    Url::parse("https://generativelanguage.googleapis.com/v1/")
41        .expect("unreachable error: failed to parse v1 base URL")
42});
43
44// ══════════════════════════════════════════════════════════════════════
45// Model enum
46// ══════════════════════════════════════════════════════════════════════
47
48/// Available Gemini model identifiers.
49///
50/// Each variant maps to a specific model version on the Gemini API.
51/// Use [`Model::Custom`] for model IDs not yet represented as variants.
52#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
53pub enum Model {
54    // ── Gemini 3.5 ──────────────────────────────────────────────
55    /// Gemini 3.5 Flash.
56    #[serde(rename = "models/gemini-3.5-flash")]
57    Gemini35Flash,
58
59    // ── Gemini 3.1 ────────────────────────────────────────────────
60    /// Gemini 3.1 Pro preview — strongest reasoning.
61    #[serde(rename = "models/gemini-3.1-pro-preview")]
62    Gemini31ProPreview,
63    /// Gemini 3.1 Flash Lite — GA, most cost-efficient for high-volume agentic tasks.
64    #[serde(rename = "models/gemini-3.1-flash-lite")]
65    Gemini31FlashLite,
66    /// Gemini 3.1 Flash Image (Nano Banana 2) — GA native image generation.
67    #[serde(rename = "models/gemini-3.1-flash-image")]
68    Gemini31FlashImage,
69
70    // ── Gemini 3 ─────────────────────────────────────────────────
71    /// Gemini 3 Pro preview.
72    #[serde(rename = "models/gemini-3-pro-preview")]
73    Gemini3ProPreview,
74    /// Gemini 3 Pro Image (Nano Banana Pro) — GA multimodal output.
75    #[serde(rename = "models/gemini-3-pro-image")]
76    Gemini3ProImage,
77    /// Gemini 3 Pro Image preview (deprecated).
78    #[deprecated(note = "Shut down June 25, 2026. Use Model::Gemini3ProImage instead.")]
79    #[serde(rename = "models/gemini-3-pro-image-preview")]
80    Gemini3ProImagePreview,
81    /// Gemini 3 Flash preview — good balance of speed and capability.
82    #[serde(rename = "models/gemini-3-flash-preview")]
83    Gemini3FlashPreview,
84
85    // ── Gemini 2.5 ───────────────────────────────────────────────
86    /// Gemini 2.5 Pro — advanced reasoning.
87    #[serde(rename = "models/gemini-2.5-pro")]
88    Gemini25Pro,
89    /// Gemini 2.5 Pro preview with TTS support.
90    #[serde(rename = "models/gemini-2.5-pro-preview-tts")]
91    Gemini25ProPreviewTts,
92    /// Gemini 2.5 Flash.
93    #[serde(rename = "models/gemini-2.5-flash")]
94    Gemini25Flash,
95    /// Gemini 2.5 Flash preview (September 2025).
96    #[serde(rename = "models/gemini-2.5-flash-preview-09-2025")]
97    Gemini25FlashPreview092025,
98    /// Gemini 2.5 Flash with image generation.
99    #[serde(rename = "models/gemini-2.5-flash-image")]
100    Gemini25FlashImage,
101    /// Deprecated: use `Gemini25FlashImage` instead.
102    #[deprecated(note = "Use Model::Gemini25FlashImage instead")]
103    #[serde(rename = "models/gemini-2.5-flash-image-preview")]
104    Gemini25FlashImagePreview,
105    /// Gemini 2.5 Flash native audio preview (December 2025).
106    #[serde(rename = "models/gemini-2.5-flash-native-audio-preview-12-2025")]
107    Gemini25FlashLive122025,
108    /// Gemini 2.5 Flash native audio preview (September 2025).
109    #[serde(rename = "models/gemini-2.5-flash-native-audio-preview-09-2025")]
110    Gemini25FlashLive092025,
111    /// Gemini 2.5 Flash preview with TTS support.
112    #[serde(rename = "models/gemini-2.5-flash-preview-tts")]
113    Gemini25FlashPreviewTts,
114    /// Gemini 2.5 Flash Lite — cost-efficient.
115    #[serde(rename = "models/gemini-2.5-flash-lite")]
116    Gemini25FlashLite,
117    /// Gemini 2.5 Flash Lite preview (September 2025).
118    #[serde(rename = "models/gemini-2.5-flash-lite-preview-09-2025")]
119    Gemini25FlashLitePreview092025,
120
121    // ── Embedding models ─────────────────────────────────────────
122    /// Gemini Embedding 2 — GA multimodal embeddings (text, image, video, audio, PDF).
123    #[serde(rename = "models/gemini-embedding-2")]
124    GeminiEmbedding2,
125    /// Gemini Embedding 001 (3072 dimensions). Replaces text-embedding-004.
126    #[serde(rename = "models/gemini-embedding-001")]
127    GeminiEmbedding001,
128
129    // ── Custom ───────────────────────────────────────────────────
130    /// A custom model identifier string (e.g. `"models/my-tuned-model"`).
131    #[serde(untagged)]
132    Custom(String),
133}
134
135impl Default for Model {
136    fn default() -> Self {
137        Self::gemini_3_7_flash()
138    }
139}
140
141impl Model {
142    /// Return the current balanced Gemini default.
143    ///
144    /// This factory avoids adding a new enum variant whenever Google releases a
145    /// model while preserving exhaustive matches for existing callers.
146    pub fn gemini_3_7_flash() -> Self {
147        Self::Custom("models/gemini-3.7-flash".to_string())
148    }
149
150    /// Return Gemini 3.6 Flash.
151    pub fn gemini_3_6_flash() -> Self {
152        Self::Custom("models/gemini-3.6-flash".to_string())
153    }
154
155    /// Return Gemini 3.5 Flash-Lite, the most cost-efficient GA model.
156    pub fn gemini_3_5_flash_lite() -> Self {
157        Self::Custom("models/gemini-3.5-flash-lite".to_string())
158    }
159
160    /// Returns the model identifier as a string slice.
161    pub fn as_str(&self) -> &str {
162        #[allow(deprecated)]
163        match self {
164            Model::Gemini35Flash => "models/gemini-3.5-flash",
165            Model::Gemini31ProPreview => "models/gemini-3.1-pro-preview",
166            Model::Gemini31FlashLite => "models/gemini-3.1-flash-lite",
167            Model::Gemini31FlashImage => "models/gemini-3.1-flash-image",
168            Model::Gemini3ProPreview => "models/gemini-3-pro-preview",
169            Model::Gemini3ProImage => "models/gemini-3-pro-image",
170            Model::Gemini3ProImagePreview => "models/gemini-3-pro-image-preview",
171            Model::Gemini3FlashPreview => "models/gemini-3-flash-preview",
172            Model::Gemini25Pro => "models/gemini-2.5-pro",
173            Model::Gemini25ProPreviewTts => "models/gemini-2.5-pro-preview-tts",
174            Model::Gemini25Flash => "models/gemini-2.5-flash",
175            Model::Gemini25FlashPreview092025 => "models/gemini-2.5-flash-preview-09-2025",
176            Model::Gemini25FlashImage => "models/gemini-2.5-flash-image",
177            Model::Gemini25FlashImagePreview => "models/gemini-2.5-flash-image-preview",
178            Model::Gemini25FlashLive122025 => {
179                "models/gemini-2.5-flash-native-audio-preview-12-2025"
180            }
181            Model::Gemini25FlashLive092025 => {
182                "models/gemini-2.5-flash-native-audio-preview-09-2025"
183            }
184            Model::Gemini25FlashPreviewTts => "models/gemini-2.5-flash-preview-tts",
185            Model::Gemini25FlashLite => "models/gemini-2.5-flash-lite",
186            Model::Gemini25FlashLitePreview092025 => "models/gemini-2.5-flash-lite-preview-09-2025",
187            Model::GeminiEmbedding2 => "models/gemini-embedding-2",
188            Model::GeminiEmbedding001 => "models/gemini-embedding-001",
189            Model::Custom(model) => model,
190        }
191    }
192
193    /// Returns the Vertex AI model resource path for this model.
194    pub fn vertex_model_path(&self, project_id: &str, location: &str) -> String {
195        #[allow(deprecated)]
196        let model_id = match self {
197            Model::Gemini35Flash => "gemini-3.5-flash",
198            Model::Gemini31ProPreview => "gemini-3.1-pro-preview",
199            Model::Gemini31FlashLite => "gemini-3.1-flash-lite",
200            Model::Gemini31FlashImage => "gemini-3.1-flash-image",
201            Model::Gemini3ProPreview => "gemini-3-pro-preview",
202            Model::Gemini3ProImage => "gemini-3-pro-image",
203            Model::Gemini3ProImagePreview => "gemini-3-pro-image-preview",
204            Model::Gemini3FlashPreview => "gemini-3-flash-preview",
205            Model::Gemini25Pro => "gemini-2.5-pro",
206            Model::Gemini25ProPreviewTts => "gemini-2.5-pro-preview-tts",
207            Model::Gemini25Flash => "gemini-2.5-flash",
208            Model::Gemini25FlashPreview092025 => "gemini-2.5-flash-preview-09-2025",
209            Model::Gemini25FlashImage => "gemini-2.5-flash-image",
210            Model::Gemini25FlashImagePreview => "gemini-2.5-flash-image-preview",
211            Model::Gemini25FlashLive122025 => "gemini-2.5-flash-native-audio-preview-12-2025",
212            Model::Gemini25FlashLive092025 => "gemini-2.5-flash-native-audio-preview-09-2025",
213            Model::Gemini25FlashPreviewTts => "gemini-2.5-flash-preview-tts",
214            Model::Gemini25FlashLite => "gemini-2.5-flash-lite",
215            Model::Gemini25FlashLitePreview092025 => "gemini-2.5-flash-lite-preview-09-2025",
216            Model::GeminiEmbedding2 => "gemini-embedding-2",
217            Model::GeminiEmbedding001 => "gemini-embedding-001",
218            Model::Custom(model) => {
219                if model.starts_with("projects/") {
220                    return model.clone();
221                }
222                if model.starts_with("publishers/") {
223                    return format!("projects/{project_id}/locations/{location}/{model}");
224                }
225                model.strip_prefix("models/").unwrap_or(model)
226            }
227        };
228        format!("projects/{project_id}/locations/{location}/publishers/google/models/{model_id}")
229    }
230}
231
232impl From<String> for Model {
233    #[allow(deprecated)]
234    fn from(model: String) -> Self {
235        // Match known model names (with or without "models/" prefix) to proper variants.
236        let bare = model.strip_prefix("models/").unwrap_or(&model);
237        match bare {
238            // Gemini 3.5 models (latest generation)
239            "gemini-3.5-flash" => Self::Gemini35Flash,
240            // Gemini 3.1 models
241            "gemini-3.1-pro-preview" => Self::Gemini31ProPreview,
242            "gemini-3.1-flash-lite" => Self::Gemini31FlashLite,
243            "gemini-3.1-flash-image" => Self::Gemini31FlashImage,
244            // Gemini 3 models
245            "gemini-3-pro-preview" => Self::Gemini3ProPreview,
246            "gemini-3-pro-image" => Self::Gemini3ProImage,
247            "gemini-3-pro-image-preview" => Self::Gemini3ProImagePreview,
248            "gemini-3-flash-preview" => Self::Gemini3FlashPreview,
249            // Gemini 2.5 models
250            "gemini-2.5-pro" => Self::Gemini25Pro,
251            "gemini-2.5-pro-preview-tts" => Self::Gemini25ProPreviewTts,
252            "gemini-2.5-flash" => Self::Gemini25Flash,
253            "gemini-2.5-flash-preview-09-2025" => Self::Gemini25FlashPreview092025,
254            "gemini-2.5-flash-image" => Self::Gemini25FlashImage,
255            "gemini-2.5-flash-image-preview" => Self::Gemini25FlashImagePreview,
256            "gemini-2.5-flash-native-audio-preview-12-2025" => Self::Gemini25FlashLive122025,
257            "gemini-2.5-flash-native-audio-preview-09-2025" => Self::Gemini25FlashLive092025,
258            "gemini-2.5-flash-preview-tts" => Self::Gemini25FlashPreviewTts,
259            "gemini-2.5-flash-lite" => Self::Gemini25FlashLite,
260            "gemini-2.5-flash-lite-preview-09-2025" => Self::Gemini25FlashLitePreview092025,
261            // Embedding models
262            "gemini-embedding-2" => Self::GeminiEmbedding2,
263            "gemini-embedding-001" => Self::GeminiEmbedding001,
264            _ => Self::Custom(model),
265        }
266    }
267}
268
269impl fmt::Display for Model {
270    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
271        #[allow(deprecated)]
272        match self {
273            Model::Custom(model) => {
274                // Ensure custom models always have the "models/" prefix for API URLs
275                if model.starts_with("models/")
276                    || model.starts_with("projects/")
277                    || model.starts_with("publishers/")
278                {
279                    write!(f, "{model}")
280                } else {
281                    write!(f, "models/{model}")
282                }
283            }
284            other => write!(f, "{}", other.as_str()),
285        }
286    }
287}
288
289// ══════════════════════════════════════════════════════════════════════
290// Error enum
291// ══════════════════════════════════════════════════════════════════════
292
293/// Errors that can occur when interacting with the Gemini API.
294#[derive(Debug, Snafu)]
295#[snafu(visibility(pub))]
296pub enum Error {
297    /// The provided API key is not a valid HTTP header value.
298    #[snafu(display("failed to parse API key"))]
299    InvalidApiKey {
300        /// The underlying header value error.
301        source: InvalidHeaderValue,
302    },
303
304    /// Failed to construct a request URL (likely an incorrect model name).
305    #[snafu(display("failed to construct URL (probably incorrect model name): {suffix}"))]
306    ConstructUrl {
307        /// The underlying URL parse error.
308        source: url::ParseError,
309        /// The URL suffix that failed to parse.
310        suffix: String,
311    },
312
313    /// An HTTP request failed before reaching the server.
314    #[snafu(display("failed to perform request: {source}"))]
315    PerformRequestNew {
316        /// The underlying reqwest error.
317        source: reqwest::Error,
318    },
319
320    /// An HTTP request to a specific URL failed.
321    #[snafu(display("failed to perform request to '{url}'"))]
322    PerformRequest {
323        /// The underlying reqwest error.
324        source: reqwest::Error,
325        /// The URL that was being requested.
326        url: Url,
327    },
328
329    /// The server returned a non-success HTTP status code.
330    #[snafu(display("bad response from server; code {code}; description: {}", description.as_deref().unwrap_or("none")))]
331    BadResponse {
332        /// The HTTP status code.
333        code: u16,
334        /// An optional error description from the server.
335        description: Option<String>,
336    },
337
338    /// A required response header was missing.
339    MissingResponseHeader {
340        /// The name of the missing header.
341        header: String,
342    },
343
344    /// Failed to read a Server-Sent Events (SSE) part from the stream.
345    #[snafu(display("failed to obtain stream SSE part"))]
346    BadPart {
347        /// The underlying event stream error.
348        source: EventStreamError<reqwest::Error>,
349    },
350
351    /// Failed to deserialize a JSON response body.
352    #[snafu(display("failed to deserialize JSON response"))]
353    Deserialize {
354        /// The underlying serde_json error.
355        source: serde_json::Error,
356    },
357
358    /// Failed to decode the response body.
359    #[snafu(display("failed to generate content"))]
360    DecodeResponse {
361        /// The underlying reqwest error.
362        source: reqwest::Error,
363    },
364
365    /// Failed to parse a URL string.
366    #[snafu(display("failed to parse URL"))]
367    UrlParse {
368        /// The underlying URL parse error.
369        source: url::ParseError,
370    },
371
372    /// Failed to build Google Cloud credentials.
373    #[snafu(display("failed to build google cloud credentials"))]
374    #[cfg(feature = "vertex")]
375    GoogleCloudAuth {
376        /// The underlying credentials build error.
377        source: google_cloud_auth::build_errors::Error,
378    },
379
380    /// Failed to obtain Google Cloud auth headers.
381    #[snafu(display("failed to obtain google cloud auth headers"))]
382    #[cfg(feature = "vertex")]
383    GoogleCloudCredentialHeaders {
384        /// The underlying credentials error.
385        source: google_cloud_auth::errors::CredentialsError,
386    },
387
388    /// Google Cloud credentials returned NotModified without cached headers.
389    #[snafu(display("google cloud credentials returned NotModified without cached headers"))]
390    GoogleCloudCredentialHeadersUnavailable,
391
392    /// Failed to parse Google Cloud credentials JSON.
393    #[snafu(display("failed to parse google cloud credentials JSON"))]
394    GoogleCloudCredentialParse {
395        /// The underlying JSON parse error.
396        source: serde_json::Error,
397    },
398
399    /// Failed to build the Google Cloud Vertex AI client.
400    #[snafu(display("failed to build google cloud vertex client"))]
401    #[cfg(feature = "vertex")]
402    GoogleCloudClientBuild {
403        /// The underlying client builder error.
404        source: google_cloud_gax::client_builder::Error,
405    },
406
407    /// Failed to send a request via the Google Cloud Vertex AI client.
408    #[snafu(display("failed to send google cloud vertex request"))]
409    #[cfg(feature = "vertex")]
410    GoogleCloudRequest {
411        /// The underlying Vertex AI error.
412        source: google_cloud_aiplatform_v1::Error,
413    },
414
415    /// Failed to serialize a Google Cloud request payload.
416    #[snafu(display("failed to serialize google cloud request"))]
417    GoogleCloudRequestSerialize {
418        /// The underlying serialization error.
419        source: serde_json::Error,
420    },
421
422    /// Failed to deserialize a Google Cloud request payload.
423    #[snafu(display("failed to deserialize google cloud request"))]
424    GoogleCloudRequestDeserialize {
425        /// The underlying deserialization error.
426        source: serde_json::Error,
427    },
428
429    /// Failed to serialize a Google Cloud response.
430    #[snafu(display("failed to serialize google cloud response"))]
431    GoogleCloudResponseSerialize {
432        /// The underlying serialization error.
433        source: serde_json::Error,
434    },
435
436    /// Failed to deserialize a Google Cloud response.
437    #[snafu(display("failed to deserialize google cloud response"))]
438    GoogleCloudResponseDeserialize {
439        /// The underlying deserialization error.
440        source: serde_json::Error,
441    },
442
443    /// The Google Cloud request payload is not a JSON object.
444    #[snafu(display("google cloud request payload is not an object"))]
445    GoogleCloudRequestNotObject,
446
447    /// Google Cloud configuration is required for this authentication mode.
448    #[snafu(display("google cloud configuration is required for this authentication mode"))]
449    MissingGoogleCloudConfig,
450
451    /// Google Cloud authentication is required for this configuration.
452    #[snafu(display("google cloud authentication is required for this configuration"))]
453    MissingGoogleCloudAuth,
454
455    /// The service account JSON is missing the `project_id` field.
456    #[snafu(display("service account JSON is missing required field 'project_id'"))]
457    MissingGoogleCloudProjectId,
458
459    /// An API key is required for this configuration.
460    #[snafu(display("api key is required for this configuration"))]
461    MissingApiKey,
462
463    /// The requested operation is not supported by the Vertex AI backend.
464    #[snafu(display(
465        "operation '{operation}' is not supported with the google cloud vertex backend (files, batch, and interactions APIs are Studio-only)"
466    ))]
467    GoogleCloudUnsupported {
468        /// The unsupported operation name.
469        operation: &'static str,
470    },
471
472    /// Failed to create a tokio runtime for the Google Cloud client.
473    #[snafu(display("failed to create tokio runtime for google cloud client"))]
474    TokioRuntime {
475        /// The underlying I/O error.
476        source: std::io::Error,
477    },
478
479    /// The Google Cloud client initialization thread panicked.
480    #[snafu(display("google cloud client initialization thread panicked"))]
481    GoogleCloudInitThreadPanicked,
482
483    /// An I/O error occurred during file operations.
484    #[snafu(display("I/O error during file operations"))]
485    Io {
486        /// The underlying I/O error.
487        source: std::io::Error,
488    },
489
490    /// The generation config is invalid.
491    #[snafu(display("invalid generation config: {message}"))]
492    InvalidGenerationConfig {
493        /// Description of the validation failure.
494        message: String,
495    },
496
497    /// A request failed client-side validation before dispatch.
498    #[snafu(display("validation error: {message}"))]
499    Validation {
500        /// Description of the validation failure.
501        message: String,
502    },
503}
504
505// ══════════════════════════════════════════════════════════════════════
506// GeminiClient — thin facade over a backend trait object
507// ══════════════════════════════════════════════════════════════════════
508
509/// Internal client for making requests to the Gemini API.
510///
511/// Delegates all operations to a [`GeminiBackend`](backend::GeminiBackend)
512/// trait object (AI Studio REST or Vertex AI).
513pub struct GeminiClient {
514    /// The model this client is configured to use.
515    pub model: Model,
516    backend: Box<dyn backend::GeminiBackend>,
517}
518
519pub(crate) fn validate_generation_config_for_model(
520    model: &Model,
521    config: &GenerationConfig,
522) -> Result<(), Error> {
523    config.validate().map_err(|message| Error::InvalidGenerationConfig { message })?;
524
525    let model = model.to_string();
526    let model = model.rsplit('/').next().unwrap_or(&model);
527    if !matches!(model, "gemini-3.6-flash" | "gemini-3.7-flash") {
528        return Ok(());
529    }
530
531    if config.temperature.is_some() || config.top_p.is_some() || config.top_k.is_some() {
532        return Err(Error::InvalidGenerationConfig {
533            message: format!(
534                "{model} does not accept temperature, top_p, or top_k; remove explicit sampling parameters"
535            ),
536        });
537    }
538    if config.candidate_count.is_some() {
539        return Err(Error::InvalidGenerationConfig {
540            message: format!(
541                "{model} does not accept candidate_count; remove it and request one candidate"
542            ),
543        });
544    }
545    if config.thinking_config.as_ref().is_some_and(|thinking| thinking.thinking_budget.is_some()) {
546        return Err(Error::InvalidGenerationConfig {
547            message: format!(
548                "{model} uses thinking levels instead of token budgets; set thinking_level and clear thinking_budget"
549            ),
550        });
551    }
552
553    Ok(())
554}
555
556#[cfg(feature = "interactions")]
557pub(crate) fn validate_interaction_generation_config_for_model(
558    model: Option<&str>,
559    config: &crate::interactions::GenerationConfig,
560) -> Result<(), Error> {
561    config.validate().map_err(|message| Error::InvalidGenerationConfig { message })?;
562
563    let Some(model) = model.map(|model| model.rsplit('/').next().unwrap_or(model)) else {
564        return Ok(());
565    };
566    if matches!(model, "gemini-3.6-flash" | "gemini-3.7-flash")
567        && (config.temperature.is_some() || config.top_p.is_some())
568    {
569        return Err(Error::InvalidGenerationConfig {
570            message: format!(
571                "{model} does not accept temperature or top_p; remove explicit sampling parameters"
572            ),
573        });
574    }
575
576    Ok(())
577}
578
579impl std::fmt::Debug for GeminiClient {
580    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
581        f.debug_struct("GeminiClient")
582            .field("model", &self.model)
583            .field("backend", &self.backend)
584            .finish()
585    }
586}
587
588impl GeminiClient {
589    /// Create a client backed by AI Studio REST.
590    fn with_studio(model: Model, studio: backend::studio::StudioBackend) -> Self {
591        Self { model, backend: Box::new(studio) }
592    }
593
594    /// Create a client backed by Vertex AI.
595    #[cfg(feature = "vertex")]
596    fn with_vertex(model: Model, vertex: backend::vertex::VertexBackend) -> Self {
597        Self { model, backend: Box::new(vertex) }
598    }
599
600    fn validate_generation_config(&self, config: &GenerationConfig) -> Result<(), Error> {
601        validate_generation_config_for_model(&self.model, config)
602    }
603
604    // ── Delegating methods ──────────────────────────────────────────────
605
606    #[instrument(skip_all, fields(
607        model,
608        messages.parts.count = request.contents.len(),
609        tools.present = request.tools.is_some(),
610        system.instruction.present = request.system_instruction.is_some(),
611        cached.content.present = request.cached_content.is_some(),
612        usage.prompt_tokens,
613        usage.candidates_tokens,
614        usage.thoughts_tokens,
615        usage.cached_content_tokens,
616        usage.total_tokens,
617    ), ret(level = Level::TRACE), err)]
618    pub(crate) async fn generate_content_raw(
619        &self,
620        request: GenerateContentRequest,
621    ) -> Result<GenerationResponse, Error> {
622        if let Some(ref gc) = request.generation_config {
623            self.validate_generation_config(gc)?;
624        }
625
626        let response = self.backend.generate_content(request).await?;
627
628        if let Some(usage) = &response.usage_metadata {
629            #[rustfmt::skip]
630            Span::current()
631                .record("usage.prompt_tokens", usage.prompt_token_count)
632                .record("usage.candidates_tokens", usage.candidates_token_count)
633                .record("usage.thoughts_tokens", usage.thoughts_token_count)
634                .record("usage.cached_content_tokens", usage.cached_content_token_count)
635                .record("usage.total_tokens", usage.total_token_count);
636            tracing::debug!("generation usage evaluated");
637        }
638
639        Ok(response)
640    }
641
642    #[instrument(skip_all, fields(
643        model,
644        messages.parts.count = request.contents.len(),
645        tools.present = request.tools.is_some(),
646        system.instruction.present = request.system_instruction.is_some(),
647        cached.content.present = request.cached_content.is_some(),
648    ), err)]
649    pub(crate) async fn generate_content_stream(
650        &self,
651        request: GenerateContentRequest,
652    ) -> Result<backend::BackendStream<GenerationResponse>, Error> {
653        if let Some(ref gc) = request.generation_config {
654            self.validate_generation_config(gc)?;
655        }
656
657        self.backend.generate_content_stream(request).await
658    }
659
660    #[instrument(skip_all, fields(
661        model,
662        task.type = request.task_type.as_ref().map(|t| format!("{:?}", t)),
663        task.title = request.title,
664        task.output.dimensionality = request.output_dimensionality,
665    ))]
666    pub(crate) async fn embed_content(
667        &self,
668        request: EmbedContentRequest,
669    ) -> Result<ContentEmbeddingResponse, Error> {
670        self.backend.embed_content(request).await
671    }
672
673    #[instrument(skip_all, fields(batch.size = request.requests.len()))]
674    pub(crate) async fn embed_content_batch(
675        &self,
676        request: BatchEmbedContentsRequest,
677    ) -> Result<BatchContentEmbeddingResponse, Error> {
678        self.backend.batch_embed_contents(request).await
679    }
680
681    #[instrument(skip_all, fields(
682        batch.display_name = request.batch.display_name,
683        batch.size = request.batch.input_config.batch_size(),
684    ))]
685    pub(crate) async fn batch_generate_content(
686        &self,
687        request: BatchGenerateContentRequest,
688    ) -> Result<BatchGenerateContentResponse, Error> {
689        self.backend.batch_generate_content(request).await
690    }
691
692    #[instrument(skip_all, fields(operation.name = name))]
693    pub(crate) async fn get_batch_operation<T: serde::de::DeserializeOwned>(
694        &self,
695        name: &str,
696    ) -> Result<T, Error> {
697        let value = self.backend.get_batch_operation(name).await?;
698        serde_json::from_value(value).context(DeserializeSnafu)
699    }
700
701    #[instrument(skip_all, fields(page.size = page_size, page.token.present = page_token.is_some()))]
702    pub(crate) async fn list_batch_operations(
703        &self,
704        page_size: Option<u32>,
705        page_token: Option<String>,
706    ) -> Result<ListBatchesResponse, Error> {
707        self.backend.list_batch_operations(page_size, page_token).await
708    }
709
710    #[instrument(skip_all, fields(page.size = page_size, page.token.present = page_token.is_some()))]
711    pub(crate) async fn list_files(
712        &self,
713        page_size: Option<u32>,
714        page_token: Option<String>,
715    ) -> Result<ListFilesResponse, Error> {
716        self.backend.list_files(page_size, page_token).await
717    }
718
719    #[instrument(skip_all, fields(operation.name = name))]
720    pub(crate) async fn cancel_batch_operation(&self, name: &str) -> Result<(), Error> {
721        self.backend.cancel_batch_operation(name).await
722    }
723
724    #[instrument(skip_all, fields(operation.name = name))]
725    pub(crate) async fn delete_batch_operation(&self, name: &str) -> Result<(), Error> {
726        self.backend.delete_batch_operation(name).await
727    }
728
729    #[instrument(skip_all, fields(
730        file.size = file_bytes.len(),
731        mime.type = mime_type.to_string(),
732        file.display_name = display_name.as_deref(),
733    ))]
734    pub(crate) async fn upload_file(
735        &self,
736        display_name: Option<String>,
737        file_bytes: Vec<u8>,
738        mime_type: Mime,
739    ) -> Result<File, Error> {
740        self.backend.upload_file(display_name, file_bytes, mime_type).await
741    }
742
743    #[instrument(skip_all, fields(file.name = name))]
744    pub(crate) async fn get_file(&self, name: &str) -> Result<File, Error> {
745        self.backend.get_file(name).await
746    }
747
748    #[instrument(skip_all, fields(file.name = name))]
749    pub(crate) async fn delete_file(&self, name: &str) -> Result<(), Error> {
750        self.backend.delete_file(name).await
751    }
752
753    #[instrument(skip_all, fields(file.name = name))]
754    pub(crate) async fn download_file(&self, name: &str) -> Result<Vec<u8>, Error> {
755        self.backend.download_file(name).await
756    }
757
758    pub(crate) async fn create_cached_content(
759        &self,
760        cached_content: CreateCachedContentRequest,
761    ) -> Result<CachedContent, Error> {
762        self.backend.create_cached_content(cached_content).await
763    }
764
765    pub(crate) async fn get_cached_content(&self, name: &str) -> Result<CachedContent, Error> {
766        self.backend.get_cached_content(name).await
767    }
768
769    pub(crate) async fn update_cached_content(
770        &self,
771        name: &str,
772        expiration: CacheExpirationRequest,
773    ) -> Result<CachedContent, Error> {
774        self.backend.update_cached_content(name, expiration).await
775    }
776
777    pub(crate) async fn delete_cached_content(&self, name: &str) -> Result<(), Error> {
778        self.backend.delete_cached_content(name).await
779    }
780
781    pub(crate) async fn list_cached_contents(
782        &self,
783        page_size: Option<i32>,
784        page_token: Option<String>,
785    ) -> Result<ListCachedContentsResponse, Error> {
786        self.backend.list_cached_contents(page_size, page_token).await
787    }
788
789    // ── Model discovery ─────────────────────────────────────────────────
790
791    #[instrument(skip_all, fields(page.size = page_size, page.token.present = page_token.is_some()))]
792    pub(crate) async fn list_models(
793        &self,
794        page_size: Option<u32>,
795        page_token: Option<String>,
796    ) -> Result<crate::model_info::ListModelsResponse, Error> {
797        self.backend.list_models(page_size, page_token).await
798    }
799
800    #[instrument(skip_all, fields(model.name = name))]
801    pub(crate) async fn get_model(
802        &self,
803        name: &str,
804    ) -> Result<crate::model_info::ModelInfo, Error> {
805        self.backend.get_model(name).await
806    }
807
808    // ── Interactions API (Beta) ─────────────────────────────────────────
809
810    #[cfg(feature = "interactions")]
811    #[instrument(skip_all, fields(
812        model = request.model.as_deref(),
813        agent = request.agent.as_deref(),
814        tools.count = request.tools.len(),
815        background = request.background,
816        previous.interaction.present = request.previous_interaction_id.is_some(),
817    ), err)]
818    pub(crate) async fn create_interaction(
819        &self,
820        request: crate::interactions::CreateInteractionRequest,
821    ) -> Result<crate::interactions::Interaction, Error> {
822        if let Some(ref gc) = request.generation_config {
823            validate_interaction_generation_config_for_model(request.model.as_deref(), gc)?;
824        }
825        self.backend.create_interaction(request).await
826    }
827
828    #[cfg(feature = "interactions")]
829    #[instrument(skip_all, fields(
830        model = request.model.as_deref(),
831        agent = request.agent.as_deref(),
832    ), err)]
833    pub(crate) async fn create_interaction_stream(
834        &self,
835        request: crate::interactions::CreateInteractionRequest,
836    ) -> Result<backend::BackendStream<crate::interactions::InteractionSseEvent>, Error> {
837        if let Some(ref gc) = request.generation_config {
838            validate_interaction_generation_config_for_model(request.model.as_deref(), gc)?;
839        }
840        self.backend.create_interaction_stream(request).await
841    }
842
843    #[cfg(feature = "interactions")]
844    #[instrument(skip_all, fields(interaction.id = id, include.input = include_input))]
845    pub(crate) async fn get_interaction(
846        &self,
847        id: &str,
848        include_input: bool,
849    ) -> Result<crate::interactions::Interaction, Error> {
850        self.backend.get_interaction(id, include_input).await
851    }
852
853    #[cfg(feature = "interactions")]
854    #[instrument(skip_all, fields(interaction.id = id))]
855    pub(crate) async fn delete_interaction(&self, id: &str) -> Result<(), Error> {
856        self.backend.delete_interaction(id).await
857    }
858
859    #[cfg(feature = "interactions")]
860    #[instrument(skip_all, fields(interaction.id = id))]
861    pub(crate) async fn cancel_interaction(
862        &self,
863        id: &str,
864    ) -> Result<crate::interactions::Interaction, Error> {
865        self.backend.cancel_interaction(id).await
866    }
867
868    #[cfg(feature = "interactions")]
869    #[instrument(skip_all, fields(agent.id = request.id))]
870    pub(crate) async fn create_agent(
871        &self,
872        request: crate::interactions::managed_agent::CreateAgentRequest,
873    ) -> Result<crate::interactions::managed_agent::SavedAgent, Error> {
874        self.backend.create_agent(request).await
875    }
876}
877
878// ══════════════════════════════════════════════════════════════════════
879// Auth helpers & builder infrastructure
880// ══════════════════════════════════════════════════════════════════════
881
882#[cfg(feature = "vertex")]
883#[derive(Debug, Clone)]
884enum GoogleCloudAuth {
885    ApiKey(String),
886    Credentials(Credentials),
887}
888
889#[cfg(feature = "vertex")]
890impl GoogleCloudAuth {
891    fn credentials(&self) -> Result<Credentials, Error> {
892        match self {
893            GoogleCloudAuth::ApiKey(api_key) => {
894                Ok(credentials::api_key_credentials::Builder::new(api_key).build())
895            }
896            GoogleCloudAuth::Credentials(credentials) => Ok(credentials.clone()),
897        }
898    }
899}
900
901#[cfg(feature = "vertex")]
902#[derive(Debug, Clone)]
903struct GoogleCloudConfig {
904    project_id: String,
905    location: String,
906}
907
908#[cfg(feature = "vertex")]
909impl GoogleCloudConfig {
910    fn endpoint(&self) -> String {
911        if self.location == "global" {
912            "https://aiplatform.googleapis.com".to_string()
913        } else {
914            format!("https://{}-aiplatform.googleapis.com", self.location)
915        }
916    }
917}
918
919#[cfg(feature = "vertex")]
920fn extract_service_account_project_id(service_account_json: &str) -> Result<String, Error> {
921    let value: serde_json::Value =
922        serde_json::from_str(service_account_json).context(GoogleCloudCredentialParseSnafu)?;
923
924    let project_id = value
925        .get("project_id")
926        .and_then(serde_json::Value::as_str)
927        .map(str::trim)
928        .filter(|s| !s.is_empty())
929        .ok_or(Error::MissingGoogleCloudProjectId)?;
930
931    Ok(project_id.to_string())
932}
933
934#[cfg(feature = "vertex")]
935fn build_vertex_prediction_service(
936    endpoint: String,
937    credentials: Credentials,
938) -> Result<PredictionService, Error> {
939    let build_in_runtime =
940        |endpoint: String, credentials: Credentials| -> Result<PredictionService, Error> {
941            let runtime = tokio::runtime::Runtime::new().context(TokioRuntimeSnafu)?;
942            runtime
943                .block_on(
944                    PredictionService::builder()
945                        .with_endpoint(endpoint)
946                        .with_credentials(credentials)
947                        .build(),
948                )
949                .context(GoogleCloudClientBuildSnafu)
950        };
951
952    if tokio::runtime::Handle::try_current().is_ok() {
953        let worker = std::thread::Builder::new()
954            .name("adk-gemini-vertex-init".to_string())
955            .spawn(move || build_in_runtime(endpoint, credentials))
956            .map_err(|source| Error::TokioRuntime { source })?;
957
958        return worker.join().map_err(|_| Error::GoogleCloudInitThreadPanicked)?;
959    }
960
961    build_in_runtime(endpoint, credentials)
962}
963
964// ══════════════════════════════════════════════════════════════════════
965// GeminiBuilder
966// ══════════════════════════════════════════════════════════════════════
967
968/// A builder for the `Gemini` client.
969///
970/// # Examples
971///
972/// ## Basic usage
973///
974/// ```no_run
975/// use adk_gemini::{GeminiBuilder, Model};
976///
977/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
978/// let gemini = GeminiBuilder::new("YOUR_API_KEY")
979///     .with_model(Model::Gemini25Pro)
980///     .build()?;
981/// # Ok(())
982/// # }
983/// ```
984pub struct GeminiBuilder {
985    model: Model,
986    client_builder: ClientBuilder,
987    base_url: Url,
988    #[cfg(feature = "vertex")]
989    google_cloud: Option<GoogleCloudConfig>,
990    api_key: Option<String>,
991    #[cfg(feature = "vertex")]
992    google_cloud_auth: Option<GoogleCloudAuth>,
993}
994
995impl GeminiBuilder {
996    /// Create a new builder with the given API key.
997    pub fn new<K: Into<String>>(key: K) -> Self {
998        Self {
999            model: Model::default(),
1000            client_builder: ClientBuilder::default(),
1001            base_url: DEFAULT_BASE_URL.clone(),
1002            #[cfg(feature = "vertex")]
1003            google_cloud: None,
1004            api_key: Some(key.into()),
1005            #[cfg(feature = "vertex")]
1006            google_cloud_auth: None,
1007        }
1008    }
1009
1010    /// Set the model to use.
1011    pub fn with_model<M: Into<Model>>(mut self, model: M) -> Self {
1012        self.model = model.into();
1013        self
1014    }
1015
1016    /// Set a custom HTTP client builder.
1017    pub fn with_http_client(mut self, client_builder: ClientBuilder) -> Self {
1018        self.client_builder = client_builder;
1019        self
1020    }
1021
1022    /// Set a custom base URL (overrides Google Cloud config).
1023    pub fn with_base_url(mut self, base_url: Url) -> Self {
1024        self.base_url = base_url;
1025        #[cfg(feature = "vertex")]
1026        {
1027            self.google_cloud = None;
1028            self.google_cloud_auth = None;
1029        }
1030        self
1031    }
1032
1033    /// Authenticate with a service account JSON key (Vertex AI).
1034    #[cfg(feature = "vertex")]
1035    pub fn with_service_account_json(mut self, service_account_json: &str) -> Result<Self, Error> {
1036        let value =
1037            serde_json::from_str(service_account_json).context(GoogleCloudCredentialParseSnafu)?;
1038        let credentials = google_cloud_auth::credentials::service_account::Builder::new(value)
1039            .build()
1040            .context(GoogleCloudAuthSnafu)?;
1041        self.google_cloud_auth = Some(GoogleCloudAuth::Credentials(credentials));
1042        Ok(self)
1043    }
1044
1045    /// Set the Google Cloud project and location for Vertex AI.
1046    #[cfg(feature = "vertex")]
1047    pub fn with_google_cloud<P: Into<String>, L: Into<String>>(
1048        mut self,
1049        project_id: P,
1050        location: L,
1051    ) -> Self {
1052        self.google_cloud =
1053            Some(GoogleCloudConfig { project_id: project_id.into(), location: location.into() });
1054        self
1055    }
1056
1057    /// Use Application Default Credentials (ADC) for Vertex AI authentication.
1058    #[cfg(feature = "vertex")]
1059    pub fn with_google_cloud_adc(mut self) -> Result<Self, Error> {
1060        let credentials = google_cloud_auth::credentials::Builder::default()
1061            .build()
1062            .context(GoogleCloudAuthSnafu)?;
1063        self.google_cloud_auth = Some(GoogleCloudAuth::Credentials(credentials));
1064        Ok(self)
1065    }
1066
1067    /// Use Workload Identity Federation JSON for Vertex AI authentication.
1068    #[cfg(feature = "vertex")]
1069    pub fn with_google_cloud_wif_json(mut self, wif_json: &str) -> Result<Self, Error> {
1070        let value = serde_json::from_str(wif_json).context(GoogleCloudCredentialParseSnafu)?;
1071        let credentials = google_cloud_auth::credentials::external_account::Builder::new(value)
1072            .build()
1073            .context(GoogleCloudAuthSnafu)?;
1074        self.google_cloud_auth = Some(GoogleCloudAuth::Credentials(credentials));
1075        Ok(self)
1076    }
1077
1078    /// Builds the `Gemini` client.
1079    pub fn build(self) -> Result<Gemini, Error> {
1080        #[cfg(feature = "vertex")]
1081        {
1082            if self.google_cloud.is_none() && self.google_cloud_auth.is_some() {
1083                return MissingGoogleCloudConfigSnafu.fail();
1084            }
1085
1086            // ── Vertex AI path ──────────────────────────────────────────────
1087            if let Some(config) = &self.google_cloud {
1088                let model = Model::Custom(
1089                    self.model.vertex_model_path(&config.project_id, &config.location),
1090                );
1091                let google_cloud_auth = match self.google_cloud_auth {
1092                    Some(auth) => auth,
1093                    None => match self.api_key {
1094                        Some(api_key) if !api_key.is_empty() => GoogleCloudAuth::ApiKey(api_key),
1095                        _ => return MissingGoogleCloudAuthSnafu.fail(),
1096                    },
1097                };
1098                let credentials = google_cloud_auth.credentials()?;
1099                let endpoint = config.endpoint();
1100                let prediction =
1101                    build_vertex_prediction_service(endpoint.clone(), credentials.clone())?;
1102
1103                let vertex = backend::vertex::VertexBackend::new(
1104                    model.clone(),
1105                    prediction,
1106                    credentials,
1107                    endpoint,
1108                );
1109
1110                return Ok(Gemini { client: Arc::new(GeminiClient::with_vertex(model, vertex)) });
1111            }
1112        }
1113
1114        // ── AI Studio REST path ─────────────────────────────────────────
1115        let api_key = self.api_key.ok_or(Error::MissingApiKey)?;
1116        if api_key.is_empty() {
1117            return MissingApiKeySnafu.fail();
1118        }
1119
1120        let studio =
1121            backend::studio::StudioBackend::new(&api_key, self.model.clone(), self.base_url)?;
1122
1123        Ok(Gemini { client: Arc::new(GeminiClient::with_studio(self.model, studio)) })
1124    }
1125}
1126
1127// ══════════════════════════════════════════════════════════════════════
1128// Gemini — the main public-facing client
1129// ══════════════════════════════════════════════════════════════════════
1130
1131/// The main public-facing Gemini API client.
1132///
1133/// Provides methods for content generation, embeddings, batch processing,
1134/// file management, caching, and model discovery.
1135pub struct Gemini {
1136    client: Arc<GeminiClient>,
1137}
1138
1139impl Gemini {
1140    /// Create a new client with the specified API key
1141    pub fn new<K: AsRef<str>>(api_key: K) -> Result<Self, Error> {
1142        Self::with_model(api_key, Model::default())
1143    }
1144
1145    /// Create a new client for the Gemini 3.1 Pro model
1146    pub fn pro<K: AsRef<str>>(api_key: K) -> Result<Self, Error> {
1147        Self::with_model(api_key, Model::Gemini31ProPreview)
1148    }
1149
1150    /// Create a new client with the specified API key and model
1151    pub fn with_model<K: AsRef<str>, M: Into<Model>>(api_key: K, model: M) -> Result<Self, Error> {
1152        Self::with_model_and_base_url(api_key, model, DEFAULT_BASE_URL.clone())
1153    }
1154
1155    /// Create a new client with the specified API key using the v1 (stable) API.
1156    pub fn with_v1<K: AsRef<str>>(api_key: K) -> Result<Self, Error> {
1157        Self::with_model_and_base_url(api_key, Model::default(), V1_BASE_URL.clone())
1158    }
1159
1160    /// Create a new client with the specified API key and model using the v1 (stable) API.
1161    pub fn with_model_v1<K: AsRef<str>, M: Into<Model>>(
1162        api_key: K,
1163        model: M,
1164    ) -> Result<Self, Error> {
1165        Self::with_model_and_base_url(api_key, model, V1_BASE_URL.clone())
1166    }
1167
1168    /// Create a new client with custom base URL
1169    pub fn with_base_url<K: AsRef<str>>(api_key: K, base_url: Url) -> Result<Self, Error> {
1170        Self::with_model_and_base_url(api_key, Model::default(), base_url)
1171    }
1172
1173    /// Create a new client using Vertex AI (Google Cloud) endpoints.
1174    #[cfg(feature = "vertex")]
1175    pub fn with_google_cloud<K: AsRef<str>, P: AsRef<str>, L: AsRef<str>>(
1176        api_key: K,
1177        project_id: P,
1178        location: L,
1179    ) -> Result<Self, Error> {
1180        Self::with_google_cloud_model(api_key, project_id, location, Model::default())
1181    }
1182
1183    /// Create a new client using Vertex AI (Google Cloud) endpoints and a specific model.
1184    #[cfg(feature = "vertex")]
1185    pub fn with_google_cloud_model<K: AsRef<str>, P: AsRef<str>, L: AsRef<str>, M: Into<Model>>(
1186        api_key: K,
1187        project_id: P,
1188        location: L,
1189        model: M,
1190    ) -> Result<Self, Error> {
1191        GeminiBuilder::new(api_key.as_ref())
1192            .with_model(model)
1193            .with_google_cloud(project_id.as_ref(), location.as_ref())
1194            .build()
1195    }
1196
1197    /// Create a new client using Vertex AI (Google Cloud) endpoints with Application Default Credentials (ADC).
1198    #[cfg(feature = "vertex")]
1199    pub fn with_google_cloud_adc<P: AsRef<str>, L: AsRef<str>>(
1200        project_id: P,
1201        location: L,
1202    ) -> Result<Self, Error> {
1203        Self::with_google_cloud_adc_model(project_id, location, Model::default())
1204    }
1205
1206    /// Create a new client using Vertex AI (Google Cloud) endpoints and a specific model with ADC.
1207    #[cfg(feature = "vertex")]
1208    pub fn with_google_cloud_adc_model<P: AsRef<str>, L: AsRef<str>, M: Into<Model>>(
1209        project_id: P,
1210        location: L,
1211        model: M,
1212    ) -> Result<Self, Error> {
1213        GeminiBuilder::new("")
1214            .with_model(model)
1215            .with_google_cloud(project_id.as_ref(), location.as_ref())
1216            .with_google_cloud_adc()?
1217            .build()
1218    }
1219
1220    /// Create a new client using Vertex AI (Google Cloud) endpoints and Workload Identity Federation JSON.
1221    #[cfg(feature = "vertex")]
1222    pub fn with_google_cloud_wif_json<P: AsRef<str>, L: AsRef<str>, M: Into<Model>>(
1223        wif_json: &str,
1224        project_id: P,
1225        location: L,
1226        model: M,
1227    ) -> Result<Self, Error> {
1228        GeminiBuilder::new("")
1229            .with_model(model)
1230            .with_google_cloud(project_id.as_ref(), location.as_ref())
1231            .with_google_cloud_wif_json(wif_json)?
1232            .build()
1233    }
1234
1235    /// Create a new client using a service account JSON key.
1236    #[cfg(feature = "vertex")]
1237    pub fn with_service_account_json(service_account_json: &str) -> Result<Self, Error> {
1238        Self::with_service_account_json_model(service_account_json, Model::default())
1239    }
1240
1241    /// Create a new client using a service account JSON key and a specific model.
1242    #[cfg(feature = "vertex")]
1243    pub fn with_service_account_json_model<M: Into<Model>>(
1244        service_account_json: &str,
1245        model: M,
1246    ) -> Result<Self, Error> {
1247        let project_id = extract_service_account_project_id(service_account_json)?;
1248        GeminiBuilder::new("")
1249            .with_model(model)
1250            .with_service_account_json(service_account_json)?
1251            .with_google_cloud(project_id, "us-central1")
1252            .build()
1253    }
1254
1255    /// Create a new client using Vertex AI (Google Cloud) endpoints and a service account JSON key.
1256    #[cfg(feature = "vertex")]
1257    pub fn with_google_cloud_service_account_json<M: Into<Model>>(
1258        service_account_json: &str,
1259        project_id: &str,
1260        location: &str,
1261        model: M,
1262    ) -> Result<Self, Error> {
1263        GeminiBuilder::new("")
1264            .with_model(model)
1265            .with_service_account_json(service_account_json)?
1266            .with_google_cloud(project_id, location)
1267            .build()
1268    }
1269
1270    /// Create a new client with the specified API key, model, and base URL
1271    pub fn with_model_and_base_url<K: AsRef<str>, M: Into<Model>>(
1272        api_key: K,
1273        model: M,
1274        base_url: Url,
1275    ) -> Result<Self, Error> {
1276        let model = model.into();
1277        let studio =
1278            backend::studio::StudioBackend::new(api_key.as_ref(), model.clone(), base_url)?;
1279        Ok(Self { client: Arc::new(GeminiClient::with_studio(model, studio)) })
1280    }
1281
1282    /// Start building a content generation request
1283    pub fn generate_content(&self) -> ContentBuilder {
1284        ContentBuilder::new(self.client.clone())
1285    }
1286
1287    /// Start building a content embedding request
1288    pub fn embed_content(&self) -> EmbedBuilder {
1289        EmbedBuilder::new(self.client.clone())
1290    }
1291
1292    /// Start building a batch content generation request
1293    pub fn batch_generate_content(&self) -> BatchBuilder {
1294        BatchBuilder::new(self.client.clone())
1295    }
1296
1297    /// Get a handle to a batch operation by its name.
1298    pub fn get_batch(&self, name: &str) -> BatchHandle {
1299        BatchHandle::new(name.to_string(), self.client.clone())
1300    }
1301
1302    /// Lists batch operations.
1303    pub fn list_batches(
1304        &self,
1305        page_size: impl Into<Option<u32>>,
1306    ) -> impl Stream<Item = Result<BatchOperation, Error>> + Send {
1307        let client = self.client.clone();
1308        let page_size = page_size.into();
1309        async_stream::try_stream! {
1310            let mut page_token: Option<String> = None;
1311            loop {
1312                let response = client
1313                    .list_batch_operations(page_size, page_token.clone())
1314                    .await?;
1315
1316                for operation in response.operations {
1317                    yield operation;
1318                }
1319
1320                if let Some(next_page_token) = response.next_page_token {
1321                    page_token = Some(next_page_token);
1322                } else {
1323                    break;
1324                }
1325            }
1326        }
1327    }
1328
1329    /// Create cached content with a fluent API.
1330    pub fn create_cache(&self) -> CacheBuilder {
1331        CacheBuilder::new(self.client.clone())
1332    }
1333
1334    /// Get a handle to cached content by its name.
1335    pub fn get_cached_content(&self, name: &str) -> CachedContentHandle {
1336        CachedContentHandle::new(name.to_string(), self.client.clone())
1337    }
1338
1339    /// Lists cached contents.
1340    pub fn list_cached_contents(
1341        &self,
1342        page_size: impl Into<Option<i32>>,
1343    ) -> impl Stream<Item = Result<CachedContentSummary, Error>> + Send {
1344        let client = self.client.clone();
1345        let page_size = page_size.into();
1346        async_stream::try_stream! {
1347            let mut page_token: Option<String> = None;
1348            loop {
1349                let response = client
1350                    .list_cached_contents(page_size, page_token.clone())
1351                    .await?;
1352
1353                for cached_content in response.cached_contents {
1354                    yield cached_content;
1355                }
1356
1357                if let Some(next_page_token) = response.next_page_token {
1358                    page_token = Some(next_page_token);
1359                } else {
1360                    break;
1361                }
1362            }
1363        }
1364    }
1365
1366    /// Start building a file resource
1367    pub fn create_file<B: Into<Vec<u8>>>(&self, bytes: B) -> crate::files::builder::FileBuilder {
1368        crate::files::builder::FileBuilder::new(self.client.clone(), bytes)
1369    }
1370
1371    /// Get a handle to a file by its name.
1372    pub async fn get_file(&self, name: &str) -> Result<FileHandle, Error> {
1373        let file = self.client.get_file(name).await?;
1374        Ok(FileHandle::new(self.client.clone(), file))
1375    }
1376
1377    /// Lists files.
1378    pub fn list_files(
1379        &self,
1380        page_size: impl Into<Option<u32>>,
1381    ) -> impl Stream<Item = Result<FileHandle, Error>> + Send {
1382        let client = self.client.clone();
1383        let page_size = page_size.into();
1384        async_stream::try_stream! {
1385            let mut page_token: Option<String> = None;
1386            loop {
1387                let response = client
1388                    .list_files(page_size, page_token.clone())
1389                    .await?;
1390
1391                for file in response.files {
1392                    yield FileHandle::new(client.clone(), file);
1393                }
1394
1395                if let Some(next_page_token) = response.next_page_token {
1396                    page_token = Some(next_page_token);
1397                } else {
1398                    break;
1399                }
1400            }
1401        }
1402    }
1403
1404    // ── Model discovery ─────────────────────────────────────────────────
1405
1406    /// Lists available Gemini models with pagination.
1407    ///
1408    /// Returns a stream of [`ModelInfo`](crate::model_info::ModelInfo) items.
1409    /// This is useful for discovering which models are available and their
1410    /// capabilities (token limits, supported methods, etc.).
1411    ///
1412    /// # Example
1413    ///
1414    /// ```rust,ignore
1415    /// use futures::StreamExt;
1416    ///
1417    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1418    /// let mut models = gemini.list_models(None);
1419    /// while let Some(model) = models.next().await {
1420    ///     let model = model?;
1421    ///     println!("{}: {}", model.name, model.display_name);
1422    /// }
1423    /// ```
1424    pub fn list_models(
1425        &self,
1426        page_size: impl Into<Option<u32>>,
1427    ) -> impl Stream<Item = Result<crate::model_info::ModelInfo, Error>> + Send {
1428        let client = self.client.clone();
1429        let page_size = page_size.into();
1430        async_stream::try_stream! {
1431            let mut page_token: Option<String> = None;
1432            loop {
1433                let response = client
1434                    .list_models(page_size, page_token.clone())
1435                    .await?;
1436
1437                for model in response.models {
1438                    yield model;
1439                }
1440
1441                if let Some(next_page_token) = response.next_page_token {
1442                    page_token = Some(next_page_token);
1443                } else {
1444                    break;
1445                }
1446            }
1447        }
1448    }
1449
1450    /// Get metadata for a specific model by name.
1451    ///
1452    /// The name can be provided with or without the `models/` prefix
1453    /// (e.g. both `"gemini-2.5-flash"` and `"models/gemini-2.5-flash"` work).
1454    ///
1455    /// # Example
1456    ///
1457    /// ```rust,ignore
1458    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1459    /// let info = gemini.get_model("gemini-2.5-flash").await?;
1460    /// println!("Input limit: {} tokens", info.input_token_limit);
1461    /// println!("Output limit: {} tokens", info.output_token_limit);
1462    /// ```
1463    pub async fn get_model(&self, name: &str) -> Result<crate::model_info::ModelInfo, Error> {
1464        self.client.get_model(name).await
1465    }
1466
1467    // ── Interactions API (Beta) ─────────────────────────────────────────
1468
1469    /// Start building an interaction using the Interactions API (Beta).
1470    ///
1471    /// The Interactions API is Google's new direction for the Gemini API. It
1472    /// provides server-side history, a typed step timeline, and native agentic
1473    /// workflows. See [`crate::interactions`] for details.
1474    ///
1475    /// # Example
1476    ///
1477    /// ```rust,ignore
1478    /// use adk_gemini::{Gemini, Model};
1479    ///
1480    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1481    /// let interaction = gemini
1482    ///     .create_interaction()
1483    ///     .model(Model::Gemini35Flash)
1484    ///     .input_text("Hello!")
1485    ///     .send()
1486    ///     .await?;
1487    /// println!("{}", interaction.output_text().unwrap_or_default());
1488    /// ```
1489    #[cfg(feature = "interactions")]
1490    pub fn create_interaction(&self) -> crate::interactions::InteractionBuilder {
1491        crate::interactions::InteractionBuilder::new(self.client.clone())
1492    }
1493
1494    /// Retrieve a stored interaction by ID.
1495    ///
1496    /// Set `include_input` to also return the `user_input` steps.
1497    #[cfg(feature = "interactions")]
1498    pub async fn get_interaction(
1499        &self,
1500        id: &str,
1501        include_input: bool,
1502    ) -> Result<crate::interactions::Interaction, Error> {
1503        self.client.get_interaction(id, include_input).await
1504    }
1505
1506    /// Delete a stored interaction by ID.
1507    #[cfg(feature = "interactions")]
1508    pub async fn delete_interaction(&self, id: &str) -> Result<(), Error> {
1509        self.client.delete_interaction(id).await
1510    }
1511
1512    /// Cancel a running background interaction by ID.
1513    #[cfg(feature = "interactions")]
1514    pub async fn cancel_interaction(
1515        &self,
1516        id: &str,
1517    ) -> Result<crate::interactions::Interaction, Error> {
1518        self.client.cancel_interaction(id).await
1519    }
1520
1521    /// Send a pre-built [`CreateInteractionRequest`](crate::interactions::CreateInteractionRequest)
1522    /// (non-streaming).
1523    ///
1524    /// The fluent [`create_interaction`](Self::create_interaction) builder is
1525    /// the ergonomic entry point for most callers. This method exists for
1526    /// callers that assemble a [`CreateInteractionRequest`](crate::interactions::CreateInteractionRequest)
1527    /// directly (e.g. a conversion layer that maps another request type onto the
1528    /// Interactions wire format) and need to dispatch it without rebuilding it
1529    /// field-by-field through the builder.
1530    #[cfg(feature = "interactions")]
1531    pub async fn send_interaction(
1532        &self,
1533        request: crate::interactions::CreateInteractionRequest,
1534    ) -> Result<crate::interactions::Interaction, Error> {
1535        self.client.create_interaction(request).await
1536    }
1537
1538    /// Send a pre-built [`CreateInteractionRequest`](crate::interactions::CreateInteractionRequest)
1539    /// as an SSE stream.
1540    ///
1541    /// Forces `stream = true` on the request before dispatching, mirroring the
1542    /// builder's [`stream`](crate::interactions::InteractionBuilder::stream)
1543    /// entry point. Like [`send_interaction`](Self::send_interaction), this is
1544    /// for callers that assemble the request directly rather than through the
1545    /// fluent builder.
1546    #[cfg(feature = "interactions")]
1547    pub async fn send_interaction_stream(
1548        &self,
1549        mut request: crate::interactions::CreateInteractionRequest,
1550    ) -> Result<backend::BackendStream<crate::interactions::InteractionSseEvent>, Error> {
1551        request.stream = Some(true);
1552        self.client.create_interaction_stream(request).await
1553    }
1554
1555    // ── Managed Agents ──────────────────────────────────────────────────
1556
1557    /// Create a managed-agent configuration on the server.
1558    ///
1559    /// Returns a [`ManagedAgentBuilder`](crate::interactions::managed_agent::ManagedAgentBuilder)
1560    /// that accumulates configuration and saves the agent when
1561    /// [`build_and_save()`](crate::interactions::managed_agent::ManagedAgentBuilder::build_and_save)
1562    /// is called.
1563    ///
1564    /// This is a direct-client capability and is not wired into the
1565    /// `adk-runner` `Agent` trait.
1566    ///
1567    /// # Example
1568    ///
1569    /// ```rust,ignore
1570    /// use adk_gemini::Gemini;
1571    ///
1572    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
1573    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1574    ///
1575    /// let agent = gemini.create_agent()
1576    ///     .id("my-coding-agent")
1577    ///     .base_agent("antigravity-preview-05-2026")
1578    ///     .system_instruction("You are a Rust expert.")
1579    ///     .build_and_save()
1580    ///     .await?;
1581    ///
1582    /// println!("Created agent: {:?}", agent.id);
1583    /// # Ok(())
1584    /// # }
1585    /// ```
1586    #[cfg(feature = "interactions")]
1587    #[instrument(skip_all)]
1588    pub fn create_agent(&self) -> crate::interactions::managed_agent::ManagedAgentBuilder {
1589        crate::interactions::managed_agent::ManagedAgentBuilder::new(self.client.clone())
1590    }
1591
1592    /// List saved managed-agent configurations.
1593    ///
1594    /// Returns a [`ListAgentsResponse`](crate::interactions::managed_agent::ListAgentsResponse)
1595    /// containing the first page of saved agents. Use the `next_page_token`
1596    /// field to fetch subsequent pages by calling the backend directly.
1597    ///
1598    /// This is a direct-client capability and is not wired into the
1599    /// `adk-runner` `Agent` trait.
1600    ///
1601    /// # Example
1602    ///
1603    /// ```rust,ignore
1604    /// use adk_gemini::Gemini;
1605    ///
1606    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
1607    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1608    ///
1609    /// let response = gemini.list_agents().await?;
1610    /// for agent in &response.agents {
1611    ///     println!("Agent: {:?}", agent.id);
1612    /// }
1613    /// # Ok(())
1614    /// # }
1615    /// ```
1616    #[cfg(feature = "interactions")]
1617    #[instrument(skip_all)]
1618    pub async fn list_agents(
1619        &self,
1620    ) -> Result<crate::interactions::managed_agent::ListAgentsResponse, Error> {
1621        self.client.backend.list_agents(None, None).await
1622    }
1623
1624    /// Get a saved managed-agent configuration by ID.
1625    ///
1626    /// Returns the [`SavedAgent`](crate::interactions::managed_agent::SavedAgent)
1627    /// with the given identifier.
1628    ///
1629    /// This is a direct-client capability and is not wired into the
1630    /// `adk-runner` `Agent` trait.
1631    ///
1632    /// # Example
1633    ///
1634    /// ```rust,ignore
1635    /// use adk_gemini::Gemini;
1636    ///
1637    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
1638    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1639    ///
1640    /// let agent = gemini.get_agent("my-coding-agent").await?;
1641    /// println!("Base agent: {:?}", agent.base_agent);
1642    /// # Ok(())
1643    /// # }
1644    /// ```
1645    #[cfg(feature = "interactions")]
1646    #[instrument(skip_all)]
1647    pub async fn get_agent(
1648        &self,
1649        id: &str,
1650    ) -> Result<crate::interactions::managed_agent::SavedAgent, Error> {
1651        self.client.backend.get_agent(id).await
1652    }
1653
1654    /// Delete a saved managed-agent configuration by ID.
1655    ///
1656    /// Removes the agent from the server. This operation is irreversible.
1657    ///
1658    /// This is a direct-client capability and is not wired into the
1659    /// `adk-runner` `Agent` trait.
1660    ///
1661    /// # Example
1662    ///
1663    /// ```rust,ignore
1664    /// use adk_gemini::Gemini;
1665    ///
1666    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
1667    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1668    ///
1669    /// gemini.delete_agent("my-coding-agent").await?;
1670    /// println!("Agent deleted.");
1671    /// # Ok(())
1672    /// # }
1673    /// ```
1674    #[cfg(feature = "interactions")]
1675    #[instrument(skip_all)]
1676    pub async fn delete_agent(&self, id: &str) -> Result<(), Error> {
1677        self.client.backend.delete_agent(id).await
1678    }
1679
1680    /// Download an environment snapshot as a tar archive.
1681    ///
1682    /// Given an environment ID (returned as `environment_id` on an
1683    /// [`Interaction`](crate::interactions::Interaction) response), downloads
1684    /// the full workspace snapshot as raw bytes.
1685    ///
1686    /// This is a direct-client capability and is not wired into the
1687    /// `adk-runner` `Agent` trait.
1688    ///
1689    /// # Example
1690    ///
1691    /// ```rust,ignore
1692    /// use adk_gemini::Gemini;
1693    ///
1694    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
1695    /// let gemini = Gemini::new("YOUR_API_KEY")?;
1696    ///
1697    /// let snapshot = gemini.download_environment("env_abc123").await?;
1698    /// std::fs::write("workspace.tar", &snapshot)?;
1699    /// println!("Downloaded {} bytes", snapshot.len());
1700    /// # Ok(())
1701    /// # }
1702    /// ```
1703    #[cfg(feature = "interactions")]
1704    #[instrument(skip_all)]
1705    pub async fn download_environment(&self, env_id: &str) -> Result<Vec<u8>, Error> {
1706        self.client.backend.download_environment(env_id).await
1707    }
1708}
1709
1710// ══════════════════════════════════════════════════════════════════════
1711// Tests
1712// ══════════════════════════════════════════════════════════════════════
1713
1714#[cfg(test)]
1715#[cfg(feature = "vertex")]
1716mod client_tests {
1717    use super::{Error, GoogleCloudConfig, extract_service_account_project_id};
1718    use crate::backend::vertex::VertexBackend;
1719
1720    #[test]
1721    fn extract_service_account_project_id_reads_project_id() {
1722        let json = r#"{
1723            "type": "service_account",
1724            "project_id": "test-project-123",
1725            "private_key_id": "key-id"
1726        }"#;
1727
1728        let project_id = extract_service_account_project_id(json).expect("project id should parse");
1729        assert_eq!(project_id, "test-project-123");
1730    }
1731
1732    #[test]
1733    fn extract_service_account_project_id_missing_field_errors() {
1734        let json = r#"{
1735            "type": "service_account",
1736            "private_key_id": "key-id"
1737        }"#;
1738
1739        let err =
1740            extract_service_account_project_id(json).expect_err("missing project_id should fail");
1741        assert!(matches!(err, Error::MissingGoogleCloudProjectId));
1742    }
1743
1744    #[test]
1745    fn extract_service_account_project_id_invalid_json_errors() {
1746        let err =
1747            extract_service_account_project_id("not-json").expect_err("invalid json should fail");
1748        assert!(matches!(err, Error::GoogleCloudCredentialParse { .. }));
1749    }
1750
1751    #[test]
1752    fn vertex_transport_error_detection_matches_http2_failure() {
1753        assert!(VertexBackend::is_transport_error(
1754            "the transport reports an error: client error (SendRequest): http2 error"
1755        ));
1756        assert!(!VertexBackend::is_transport_error("permission denied"));
1757    }
1758
1759    #[test]
1760    fn vertex_regional_endpoint_uses_location_prefix() {
1761        let config = GoogleCloudConfig {
1762            project_id: "my-project".to_string(),
1763            location: "us-central1".to_string(),
1764        };
1765        assert_eq!(config.endpoint(), "https://us-central1-aiplatform.googleapis.com");
1766    }
1767
1768    #[test]
1769    fn vertex_global_endpoint_omits_location_prefix() {
1770        let config = GoogleCloudConfig {
1771            project_id: "my-project".to_string(),
1772            location: "global".to_string(),
1773        };
1774        assert_eq!(config.endpoint(), "https://aiplatform.googleapis.com");
1775    }
1776
1777    #[test]
1778    fn vertex_other_regional_endpoint_formats_correctly() {
1779        let config = GoogleCloudConfig {
1780            project_id: "my-project".to_string(),
1781            location: "europe-west4".to_string(),
1782        };
1783        assert_eq!(config.endpoint(), "https://europe-west4-aiplatform.googleapis.com");
1784    }
1785}
adk_gemini/client.rs

adk_gemini/
client.rs
