litellm_rust/providers/

gemini.rs

use crate::config::ProviderConfig;
use crate::error::{LiteLLMError, Result};
use crate::http::send_json;
use crate::providers::{merge_extra_body, resolve_api_key};
use crate::stream::{parse_gemini_sse_stream, ChatStream};
use crate::types::{
    ChatContentPart, ChatContentPartFile, ChatContentPartImageUrl, ChatContentPartInputAudio,
    ChatContentPartText, ChatFile, ChatImageUrl, ChatInputAudio, ChatMessage, ChatMessageContent,
    ChatRequest, ChatResponse, EmbeddingRequest, EmbeddingResponse, ImageData, ImageEditRequest,
    ImageInputData, ImageRequest, ImageResponse, Reasoning, Usage, VideoRequest, VideoResponse,
};
use base64::{engine::general_purpose, Engine as _};
use mime_guess::MimeGuess;
use reqwest::header::CONTENT_TYPE;
use reqwest::Client;
use serde_json::{Map, Value};
use std::env;
use tokio::time::{sleep, Duration};

/// Default maximum polling attempts for video generation (240 * 5s = 20 minutes)
pub const DEFAULT_VIDEO_MAX_POLL_ATTEMPTS: u32 = 240;
/// Default polling interval for video generation status checks
pub const DEFAULT_VIDEO_POLL_INTERVAL_SECS: u64 = 5;

pub async fn chat(client: &Client, cfg: &ProviderConfig, req: ChatRequest) -> Result<ChatResponse> {
    let base = cfg
        .base_url
        .clone()
        .ok_or_else(|| LiteLLMError::Config("base_url required".into()))?;
    let key = resolve_api_key(cfg)?
        .ok_or_else(|| LiteLLMError::MissingApiKey("GEMINI_API_KEY".into()))?;
    let model = req.model.trim_start_matches("models/").to_string();
    let url = format!(
        "{}/models/{}:generateContent",
        base.trim_end_matches('/'),
        model
    );
    let body = build_chat_body(client, req, &model).await?;

    let mut builder = client.post(url).header("x-goog-api-key", key).json(&body);
    for (k, v) in &cfg.extra_headers {
        builder = builder.header(k, v);
    }

    let (resp, _headers) = send_json::<Value>(builder).await?;
    let (content, reasoning) = split_text_and_reasoning(&resp);
    let usage = parse_usage(&resp);
    let debug = env::var("LITELLM_GEMINI_DEBUG").ok().as_deref() == Some("1")
        || env::var("YOETZ_GEMINI_DEBUG").ok().as_deref() == Some("1");
    if debug {
        if let Ok(pretty) = serde_json::to_string_pretty(&resp) {
            eprintln!("litellm-rust gemini raw response:\n{pretty}");
        }
    }

    let tool_calls = extract_gemini_tool_calls(&resp);
    Ok(ChatResponse {
        content,
        reasoning,
        tool_calls,
        usage,
        response_id: None,
        header_cost: None,
        raw: if debug { Some(resp) } else { None },
    })
}

/// Stream a chat completion through `streamGenerateContent?alt=sse`.
pub async fn chat_stream(
    client: &Client,
    cfg: &ProviderConfig,
    req: ChatRequest,
) -> Result<ChatStream> {
    let base = cfg
        .base_url
        .clone()
        .ok_or_else(|| LiteLLMError::Config("base_url required".into()))?;
    let key = resolve_api_key(cfg)?
        .ok_or_else(|| LiteLLMError::MissingApiKey("GEMINI_API_KEY".into()))?;
    let model = req.model.trim_start_matches("models/").to_string();
    let url = format!(
        "{}/models/{}:streamGenerateContent",
        base.trim_end_matches('/'),
        model
    );
    let body = build_chat_body(client, req, &model).await?;

    let mut builder = client
        .post(url)
        .query(&[("alt", "sse")])
        .header("x-goog-api-key", key)
        .header("accept", "text/event-stream")
        .json(&body);
    for (k, v) in &cfg.extra_headers {
        builder = builder.header(k, v);
    }

    let resp = builder.send().await.map_err(LiteLLMError::from)?;
    let status = resp.status();
    if !status.is_success() {
        let text = resp.text().await.map_err(LiteLLMError::from)?;
        return Err(LiteLLMError::http(format!(
            "http {}: {}",
            status.as_u16(),
            text
        )));
    }

    Ok(parse_gemini_sse_stream(resp.bytes_stream()))
}

/// Build the `generateContent` / `streamGenerateContent` request body. Both
/// endpoints take the same `GenerateContentRequest` shape.
async fn build_chat_body(client: &Client, req: ChatRequest, model: &str) -> Result<Value> {
    let generation_config = build_generation_config(&req, model)?;

    let mut messages = req.messages;
    let system_instruction = extract_system_instruction(&mut messages)?;
    if messages.is_empty() && !system_instruction.is_empty() {
        messages.push(ChatMessage {
            role: "user".to_string(),
            content: ChatMessageContent::Text(" ".to_string()),
            name: None,
            tool_call_id: None,
            tool_calls: None,
            function_call: None,
            provider_specific_fields: None,
        });
    }

    let contents = gemini_contents_from_messages(client, messages, model).await?;

    let mut body = serde_json::json!({ "contents": contents });
    if !system_instruction.is_empty() {
        body["system_instruction"] = serde_json::json!({ "parts": system_instruction });
    }
    if !generation_config.is_empty() {
        body["generationConfig"] = Value::Object(generation_config);
    }
    merge_extra_body(
        &mut body,
        req.extra_body.as_ref(),
        &["contents", "system_instruction", "generationConfig"],
    )?;
    Ok(body)
}

