agent_sdk_providers/impls/

gemini.rs

//! Google Gemini API provider implementation.
//!
//! This module provides an implementation of `LlmProvider` for the Google Gemini
//! API (`generativelanguage.googleapis.com`).

pub(crate) mod data;

use crate::attachments::validate_request_attachments;
use crate::provider::LlmProvider;
use crate::streaming::{StreamBox, StreamDelta, StreamErrorKind};
use agent_sdk_foundation::llm::{ChatOutcome, ChatRequest, ChatResponse, ThinkingConfig};
use anyhow::Result;
use async_trait::async_trait;
use data::{
    ApiContent, ApiFunctionCallingConfig, ApiGenerateContentRequest, ApiGenerateContentResponse,
    ApiGenerationConfig, ApiPart, ApiUsageMetadata, build_api_contents, build_content_blocks,
    convert_tools_to_config, gemini_response_schema, map_finish_reason, map_thinking_config,
};
use reqwest::StatusCode;

const API_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";

/// Connect timeout for the HTTP client (matches Anthropic/Vertex).
const CONNECT_TIMEOUT_SECS: u64 = 30;
/// TCP keepalive interval to keep long streaming connections from dropping.
const TCP_KEEPALIVE_SECS: u64 = 30;
/// Per-request read timeout for the **non-streaming** `chat()` path. Bounds a
/// black-holed endpoint so a single turn cannot hang the agent loop forever.
/// Streaming requests intentionally have no overall timeout.
const CHAT_READ_TIMEOUT_SECS: u64 = 300;

/// Build the shared HTTP client with connect + keepalive timeouts, falling back
/// to a default client (with a logged warning) if the builder fails.
fn build_http_client() -> reqwest::Client {
    reqwest::Client::builder()
        .connect_timeout(std::time::Duration::from_secs(CONNECT_TIMEOUT_SECS))
        .tcp_keepalive(std::time::Duration::from_secs(TCP_KEEPALIVE_SECS))
        .build()
        .unwrap_or_else(|error| {
            log::warn!(
                "failed to build Gemini HTTP client with timeouts ({error}); using default client"
            );
            reqwest::Client::new()
        })
}

// Gemini 3.1 series
pub const MODEL_GEMINI_31_PRO: &str = "gemini-3.1-pro-preview";
pub const MODEL_GEMINI_31_FLASH_LITE: &str = "gemini-3.1-flash-lite-preview";

// Gemini 3 series
pub const MODEL_GEMINI_3_FLASH: &str = "gemini-3-flash-preview";

// Legacy Gemini 3.0 Pro model kept for explicit opt-in.
pub const MODEL_GEMINI_3_PRO: &str = "gemini-3.0-pro";

// Gemini 2.5 series
pub const MODEL_GEMINI_25_FLASH: &str = "gemini-2.5-flash";
pub const MODEL_GEMINI_25_PRO: &str = "gemini-2.5-pro";

// Gemini 2.0 series
pub const MODEL_GEMINI_2_FLASH: &str = "gemini-2.0-flash";
pub const MODEL_GEMINI_2_FLASH_LITE: &str = "gemini-2.0-flash-lite";

/// Google Gemini LLM provider.
#[derive(Clone)]
pub struct GeminiProvider {
    client: reqwest::Client,
    api_key: String,
    model: String,
    base_url: String,
    thinking: Option<ThinkingConfig>,
    /// When true, send the API key via `x-goog-api-key` header instead of a
    /// query parameter. Required when routing through proxies.
    use_header_auth: bool,
    /// Extra headers applied to every request (e.g. for gateway authentication).
    extra_headers: Vec<(String, String)>,
}

impl GeminiProvider {
    /// The conventional environment variable holding the Gemini API key.
    pub const API_KEY_ENV: &'static str = "GEMINI_API_KEY";

