ic-rig 0.2.0

//! Google Gemini `generateContent` API provider.
//!
//! # Usage
//!
//! ```rust,ignore
//! use irig::providers::gemini::{Client, GEMINI_3_5_FLASH};
//!
//! let client = Client::new(my_http, "AIza...");
//! let model  = client.model(GEMINI_3_5_FLASH);
//!
//! let agent = irig::Agent::builder(model)
//!     .preamble("You are a helpful assistant.")
//!     .build();
//!
//! let reply = agent.prompt("Hello!").await?;
//! ```

use crate::{
    completion::{CompletionError, CompletionModel, CompletionRequest, CompletionResponse, ModelChoice, Usage},
    embeddings::{Embedding, EmbeddingError, EmbeddingModel},
    http::{HttpClient, HttpRequest},
    message::{AssistantContent, Message, ToolCall, UserContent},
    tool::ToolDefinition,
};
use serde::{Deserialize, Serialize};

// ── Model constants ───────────────────────────────────────────────────────────

// ── Completion model constants ────────────────────────────────────────────────

// Gemini 3 (current generation, recommended).
pub const GEMINI_3_1_PRO_PREVIEW: &str = "gemini-3.1-pro-preview";
pub const GEMINI_3_7_FLASH: &str = "gemini-3.7-flash";
pub const GEMINI_3_6_FLASH: &str = "gemini-3.6-flash";
pub const GEMINI_3_5_FLASH: &str = "gemini-3.5-flash";
pub const GEMINI_3_5_FLASH_LITE: &str = "gemini-3.5-flash-lite";
pub const GEMINI_3_1_FLASH_LITE: &str = "gemini-3.1-flash-lite";

// Gemini 2.5 (previous generation; GA-stable, Google's safe fallback baseline
// until it goes off GA on 2026-10-16).
/// Was previously pinned to the dated preview ID `gemini-2.5-pro-preview-05-06`,
/// which Google now redirects to this stable ID — updated 2026-08-27.
pub const GEMINI_2_5_PRO: &str = "gemini-2.5-pro";
pub const GEMINI_2_5_FLASH: &str = "gemini-2.5-flash";
pub const GEMINI_2_5_FLASH_LITE: &str = "gemini-2.5-flash-lite";

/// Shut down by Google on 2026-06-01; kept only so old code still compiles.
#[deprecated(note = "shut down by Google on 2026-06-01; use GEMINI_3_5_FLASH or GEMINI_2_5_FLASH instead")]
pub const GEMINI_2_0_FLASH: &str = "gemini-2.0-flash";
/// Shut down by Google on 2026-06-01; kept only so old code still compiles.
#[deprecated(note = "shut down by Google on 2026-06-01; use GEMINI_3_1_FLASH_LITE or GEMINI_2_5_FLASH_LITE instead")]
pub const GEMINI_2_0_FLASH_LITE: &str = "gemini-2.0-flash-lite";
/// Shut down by Google on 2025-09-29; kept only so old code still compiles.
#[deprecated(note = "shut down by Google on 2025-09-29; use GEMINI_3_1_PRO_PREVIEW or GEMINI_2_5_PRO instead")]
pub const GEMINI_1_5_PRO: &str = "gemini-1.5-pro";
/// Shut down by Google on 2025-09-29; kept only so old code still compiles.
#[deprecated(note = "shut down by Google on 2025-09-29; use GEMINI_3_5_FLASH or GEMINI_2_5_FLASH instead")]
pub const GEMINI_1_5_FLASH: &str = "gemini-1.5-flash";

// ── Embedding model constants ─────────────────────────────────────────────────

pub const GEMINI_EMBEDDING_001: &str = "gemini-embedding-001";
/// Multimodal embedding model.
pub const GEMINI_EMBEDDING_2_PREVIEW: &str = "gemini-embedding-2-preview";

/// Shut down by Google on 2026-01-14; kept only so old code still compiles.
#[deprecated(note = "shut down by Google on 2026-01-14; use GEMINI_EMBEDDING_001 instead")]
pub const TEXT_EMBEDDING_004: &str = "text-embedding-004";
/// Shut down by Google in October 2025; kept only so old code still compiles.
#[deprecated(note = "shut down by Google in October 2025; use GEMINI_EMBEDDING_001 instead")]
pub const EMBEDDING_001: &str = "embedding-001";

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

