flow_adapter_ai/providers/

gemini.rs

use async_trait::async_trait;
use futures_util::StreamExt;
use serde::Deserialize;
use serde_json::{json, Value};
use std::time::Instant;

use crate::error::{redact_error_body, CloudAiError};
use crate::registry::{CloudAiProvider, ModelCapabilityInfo};
use crate::request::{
    CloudAiRequest, CloudAiResponse, EmbeddingResponse, ToolDispatcher, ToolSpec,
};
use crate::stream::{LlmStreamEvent, LlmStreamSink};

const NAME: &str = "gemini";
const ENV_VAR: &str = "GOOGLE_API_KEY";
const ENDPOINT_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/models";

// Keep this list aligned with what the Google Generative Language API
// actually serves under v1beta. As of mid-2026, `gemini-2.0-pro` no longer
// resolves (404 from generativelanguage.googleapis.com), so the current
// generation models are 2.5 with the previous-gen 2.0-flash + 1.5-flash
// kept around as cheaper fallbacks.
const DEFAULT_MODELS: &[&str] = &[
    "gemini-2.5-pro",
    "gemini-2.5-flash",
    "gemini-2.0-flash",
    "gemini-1.5-flash",
];

pub struct GeminiProvider {
    client: reqwest::Client,
}

impl GeminiProvider {
    pub fn new() -> Self {
        Self {
            client: reqwest::Client::new(),
        }
    }
}

impl Default for GeminiProvider {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl CloudAiProvider for GeminiProvider {
    fn name(&self) -> &str {
        NAME
    }

    fn env_var(&self) -> &str {
        ENV_VAR
    }

    fn default_models(&self) -> &[&str] {
        DEFAULT_MODELS
    }

    fn model_capabilities(&self) -> Vec<ModelCapabilityInfo> {
        vec![
            ModelCapabilityInfo::new(
                "gemini-2.5-pro",
                &["reasoning", "vision", "tool_use", "code"],
            ),
            ModelCapabilityInfo::new(
                "gemini-2.5-flash",
                &["reasoning", "vision", "tool_use", "code"],
            ),
            ModelCapabilityInfo::new("gemini-2.0-flash", &["vision", "tool_use", "code"]),
            ModelCapabilityInfo::new("gemini-1.5-flash", &["vision", "tool_use", "code"]),
        ]
    }

    async fn invoke(&self, req: &CloudAiRequest) -> Result<CloudAiResponse, CloudAiError> {
        // Send the API key via the `x-goog-api-key` header rather than the
        // `?key=` query string. `reqwest` error messages on transport failures
        // (DNS, TLS, redirect, timeout) routinely include the full URL - with
        // the key in the query string those error strings, surfaced through
        // CloudAiError::Http to logs and the renderer, would leak the key.
        // Header values are not echoed by reqwest's Display impl.
        let url = format!("{ENDPOINT_BASE}/{}:generateContent", req.model);
        let body = build_body(req);

        let started = Instant::now();
        let resp = self
            .client
            .post(&url)
            .header("content-type", "application/json")
            .header("x-goog-api-key", &req.api_key)
            .json(&body)
            .send()
            .await
            .map_err(|e| CloudAiError::Http {
                provider: NAME.into(),
                reason: e.to_string(),
            })?;

        let status = resp.status();
        let raw_body = resp.text().await.map_err(|e| CloudAiError::Http {
            provider: NAME.into(),
            reason: e.to_string(),
        })?;
        let latency_ms = started.elapsed().as_millis() as u64;

        if !status.is_success() {
            return Err(CloudAiError::Status {
                provider: NAME.into(),
                status: status.as_u16(),
                body: redact_error_body(&raw_body),
            });
        }

        parse_response(&raw_body, &req.model, latency_ms)
    }