    /// Create a new Gemini provider with the specified API key and model.
    #[must_use]
    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            client: build_http_client(),
            api_key: api_key.into(),
            model: model.into(),
            base_url: API_BASE_URL.to_owned(),
            thinking: None,
            use_header_auth: true,
            extra_headers: Vec::new(),
        }
    }

    /// Effective output-token budget for a request.
    ///
    /// Mirrors the Anthropic provider: when the caller did not explicitly set
    /// `max_tokens`, substitute the provider/model default
    /// ([`default_max_tokens`](LlmProvider::default_max_tokens)) instead of
    /// silently capping at `ChatRequest::DEFAULT_MAX_TOKENS`.
    fn effective_max_tokens(&self, request: &ChatRequest) -> u32 {
        if request.max_tokens_explicit {
            request.max_tokens
        } else {
            self.default_max_tokens()
        }
    }

    /// Create a provider using Gemini Flash, reading the API key from the
    /// conventional [`GEMINI_API_KEY`](Self::API_KEY_ENV) environment variable.
    ///
    /// # Panics
    ///
    /// Panics if `GEMINI_API_KEY` is not set. Prefer
    /// [`try_from_env`](Self::try_from_env) outside of examples/tests.
    #[must_use]
    pub fn from_env() -> Self {
        Self::try_from_env().unwrap_or_else(|e| panic!("{e}"))
    }

    /// Create a provider using Gemini Flash, reading the API key from the
    /// conventional [`GEMINI_API_KEY`](Self::API_KEY_ENV) environment variable.
    ///
    /// # Errors
    ///
    /// Returns an error if `GEMINI_API_KEY` is unset or not valid UTF-8.
    pub fn try_from_env() -> Result<Self> {
        let api_key = std::env::var(Self::API_KEY_ENV).map_err(|_| {
            anyhow::anyhow!("environment variable `{}` is not set", Self::API_KEY_ENV)
        })?;
        Ok(Self::flash(api_key))
    }

    /// Create a provider using Gemini 3 Flash Preview (fast and capable, current default).
    #[must_use]
    pub fn flash(api_key: impl Into<String>) -> Self {
        Self::new(api_key, MODEL_GEMINI_3_FLASH)
    }

    /// Create a provider using Gemini 3.1 Flash Lite Preview.
    #[must_use]
    pub fn flash_lite_31(api_key: String) -> Self {
        Self::new(api_key, MODEL_GEMINI_31_FLASH_LITE.to_owned())
    }

    /// Create a provider using Gemini 2.0 Flash Lite (fastest, most cost-effective).
    #[must_use]
    pub fn flash_lite(api_key: String) -> Self {
        Self::new(api_key, MODEL_GEMINI_2_FLASH_LITE.to_owned())
    }

    /// Create a provider using Gemini 3.1 Pro Preview.
    #[must_use]
    pub fn pro_31(api_key: String) -> Self {
        Self::new(api_key, MODEL_GEMINI_31_PRO.to_owned())
    }

    /// Create a provider using Gemini 3.1 Pro Preview (current recommended pro model).
    #[must_use]
    pub fn pro(api_key: String) -> Self {
        Self::new(api_key, MODEL_GEMINI_31_PRO.to_owned())
    }

    /// Set the provider-owned thinking configuration for this model.
    #[must_use]
    pub const fn with_thinking(mut self, thinking: ThinkingConfig) -> Self {
        self.thinking = Some(thinking);
        self
    }

    /// Override the base URL.
    #[must_use]
    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
        self.base_url = base_url.into();
        self
    }

    /// Send the API key via `x-goog-api-key` header instead of `?key=` query
    /// parameter. Required when routing through proxies.
    #[must_use]
    pub const fn with_header_auth(mut self) -> Self {
        self.use_header_auth = true;
        self
    }

    /// Add extra HTTP headers applied to every request.
    #[must_use]
    pub fn with_extra_headers(mut self, headers: Vec<(String, String)>) -> Self {
        self.extra_headers = headers;
        self
    }

    /// Apply auth + extra headers. Skips provider auth when `api_key` is
    /// empty (BYOK gateway mode).
    fn apply_auth(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
        let builder = if self.api_key.is_empty() {
            builder
        } else if self.use_header_auth {
            builder.header("x-goog-api-key", &self.api_key)
        } else {
            builder.query(&[("key", &self.api_key)])
        };
        self.extra_headers
            .iter()
            .fold(builder, |b, (k, v)| b.header(k.as_str(), v.as_str()))
    }
}

