gemini_adk_rs/llm/
gemini.rs

1//! Concrete Gemini LLM implementation using gemini-live `Client`.
2//!
3//! The [`GeminiLlm`] struct is always available for type references and registry
4//! wiring. Actual HTTP generation requires the `gemini-llm` feature flag, which
5//! pulls in `gemini-live/http` and `gemini-live/generate`.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use regex::Regex;
12use std::sync::LazyLock;
13
14#[cfg(feature = "gemini-llm")]
15use crate::llm::TokenUsage;
16use crate::llm::{
17    BaseLlm, EnvTokenProvider, GcloudTokenProvider, LlmError, LlmRequest, LlmResponse,
18    TokenProvider,
19};
20use crate::utils::variant::{GoogleLlmVariant, get_google_llm_variant};
21
22/// Parameters for constructing a [`GeminiLlm`].
23#[derive(Default)]
24pub struct GeminiLlmParams {
25    /// Model name. Defaults to the `GEMINI_MODEL` env var if set, else
26    /// `gemini-flash-latest` on Google AI (a rolling alias the catalog
27    /// keeps serving) or `gemini-2.5-flash` on Vertex AI.
28    pub model: Option<String>,
29    /// API key for Gemini API (non-Vertex).
30    pub api_key: Option<String>,
31    /// Whether to use Vertex AI backend.
32    pub vertexai: Option<bool>,
33    /// Google Cloud project ID (Vertex AI only).
34    pub project: Option<String>,
35    /// Google Cloud region (Vertex AI only, defaults to "us-central1").
36    pub location: Option<String>,
37    /// Custom HTTP headers for requests.
38    pub headers: Option<HashMap<String, String>>,
39    /// Custom token provider for VertexAI. Defaults to reading `GOOGLE_ACCESS_TOKEN` env var.
40    pub token_provider: Option<Arc<dyn TokenProvider>>,
41}
42
43/// Concrete Gemini LLM implementation using gemini-live `Client`.
44///
45/// The gemini-live `Client` is created once at construction time and reused for
46/// all `generate()` calls, matching the JS GenAI SDK pattern where a single
47/// `GoogleGenAI` instance is shared across requests.
48pub struct GeminiLlm {
49    model: String,
50    variant: GoogleLlmVariant,
51    /// Stored for constructing the gemini-live `Client` when `gemini-llm` is enabled.
52    #[allow(dead_code)]
53    params: GeminiLlmParams,
54    /// Token provider for VertexAI token refresh.
55    #[allow(dead_code)]
56    token_provider: Arc<dyn TokenProvider>,
57    /// Cached gemini-live Client, created once at construction time.
58    #[cfg(feature = "gemini-llm")]
59    client: gemini_genai_rs::Client,
60}
61
62static SUPPORTED_PATTERNS: LazyLock<Vec<Regex>> = LazyLock::new(|| {
63    vec![
64        Regex::new(r"^gemini-.*$").unwrap(),
65        Regex::new(r"^projects/.*/endpoints/.*$").unwrap(),
66        Regex::new(r"^projects/.*/models/gemini.*$").unwrap(),
67    ]
68});
69
70impl GeminiLlm {
71    /// Create a new `GeminiLlm` from parameters.
72    ///
73    /// Resolves defaults for model, variant, API key, project, and location
74    /// from parameters first, then falls back to environment variables.
75    /// The gemini-live `Client` is created once here and reused for all calls.
76    pub fn new(mut params: GeminiLlmParams) -> Self {
77        // Resolve variant from params or env
78        let variant = if let Some(true) = params.vertexai {
79            GoogleLlmVariant::VertexAi
80        } else if let Some(false) = params.vertexai {
81            GoogleLlmVariant::GeminiApi
82        } else {
83            get_google_llm_variant()
84        };
85
86        // Resolve model: params, then GEMINI_TEXT_MODEL, then the shared
87        // GEMINI_MODEL, then a per-variant default. Google AI retires dated
88        // names but serves the rolling `gemini-flash-latest` alias; Vertex AI
89        // keeps versioned GA names and does not carry the alias.
90        let model = params
91            .model
92            .clone()
93            .or_else(|| {
94                ["GEMINI_TEXT_MODEL", "GEMINI_MODEL"]
95                    .iter()
96                    .find_map(|k| std::env::var(k).ok())
97                    .filter(|m| !m.trim().is_empty())
98            })
99            .unwrap_or_else(|| match variant {
100                GoogleLlmVariant::GeminiApi => "gemini-flash-latest".to_string(),
101                GoogleLlmVariant::VertexAi => "gemini-2.5-flash".to_string(),
102            });
103        if model.contains("native-audio") || model.contains("-live-") {
104            tracing::warn!(
105                model = %model,
106                "GeminiLlm resolved a Live (bidi) model name for generateContent; set \
107                 GEMINI_TEXT_MODEL (or GeminiLlmParams::model) to a text model — a shared \
108                 GEMINI_MODEL pointing at the Live model 404s here"
109            );
110        }
111
112        // Resolve API key from params or env
113        if params.api_key.is_none() && variant == GoogleLlmVariant::GeminiApi {
114            // Same acceptance chain as the Live connect path, so one exported
115            // variable works for both halves of the stack.
116            params.api_key = std::env::var("GOOGLE_GENAI_API_KEY")
117                .or_else(|_| std::env::var("GEMINI_API_KEY"))
118                .or_else(|_| std::env::var("GOOGLE_API_KEY"))
119                .ok();
120        }
121
122        // Resolve project/location from env for Vertex AI
123        if variant == GoogleLlmVariant::VertexAi {
124            if params.project.is_none() {
125                params.project = std::env::var("GOOGLE_CLOUD_PROJECT").ok();
126            }
127            if params.location.is_none() {
128                params.location = std::env::var("GOOGLE_CLOUD_LOCATION").ok();
129            }
130        }
131
132        // Resolve token provider for VertexAI.
133        // Default to GcloudTokenProvider (env var -> gcloud CLI fallback) for VertexAI,
134        // matching the auth resolution in build_session_config(). For GeminiApi, use
135        // EnvTokenProvider since API key auth doesn't need token refresh.
136        let token_provider: Arc<dyn TokenProvider> =
137            params.token_provider.take().unwrap_or_else(|| {
138                if variant == GoogleLlmVariant::VertexAi {
139                    Arc::new(GcloudTokenProvider::new(std::time::Duration::from_secs(
140                        45 * 60,
141                    )))
142                } else {
143                    Arc::new(EnvTokenProvider)
144                }
145            });
146
147        // Create the gemini-live Client once, reuse across generate() calls.
148        // For VertexAI, use from_vertex_refreshable() so the token is dynamically
149        // refreshed on every REST API call (via auth_headers()), preventing 401
150        // errors from stale tokens during long-running sessions.
151        #[cfg(feature = "gemini-llm")]
152        let client = {
153            use gemini_genai_rs::{Client, prelude::ModelId};
154            match variant {
155                GoogleLlmVariant::GeminiApi => {
156                    let api_key = params.api_key.as_deref().unwrap_or("");
157                    Client::from_api_key(api_key).model(ModelId::new(model.clone()))
158                }
159                GoogleLlmVariant::VertexAi => {
160                    let project = params.project.as_deref().unwrap_or("").to_string();
161                    let location = params
162                        .location
163                        .as_deref()
164                        .unwrap_or("us-central1")
165                        .to_string();
166                    let tp = token_provider.clone();
167                    Client::from_vertex_refreshable(project, location, move || tp.token())
168                        .model(ModelId::new(model.clone()))
169                }
170            }
171        };
172
173        Self {
174            model,
175            variant,
176            params,
177            token_provider,
178            #[cfg(feature = "gemini-llm")]
179            client,
180        }
181    }
182
183    /// Check if a model name is supported by `GeminiLlm`.
184    pub fn is_supported(model: &str) -> bool {
185        SUPPORTED_PATTERNS.iter().any(|re| re.is_match(model))
186    }
187
188    /// Get the variant (VertexAI vs GeminiApi).
189    pub fn variant(&self) -> GoogleLlmVariant {
190        self.variant
191    }
192
193    /// Preprocess request: remove labels and displayName for non-Vertex (Gemini API).
194    fn preprocess_request(_request: &mut LlmRequest) {
195        // For Gemini API backend: remove labels and displayName from tools.
196        // This is a no-op for now since LlmRequest doesn't have those fields yet.
197        // In a full implementation, this would strip Vertex-only fields.
198    }
199}
200
201#[async_trait]
202impl BaseLlm for GeminiLlm {
203    fn model_id(&self) -> &str {
204        &self.model
205    }
206
207    async fn generate(&self, mut request: LlmRequest) -> Result<LlmResponse, LlmError> {
208        Self::preprocess_request(&mut request);
209
210        // Feature-gate the actual HTTP call behind gemini-live's generate + http features.
211        #[cfg(feature = "gemini-llm")]
212        {
213            use gemini_genai_rs::generate::GenerateContentConfig;
214            use gemini_genai_rs::prelude::*;
215
216            // Build GenerateContentConfig from LlmRequest — move, don't clone.
217            let mut config = if request.contents.is_empty() {
218                GenerateContentConfig::from_text("")
219            } else {
220                GenerateContentConfig::from_contents(std::mem::take(&mut request.contents))
221            };
222
223            if let Some(sys) = request.system_instruction.take() {
224                config = config.system_instruction(&sys);
225            }
226            if !request.tools.is_empty() {
227                config.tools = std::mem::take(&mut request.tools);
228            }
229            if let Some(temp) = request.temperature {
230                config = config.temperature(temp);
231            }
232            if let Some(max) = request.max_output_tokens {
233                config = config.max_output_tokens(max);
234            }
235            if request.response_mime_type.is_some() || request.response_json_schema.is_some() {
236                let gc = config
237                    .generation_config
238                    .get_or_insert_with(gemini_genai_rs::prelude::GenerationConfig::default);
239                if let Some(mime) = request.response_mime_type.take() {
240                    gc.response_mime_type = Some(mime);
241                }
242                if let Some(schema) = request.response_json_schema.take() {
243                    gc.response_json_schema = Some(schema);
244                }
245            }
246
247            let response = self
248                .client
249                .generate_content_with(config, None)
250                .await
251                .map_err(|e| LlmError::RequestFailed(e.to_string()))?;
252
253            let content = response
254                .candidates
255                .first()
256                .and_then(|c| c.content.clone())
257                .unwrap_or_else(|| Content {
258                    role: Some(Role::Model),
259                    parts: vec![],
260                });
261
262            let finish_reason = response
263                .candidates
264                .first()
265                .and_then(|c| c.finish_reason)
266                .map(|r| format!("{r:?}"));
267
268            let usage = response.usage_metadata.map(|u| TokenUsage {
269                prompt_tokens: u.prompt_token_count.unwrap_or(0),
270                completion_tokens: u.response_token_count.unwrap_or(0),
271                total_tokens: u.total_token_count.unwrap_or(0),
272            });
273
274            Ok(LlmResponse {
275                content,
276                finish_reason,
277                usage,
278            })
279        }
280
281        #[cfg(not(feature = "gemini-llm"))]
282        {
283            // Suppress unused-variable warnings when the feature is disabled.
284            let _ = request;
285            Err(LlmError::RequestFailed(
286                "GeminiLlm requires the 'gemini-llm' feature flag \
287                 (depends on gemini-live HTTP client)"
288                    .into(),
289            ))
290        }
291    }
292
293    /// Pre-warm the HTTP connection pool by making a lightweight request.
294    ///
295    /// Establishes the TCP+TLS connection so the first real `generate()`
296    /// call doesn't pay the ~100-300ms handshake penalty. reqwest's
297    /// connection pool keeps it alive for subsequent calls.
298    async fn warm_up(&self) -> Result<(), LlmError> {
299        #[cfg(feature = "gemini-llm")]
300        {
301            use gemini_genai_rs::generate::GenerateContentConfig;
302            let config = GenerateContentConfig::from_text(".").max_output_tokens(1);
303            let _ = self.client.generate_content_with(config, None).await;
304        }
305        Ok(())
306    }
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312
313    #[test]
314    fn default_model_is_the_rolling_flash_alias() {
315        let llm = GeminiLlm::new(GeminiLlmParams {
316            vertexai: Some(false),
317            ..Default::default()
318        });
319        let expected = std::env::var("GEMINI_MODEL")
320            .ok()
321            .filter(|m| !m.trim().is_empty())
322            .unwrap_or_else(|| "gemini-flash-latest".to_string());
323        assert_eq!(llm.model_id(), expected);
324    }
325
326    #[test]
327    fn default_model_on_vertex_is_versioned() {
328        if std::env::var("GEMINI_MODEL").is_ok_and(|m| !m.trim().is_empty()) {
329            return; // env override wins by design; nothing to assert here
330        }
331        let llm = GeminiLlm::new(GeminiLlmParams {
332            vertexai: Some(true),
333            ..Default::default()
334        });
335        assert_eq!(llm.model_id(), "gemini-2.5-flash");
336    }
337
338    #[test]
339    fn explicit_model() {
340        let llm = GeminiLlm::new(GeminiLlmParams {
341            model: Some("gemini-2.0-pro".into()),
342            ..Default::default()
343        });
344        assert_eq!(llm.model_id(), "gemini-2.0-pro");
345    }
346
347    #[test]
348    fn variant_from_params_vertex() {
349        let llm = GeminiLlm::new(GeminiLlmParams {
350            vertexai: Some(true),
351            ..Default::default()
352        });
353        assert_eq!(llm.variant(), GoogleLlmVariant::VertexAi);
354    }
355
356    #[test]
357    fn variant_from_params_gemini_api() {
358        let llm = GeminiLlm::new(GeminiLlmParams {
359            vertexai: Some(false),
360            ..Default::default()
361        });
362        assert_eq!(llm.variant(), GoogleLlmVariant::GeminiApi);
363    }
364
365    #[test]
366    fn is_supported_gemini_models() {
367        assert!(GeminiLlm::is_supported("gemini-2.5-flash"));
368        assert!(GeminiLlm::is_supported("gemini-2.0-pro"));
369        assert!(GeminiLlm::is_supported("gemini-1.5-pro-001"));
370    }
371
372    #[test]
373    fn is_supported_non_gemini_models() {
374        assert!(!GeminiLlm::is_supported("gpt-4"));
375        assert!(!GeminiLlm::is_supported("claude-3-opus"));
376        assert!(!GeminiLlm::is_supported("llama-3"));
377    }
378
379    #[test]
380    fn is_supported_vertex_ai_resource_paths() {
381        assert!(GeminiLlm::is_supported(
382            "projects/my-project/endpoints/12345"
383        ));
384        assert!(GeminiLlm::is_supported(
385            "projects/my-project/models/gemini-2.5-flash"
386        ));
387    }
388
389    #[test]
390    fn model_id_returns_correct_string() {
391        let llm = GeminiLlm::new(GeminiLlmParams {
392            model: Some("gemini-2.5-flash-preview-04-17".into()),
393            ..Default::default()
394        });
395        assert_eq!(llm.model_id(), "gemini-2.5-flash-preview-04-17");
396    }
397
398    #[test]
399    fn base_llm_is_object_safe() {
400        fn _assert_object_safe(_: &dyn BaseLlm) {}
401    }
402
403    #[test]
404    fn gemini_llm_is_send_sync() {
405        fn _assert_send_sync<T: Send + Sync>() {}
406        _assert_send_sync::<GeminiLlm>();
407    }
408}
gemini_adk_rs/llm/gemini.rs

gemini_adk_rs/llm/
gemini.rs