fn build_generation_config(req: &ChatRequest, model: &str) -> Result<Map<String, Value>> {
    let mut generation_config = Map::new();
    if let Some(temp) = req.temperature {
        generation_config.insert("temperature".to_string(), serde_json::json!(temp));
    }
    if let Some(max_tokens) = req.max_tokens {
        generation_config.insert("maxOutputTokens".to_string(), serde_json::json!(max_tokens));
    }
    if let Some(thinking_config) =
        build_thinking_config(req.thinking.as_ref(), req.reasoning_effort.as_ref(), model)?
    {
        generation_config.insert("thinkingConfig".to_string(), thinking_config);
    }
    Ok(generation_config)
}

/// Map `thinking` / `reasoning_effort` onto Gemini's
/// `generationConfig.thinkingConfig` (`includeThoughts`, `thinkingBudget`,
/// `thinkingLevel`; see <https://ai.google.dev/api/generate-content#ThinkingConfig>).
///
/// - `thinking` as an object with a `type` key is the LiteLLM/Anthropic shape
///   (`{"type": "enabled", "budget_tokens": 1024}`). `type` sets
///   `includeThoughts`; `budget_tokens` becomes `thinkingBudget` on models
///   before Gemini 3 only, as in Python LiteLLM.
/// - `thinking` as an object without `type` is Gemini-native (for example
///   `{"thinkingBudget": 1024}`) and is sent as is.
/// - `thinking` as a bool sets `includeThoughts`; as a number it sets
///   `thinkingBudget`.
/// - `reasoning_effort` is used only when `thinking` did not already set a
///   budget or a level, because Gemini rejects both in one request.
///   - Gemini 3 and newer take a `thinkingLevel`: `low`, `medium`, `high`
///     become `LOW`, `MEDIUM`, `HIGH`; `minimal`, `none` and `disable` become
///     `LOW`. Gemini 3 cannot turn thinking off, and not every Gemini 3 model
///     accepts `MINIMAL`, so `LOW` is the lowest level that all of them take.
///   - Older models reject `thinkingLevel`, so the effort becomes a
///     `thinkingBudget` with Python LiteLLM's token values (`none` /
///     `disable` 0, `minimal` 128, `low` 1024, `medium` 2048, `high` 4096),
///     raised to the model's minimum: `gemini-2.5-pro` cannot turn thinking
///     off, so `none` / `disable` send 128; Flash-Lite models accept 0 or
///     512 and more, so `minimal` sends 512.
///   - Unless `thinking` set it, `includeThoughts` is `true`, or `false` for
///     `none` / `disable`. Other effort values are a config error.
fn build_thinking_config(
    thinking: Option<&Value>,
    reasoning_effort: Option<&Value>,
    model: &str,
) -> Result<Option<Value>> {
    let pre_gemini_3 = is_pre_gemini_3(model);
    let mut cfg = Map::new();
    match thinking {
        Some(Value::Object(obj)) if obj.contains_key("type") => {
            let budget = obj.get("budget_tokens");
            let include_thoughts = match obj.get("type").and_then(Value::as_str) {
                Some("enabled") => budget.and_then(Value::as_i64) != Some(0),
                Some("disabled") => false,
                _ => {
                    return Err(LiteLLMError::Config(
                        "thinking.type must be \"enabled\" or \"disabled\"".into(),
                    ))
                }
            };
            cfg.insert("includeThoughts".to_string(), include_thoughts.into());
            if let (true, Some(budget)) = (pre_gemini_3, budget) {
                let budget = budget.as_i64().ok_or_else(|| {
                    LiteLLMError::Config("thinking.budget_tokens must be an integer".into())
                })?;
                cfg.insert("thinkingBudget".to_string(), budget.into());
            }
        }
        Some(Value::Object(obj)) => cfg.extend(obj.clone()),
        Some(Value::Bool(b)) => {
            cfg.insert("includeThoughts".to_string(), Value::Bool(*b));
        }
        Some(Value::Number(n)) => {
            cfg.insert("thinkingBudget".to_string(), Value::Number(n.clone()));
        }
        _ => {}
    }
    let has_level_or_budget =
        cfg.contains_key("thinkingLevel") || cfg.contains_key("thinkingBudget");
    if let (false, Some(effort)) = (has_level_or_budget, reasoning_effort) {
        let effort = effort.as_str().ok_or_else(|| {
            LiteLLMError::Config("gemini reasoning_effort must be a string".into())
        })?;
        let (key, value) = if pre_gemini_3 {
            let budget = effort_to_thinking_budget(effort, model)
                .ok_or_else(|| unsupported_reasoning_effort(effort))?;
            ("thinkingBudget", Value::from(budget))
        } else {
            let level = effort_to_thinking_level(effort)
                .ok_or_else(|| unsupported_reasoning_effort(effort))?;
            ("thinkingLevel", Value::from(level))
        };
        cfg.insert(key.to_string(), value);
        let include_thoughts = !matches!(effort, "none" | "disable");
        cfg.entry("includeThoughts")
            .or_insert(Value::Bool(include_thoughts));
    }
    Ok((!cfg.is_empty()).then_some(Value::Object(cfg)))
}

fn unsupported_reasoning_effort(effort: &str) -> LiteLLMError {
    LiteLLMError::Config(format!(
        "unsupported gemini reasoning_effort `{effort}`; expected none, disable, minimal, low, medium, or high"
    ))
}

/// Token budgets used by Python LiteLLM for pre-Gemini-3 `reasoning_effort`,
/// raised to the minimum the model accepts.
fn effort_to_thinking_budget(effort: &str, model: &str) -> Option<u32> {
    let name = model.rsplit('/').next().unwrap_or(model).to_lowercase();
    let budget = match effort {
        "none" | "disable" if name.contains("gemini-2.5-pro") => 128,
        "none" | "disable" => 0,
        "minimal" if name.contains("flash-lite") => 512,
        "minimal" => 128,
        "low" => 1024,
        "medium" => 2048,
        "high" => 4096,
        _ => return None,
    };
    Some(budget)
}

