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    /// Scheme, host and port to send requests to instead of the platform's,
42    /// such as a gateway or a local stand-in (see `Client::base_url`).
43    pub base_url: Option<String>,
44}
45
46/// Concrete Gemini LLM implementation using gemini-live `Client`.
47///
48/// The gemini-live `Client` is created once at construction time and reused for
49/// all `generate()` calls, matching the JS GenAI SDK pattern where a single
50/// `GoogleGenAI` instance is shared across requests.
51pub struct GeminiLlm {
52    model: String,
53    variant: GoogleLlmVariant,
54    /// Stored for constructing the gemini-live `Client` when `gemini-llm` is enabled.
55    #[allow(dead_code)]
56    params: GeminiLlmParams,
57    /// Token provider for VertexAI token refresh.
58    #[allow(dead_code)]
59    token_provider: Arc<dyn TokenProvider>,
60    /// Cached gemini-live Client, created once at construction time.
61    #[cfg(feature = "gemini-llm")]
62    client: gemini_genai_rs::Client,
63}
64
65/// The name the API uses for an enum value (`"MAX_TOKENS"`, not `MaxTokens`).
66#[cfg(feature = "gemini-llm")]
67fn wire_name(value: &impl serde::Serialize) -> String {
68    serde_json::to_value(value)
69        .ok()
70        .and_then(|v| v.as_str().map(str::to_owned))
71        .unwrap_or_default()
72}
73
74/// Keep what the transport knew: the status, whether credentials failed,
75/// whether the provider was reachable at all.
76#[cfg(feature = "gemini-llm")]
77fn llm_error(error: gemini_genai_rs::generate::GenerateError) -> LlmError {
78    use gemini_genai_rs::client::http::HttpError;
79    use gemini_genai_rs::generate::GenerateError;
80
81    match error {
82        GenerateError::Http(HttpError::ApiError {
83            status, message, ..
84        }) => LlmError::Api { status, message },
85        GenerateError::Http(HttpError::Request(e)) => LlmError::Transport(e.to_string()),
86        GenerateError::Http(e @ HttpError::RetriesExhausted { .. }) => {
87            LlmError::Transport(e.to_string())
88        }
89        GenerateError::Http(HttpError::Auth(e)) | GenerateError::Auth(e) => {
90            LlmError::Auth(e.to_string())
91        }
92        GenerateError::SafetyBlocked { reason } | GenerateError::PromptBlocked { reason } => {
93            LlmError::ContentFiltered(wire_name(&reason))
94        }
95        other => LlmError::RequestFailed(other.to_string()),
96    }
97}
98
99static SUPPORTED_PATTERNS: LazyLock<Vec<Regex>> = LazyLock::new(|| {
100    vec![
101        Regex::new(r"^gemini-.*$").unwrap(),
102        Regex::new(r"^projects/.*/endpoints/.*$").unwrap(),
103        Regex::new(r"^projects/.*/models/gemini.*$").unwrap(),
104    ]
105});
106
107impl GeminiLlm {
108    /// Create a new `GeminiLlm` from parameters.
109    ///
110    /// Resolves defaults for model, variant, API key, project, and location
111    /// from parameters first, then falls back to environment variables.
112    /// The gemini-live `Client` is created once here and reused for all calls.
113    pub fn new(mut params: GeminiLlmParams) -> Self {
114        // Resolve variant from params or env
115        let variant = if let Some(true) = params.vertexai {
116            GoogleLlmVariant::VertexAi
117        } else if let Some(false) = params.vertexai {
118            GoogleLlmVariant::GeminiApi
119        } else {
120            get_google_llm_variant()
121        };
122
123        // Resolve model: params, then GEMINI_TEXT_MODEL, then the shared
124        // GEMINI_MODEL, then a per-variant default. Google AI retires dated
125        // names but serves the rolling `gemini-flash-latest` alias; Vertex AI
126        // keeps versioned GA names and does not carry the alias.
127        let model = params
128            .model
129            .clone()
130            .or_else(|| {
131                ["GEMINI_TEXT_MODEL", "GEMINI_MODEL"]
132                    .iter()
133                    .find_map(|k| std::env::var(k).ok())
134                    .filter(|m| !m.trim().is_empty())
135            })
136            .unwrap_or_else(|| match variant {
137                GoogleLlmVariant::GeminiApi => "gemini-flash-latest".to_string(),
138                GoogleLlmVariant::VertexAi => "gemini-2.5-flash".to_string(),
139            });
140        if model.contains("native-audio") || model.contains("-live-") {
141            tracing::warn!(
142                model = %model,
143                "GeminiLlm resolved a Live (bidi) model name for generateContent; set \
144                 GEMINI_TEXT_MODEL (or GeminiLlmParams::model) to a text model — a shared \
145                 GEMINI_MODEL pointing at the Live model 404s here"
146            );
147        }
148
149        // Resolve API key from params or env
150        if params.api_key.is_none() && variant == GoogleLlmVariant::GeminiApi {
151            // Same acceptance chain as the Live connect path, so one exported
152            // variable works for both halves of the stack.
153            params.api_key = std::env::var("GOOGLE_GENAI_API_KEY")
154                .or_else(|_| std::env::var("GEMINI_API_KEY"))
155                .or_else(|_| std::env::var("GOOGLE_API_KEY"))
156                .ok();
157        }
158
159        // Resolve project/location from env for Vertex AI
160        if variant == GoogleLlmVariant::VertexAi {
161            if params.project.is_none() {
162                params.project = std::env::var("GOOGLE_CLOUD_PROJECT").ok();
163            }
164            if params.location.is_none() {
165                params.location = std::env::var("GOOGLE_CLOUD_LOCATION").ok();
166            }
167        }
168
169        // Resolve token provider for VertexAI.
170        // Default to GcloudTokenProvider (env var -> gcloud CLI fallback) for VertexAI,
171        // matching the auth resolution in build_session_config(). For GeminiApi, use
172        // EnvTokenProvider since API key auth doesn't need token refresh.
173        let token_provider: Arc<dyn TokenProvider> =
174            params.token_provider.take().unwrap_or_else(|| {
175                if variant == GoogleLlmVariant::VertexAi {
176                    Arc::new(GcloudTokenProvider::new(std::time::Duration::from_secs(
177                        45 * 60,
178                    )))
179                } else {
180                    Arc::new(EnvTokenProvider)
181                }
182            });
183
184        // Create the gemini-live Client once, reuse across generate() calls.
185        // For VertexAI, use from_vertex_refreshable() so the token is dynamically
186        // refreshed on every REST API call (via auth_headers()), preventing 401
187        // errors from stale tokens during long-running sessions.
188        #[cfg(feature = "gemini-llm")]
189        let client = {
190            use gemini_genai_rs::{Client, prelude::ModelId};
191            let client = match variant {
192                GoogleLlmVariant::GeminiApi => {
193                    let api_key = params.api_key.as_deref().unwrap_or("");
194                    Client::from_api_key(api_key).model(ModelId::new(model.clone()))
195                }
196                GoogleLlmVariant::VertexAi => {
197                    let project = params.project.as_deref().unwrap_or("").to_string();
198                    let location = params
199                        .location
200                        .as_deref()
201                        .unwrap_or("us-central1")
202                        .to_string();
203                    let tp = token_provider.clone();
204                    Client::from_vertex_refreshable(project, location, move || tp.token())
205                        .model(ModelId::new(model.clone()))
206                }
207            };
208            match &params.base_url {
209                Some(base) => client.base_url(base.clone()),
210                None => client,
211            }
212        };
213
214        Self {
215            model,
216            variant,
217            params,
218            token_provider,
219            #[cfg(feature = "gemini-llm")]
220            client,
221        }
222    }
223
224    /// A `GeminiLlm` configured from the environment, checked before any
225    /// request is made.
226    ///
227    /// Reads the same variables as [`new`](Self::new) —
228    /// `GEMINI_API_KEY` (or `GOOGLE_GENAI_API_KEY`, `GOOGLE_API_KEY`) for Google
229    /// AI; `GOOGLE_GENAI_USE_VERTEXAI=true` with `GOOGLE_CLOUD_PROJECT` and
230    /// optionally `GOOGLE_CLOUD_LOCATION` for Vertex AI; `GEMINI_TEXT_MODEL` or
231    /// `GEMINI_MODEL` for the model — and fails with a message naming the
232    /// variable to set when one is missing, or when the model is a Live model
233    /// the text API does not serve. [`new`](Self::new) accepts the same
234    /// configuration and fails on the first request instead.
235    ///
236    /// ```no_run
237    /// use gemini_adk_rs::llm::GeminiLlm;
238    ///
239    /// let llm = GeminiLlm::from_env()?;
240    /// # Ok::<(), gemini_adk_rs::llm::LlmError>(())
241    /// ```
242    pub fn from_env() -> Result<Self, LlmError> {
243        Self::try_new(GeminiLlmParams::default())
244    }
245
246    /// Like [`new`](Self::new), but checks the resolved configuration first;
247    /// see [`from_env`](Self::from_env).
248    pub fn try_new(params: GeminiLlmParams) -> Result<Self, LlmError> {
249        let llm = Self::new(params);
250        llm.check()?;
251        Ok(llm)
252    }
253
254    /// What would make every request fail, found before sending one.
255    fn check(&self) -> Result<(), LlmError> {
256        if super::ModelCapabilities::infer_from_id(&self.model).live_bidi {
257            return Err(LlmError::Config(format!(
258                "`{}` is a Live model, which the text API does not serve. Set \
259                 GEMINI_TEXT_MODEL (or `GeminiLlmParams::model`) to a text model such as \
260                 `gemini-flash-latest`; keep GEMINI_MODEL for the Live session",
261                self.model
262            )));
263        }
264        match self.variant {
265            GoogleLlmVariant::GeminiApi
266                if self
267                    .params
268                    .api_key
269                    .as_deref()
270                    .is_none_or(|k| k.trim().is_empty()) =>
271            {
272                Err(LlmError::Auth(
273                    "no Gemini API key. Set GEMINI_API_KEY (create one at \
274                     https://aistudio.google.com/apikey), or use Vertex AI with \
275                     GOOGLE_GENAI_USE_VERTEXAI=true and GOOGLE_CLOUD_PROJECT"
276                        .into(),
277                ))
278            }
279            GoogleLlmVariant::VertexAi
280                if self
281                    .params
282                    .project
283                    .as_deref()
284                    .is_none_or(|p| p.trim().is_empty()) =>
285            {
286                Err(LlmError::Config(
287                    "Vertex AI needs a project. Set GOOGLE_CLOUD_PROJECT (and optionally \
288                     GOOGLE_CLOUD_LOCATION), or unset GOOGLE_GENAI_USE_VERTEXAI to use a \
289                     Gemini API key"
290                        .into(),
291                ))
292            }
293            _ => Ok(()),
294        }
295    }
296
297    /// Turn a wire response into an [`LlmResponse`], or the error it carries:
298    /// a blocked prompt or a reply withheld by content safety is an error, not
299    /// an empty answer.
300    #[cfg(feature = "gemini-llm")]
301    fn from_generate_response(
302        response: gemini_genai_rs::generate::GenerateContentResponse,
303    ) -> Result<LlmResponse, LlmError> {
304        use gemini_genai_rs::prelude::{Content, FinishReason, Role};
305
306        if let Some(reason) = response
307            .prompt_feedback
308            .as_ref()
309            .and_then(|f| f.block_reason)
310        {
311            return Err(LlmError::ContentFiltered(wire_name(&reason)));
312        }
313        let candidate = response.candidates.into_iter().next();
314        let finish = candidate.as_ref().and_then(|c| c.finish_reason);
315        let content = candidate.and_then(|c| c.content).unwrap_or(Content {
316            role: Some(Role::Model),
317            parts: vec![],
318        });
319        if content.parts.is_empty()
320            && let Some(
321                reason @ (FinishReason::Safety
322                | FinishReason::Recitation
323                | FinishReason::Blocklist
324                | FinishReason::ProhibitedContent
325                | FinishReason::Spii),
326            ) = finish
327        {
328            return Err(LlmError::ContentFiltered(wire_name(&reason)));
329        }
330
331        // Thinking tokens are billed as output, so they count as completion.
332        let usage = response.usage_metadata.map(|u| {
333            let prompt = u.prompt_token_count.unwrap_or(0);
334            let completion = u
335                .response_token_count
336                .unwrap_or(0)
337                .saturating_add(u.thoughts_token_count.unwrap_or(0));
338            TokenUsage {
339                prompt_tokens: prompt,
340                completion_tokens: completion,
341                total_tokens: u
342                    .total_token_count
343                    .unwrap_or(prompt.saturating_add(completion)),
344            }
345        });
346
347        Ok(LlmResponse {
348            content,
349            finish_reason: finish.map(|r| wire_name(&r)),
350            usage,
351        })
352    }
353
354    /// Check if a model name is supported by `GeminiLlm`.
355    pub fn is_supported(model: &str) -> bool {
356        SUPPORTED_PATTERNS.iter().any(|re| re.is_match(model))
357    }
358
359    /// Get the variant (VertexAI vs GeminiApi).
360    pub fn variant(&self) -> GoogleLlmVariant {
361        self.variant
362    }
363
364    /// Map every field of an [`LlmRequest`] onto the wire request, and the
365    /// per-request model override onto the model to call.
366    #[cfg(feature = "gemini-llm")]
367    fn to_generate_config(
368        mut request: LlmRequest,
369    ) -> (
370        gemini_genai_rs::generate::GenerateContentConfig,
371        Option<gemini_genai_rs::prelude::ModelId>,
372    ) {
373        use gemini_genai_rs::generate::GenerateContentConfig;
374        use gemini_genai_rs::prelude::{GenerationConfig, ModelId, ThinkingConfig};
375
376        let mut config = if request.contents.is_empty() {
377            GenerateContentConfig::from_text("")
378        } else {
379            GenerateContentConfig::from_contents(std::mem::take(&mut request.contents))
380        };
381        if let Some(sys) = request.system_instruction.take() {
382            config = config.system_instruction(&sys);
383        }
384        config.tools = std::mem::take(&mut request.tools);
385
386        let generation = GenerationConfig {
387            temperature: request.temperature,
388            max_output_tokens: request.max_output_tokens,
389            top_p: request.top_p,
390            top_k: request.top_k,
391            stop_sequences: (!request.stop_sequences.is_empty())
392                .then(|| std::mem::take(&mut request.stop_sequences)),
393            thinking_config: request.thinking_budget.map(|budget| ThinkingConfig {
394                thinking_budget: Some(budget),
395                ..ThinkingConfig::default()
396            }),
397            response_mime_type: request.response_mime_type.take(),
398            response_json_schema: request.response_json_schema.take(),
399            ..GenerationConfig::default()
400        };
401        config.generation_config =
402            (generation != GenerationConfig::default()).then_some(generation);
403
404        (config, request.model.take().map(ModelId::new))
405    }
406}
407
408#[async_trait]
409impl BaseLlm for GeminiLlm {
410    fn model_id(&self) -> &str {
411        &self.model
412    }
413
414    async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError> {
415        // Feature-gate the actual HTTP call behind gemini-live's generate + http features.
416        #[cfg(feature = "gemini-llm")]
417        {
418            let (config, model) = Self::to_generate_config(request);
419
420            let response = self
421                .client
422                .generate_content_with(config, model.as_ref())
423                .await
424                .map_err(llm_error)?;
425            Self::from_generate_response(response)
426        }
427
428        #[cfg(not(feature = "gemini-llm"))]
429        {
430            // Suppress unused-variable warnings when the feature is disabled.
431            let _ = request;
432            Err(LlmError::RequestFailed(
433                "GeminiLlm requires the 'gemini-llm' feature flag \
434                 (depends on gemini-live HTTP client)"
435                    .into(),
436            ))
437        }
438    }
439
440    #[cfg(feature = "gemini-llm")]
441    async fn generate_stream(&self, request: LlmRequest) -> Result<super::LlmStream, LlmError> {
442        use futures_util::StreamExt;
443
444        let (config, model) = Self::to_generate_config(request);
445        let chunks = self
446            .client
447            .stream_generate_content_with(config, model.as_ref())
448            .await
449            .map_err(llm_error)?;
450        Ok(chunks
451            .map(|chunk| {
452                chunk
453                    .map_err(llm_error)
454                    .and_then(Self::from_generate_response)
455            })
456            .boxed())
457    }
458
459    /// Pre-warm the HTTP connection pool by making a lightweight request.
460    ///
461    /// Establishes the TCP+TLS connection so the first real `generate()`
462    /// call doesn't pay the ~100-300ms handshake penalty. reqwest's
463    /// connection pool keeps it alive for subsequent calls.
464    async fn warm_up(&self) -> Result<(), LlmError> {
465        #[cfg(feature = "gemini-llm")]
466        {
467            use gemini_genai_rs::generate::GenerateContentConfig;
468            let config = GenerateContentConfig::from_text(".").max_output_tokens(1);
469            let _ = self.client.generate_content_with(config, None).await;
470        }
471        Ok(())
472    }
473}
474
475#[cfg(test)]
476mod tests {
477    use super::*;
478
479    /// Every setting an agent can make must reach the wire body; a field that
480    /// is accepted and dropped is the bug this guards against.
481    #[cfg(feature = "gemini-llm")]
482    #[test]
483    fn every_request_field_reaches_the_wire() {
484        let request = LlmRequest {
485            model: Some("gemini-2.5-pro".into()),
486            system_instruction: Some("Be brief.".into()),
487            tools: vec![gemini_genai_rs::prelude::Tool::google_search()],
488            temperature: Some(0.2),
489            max_output_tokens: Some(64),
490            top_p: Some(0.9),
491            top_k: Some(20),
492            stop_sequences: vec!["END".into()],
493            thinking_budget: Some(512),
494            response_mime_type: Some("application/json".into()),
495            response_json_schema: Some(serde_json::json!({ "type": "object" })),
496            ..LlmRequest::from_text("hi")
497        };
498        let (config, model) = GeminiLlm::to_generate_config(request);
499        assert_eq!(
500            model
501                .as_ref()
502                .map(gemini_genai_rs::prelude::ModelId::as_str),
503            Some("gemini-2.5-pro")
504        );
505
506        let body = config.to_request_body();
507        let gc = &body["generationConfig"];
508        assert_eq!(gc["temperature"], 0.2_f32 as f64);
509        assert_eq!(gc["maxOutputTokens"], 64);
510        assert_eq!(gc["topP"], 0.9_f32 as f64);
511        assert_eq!(gc["topK"], 20);
512        assert_eq!(gc["stopSequences"], serde_json::json!(["END"]));
513        assert_eq!(gc["thinkingConfig"]["thinkingBudget"], 512);
514        assert_eq!(gc["responseMimeType"], "application/json");
515        assert_eq!(gc["responseJsonSchema"]["type"], "object");
516        assert!(body["tools"][0].get("googleSearch").is_some(), "{body}");
517        assert!(body["systemInstruction"].is_object(), "{body}");
518    }
519
520    #[cfg(feature = "gemini-llm")]
521    fn wire(json: serde_json::Value) -> Result<LlmResponse, LlmError> {
522        GeminiLlm::from_generate_response(serde_json::from_value(json).unwrap())
523    }
524
525    /// A blocked prompt used to come back as an empty answer.
526    #[cfg(feature = "gemini-llm")]
527    #[test]
528    fn a_blocked_prompt_is_an_error_with_its_reason() {
529        let err =
530            wire(serde_json::json!({ "promptFeedback": { "blockReason": "SAFETY" } })).unwrap_err();
531        assert!(
532            matches!(&err, LlmError::ContentFiltered(r) if r == "SAFETY"),
533            "{err}"
534        );
535    }
536
537    #[cfg(feature = "gemini-llm")]
538    #[test]
539    fn a_withheld_reply_is_an_error_but_a_truncated_one_is_not() {
540        let withheld = wire(serde_json::json!({
541            "candidates": [{ "finishReason": "PROHIBITED_CONTENT" }]
542        }));
543        assert!(matches!(withheld, Err(LlmError::ContentFiltered(r)) if r == "PROHIBITED_CONTENT"));
544
545        let truncated = wire(serde_json::json!({
546            "candidates": [{
547                "content": { "role": "model", "parts": [{ "text": "Once upon" }] },
548                "finishReason": "MAX_TOKENS"
549            }]
550        }))
551        .unwrap();
552        assert_eq!(truncated.text(), "Once upon");
553        assert_eq!(truncated.finish_reason.as_deref(), Some("MAX_TOKENS"));
554    }
555
556    /// `generateContent` names the output count `candidatesTokenCount`; it was
557    /// read as zero. Thinking tokens are billed as output.
558    #[cfg(feature = "gemini-llm")]
559    #[test]
560    fn usage_reads_the_rest_field_names() {
561        let response = wire(serde_json::json!({
562            "candidates": [{ "content": { "role": "model", "parts": [{ "text": "hi" }] } }],
563            "usageMetadata": {
564                "promptTokenCount": 7,
565                "candidatesTokenCount": 3,
566                "thoughtsTokenCount": 5,
567                "totalTokenCount": 15
568            }
569        }))
570        .unwrap();
571        assert_eq!(
572            response.usage,
573            Some(TokenUsage {
574                prompt_tokens: 7,
575                completion_tokens: 8,
576                total_tokens: 15,
577            })
578        );
579    }
580
581    #[cfg(feature = "gemini-llm")]
582    #[test]
583    fn transport_errors_keep_their_status() {
584        use gemini_genai_rs::client::http::HttpError;
585        let err = llm_error(gemini_genai_rs::generate::GenerateError::Http(
586            HttpError::ApiError {
587                status: 429,
588                message: "Resource exhausted".into(),
589                body: None,
590            },
591        ));
592        assert!(err.is_rate_limited(), "{err}");
593        assert!(err.to_string().contains("Resource exhausted"), "{err}");
594    }
595
596    #[test]
597    fn check_names_the_missing_key() {
598        let err = GeminiLlm::try_new(GeminiLlmParams {
599            vertexai: Some(false),
600            api_key: Some("  ".into()),
601            model: Some("gemini-flash-latest".into()),
602            ..Default::default()
603        })
604        .err()
605        .expect("a blank key cannot work");
606        assert!(err.is_auth(), "{err}");
607        assert!(err.to_string().contains("GEMINI_API_KEY"), "{err}");
608    }
609
610    #[test]
611    fn check_rejects_a_live_model_for_text() {
612        let err = GeminiLlm::try_new(GeminiLlmParams {
613            vertexai: Some(false),
614            api_key: Some("key".into()),
615            model: Some("gemini-2.5-flash-native-audio-preview-12-2025".into()),
616            ..Default::default()
617        })
618        .err()
619        .expect("a Live model cannot serve generateContent");
620        assert!(err.to_string().contains("GEMINI_TEXT_MODEL"), "{err}");
621    }
622
623    #[test]
624    fn check_needs_a_vertex_project() {
625        let err = GeminiLlm::try_new(GeminiLlmParams {
626            vertexai: Some(true),
627            project: Some(String::new()),
628            model: Some("gemini-2.5-flash".into()),
629            ..Default::default()
630        })
631        .err()
632        .expect("Vertex AI without a project cannot work");
633        assert!(err.to_string().contains("GOOGLE_CLOUD_PROJECT"), "{err}");
634    }
635
636    #[test]
637    fn check_accepts_a_complete_configuration() {
638        assert!(
639            GeminiLlm::try_new(GeminiLlmParams {
640                vertexai: Some(false),
641                api_key: Some("key".into()),
642                model: Some("gemini-flash-latest".into()),
643                ..Default::default()
644            })
645            .is_ok()
646        );
647    }
648
649    /// A request with no settings sends no `generationConfig` at all.
650    #[cfg(feature = "gemini-llm")]
651    #[test]
652    fn an_unconfigured_request_sends_no_generation_config() {
653        let (config, model) = GeminiLlm::to_generate_config(LlmRequest::from_text("hi"));
654        assert!(model.is_none());
655        assert!(config.to_request_body().get("generationConfig").is_none());
656    }
657
658    #[test]
659    fn default_model_is_the_rolling_flash_alias() {
660        let llm = GeminiLlm::new(GeminiLlmParams {
661            vertexai: Some(false),
662            ..Default::default()
663        });
664        let expected = std::env::var("GEMINI_MODEL")
665            .ok()
666            .filter(|m| !m.trim().is_empty())
667            .unwrap_or_else(|| "gemini-flash-latest".to_string());
668        assert_eq!(llm.model_id(), expected);
669    }
670
671    #[test]
672    fn default_model_on_vertex_is_versioned() {
673        if std::env::var("GEMINI_MODEL").is_ok_and(|m| !m.trim().is_empty()) {
674            return; // env override wins by design; nothing to assert here
675        }
676        let llm = GeminiLlm::new(GeminiLlmParams {
677            vertexai: Some(true),
678            ..Default::default()
679        });
680        assert_eq!(llm.model_id(), "gemini-2.5-flash");
681    }
682
683    #[test]
684    fn explicit_model() {
685        let llm = GeminiLlm::new(GeminiLlmParams {
686            model: Some("gemini-2.0-pro".into()),
687            ..Default::default()
688        });
689        assert_eq!(llm.model_id(), "gemini-2.0-pro");
690    }
691
692    #[test]
693    fn variant_from_params_vertex() {
694        let llm = GeminiLlm::new(GeminiLlmParams {
695            vertexai: Some(true),
696            ..Default::default()
697        });
698        assert_eq!(llm.variant(), GoogleLlmVariant::VertexAi);
699    }
700
701    #[test]
702    fn variant_from_params_gemini_api() {
703        let llm = GeminiLlm::new(GeminiLlmParams {
704            vertexai: Some(false),
705            ..Default::default()
706        });
707        assert_eq!(llm.variant(), GoogleLlmVariant::GeminiApi);
708    }
709
710    #[test]
711    fn is_supported_gemini_models() {
712        assert!(GeminiLlm::is_supported("gemini-2.5-flash"));
713        assert!(GeminiLlm::is_supported("gemini-2.0-pro"));
714        assert!(GeminiLlm::is_supported("gemini-1.5-pro-001"));
715    }
716
717    #[test]
718    fn is_supported_non_gemini_models() {
719        assert!(!GeminiLlm::is_supported("gpt-4"));
720        assert!(!GeminiLlm::is_supported("claude-3-opus"));
721        assert!(!GeminiLlm::is_supported("llama-3"));
722    }
723
724    #[test]
725    fn is_supported_vertex_ai_resource_paths() {
726        assert!(GeminiLlm::is_supported(
727            "projects/my-project/endpoints/12345"
728        ));
729        assert!(GeminiLlm::is_supported(
730            "projects/my-project/models/gemini-2.5-flash"
731        ));
732    }
733
734    #[test]
735    fn model_id_returns_correct_string() {
736        let llm = GeminiLlm::new(GeminiLlmParams {
737            model: Some("gemini-2.5-flash-preview-04-17".into()),
738            ..Default::default()
739        });
740        assert_eq!(llm.model_id(), "gemini-2.5-flash-preview-04-17");
741    }
742
743    #[test]
744    fn base_llm_is_object_safe() {
745        fn _assert_object_safe(_: &dyn BaseLlm) {}
746    }
747
748    #[test]
749    fn gemini_llm_is_send_sync() {
750        fn _assert_send_sync<T: Send + Sync>() {}
751        _assert_send_sync::<GeminiLlm>();
752    }
753}
gemini_adk_rs/llm/gemini.rs

gemini_adk_rs/llm/
gemini.rs