#[async_trait]
#[allow(clippy::too_many_lines)]
impl LlmProvider for GeminiProvider {
    async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome> {
        let thinking = match self.resolve_thinking_config(request.thinking.as_ref()) {
            Ok(thinking) => thinking,
            Err(error) => return Ok(ChatOutcome::InvalidRequest(error.to_string())),
        };
        if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
            return Ok(ChatOutcome::InvalidRequest(error.to_string()));
        }
        let contents = build_api_contents(&request.messages);
        let tools = request
            .tools
            .as_ref()
            .map(|t| convert_tools_to_config(t.clone()));
        let tool_config = request
            .tool_choice
            .as_ref()
            .map(ApiFunctionCallingConfig::from_tool_choice);
        let system_instruction = if request.system.is_empty() {
            None
        } else {
            Some(ApiContent {
                role: None,
                parts: vec![ApiPart::Text {
                    text: request.system.clone(),
                    thought_signature: None,
                }],
            })
        };

        let thinking_config = thinking.as_ref().map(map_thinking_config);
        let (response_mime_type, response_schema) =
            request.response_format.as_ref().map_or((None, None), |rf| {
                (
                    Some("application/json"),
                    Some(gemini_response_schema(&rf.schema)),
                )
            });

        let max_tokens = self.effective_max_tokens(&request);
        let api_request = ApiGenerateContentRequest {
            contents: &contents,
            system_instruction: system_instruction.as_ref(),
            tools: tools.as_ref().map(std::slice::from_ref),
            tool_config,
            generation_config: Some(ApiGenerationConfig {
                max_output_tokens: Some(max_tokens),
                thinking_config,
                response_mime_type,
                response_schema,
            }),
            cached_content: request.cached_content.as_deref(),
        };

        log::debug!(
            "Gemini LLM request model={} max_tokens={}",
            self.model,
            max_tokens
        );