/// `ThinkingLevel` enum values for Gemini 3 and newer. `LOW` is the lowest
/// level every Gemini 3 model accepts.
fn effort_to_thinking_level(effort: &str) -> Option<&'static str> {
    match effort {
        "none" | "disable" | "minimal" | "low" => Some("LOW"),
        "medium" => Some("MEDIUM"),
        "high" => Some("HIGH"),
        _ => None,
    }
}

/// True for every model that must not get `thinkingLevel`: `gemini-1.x`,
/// `gemini-2.x`, `gemini-exp*`, the old `gemini-pro` / `gemini-flash` names,
/// and names without the `gemini-` prefix (Gemma, tuned models). The
/// `*-latest` aliases point at newer models. This is Python LiteLLM's
/// `_is_gemini_3_or_newer`, inverted.
fn is_pre_gemini_3(model: &str) -> bool {
    let name = model.rsplit('/').next().unwrap_or(model).to_lowercase();
    let Some(rest) = name.strip_prefix("gemini-") else {
        return true;
    };
    let head = rest.split('-').next().unwrap_or_default();
    match head {
        "exp" => true,
        "pro" | "flash" => !matches!(
            rest,
            "pro-latest" | "flash-latest" | "pro-lite-latest" | "flash-lite-latest"
        ),
        _ => {
            let (major, minor) = head.split_once('.').unwrap_or((head, "0"));
            matches!(major, "1" | "2")
                && !minor.is_empty()
                && minor.chars().all(|c| c.is_ascii_digit())
        }
    }
}

/// Embed text through `embedContent` (one input) or `batchEmbedContents`
/// (an array of inputs).
pub async fn embeddings(
    client: &Client,
    cfg: &ProviderConfig,
    req: EmbeddingRequest,
) -> Result<EmbeddingResponse> {
    let base = cfg
        .base_url
        .clone()
        .ok_or_else(|| LiteLLMError::Config("base_url required".into()))?;
    let key = resolve_api_key(cfg)?
        .ok_or_else(|| LiteLLMError::MissingApiKey("GEMINI_API_KEY".into()))?;
    let model = req.model.trim_start_matches("models/").to_string();
    let inputs = embedding_inputs(&req.input)?;
    if inputs.is_empty() {
        return Err(LiteLLMError::Config(
            "embeddings input must not be empty".into(),
        ));
    }
    let content = |text: &str| serde_json::json!({ "parts": [{ "text": text }] });
    let extra_body = req.extra_body.as_ref();

    // `extra_body` holds per-request `EmbedContentRequest` fields such as
    // `taskType` or `outputDimensionality`. A batch carries them on each item.
    let single = inputs.len() == 1;
    let (method, body) = if single {
        let mut body = serde_json::json!({ "content": content(&inputs[0]) });
        merge_extra_body(&mut body, extra_body, &["model", "content"])?;
        ("embedContent", body)
    } else {
        let requests = inputs
            .iter()
            .map(|text| {
                let mut item = serde_json::json!({
                    "model": format!("models/{model}"),
                    "content": content(text),
                });
                merge_extra_body(&mut item, extra_body, &["model", "content"])?;
                Ok(item)
            })
            .collect::<Result<Vec<Value>>>()?;
        (
            "batchEmbedContents",
            serde_json::json!({ "requests": requests }),
        )
    };
    let url = format!("{}/models/{}:{}", base.trim_end_matches('/'), model, method);

    let mut builder = client.post(url).header("x-goog-api-key", key).json(&body);
    for (k, v) in &cfg.extra_headers {
        builder = builder.header(k, v);
    }

    let (resp, _headers) = send_json::<Value>(builder).await?;
    let vectors = if single {
        vec![embedding_values(resp.get("embedding"))?]
    } else {
        resp.get("embeddings")
            .and_then(Value::as_array)
            .ok_or_else(|| LiteLLMError::Parse("gemini batch response has no embeddings".into()))?
            .iter()
            .map(|item| embedding_values(Some(item)))
            .collect::<Result<Vec<_>>>()?
    };
    if vectors.len() != inputs.len() {
        return Err(LiteLLMError::Parse(format!(
            "gemini returned {} embeddings for {} inputs",
            vectors.len(),
            inputs.len()
        )));
    }
    // Embedding usage only reports `promptTokenCount`.
    let usage = Usage {
        prompt_tokens: resp
            .pointer("/usageMetadata/promptTokenCount")
            .and_then(Value::as_u64),
        ..Usage::default()
    };

    Ok(EmbeddingResponse {
        vectors,
        usage,
        raw: None,
    })
}

/// Read `values` from a Gemini `ContentEmbedding`.
fn embedding_values(embedding: Option<&Value>) -> Result<Vec<f32>> {
    embedding
        .and_then(|e| e.get("values"))
        .and_then(Value::as_array)
        .ok_or_else(|| LiteLLMError::Parse("gemini embedding has no values array".into()))?
        .iter()
        .map(|v| {
            v.as_f64()
                .map(|f| f as f32)
                .ok_or_else(|| LiteLLMError::Parse("gemini embedding value is not a number".into()))
        })
        .collect()
}

/// Accepts the OpenAI-style embeddings `input`: one string or an array of
/// strings. Gemini embedding requests here are text only.
fn embedding_inputs(input: &Value) -> Result<Vec<String>> {
    match input {
        Value::String(s) => Ok(vec![s.clone()]),
        Value::Array(items) => items
            .iter()
            .map(|item| {
                item.as_str().map(str::to_string).ok_or_else(|| {
                    LiteLLMError::Config(
                        "gemini embeddings input array must contain strings".into(),
                    )
                })
            })
            .collect(),
        _ => Err(LiteLLMError::Config(
            "gemini embeddings input must be a string or array of strings".into(),
        )),
    }
}

