ic-rig 0.1.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")]
enum ApiPart {
    /// Plain text.
    Text(String),
    /// A tool call emitted by the model (`role: "model"`).
    FunctionCall(ApiFunctionCall),
    /// The result of a tool call (`role: "user"`).
    FunctionResponse(ApiFunctionResponse),
}

#[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>,
}

#[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 generation_config =
        if req.temperature.is_some() || req.max_tokens.is_some() {
            Some(GenerationConfig {
                temperature: req.temperature,
                max_output_tokens: req.max_tokens,
            })
        } 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 tool_calls: Vec<ToolCall> = Vec::new();

    for part in candidate.content.parts {
        match part {
            ApiPart::Text(t) => text = Some(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, 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())
    }
}