// ── Client ────────────────────────────────────────────────────────────────────

/// Gemini API client. Use [`model`](Client::model) to get a [`CompletionModel`].
pub struct Client<H> {
    http: H,
    api_key: String,
    base_url: String,
}

impl<H: HttpClient + Clone> Client<H> {
    pub fn new(http: H, api_key: impl Into<String>) -> Self {
        Self { http, api_key: api_key.into(), base_url: BASE_URL.to_owned() }
    }

    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = url.into();
        self
    }

    pub fn model(&self, model: impl Into<String>) -> Model<H> {
        Model {
            http: self.http.clone(),
            api_key: self.api_key.clone(),
            base_url: self.base_url.clone(),
            model: model.into(),
        }
    }

    /// Produce a Gemini [`GeminiEmbeddingModel`].
    pub fn embedding_model(&self, model: impl Into<String>) -> GeminiEmbeddingModel<H> {
        GeminiEmbeddingModel {
            http: self.http.clone(),
            api_key: self.api_key.clone(),
            base_url: self.base_url.clone(),
            model: model.into(),
        }
    }
}

// ── Model ─────────────────────────────────────────────────────────────────────

pub struct Model<H> {
    http: H,
    api_key: String,
    base_url: String,
    model: String,
}

impl<H: HttpClient> CompletionModel for Model<H> {
    type Error = CompletionError;

    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, CompletionError> {
        let body = build_request(request)?;
        let bytes = serde_json::to_vec(&body)?;

        // Gemini passes the API key as a query parameter, not a header.
        let url = format!(
            "{}/{}:generateContent?key={}",
            self.base_url, self.model, self.api_key
        );

        let http_req = HttpRequest::new(url).json_body(bytes);

        let resp = self.http.post(http_req).await
            .map_err(|e| CompletionError::Http(e.to_string()))?;

        if !resp.is_success() {
            let message = String::from_utf8_lossy(&resp.body).into_owned();
            return Err(CompletionError::Provider { status: resp.status, message });
        }

        let api_resp: ApiResponse = resp.json()?;
        parse_response(api_resp)
    }
}

// ── Wire types (Gemini JSON format) ──────────────────────────────────────────

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ApiRequest {
    contents: Vec<ApiContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system_instruction: Option<ApiSystemInstruction>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<ApiTools>,
    #[serde(skip_serializing_if = "Option::is_none")]
    generation_config: Option<GenerationConfig>,
}

#[derive(Serialize)]
struct ApiSystemInstruction {
    parts: Vec<ApiPart>,
}

/// Gemini groups all tool declarations under a single `tools` object.
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ApiTools {
    function_declarations: Vec<ApiFunctionDeclaration>,
}

#[derive(Serialize)]
struct ApiFunctionDeclaration {
    name: String,
    description: String,
    parameters: serde_json::Value,
}

#[derive(Serialize, Deserialize)]
struct ApiContent {
    role: String,
    parts: Vec<ApiPart>,
}

#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(try_from = "RawPart")]
enum ApiPart {
    /// Plain text.
    Text(String),
    /// Thinking / chain-of-thought trace — a text part with `"thought":
    /// true` alongside it. Only appears when thinking is enabled. Never
    /// treated as the answer — collected into
    /// [`CompletionResponse::reasoning`](crate::completion::CompletionResponse::reasoning)
    /// instead.
    Thought(String),
    /// A tool call emitted by the model (`role: "model"`).
    FunctionCall(ApiFunctionCall),
    /// The result of a tool call (`role: "user"`).
    FunctionResponse(ApiFunctionResponse),
}

/// Incoming Gemini parts carry `text` alongside a sibling `thought: bool`
/// flag rather than a distinct wire "type", which doesn't fit a plain
/// externally-tagged enum — so we deserialize this loose shape first and
/// resolve it into an [`ApiPart`] via [`TryFrom`].
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct RawPart {
    #[serde(default)]
    text: Option<String>,
    #[serde(default)]
    thought: bool,
    #[serde(default)]
    function_call: Option<ApiFunctionCall>,
    #[serde(default)]
    function_response: Option<ApiFunctionResponse>,
}