/// Map Gemini `functionCall` parts to the OpenAI tool_calls shape. Gemini emits
/// `candidates[0].content.parts[].functionCall = { name, args }`; the OpenAI
/// shape is `{ id, type: "function", function: { name, arguments } }` where
/// `arguments` is a JSON string. Gemini doesn't give per-call ids, so we
/// synthesize `call_<idx>` to keep IDs stable within a response.
fn extract_gemini_tool_calls(resp: &serde_json::Value) -> Option<serde_json::Value> {
    let parts = resp
        .get("candidates")?
        .as_array()?
        .first()?
        .get("content")?
        .get("parts")?
        .as_array()?;
    let mut calls: Vec<serde_json::Value> = Vec::new();
    for (idx, part) in parts.iter().enumerate() {
        let Some(fc) = part.get("functionCall") else {
            continue;
        };
        let name = fc.get("name").cloned().unwrap_or(serde_json::Value::Null);
        let args = fc
            .get("args")
            .cloned()
            .unwrap_or_else(|| serde_json::json!({}));
        let arguments = serde_json::to_string(&args).unwrap_or_else(|_| "{}".into());
        calls.push(serde_json::json!({
            "id": format!("call_{idx}"),
            "type": "function",
            "function": { "name": name, "arguments": arguments },
        }));
    }
    if calls.is_empty() {
        None
    } else {
        Some(serde_json::Value::Array(calls))
    }
}

/// Video generation options for configurable timeouts.
#[derive(Debug, Clone)]
pub struct VideoGenerationOptions {
    /// Maximum number of polling attempts
    pub max_poll_attempts: u32,
    /// Interval between polling attempts in seconds
    pub poll_interval_secs: u64,
}

impl Default for VideoGenerationOptions {
    fn default() -> Self {
        Self {
            max_poll_attempts: DEFAULT_VIDEO_MAX_POLL_ATTEMPTS,
            poll_interval_secs: DEFAULT_VIDEO_POLL_INTERVAL_SECS,
        }
    }
}

pub async fn video_generation(
    client: &Client,
    cfg: &ProviderConfig,
    req: VideoRequest,
) -> Result<VideoResponse> {
    video_generation_with_options(client, cfg, req, VideoGenerationOptions::default()).await
}

pub async fn video_generation_with_options(
    client: &Client,
    cfg: &ProviderConfig,
    req: VideoRequest,
    options: VideoGenerationOptions,
) -> Result<VideoResponse> {
    let base = cfg
        .base_url
        .clone()
        .ok_or_else(|| LiteLLMError::Config("base_url required".into()))?;
    let key = resolve_api_key(cfg)?
        .ok_or_else(|| LiteLLMError::MissingApiKey("GEMINI_API_KEY".into()))?;
    let model = req.model.trim_start_matches("models/");

    let url = format!(
        "{}/models/{}:predictLongRunning",
        base.trim_end_matches('/'),
        model
    );

    let mut parameters = serde_json::Map::new();
    if let Some(seconds) = req.seconds {
        parameters.insert("durationSeconds".to_string(), serde_json::json!(seconds));
    }
    if let Some(size) = req.size {
        parameters.insert("resolution".to_string(), serde_json::json!(size));
    }

    let body = serde_json::json!({
        "instances": [{ "prompt": req.prompt }],
        "parameters": Value::Object(parameters),
    });

    let mut builder = client.post(url).header("x-goog-api-key", &key).json(&body);
    for (k, v) in &cfg.extra_headers {
        builder = builder.header(k, v);
    }

    let (resp, _headers) = send_json::<Value>(builder).await?;
    let op_name = resp
        .get("name")
        .and_then(|v| v.as_str())
        .ok_or_else(|| LiteLLMError::Parse("missing operation name".into()))?;

    let op_url = if op_name.starts_with("http") {
        op_name.to_string()
    } else {
        format!("{}/{}", base.trim_end_matches('/'), op_name)
    };

    let poll_interval = Duration::from_secs(options.poll_interval_secs);

    for attempt in 0..options.max_poll_attempts {
        let mut poll = client.get(&op_url).header("x-goog-api-key", &key);
        for (k, v) in &cfg.extra_headers {
            poll = poll.header(k, v);
        }
        let (op_resp, _headers) = send_json::<Value>(poll).await?;

        if op_resp.get("done").and_then(|v| v.as_bool()) == Some(true) {
            if op_resp.get("error").is_some() {
                return Err(LiteLLMError::http("video generation failed"));
            }

            let uri = extract_video_uri(&op_resp)
                .ok_or_else(|| LiteLLMError::Parse("missing video uri".into()))?;

            return Ok(VideoResponse {
                video_url: Some(uri),
                raw: None,
            });
        }

        if attempt + 1 >= options.max_poll_attempts {
            return Err(LiteLLMError::http(format!(
                "video generation timed out after {} attempts",
                options.max_poll_attempts
            )));
        }
        sleep(poll_interval).await;
    }

    Err(LiteLLMError::http("video generation timed out"))
}