    async fn invoke_stream(
        &self,
        req: &CloudAiRequest,
        sink: &dyn LlmStreamSink,
    ) -> Result<CloudAiResponse, CloudAiError> {
        let call_id = req.call_id.clone().unwrap_or_default();
        // `alt=sse` switches Gemini's stream from chunked JSON arrays to
        // OpenAI-style `data: <json>` frames, which is far easier to
        // parse and matches the path the openai_compat helper takes.
        let url = format!(
            "{ENDPOINT_BASE}/{}:streamGenerateContent?alt=sse",
            req.model
        );
        let body = build_body(req);
        let started = Instant::now();
        let resp = self
            .client
            .post(&url)
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .header("x-goog-api-key", &req.api_key)
            .json(&body)
            .send()
            .await
            .map_err(|e| CloudAiError::Http {
                provider: NAME.into(),
                reason: e.to_string(),
            })?;

        let status = resp.status();
        if !status.is_success() {
            let raw = resp.text().await.unwrap_or_default();
            let err = CloudAiError::Status {
                provider: NAME.into(),
                status: status.as_u16(),
                body: redact_error_body(&raw),
            };
            sink.emit(LlmStreamEvent::error(&call_id, NAME, err.to_string()))
                .await;
            return Err(err);
        }

        let mut stream = resp.bytes_stream();
        let mut buffer = String::new();
        let mut full_text = String::new();
        let mut finish_reason = String::from("unknown");
        let mut input_tokens: u32 = 0;
        let mut output_tokens: u32 = 0;
        let mut errored: Option<String> = None;

        while let Some(chunk) = stream.next().await {
            let bytes = match chunk {
                Ok(b) => b,
                Err(e) => {
                    errored = Some(e.to_string());
                    break;
                }
            };
            buffer.push_str(&String::from_utf8_lossy(&bytes));
            while let Some(idx) = buffer.find("\n\n") {
                let frame = buffer[..idx].to_string();
                buffer.drain(..idx + 2);
                for raw_line in frame.split('\n') {
                    let line = raw_line.trim_start();
                    let Some(payload) = line.strip_prefix("data:") else {
                        continue;
                    };
                    let payload = payload.trim();
                    if payload.is_empty() {
                        continue;
                    }
                    let value: serde_json::Value = match serde_json::from_str(payload) {
                        Ok(v) => v,
                        Err(_) => continue,
                    };
                    if let Some(usage) = value.get("usageMetadata") {
                        input_tokens = usage
                            .get("promptTokenCount")
                            .and_then(|v| v.as_u64())
                            .map(|v| v as u32)
                            .unwrap_or(input_tokens);
                        output_tokens = usage
                            .get("candidatesTokenCount")
                            .and_then(|v| v.as_u64())
                            .map(|v| v as u32)
                            .unwrap_or(output_tokens);
                    }
                    if let Some(candidate) = value.get("candidates").and_then(|c| c.get(0)) {
                        if let Some(reason) =
                            candidate.get("finishReason").and_then(|r| r.as_str())
                        {
                            finish_reason = reason.to_string();
                        }
                        if let Some(parts) = candidate
                            .get("content")
                            .and_then(|c| c.get("parts"))
                            .and_then(|p| p.as_array())
                        {
                            for part in parts {
                                if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
                                    if !text.is_empty() {
                                        full_text.push_str(text);
                                        sink.emit(LlmStreamEvent {
                                            call_id: call_id.clone(),
                                            provider: NAME.to_string(),
                                            model: req.model.clone(),
                                            delta: text.to_string(),
                                            done: false,
                                            error: None,
                                        })
                                        .await;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        if let Some(err) = errored {
            sink.emit(LlmStreamEvent::error(&call_id, NAME, err.clone()))
                .await;
            return Err(CloudAiError::Http {
                provider: NAME.into(),
                reason: err,
            });
        }
        sink.emit(LlmStreamEvent {
            call_id: call_id.clone(),
            provider: NAME.to_string(),
            model: req.model.clone(),
            delta: String::new(),
            done: true,
            error: None,
        })
        .await;
        if full_text.is_empty() {
            return Err(CloudAiError::Shape {
                provider: NAME.into(),
                detail: "stream closed before any text was emitted".into(),
            });
        }
        let latency_ms = started.elapsed().as_millis() as u64;
        Ok(CloudAiResponse {
            provider: NAME.to_string(),
            model: req.model.clone(),
            text: full_text,
            finish_reason,
            input_tokens,
            output_tokens,
            latency_ms,
        })
    }

    async fn invoke_tools(
        &self,
        req: &CloudAiRequest,
        tools: &[ToolSpec],
        dispatcher: &dyn ToolDispatcher,
        max_iters: usize,
    ) -> Result<CloudAiResponse, CloudAiError> {
        let started = Instant::now();
        let url = format!("{ENDPOINT_BASE}/{}:generateContent", req.model);
        let mut contents = vec![json!({ "role": "user", "parts": user_parts(req) })];
        let fn_decls: Vec<_> = tools
            .iter()
            .map(|t| {
                json!({ "name": t.name, "description": t.description, "parameters": t.parameters })
            })
            .collect();
        let tools_body = json!([{ "functionDeclarations": fn_decls }]);
        let mut input_tokens: u32 = 0;
        let mut output_tokens: u32 = 0;

        for _ in 0..max_iters.max(1) {
            let mut body = json!({ "contents": contents, "tools": tools_body });
            if let Some(s) = req.system.as_deref().filter(|s| !s.is_empty()) {
                body["system_instruction"] = json!({ "parts": [{ "text": s }] });
            }
            if let Some(cfg) = generation_config(req) {
                body["generationConfig"] = cfg;
            }

            let resp = self
                .client
                .post(&url)
                .header("content-type", "application/json")
                .header("x-goog-api-key", &req.api_key)
                .json(&body)
                .send()
                .await
                .map_err(|e| CloudAiError::Http {
                    provider: NAME.into(),
                    reason: e.to_string(),
                })?;
            let status = resp.status();
            let raw = resp.text().await.unwrap_or_default();
            if !status.is_success() {
                return Err(CloudAiError::Status {
                    provider: NAME.into(),
                    status: status.as_u16(),
                    body: redact_error_body(&raw),
                });
            }
            let v: Value = serde_json::from_str(&raw).map_err(|e| CloudAiError::Parse {
                provider: NAME.into(),
                reason: e.to_string(),
            })?;
            if let Some(u) = v.get("usageMetadata") {
                input_tokens = input_tokens.saturating_add(
                    u.get("promptTokenCount").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
                );
                output_tokens = output_tokens.saturating_add(
                    u.get("candidatesTokenCount").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
                );
            }
            let cand = v.get("candidates").and_then(|c| c.get(0));
            let finish = cand
                .and_then(|c| c.get("finishReason"))
                .and_then(|r| r.as_str())
                .unwrap_or("STOP")
                .to_string();
            let parts = cand
                .and_then(|c| c.get("content"))
                .and_then(|c| c.get("parts"))
                .and_then(|p| p.as_array())
                .cloned()
                .unwrap_or_default();

            let mut text = String::new();
            let mut calls: Vec<(String, Value)> = Vec::new();
            for part in &parts {
                if let Some(t) = part.get("text").and_then(|t| t.as_str()) {
                    text.push_str(t);
                }
                if let Some(fc) = part.get("functionCall") {
                    calls.push((
                        fc.get("name").and_then(|n| n.as_str()).unwrap_or("").to_string(),
                        fc.get("args").cloned().unwrap_or_else(|| json!({})),
                    ));
                }
            }

            if calls.is_empty() {
                if text.is_empty() {
                    return Err(CloudAiError::Shape {
                        provider: NAME.into(),
                        detail: "empty candidate content".into(),
                    });
                }
                return Ok(CloudAiResponse {
                    provider: NAME.into(),
                    model: req.model.clone(),
                    text,
                    finish_reason: finish,
                    input_tokens,
                    output_tokens,
                    latency_ms: started.elapsed().as_millis() as u64,
                });
            }

            // Echo the model's functionCall turn, then a user turn carrying each
            // functionResponse (errors fed back so the model can recover).
            contents.push(json!({ "role": "model", "parts": parts }));
            let mut resp_parts = Vec::new();
            for (name, args) in &calls {
                let out = match dispatcher.call(name, args).await {
                    Ok(v) => v,
                    Err(e) => json!({ "error": e }),
                };
                resp_parts.push(json!({
                    "functionResponse": { "name": name, "response": { "result": out } }
                }));
            }
            contents.push(json!({ "role": "user", "parts": resp_parts }));
        }

        Err(CloudAiError::Shape {
            provider: NAME.into(),
            detail: format!("tool loop did not converge within {max_iters} iterations"),
        })
    }

    async fn embed(&self, req: &CloudAiRequest) -> Result<EmbeddingResponse, CloudAiError> {
        let url = format!("{ENDPOINT_BASE}/{}:embedContent", req.model);
        let body = json!({
            "model": format!("models/{}", req.model),
            "content": { "parts": [{ "text": req.prompt }] },
        });
        let started = Instant::now();
        let resp = self
            .client
            .post(&url)
            .header("content-type", "application/json")
            .header("x-goog-api-key", &req.api_key)
            .json(&body)
            .send()
            .await
            .map_err(|e| CloudAiError::Http {
                provider: NAME.into(),
                reason: e.to_string(),
            })?;
        let status = resp.status();
        let raw = resp.text().await.unwrap_or_default();
        let latency_ms = started.elapsed().as_millis() as u64;
        if !status.is_success() {
            return Err(CloudAiError::Status {
                provider: NAME.into(),
                status: status.as_u16(),
                body: redact_error_body(&raw),
            });
        }
        let v: Value = serde_json::from_str(&raw).map_err(|e| CloudAiError::Parse {
            provider: NAME.into(),
            reason: e.to_string(),
        })?;
        let embedding: Vec<f32> = v
            .get("embedding")
            .and_then(|e| e.get("values"))
            .and_then(|vals| vals.as_array())
            .map(|a| a.iter().filter_map(|x| x.as_f64().map(|f| f as f32)).collect())
            .unwrap_or_default();
        if embedding.is_empty() {
            return Err(CloudAiError::Shape {
                provider: NAME.into(),
                detail: "no embedding values in response".into(),
            });
        }
        let dims = embedding.len();
        Ok(EmbeddingResponse {
            provider: NAME.into(),
            model: req.model.clone(),
            embedding,
            dims,
            latency_ms,
        })
    }
}

/// User-content `parts`: the prompt text plus an image part per request image
/// (`inlineData` for `data:` URLs, `fileData` for `http(s)` URLs).
fn user_parts(req: &CloudAiRequest) -> Vec<serde_json::Value> {
    let mut parts = vec![json!({ "text": req.prompt })];
    for img in &req.images {
        parts.push(image_part(img));
    }
    parts
}

fn image_part(img: &str) -> serde_json::Value {
    if let Some(rest) = img.strip_prefix("data:") {
        if let Some((mime, data)) = rest.split_once(";base64,") {
            return json!({ "inlineData": { "mimeType": mime, "data": data } });
        }
    }
    json!({ "fileData": { "fileUri": img } })
}

/// Build the `generationConfig` block (sampling + optional thinking budget), or
/// `None` when empty.
fn generation_config(req: &CloudAiRequest) -> Option<serde_json::Value> {
    let mut config = serde_json::Map::new();
    if let Some(m) = req.max_tokens {
        config.insert("maxOutputTokens".into(), json!(m));
    }
    if let Some(t) = req.temperature {
        config.insert("temperature".into(), json!(t));
    }
    if let Some(think) = req.reasoning {
        // `thinkingBudget: 0` disables reasoning; a positive budget enables it
        // (Gemini 2.5 thinking models).
        config.insert(
            "thinkingConfig".into(),
            json!({ "thinkingBudget": if think { 1024 } else { 0 } }),
        );
    }
    if config.is_empty() {
        None
    } else {
        Some(serde_json::Value::Object(config))
    }
}

pub fn build_body(req: &CloudAiRequest) -> serde_json::Value {
    let mut payload = json!({
        "contents": [{ "parts": user_parts(req) }],
    });
    if let Some(s) = req.system.as_deref().filter(|s| !s.is_empty()) {
        // Gemini's v1beta API takes the system prompt as a top-level
        // `system_instruction` containing the same `parts`/`text`
        // shape used for `contents`. Sending it via this field (rather
        // than concatenating into `contents`) is what unlocks the
        // model's instruction-following behavior.
        payload["system_instruction"] = json!({
            "parts": [{ "text": s }]
        });
    }
    if let Some(cfg) = generation_config(req) {
        payload["generationConfig"] = cfg;
    }
    payload
}

#[derive(Debug, Deserialize)]
struct RawResponse {
    candidates: Vec<Candidate>,
    #[serde(rename = "usageMetadata", default)]
    usage_metadata: Option<UsageMetadata>,
}

#[derive(Debug, Deserialize)]
struct Candidate {
    content: Option<Content>,
    #[serde(rename = "finishReason", default)]
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
struct Content {
    parts: Vec<Part>,
}

#[derive(Debug, Deserialize)]
struct Part {
    #[serde(default)]
    text: Option<String>,
}

#[derive(Debug, Deserialize)]
struct UsageMetadata {
    #[serde(rename = "promptTokenCount", default)]
    prompt_token_count: u32,
    #[serde(rename = "candidatesTokenCount", default)]
    candidates_token_count: u32,
}

pub fn parse_response(
    body: &str,
    model: &str,
    latency_ms: u64,
) -> Result<CloudAiResponse, CloudAiError> {
    let raw: RawResponse = serde_json::from_str(body).map_err(|e| CloudAiError::Parse {
        provider: NAME.into(),
        reason: e.to_string(),
    })?;

    let candidate = raw
        .candidates
        .into_iter()
        .next()
        .ok_or(CloudAiError::Shape {
            provider: NAME.into(),
            detail: "no candidates in response".into(),
        })?;

    let text = candidate
        .content
        .map(|c| {
            c.parts
                .into_iter()
                .filter_map(|p| p.text)
                .collect::<Vec<_>>()
                .join("")
        })
        .unwrap_or_default();

    if text.is_empty() {
        return Err(CloudAiError::Shape {
            provider: NAME.into(),
            detail: "empty candidate content".into(),
        });
    }

    let (input, output) = match raw.usage_metadata {
        Some(u) => (u.prompt_token_count, u.candidates_token_count),
        None => (0, 0),
    };

    Ok(CloudAiResponse {
        provider: NAME.into(),
        model: model.to_string(),
        text,
        finish_reason: candidate.finish_reason.unwrap_or_else(|| "unknown".into()),
        input_tokens: input,
        output_tokens: output,
        latency_ms,
    })
}