        let builder = self
            .client
            .post(format!(
                "{}/models/{}:generateContent",
                self.base_url, self.model
            ))
            .header("Content-Type", "application/json")
            .timeout(std::time::Duration::from_secs(CHAT_READ_TIMEOUT_SECS));
        let response = self
            .apply_auth(builder)
            .json(&api_request)
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("request failed: {e}"))?;

        let status = response.status();
        let bytes = response
            .bytes()
            .await
            .map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?;

        log::debug!(
            "Gemini LLM response status={} body_len={}",
            status,
            bytes.len()
        );

        if status == StatusCode::TOO_MANY_REQUESTS {
            return Ok(ChatOutcome::RateLimited);
        }

        if status.is_server_error() {
            let body = String::from_utf8_lossy(&bytes);
            log::error!("Gemini server error status={status} body={body}");
            return Ok(ChatOutcome::ServerError(body.into_owned()));
        }

        if status.is_client_error() {
            let body = String::from_utf8_lossy(&bytes);
            log::warn!("Gemini client error status={status} body={body}");
            return Ok(ChatOutcome::InvalidRequest(body.into_owned()));
        }

        let api_response: ApiGenerateContentResponse = serde_json::from_slice(&bytes)
            .map_err(|e| anyhow::anyhow!("failed to parse response: {e}"))?;

        let candidate = api_response
            .candidates
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("no candidates in response"))?;

        let content = build_content_blocks(&candidate.content);

        if content.is_empty() && !candidate.content.parts.is_empty() {
            log::warn!(
                "Gemini parts not converted to content blocks raw_parts={:?}",
                candidate.content.parts
            );
        }

        let has_tool_calls = content
            .iter()
            .any(|b| matches!(b, agent_sdk_foundation::llm::ContentBlock::ToolUse { .. }));

        let stop_reason = candidate
            .finish_reason
            .as_ref()
            .map(|r| map_finish_reason(r, has_tool_calls));

        let usage = api_response
            .usage_metadata
            .unwrap_or(ApiUsageMetadata {
                prompt: 0,
                candidates: 0,
                cached_content: 0,
            })
            .into_usage();

        Ok(ChatOutcome::Success(ChatResponse {
            id: String::new(),
            content,
            model: self.model.clone(),
            stop_reason,
            usage,
        }))
    }

    fn chat_stream(&self, request: ChatRequest) -> StreamBox<'_> {
        Box::pin(async_stream::stream! {
            let thinking = match self.resolve_thinking_config(request.thinking.as_ref()) {
                Ok(thinking) => thinking,
                Err(error) => {
                    yield Ok(StreamDelta::Error {
                        message: error.to_string(),
                        kind: StreamErrorKind::InvalidRequest,
                    });
                    return;
                }
            };
            if let Err(error) = validate_request_attachments(self.provider(), self.model(), &request) {
                yield Ok(StreamDelta::Error {
                    message: error.to_string(),
                    kind: StreamErrorKind::InvalidRequest,
                });
                return;
            }
            let contents = build_api_contents(&request.messages);
            let tools = request
            .tools
            .as_ref()
            .map(|t| convert_tools_to_config(t.clone()));
            let tool_config = request
                .tool_choice
                .as_ref()
                .map(ApiFunctionCallingConfig::from_tool_choice);
            let system_instruction = if request.system.is_empty() {
                None
            } else {
                Some(ApiContent {
                    role: None,
                    parts: vec![ApiPart::Text {
                        text: request.system.clone(),
                        thought_signature: None,
                    }],
                })
            };

            let thinking_config = thinking.as_ref().map(map_thinking_config);
            let (response_mime_type, response_schema) = request
                .response_format
                .as_ref()
                .map_or((None, None), |rf| {
                    (
                        Some("application/json"),
                        Some(gemini_response_schema(&rf.schema)),
                    )
                });

            let max_tokens = self.effective_max_tokens(&request);
            let api_request = ApiGenerateContentRequest {
                contents: &contents,
                system_instruction: system_instruction.as_ref(),
                tools: tools.as_ref().map(std::slice::from_ref),
                tool_config,
                generation_config: Some(ApiGenerationConfig {
                    max_output_tokens: Some(max_tokens),
                    thinking_config,
                    response_mime_type,
                    response_schema,
                }),
                cached_content: request.cached_content.as_deref(),
            };

            log::debug!(
                "Gemini streaming LLM request model={} max_tokens={}",
                self.model,
                max_tokens
            );

            let stream_builder = self
                .client
                .post(format!(
                    "{}/models/{}:streamGenerateContent",
                    self.base_url, self.model
                ))
                .header("Content-Type", "application/json")
                .query(&[("alt", "sse")]);
            let response = match self
                .apply_auth(stream_builder)
                .json(&api_request)
                .send()
                .await
            {
                Ok(r) => r,
                Err(e) => {
                    // Include the cause so 401 detection / diagnostics survive.
                    yield Err(anyhow::anyhow!("request failed: {e}"));
                    return;
                }
            };

            let status = response.status();
            if !status.is_success() {
                let body = response.text().await.unwrap_or_default();
                let kind = if status == StatusCode::TOO_MANY_REQUESTS {
                    StreamErrorKind::RateLimited
                } else if status.is_server_error() {
                    StreamErrorKind::ServerError
                } else {
                    StreamErrorKind::InvalidRequest
                };
                log::warn!("Gemini error status={status} body={body}");
                yield Ok(StreamDelta::Error {
                    message: body,
                    kind,
                });
                return;
            }

            let mut inner = data::stream_gemini_response(response);
            while let Some(item) = futures::StreamExt::next(&mut inner).await {
                yield item;
            }
        })
    }

    fn model(&self) -> &str {
        &self.model
    }

    fn provider(&self) -> &'static str {
        "gemini"
    }

    fn configured_thinking(&self) -> Option<&ThinkingConfig> {
        self.thinking.as_ref()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_new_creates_provider_with_custom_model() {
        let provider = GeminiProvider::new("test-api-key".to_string(), "custom-model".to_string());

        assert_eq!(provider.model(), "custom-model");
        assert_eq!(provider.provider(), "gemini");
    }

    #[test]
    fn test_flash_factory_creates_flash_provider() {
        let provider = GeminiProvider::flash("test-api-key".to_string());

        assert_eq!(provider.model(), MODEL_GEMINI_3_FLASH);
        assert_eq!(provider.provider(), "gemini");
    }

    #[test]
    fn test_flash_lite_factory_creates_flash_lite_provider() {
        let provider = GeminiProvider::flash_lite("test-api-key".to_string());

        assert_eq!(provider.model(), MODEL_GEMINI_2_FLASH_LITE);
        assert_eq!(provider.provider(), "gemini");
    }

    #[test]
    fn test_flash_lite_31_factory_creates_flash_lite_provider() {
        let provider = GeminiProvider::flash_lite_31("test-api-key".to_string());

        assert_eq!(provider.model(), MODEL_GEMINI_31_FLASH_LITE);
        assert_eq!(provider.provider(), "gemini");
    }

    #[test]
    fn test_pro_factory_creates_pro_provider() {
        let provider = GeminiProvider::pro("test-api-key".to_string());

        assert_eq!(provider.model(), MODEL_GEMINI_31_PRO);
        assert_eq!(provider.provider(), "gemini");
    }

    #[test]
    fn test_pro_31_factory_creates_pro_provider() {
        let provider = GeminiProvider::pro_31("test-api-key".to_string());

        assert_eq!(provider.model(), MODEL_GEMINI_31_PRO);
        assert_eq!(provider.provider(), "gemini");
    }

    #[test]
    fn test_model_constants_have_expected_values() {
        assert_eq!(MODEL_GEMINI_31_PRO, "gemini-3.1-pro-preview");
        assert_eq!(MODEL_GEMINI_31_FLASH_LITE, "gemini-3.1-flash-lite-preview");
        assert_eq!(MODEL_GEMINI_3_FLASH, "gemini-3-flash-preview");
        assert_eq!(MODEL_GEMINI_3_PRO, "gemini-3.0-pro");
        assert_eq!(MODEL_GEMINI_25_FLASH, "gemini-2.5-flash");
        assert_eq!(MODEL_GEMINI_25_PRO, "gemini-2.5-pro");
        assert_eq!(MODEL_GEMINI_2_FLASH, "gemini-2.0-flash");
        assert_eq!(MODEL_GEMINI_2_FLASH_LITE, "gemini-2.0-flash-lite");
    }

    #[test]
    fn test_gemini_20_models_reject_thinking() {
        let provider = GeminiProvider::flash_lite("test-api-key".to_string());
        let error = provider
            .validate_thinking_config(Some(&ThinkingConfig::new(10_000)))
            .unwrap_err();
        assert!(error.to_string().contains("thinking is not supported"));
    }

    #[test]
    fn test_default_uses_header_auth() {
        let provider = GeminiProvider::new("test-key".to_string(), "model".to_string());
        assert!(
            provider.use_header_auth,
            "Default should use header auth for security"
        );
    }

    #[test]
    fn test_provider_is_cloneable() {
        let provider = GeminiProvider::new("test-api-key".to_string(), "test-model".to_string());
        let cloned = provider.clone();

        assert_eq!(provider.model(), cloned.model());
        assert_eq!(provider.provider(), cloned.provider());
    }

    fn request_with_max_tokens(max_tokens: u32, explicit: bool) -> ChatRequest {
        ChatRequest {
            system: String::new(),
            messages: vec![agent_sdk_foundation::llm::Message::user("hi")],
            tools: None,
            max_tokens,
            max_tokens_explicit: explicit,
            session_id: None,
            cached_content: None,
            thinking: None,
            tool_choice: None,
            response_format: None,
        }
    }

    #[test]
    fn test_effective_max_tokens_honors_explicit_budget() {
        let provider = GeminiProvider::pro("test-api-key".to_string());
        let request = request_with_max_tokens(123, true);
        assert_eq!(provider.effective_max_tokens(&request), 123);
    }

    #[test]
    fn test_effective_max_tokens_uses_default_when_implicit() {
        // An implicit budget must fall back to the provider/model default, not
        // be silently capped at ChatRequest::DEFAULT_MAX_TOKENS.
        let provider = GeminiProvider::pro("test-api-key".to_string());
        let request = request_with_max_tokens(4096, false);
        assert_eq!(
            provider.effective_max_tokens(&request),
            provider.default_max_tokens()
        );
    }
}