pub async fn image_generation(
    client: &Client,
    cfg: &ProviderConfig,
    req: ImageRequest,
) -> Result<ImageResponse> {
    let base = cfg
        .base_url
        .clone()
        .ok_or_else(|| LiteLLMError::Config("base_url required".into()))?;
    let key = resolve_api_key(cfg)?
        .ok_or_else(|| LiteLLMError::MissingApiKey("GEMINI_API_KEY".into()))?;
    let model = req.model.trim_start_matches("models/");
    let n = req.n.unwrap_or(1);

    if model.contains("gemini") {
        // Path A: Gemini native image generation via generateContent
        let url = format!(
            "{}/models/{}:generateContent",
            base.trim_end_matches('/'),
            model
        );

        let mut body = serde_json::json!({
            "contents": [{ "parts": [{ "text": req.prompt }] }],
            "generationConfig": { "response_modalities": ["IMAGE", "TEXT"] }
        });
        merge_extra_body(
            &mut body,
            req.extra_body.as_ref(),
            &["contents", "generationConfig"],
        )?;

        let mut builder = client.post(url).header("x-goog-api-key", &key).json(&body);
        for (k, v) in &cfg.extra_headers {
            builder = builder.header(k, v);
        }

        let (resp, _headers) = send_json::<Value>(builder).await?;
        let usage = parse_usage(&resp);
        let images = extract_images_from_candidates(&resp);

        Ok(ImageResponse {
            images,
            usage,
            raw: Some(resp),
        })
    } else {
        // Path B: Imagen models via predict endpoint
        let url = format!("{}/models/{}:predict", base.trim_end_matches('/'), model);

        let mut body = serde_json::json!({
            "instances": [{ "prompt": req.prompt }],
            "parameters": { "sampleCount": n }
        });
        merge_extra_body(
            &mut body,
            req.extra_body.as_ref(),
            &["instances", "parameters"],
        )?;

        let mut builder = client.post(url).header("x-goog-api-key", &key).json(&body);
        for (k, v) in &cfg.extra_headers {
            builder = builder.header(k, v);
        }

        let (resp, _headers) = send_json::<Value>(builder).await?;
        let images = extract_images_from_predictions(&resp);

        Ok(ImageResponse {
            images,
            usage: Usage::default(),
            raw: Some(resp),
        })
    }
}

pub async fn image_editing(
    client: &Client,
    cfg: &ProviderConfig,
    req: ImageEditRequest,
) -> Result<ImageResponse> {
    let base = cfg
        .base_url
        .clone()
        .ok_or_else(|| LiteLLMError::Config("base_url required".into()))?;
    let key = resolve_api_key(cfg)?
        .ok_or_else(|| LiteLLMError::MissingApiKey("GEMINI_API_KEY".into()))?;
    let model = req.model.trim_start_matches("models/");

    let url = format!(
        "{}/models/{}:generateContent",
        base.trim_end_matches('/'),
        model
    );

    let mut parts = Vec::with_capacity(req.images.len() + 1);
    parts.push(serde_json::json!({ "text": req.prompt }));
    for image in &req.images {
        parts.push(image_input_to_part(client, image).await?);
    }

    let body = serde_json::json!({
        "contents": [{ "role": "user", "parts": parts }],
        "generationConfig": {
            "response_modalities": ["IMAGE", "TEXT"]
        }
    });

    let mut builder = client.post(url).header("x-goog-api-key", &key).json(&body);
    for (k, v) in &cfg.extra_headers {
        builder = builder.header(k, v);
    }

    let (resp, _headers) = send_json::<Value>(builder).await?;
    let usage = parse_usage(&resp);
    let images = extract_images_from_candidates(&resp);

    Ok(ImageResponse {
        images,
        usage,
        raw: Some(resp),
    })
}

async fn image_input_to_part(_client: &Client, input: &ImageInputData) -> Result<Value> {
    if let Some(b64) = &input.b64_json {
        let mime = input
            .mime_type
            .as_deref()
            .unwrap_or("image/png")
            .to_string();
        return Ok(serde_json::json!({
            "inline_data": { "mime_type": mime, "data": b64 }
        }));
    }
    if let Some(url) = &input.url {
        return process_gemini_media_url(url, input.mime_type.as_deref(), None, None, "");
    }
    Err(LiteLLMError::Config(
        "ImageInputData requires b64_json or url".into(),
    ))
}

fn extract_images_from_candidates(resp: &Value) -> Vec<ImageData> {
    let mut images = Vec::new();
    if let Some(candidates) = resp.get("candidates").and_then(|v| v.as_array()) {
        for candidate in candidates {
            if let Some(parts) = candidate
                .get("content")
                .and_then(|v| v.get("parts"))
                .and_then(|v| v.as_array())
            {
                for part in parts {
                    if let Some(data) = part
                        .get("inlineData")
                        .and_then(|v| v.get("data"))
                        .and_then(|v| v.as_str())
                    {
                        let mime = part
                            .get("inlineData")
                            .and_then(|v| v.get("mimeType"))
                            .and_then(|v| v.as_str())
                            .map(|s| s.to_string());
                        images.push(ImageData {
                            b64_json: Some(data.to_string()),
                            url: None,
                            revised_prompt: None,
                            mime_type: mime,
                        });
                    }
                }
            }
        }
    }
    images
}

fn extract_images_from_predictions(resp: &Value) -> Vec<ImageData> {
    let mut images = Vec::new();
    if let Some(predictions) = resp.get("predictions").and_then(|v| v.as_array()) {
        for prediction in predictions {
            if let Some(data) = prediction
                .get("bytesBase64Encoded")
                .and_then(|v| v.as_str())
            {
                let mime = prediction
                    .get("mimeType")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string());
                images.push(ImageData {
                    b64_json: Some(data.to_string()),
                    url: None,
                    revised_prompt: None,
                    mime_type: mime,
                });
            }
        }
    }
    images
}

/// Split the first candidate's text parts into visible content and
/// `"thought": true` parts (returned when `includeThoughts` is set) into
/// reasoning. Used for both `generateContent` responses and each
/// `streamGenerateContent` event.
pub(crate) fn split_text_and_reasoning(resp: &Value) -> (String, Option<Reasoning>) {
    let mut content = String::new();
    let mut thoughts = String::new();
    let parts = resp
        .pointer("/candidates/0/content/parts")
        .and_then(Value::as_array);
    for part in parts.into_iter().flatten() {
        let Some(text) = part.get("text").and_then(Value::as_str) else {
            continue;
        };
        if part.get("thought").and_then(Value::as_bool) == Some(true) {
            thoughts.push_str(text);
        } else {
            content.push_str(text);
        }
    }
    (content, Reasoning::from_text(thoughts))
}