impl TryFrom<RawPart> for ApiPart {
    type Error = String;

    fn try_from(raw: RawPart) -> Result<Self, Self::Error> {
        if let Some(fc) = raw.function_call {
            Ok(ApiPart::FunctionCall(fc))
        } else if let Some(fr) = raw.function_response {
            Ok(ApiPart::FunctionResponse(fr))
        } else if let Some(text) = raw.text {
            Ok(if raw.thought { ApiPart::Thought(text) } else { ApiPart::Text(text) })
        } else {
            Err("part has none of text, functionCall, functionResponse".into())
        }
    }
}

#[derive(Serialize, Deserialize)]
struct ApiFunctionCall {
    name: String,
    /// Gemini delivers arguments as a JSON **object** (not a string).
    args: serde_json::Value,
}

#[derive(Serialize, Deserialize)]
struct ApiFunctionResponse {
    name: String,
    response: serde_json::Value,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct GenerationConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_output_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    thinking_config: Option<ThinkingConfig>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ThinkingConfig {
    /// `-1` = dynamic (model decides), `0` = disabled. Some Gemini 3 models
    /// (e.g. `gemini-3.1-pro-preview`) can't fully disable thinking — Google's
    /// own docs note this — so `0` may not be honored on every model.
    thinking_budget: i32,
    /// Whether to return thought-summary parts (`"thought": true`) at all;
    /// without this, thinking may still happen internally but won't be
    /// visible in `CompletionResponse::reasoning`.
    include_thoughts: bool,
}

#[derive(Deserialize)]
struct ApiResponse {
    candidates: Vec<ApiCandidate>,
    #[serde(rename = "usageMetadata")]
    usage_metadata: Option<ApiUsageMetadata>,
}

#[derive(Deserialize)]
struct ApiCandidate {
    content: ApiContent,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct ApiUsageMetadata {
    prompt_token_count: u32,
    candidates_token_count: u32,
}

// ── Conversion helpers ────────────────────────────────────────────────────────

fn build_request(req: CompletionRequest) -> Result<ApiRequest, CompletionError> {
    let mut system_instruction: Option<ApiSystemInstruction> = None;
    let mut chat_messages: Vec<Message> = Vec::new();

    for msg in req.messages {
        match msg {
            Message::System { content } => {
                system_instruction = Some(ApiSystemInstruction {
                    parts: vec![ApiPart::Text(content)],
                });
            }
            other => chat_messages.push(other),
        }
    }

    let contents = convert_messages(chat_messages)?;

    let tools = if req.tools.is_empty() {
        Vec::new()
    } else {
        vec![ApiTools {
            function_declarations: req.tools.into_iter().map(convert_tool).collect(),
        }]
    };

    let thinking_config = req.thinking.map(|enabled| ThinkingConfig {
        thinking_budget: if enabled { -1 } else { 0 },
        include_thoughts: enabled,
    });

    let generation_config =
        if req.temperature.is_some() || req.max_tokens.is_some() || thinking_config.is_some() {
            Some(GenerationConfig {
                temperature: req.temperature,
                max_output_tokens: req.max_tokens,
                thinking_config,
            })
        } else {
            None
        };

    Ok(ApiRequest { contents, system_instruction, tools, generation_config })
}

/// Convert irig messages to Gemini's `contents` format.
///
/// Key differences:
/// - Gemini uses `"user"` and `"model"` roles (not `"assistant"`).
/// - Tool calls are `functionCall` parts inside a `"model"` turn.
/// - Tool results are `functionResponse` parts inside a `"user"` turn.
fn convert_messages(messages: Vec<Message>) -> Result<Vec<ApiContent>, CompletionError> {
    let mut out = Vec::new();

    for msg in messages {
        match msg {
            Message::System { .. } => {
                // Extracted above; ignore stragglers.
            }

            Message::User { content } => {
                let parts: Vec<ApiPart> = content
                    .into_iter()
                    .map(|part| match part {
                        UserContent::Text(t) => ApiPart::Text(t.text),
                        UserContent::ToolResult(r) => {
                            ApiPart::FunctionResponse(ApiFunctionResponse {
                                name: r.name,
                                // Wrap the content string in a JSON object so
                                // Gemini receives a structured response field.
                                response: serde_json::json!({ "content": r.content }),
                            })
                        }
                    })
                    .collect();

                out.push(ApiContent { role: "user".into(), parts });
            }

            Message::Assistant { content } => {
                let parts: Vec<ApiPart> = content
                    .into_iter()
                    .map(|part| match part {
                        AssistantContent::Text(t) => ApiPart::Text(t.text),
                        AssistantContent::ToolCall(c) => {
                            ApiPart::FunctionCall(ApiFunctionCall {
                                name: c.name,
                                args: c.arguments,
                            })
                        }
                    })
                    .collect();

                out.push(ApiContent { role: "model".into(), parts });
            }
        }
    }

    Ok(out)
}

fn convert_tool(def: ToolDefinition) -> ApiFunctionDeclaration {
    ApiFunctionDeclaration {
        name: def.name,
        description: def.description,
        parameters: def.parameters,
    }
}

fn parse_response(resp: ApiResponse) -> Result<CompletionResponse, CompletionError> {
    let candidate = resp
        .candidates
        .into_iter()
        .next()
        .ok_or_else(|| CompletionError::Response("no candidates in response".into()))?;

    let usage = resp.usage_metadata.map(|u| Usage {
        prompt_tokens: u.prompt_token_count,
        completion_tokens: u.candidates_token_count,
    });

    let mut text: Option<String> = None;
    let mut reasoning: Option<String> = None;
    let mut tool_calls: Vec<ToolCall> = Vec::new();

    for part in candidate.content.parts {
        match part {
            ApiPart::Text(t) => text = Some(t),
            ApiPart::Thought(t) => {
                reasoning = Some(match reasoning {
                    Some(existing) => format!("{existing}\n{t}"),
                    None => t,
                });
            }
            ApiPart::FunctionCall(fc) => {
                // Gemini doesn't supply a persistent call id; synthesise one
                // from the name so tool results can be correlated.
                let id = format!("call_{}", fc.name);
                tool_calls.push(ToolCall { id, name: fc.name, arguments: fc.args });
            }
            ApiPart::FunctionResponse(_) => {
                // Function responses only appear in user turns; skip if echoed.
            }
        }
    }

    let choice = if !tool_calls.is_empty() {
        ModelChoice::ToolCall(tool_calls)
    } else {
        let t = text.ok_or_else(|| CompletionError::Response("empty parts in candidate".into()))?;
        ModelChoice::Message(t)
    };

    Ok(CompletionResponse { choice, reasoning, usage })
}

// ═════════════════════════════════════════════════════════════════════════════
// Embeddings
// ═════════════════════════════════════════════════════════════════════════════

/// A Gemini embedding model that implements [`EmbeddingModel`].
pub struct GeminiEmbeddingModel<H> {
    http: H,
    api_key: String,
    base_url: String,
    model: String,
}

// ── Wire types ────────────────────────────────────────────────────────────────

/// Gemini `batchEmbedContents` request.
#[derive(Serialize)]
struct BatchEmbedRequest<'a> {
    requests: Vec<EmbedContentRequest<'a>>,
}

#[derive(Serialize)]
struct EmbedContentRequest<'a> {
    model: &'a str,
    content: EmbedContent<'a>,
}