pub(crate) fn parse_usage(resp: &Value) -> Usage {
    if let Some(meta) = resp.get("usageMetadata").and_then(|v| v.as_object()) {
        return Usage {
            prompt_tokens: meta.get("promptTokenCount").and_then(|v| v.as_u64()),
            completion_tokens: meta.get("candidatesTokenCount").and_then(|v| v.as_u64()),
            thoughts_tokens: meta.get("thoughtsTokenCount").and_then(|v| v.as_u64()),
            total_tokens: meta.get("totalTokenCount").and_then(|v| v.as_u64()),
            cost_usd: None,
            ..Default::default()
        };
    }
    Usage::default()
}

fn extract_system_instruction(messages: &mut Vec<ChatMessage>) -> Result<Vec<Value>> {
    let mut parts = Vec::new();
    let mut indices = Vec::new();
    for (idx, msg) in messages.iter().enumerate() {
        if msg.role == "system" {
            match &msg.content {
                ChatMessageContent::Text(text) => {
                    if !text.is_empty() {
                        parts.push(serde_json::json!({ "text": text }));
                    }
                }
                ChatMessageContent::Parts(content_parts) => {
                    let mut system_text = String::new();
                    for part in content_parts {
                        if let ChatContentPart::Text(ChatContentPartText { text, .. }) = part {
                            system_text.push_str(text);
                        }
                    }
                    if !system_text.is_empty() {
                        parts.push(serde_json::json!({ "text": system_text }));
                    }
                }
            }
            indices.push(idx);
        }
    }
    for idx in indices.into_iter().rev() {
        messages.remove(idx);
    }
    Ok(parts)
}

async fn gemini_contents_from_messages(
    client: &Client,
    messages: Vec<ChatMessage>,
    model: &str,
) -> Result<Vec<Value>> {
    let mut contents: Vec<Value> = Vec::new();
    let mut msg_i = 0;
    let mut tool_call_responses: Vec<Value> = Vec::new();
    let mut last_tool_calls: Vec<ToolCallInfo> = Vec::new();

    while msg_i < messages.len() {
        let role = messages[msg_i].role.as_str();
        if role == "user" {
            let mut parts: Vec<Value> = Vec::new();
            while msg_i < messages.len() && messages[msg_i].role == "user" {
                parts.extend(
                    gemini_parts_from_content(client, &messages[msg_i].content, model).await?,
                );
                msg_i += 1;
            }
            if !parts.is_empty() {
                contents.push(serde_json::json!({ "role": "user", "parts": parts }));
            }
            continue;
        }

        if role == "assistant" {
            let mut parts: Vec<Value> = Vec::new();
            while msg_i < messages.len() && messages[msg_i].role == "assistant" {
                let message = &messages[msg_i];
                parts.extend(gemini_parts_from_content(client, &message.content, model).await?);
                let (tool_parts, tool_infos) = tool_call_parts_from_message(message)?;
                if !tool_parts.is_empty() {
                    parts.extend(tool_parts);
                    if !tool_infos.is_empty() {
                        last_tool_calls = tool_infos;
                    }
                }
                msg_i += 1;
            }
            if !parts.is_empty() {
                contents.push(serde_json::json!({ "role": "model", "parts": parts }));
            }
            continue;
        }

        if role == "tool" || role == "function" {
            while msg_i < messages.len()
                && (messages[msg_i].role == "tool" || messages[msg_i].role == "function")
            {
                let response_parts =
                    tool_response_parts(client, &messages[msg_i], &last_tool_calls).await?;
                tool_call_responses.extend(response_parts);
                msg_i += 1;
            }
            if !tool_call_responses.is_empty() {
                contents.push(serde_json::json!({ "parts": tool_call_responses }));
                tool_call_responses = Vec::new();
            }
            continue;
        }

        let parts = gemini_parts_from_content(client, &messages[msg_i].content, model).await?;
        if !parts.is_empty() {
            contents.push(serde_json::json!({ "role": "user", "parts": parts }));
        }
        msg_i += 1;
    }

    if !tool_call_responses.is_empty() {
        contents.push(serde_json::json!({ "parts": tool_call_responses }));
    }
    if contents.is_empty() {
        contents.push(serde_json::json!({
            "role": "user",
            "parts": [{ "text": " " }]
        }));
    }
    Ok(contents)
}

async fn gemini_parts_from_content(
    _client: &Client,
    content: &ChatMessageContent,
    model: &str,
) -> Result<Vec<Value>> {
    match content {
        ChatMessageContent::Text(text) => Ok(vec![serde_json::json!({ "text": text })]),
        ChatMessageContent::Parts(parts) => {
            let mut out = Vec::new();
            for part in parts {
                match part {
                    ChatContentPart::Text(ChatContentPartText { text, .. }) => {
                        if !text.is_empty() {
                            out.push(serde_json::json!({ "text": text }));
                        }
                    }
                    ChatContentPart::ImageUrl(ChatContentPartImageUrl { image_url, .. }) => {
                        let detail = match image_url {
                            ChatImageUrl::Object(obj) => obj.detail.as_deref(),
                            ChatImageUrl::Url(_) => None,
                        };
                        out.push(process_gemini_media(image_url, detail, None, None, model)?);
                    }
                    ChatContentPart::InputAudio(ChatContentPartInputAudio {
                        input_audio, ..
                    }) => {
                        out.push(process_gemini_audio(input_audio)?);
                    }
                    ChatContentPart::File(ChatContentPartFile { file, .. }) => {
                        out.push(process_gemini_file(file, model)?);
                    }
                    ChatContentPart::Other(value) => {
                        return Err(LiteLLMError::Config(format!(
                            "unsupported gemini content part: {}",
                            value
                        )));
                    }
                }
            }
            Ok(out)
        }
    }
}

#[derive(Debug, Clone)]
struct ToolCallInfo {
    id: Option<String>,
    name: String,
}