#[derive(Serialize)]
struct EmbedContent<'a> {
    parts: [EmbedPart<'a>; 1],
}

#[derive(Serialize)]
struct EmbedPart<'a> {
    text: &'a str,
}

#[derive(Deserialize)]
struct BatchEmbedResponse {
    embeddings: Vec<GeminiEmbedding>,
}

#[derive(Deserialize)]
struct GeminiEmbedding {
    values: Vec<f64>,
}

// ── Trait impl ────────────────────────────────────────────────────────────────

impl<H: HttpClient> EmbeddingModel for GeminiEmbeddingModel<H> {
    /// Gemini's batch endpoint accepts up to 100 requests at a time.
    const MAX_DOCUMENTS: usize = 100;

    type Error = EmbeddingError;

    fn ndims(&self) -> usize {
        // text-embedding-004 produces 768-dim vectors.
        768
    }

    async fn embed_texts(&self, texts: Vec<String>) -> Result<Vec<Embedding>, EmbeddingError> {
        let model_path = format!("models/{}", self.model);

        let requests: Vec<EmbedContentRequest> = texts
            .iter()
            .map(|t| EmbedContentRequest {
                model: &model_path,
                content: EmbedContent { parts: [EmbedPart { text: t }] },
            })
            .collect();

        let body = serde_json::to_vec(&BatchEmbedRequest { requests })?;

        let url = format!(
            "{}/{}:batchEmbedContents?key={}",
            self.base_url, model_path, self.api_key
        );

        let req = HttpRequest::new(url).json_body(body);

        let resp = self.http.post(req).await
            .map_err(|e| EmbeddingError::Http(e.to_string()))?;

        if !resp.is_success() {
            let message = String::from_utf8_lossy(&resp.body).into_owned();
            return Err(EmbeddingError::Provider { status: resp.status, message });
        }

        let api_resp: BatchEmbedResponse = resp.json()?;

        if api_resp.embeddings.len() != texts.len() {
            return Err(EmbeddingError::Response(format!(
                "expected {} embeddings, got {}",
                texts.len(),
                api_resp.embeddings.len(),
            )));
        }

        Ok(api_resp
            .embeddings
            .into_iter()
            .zip(texts)
            .map(|(e, document)| Embedding { document, vec: e.values })
            .collect())
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    #[test]
    fn thinking_toggle_serializes_expected_shape() {
        let mut on = CompletionRequest::new(vec![Message::user("hi")]);
        on.thinking = Some(true);
        let json = serde_json::to_value(build_request(on).unwrap()).unwrap();
        assert_eq!(json["generationConfig"]["thinkingConfig"]["thinkingBudget"], -1);
        assert_eq!(json["generationConfig"]["thinkingConfig"]["includeThoughts"], true);

        let mut off = CompletionRequest::new(vec![Message::user("hi")]);
        off.thinking = Some(false);
        let json = serde_json::to_value(build_request(off).unwrap()).unwrap();
        assert_eq!(json["generationConfig"]["thinkingConfig"]["thinkingBudget"], 0);
        assert_eq!(json["generationConfig"]["thinkingConfig"]["includeThoughts"], false);

        let unset = CompletionRequest::new(vec![Message::user("hi")]);
        let json = serde_json::to_value(build_request(unset).unwrap()).unwrap();
        assert!(json.get("generationConfig").is_none());
    }

    fn make_response(json: &str) -> ApiResponse {
        serde_json::from_str(json).expect("test fixture must deserialise")
    }

    #[test]
    fn thought_part_is_kept_out_of_the_answer() {
        // Regression test: before ApiPart handled the `thought: true`
        // sibling field, a thinking-enabled response would fail to parse
        // entirely (it doesn't fit a plain externally-tagged enum).
        let resp = make_response(
            r#"{
                "candidates": [{
                    "content": {
                        "role": "model",
                        "parts": [
                            {"text": "Reasoning about the problem...", "thought": true},
                            {"text": "Final answer."}
                        ]
                    }
                }],
                "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5}
            }"#,
        );

        let result = parse_response(resp).unwrap();
        assert!(matches!(result.choice, ModelChoice::Message(ref s) if s == "Final answer."));
        assert_eq!(result.reasoning.as_deref(), Some("Reasoning about the problem..."));
    }

    #[test]
    fn plain_text_response_has_no_reasoning() {
        let resp = make_response(
            r#"{
                "candidates": [{
                    "content": {
                        "role": "model",
                        "parts": [{"text": "Hello, world!"}]
                    }
                }],
                "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2}
            }"#,
        );

        let result = parse_response(resp).unwrap();
        assert!(matches!(result.choice, ModelChoice::Message(ref s) if s == "Hello, world!"));
        assert_eq!(result.reasoning, None);
    }
}