fn tool_call_parts_from_message(message: &ChatMessage) -> Result<(Vec<Value>, Vec<ToolCallInfo>)> {
    let mut parts = Vec::new();
    let mut infos = Vec::new();

    if let Some(tool_calls) = &message.tool_calls {
        if let Some(array) = tool_calls.as_array() {
            for tool in array {
                let function = tool
                    .get("function")
                    .ok_or_else(|| LiteLLMError::Config("tool_call missing function".into()))?;
                let name = function
                    .get("name")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| LiteLLMError::Config("tool_call missing name".into()))?;
                let args_raw = function
                    .get("arguments")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                let args = if args_raw.trim().is_empty() {
                    Value::Object(serde_json::Map::new())
                } else {
                    serde_json::from_str(args_raw)
                        .map_err(|e| LiteLLMError::Parse(e.to_string()))?
                };
                parts.push(serde_json::json!({
                    "function_call": { "name": name, "args": args }
                }));
                let id = tool
                    .get("id")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string());
                infos.push(ToolCallInfo {
                    id,
                    name: name.to_string(),
                });
            }
            return Ok((parts, infos));
        }
    }

    if let Some(function_call) = &message.function_call {
        let name = function_call
            .get("name")
            .and_then(|v| v.as_str())
            .ok_or_else(|| LiteLLMError::Config("function_call missing name".into()))?;
        let args_raw = function_call
            .get("arguments")
            .and_then(|v| v.as_str())
            .unwrap_or("");
        let args = if args_raw.trim().is_empty() {
            Value::Object(serde_json::Map::new())
        } else {
            serde_json::from_str(args_raw).map_err(|e| LiteLLMError::Parse(e.to_string()))?
        };
        parts.push(serde_json::json!({
            "function_call": { "name": name, "args": args }
        }));
        infos.push(ToolCallInfo {
            id: None,
            name: name.to_string(),
        });
    }

    Ok((parts, infos))
}

async fn tool_response_parts(
    client: &Client,
    message: &ChatMessage,
    last_tool_calls: &[ToolCallInfo],
) -> Result<Vec<Value>> {
    let name = resolve_tool_name(message, last_tool_calls)?;
    let (content_str, mut inline_parts) = extract_tool_content(client, &message.content).await?;
    let response = parse_tool_response_data(&content_str);
    let function_part = serde_json::json!({
        "function_response": { "name": name, "response": response }
    });
    let mut parts = vec![function_part];
    parts.append(&mut inline_parts);
    Ok(parts)
}

fn resolve_tool_name(message: &ChatMessage, last_tool_calls: &[ToolCallInfo]) -> Result<String> {
    if let Some(name) = &message.name {
        return Ok(name.clone());
    }
    if let Some(tool_call_id) = &message.tool_call_id {
        if let Some(info) = last_tool_calls
            .iter()
            .find(|info| info.id.as_deref() == Some(tool_call_id.as_str()))
        {
            return Ok(info.name.clone());
        }
    }
    Err(LiteLLMError::Config(
        "missing tool name for tool response".into(),
    ))
}

async fn extract_tool_content(
    client: &Client,
    content: &ChatMessageContent,
) -> Result<(String, Vec<Value>)> {
    let mut text = String::new();
    let mut inline_parts = Vec::new();
    match content {
        ChatMessageContent::Text(t) => {
            text.push_str(t);
        }
        ChatMessageContent::Parts(parts) => {
            for part in parts {
                match part {
                    ChatContentPart::Text(ChatContentPartText { text: t, .. }) => {
                        text.push_str(t);
                    }
                    ChatContentPart::ImageUrl(ChatContentPartImageUrl { image_url, .. }) => {
                        inline_parts.push(inline_data_from_image_url(client, image_url).await?);
                    }
                    ChatContentPart::InputAudio(ChatContentPartInputAudio {
                        input_audio, ..
                    }) => {
                        inline_parts.push(inline_data_from_audio(input_audio)?);
                    }
                    ChatContentPart::File(ChatContentPartFile { file, .. }) => {
                        inline_parts.push(inline_data_from_file(client, file).await?);
                    }
                    ChatContentPart::Other(_) => {}
                }
            }
        }
    }
    Ok((text, inline_parts))
}

fn parse_tool_response_data(content: &str) -> Value {
    let trimmed = content.trim();
    if trimmed.starts_with('{') || trimmed.starts_with('[') {
        if let Ok(value) = serde_json::from_str(trimmed) {
            return value;
        }
    }
    serde_json::json!({ "content": content })
}

async fn inline_data_from_image_url(client: &Client, image_url: &ChatImageUrl) -> Result<Value> {
    let (url, format) = match image_url {
        ChatImageUrl::Url(url) => (url.as_str(), None),
        ChatImageUrl::Object(obj) => (obj.url.as_str(), obj.format.as_deref()),
    };
    inline_data_from_url(client, url, format).await
}

fn inline_data_from_audio(input_audio: &ChatInputAudio) -> Result<Value> {
    let format = if input_audio.format.starts_with("audio/") {
        input_audio.format.clone()
    } else {
        format!("audio/{}", input_audio.format)
    };
    Ok(serde_json::json!({
        "inline_data": { "mime_type": format, "data": input_audio.data }
    }))
}

async fn inline_data_from_file(client: &Client, file: &ChatFile) -> Result<Value> {
    let passed = file
        .file_id
        .as_ref()
        .or(file.file_data.as_ref())
        .ok_or_else(|| LiteLLMError::Config("file_id or file_data required".into()))?;
    inline_data_from_url(client, passed, file.format.as_deref()).await
}

async fn inline_data_from_url(client: &Client, url: &str, format: Option<&str>) -> Result<Value> {
    if let Some((mime, data)) = parse_data_url(url, format) {
        return Ok(serde_json::json!({
            "inline_data": { "mime_type": mime, "data": data }
        }));
    }
    if url.starts_with("http://") || url.starts_with("https://") {
        let (mime, data) = fetch_bytes_with_mime(client, url, format).await?;
        return Ok(serde_json::json!({
            "inline_data": { "mime_type": mime, "data": data }
        }));
    }
    Err(LiteLLMError::Config("unsupported inline data url".into()))
}

async fn fetch_bytes_with_mime(
    client: &Client,
    url: &str,
    format: Option<&str>,
) -> Result<(String, String)> {
    let resp = client.get(url).send().await.map_err(LiteLLMError::from)?;
    let headers = resp.headers().clone();
    let bytes = resp.bytes().await.map_err(LiteLLMError::from)?;
    let header_mime = headers
        .get(CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .map(|v| v.split(';').next().unwrap_or(v).trim().to_string());
    let mime = format
        .map(|v| v.to_string())
        .or(header_mime)
        .or_else(|| mime_type_from_url(url))
        .unwrap_or_else(|| "application/octet-stream".to_string());
    let data = general_purpose::STANDARD.encode(bytes);
    Ok((mime, data))
}

fn process_gemini_audio(input_audio: &ChatInputAudio) -> Result<Value> {
    let format = if input_audio.format.starts_with("audio/") {
        input_audio.format.clone()
    } else {
        format!("audio/{}", input_audio.format)
    };
    Ok(serde_json::json!({
        "inline_data": {
            "mime_type": format,
            "data": input_audio.data
        }
    }))
}

fn process_gemini_file(file: &ChatFile, model: &str) -> Result<Value> {
    let passed = file
        .file_id
        .as_ref()
        .or(file.file_data.as_ref())
        .ok_or_else(|| LiteLLMError::Config("file_id or file_data required".into()))?;
    process_gemini_media_url(
        passed,
        file.format.as_deref(),
        file.detail.as_deref(),
        file.video_metadata.as_ref(),
        model,
    )
}

fn process_gemini_media(
    image_url: &ChatImageUrl,
    detail: Option<&str>,
    video_metadata: Option<&Value>,
    format: Option<&str>,
    model: &str,
) -> Result<Value> {
    let (url, format) = match image_url {
        ChatImageUrl::Url(url) => (url.as_str(), format),
        ChatImageUrl::Object(obj) => (obj.url.as_str(), obj.format.as_deref().or(format)),
    };
    process_gemini_media_url(url, format, detail, video_metadata, model)
}

fn process_gemini_media_url(
    url: &str,
    format: Option<&str>,
    detail: Option<&str>,
    video_metadata: Option<&Value>,
    _model: &str,
) -> Result<Value> {
    if url.starts_with("gs://") || url.starts_with("https://") || url.starts_with("http://") {
        let mime_type = format
            .map(|v| v.to_string())
            .or_else(|| mime_type_from_url(url))
            .ok_or_else(|| LiteLLMError::Config("missing media mime type".into()))?;
        let mut part = serde_json::json!({
            "file_data": { "mime_type": mime_type, "file_uri": url }
        });
        apply_gemini_metadata(&mut part, detail, video_metadata);
        return Ok(part);
    }

    if let Some((media_type, data)) = parse_data_url(url, format) {
        let mut part = serde_json::json!({
            "inline_data": { "mime_type": media_type, "data": data }
        });
        apply_gemini_metadata(&mut part, detail, video_metadata);
        return Ok(part);
    }

    Err(LiteLLMError::Config("unsupported gemini media url".into()))
}

fn apply_gemini_metadata(part: &mut Value, detail: Option<&str>, video_metadata: Option<&Value>) {
    if let Some(detail) = detail {
        if let Some(level) = detail_to_media_resolution(detail) {
            if let Some(obj) = part.as_object_mut() {
                obj.insert(
                    "media_resolution".to_string(),
                    serde_json::json!({ "level": level }),
                );
            }
        }
    }
    if let Some(video_metadata) = video_metadata {
        if let Some(obj) = part.as_object_mut() {
            obj.insert("video_metadata".to_string(), video_metadata.clone());
        }
    }
}

fn detail_to_media_resolution(detail: &str) -> Option<&'static str> {
    match detail {
        "low" => Some("MEDIA_RESOLUTION_LOW"),
        "medium" => Some("MEDIA_RESOLUTION_MEDIUM"),
        "high" => Some("MEDIA_RESOLUTION_HIGH"),
        "ultra_high" => Some("MEDIA_RESOLUTION_ULTRA_HIGH"),
        _ => None,
    }
}

fn parse_data_url(url: &str, override_format: Option<&str>) -> Option<(String, String)> {
    if url.starts_with("data:") {
        let stripped = url.strip_prefix("data:").unwrap_or(url);
        if let Some((meta, data)) = stripped.split_once(',') {
            let mut media_type = meta.split(';').next().unwrap_or("application/octet-stream");
            if let Some(fmt) = override_format {
                media_type = fmt;
            }
            return Some((media_type.to_string(), data.to_string()));
        }
    }
    None
}

fn mime_type_from_url(url: &str) -> Option<String> {
    let path = url.split('?').next().unwrap_or(url);
    MimeGuess::from_path(path)
        .first_raw()
        .map(|m| m.to_string())
}

fn extract_video_uri(resp: &Value) -> Option<String> {
    let response = resp.get("response")?;
    if let Some(uri) = response
        .pointer("/generateVideoResponse/generatedSamples/0/video/uri")
        .and_then(|v| v.as_str())
    {
        return Some(uri.to_string());
    }
    if let Some(uri) = response
        .pointer("/generatedVideos/0/uri")
        .and_then(|v| v.as_str())
    {
        return Some(uri.to_string());
    }
    if let Some(uri) = response.pointer("/videos/0/uri").and_then(|v| v.as_str()) {
        return Some(uri.to_string());
    }
    None
}