kcode-gemini-api 0.4.0

use serde_json::{Value, json};

use crate::{Error, Result, TokenUsage};

const MAX_MODEL_CHARACTERS: usize = 128;
const MAX_AGENT_INPUT_CHARACTERS: usize = 2_000_000;
const MAX_TOOLS: usize = 64;

/// One function available to a Gemini text turn.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AgentTool {
    /// Exact function name.
    pub name: String,
    /// Human-readable function description.
    pub description: String,
    /// JSON Schema object accepted by the function.
    pub input_schema: Value,
}

/// One stateless Gemini text/tool interaction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AgentTurnRequest {
    /// Exact Gemini model identifier.
    pub model: String,
    /// Complete input for this interaction.
    pub input: String,
    /// Provider-neutral reasoning effort.
    pub reasoning_effort: String,
    /// Functions available to the model.
    pub tools: Vec<AgentTool>,
}

impl AgentTurnRequest {
    /// Constructs a tool-free request with medium reasoning effort.
    pub fn new(model: impl Into<String>, input: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            input: input.into(),
            reasoning_effort: "medium".into(),
            tools: Vec::new(),
        }
    }

    pub(crate) fn validate(&self) -> Result<()> {
        validate_model(&self.model)?;
        if self.input.trim().is_empty() || self.input.chars().count() > MAX_AGENT_INPUT_CHARACTERS {
            return Err(Error::InvalidInput(format!(
                "agent input must contain 1 through {MAX_AGENT_INPUT_CHARACTERS} characters"
            )));
        }
        if !matches!(
            self.reasoning_effort.as_str(),
            "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max"
        ) {
            return Err(Error::InvalidInput(
                "reasoning effort is not supported".into(),
            ));
        }
        if self.tools.len() > MAX_TOOLS {
            return Err(Error::InvalidInput(format!(
                "agent turn may expose at most {MAX_TOOLS} tools"
            )));
        }
        for tool in &self.tools {
            if tool.name.is_empty()
                || tool.name.chars().count() > 100
                || !tool
                    .name
                    .bytes()
                    .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
            {
                return Err(Error::InvalidInput(
                    "agent tool name must be a safe non-empty identifier".into(),
                ));
            }
            if tool.description.trim().is_empty()
                || tool.description.chars().count() > 4_000
                || !tool.input_schema.is_object()
            {
                return Err(Error::InvalidInput(
                    "agent tools require a bounded description and object JSON Schema".into(),
                ));
            }
        }
        Ok(())
    }

    pub(crate) fn payload(&self) -> Value {
        let mut payload = json!({
            "model": self.model,
            "input": self.input,
            "generation_config": {
                "thinking_level": thinking_level(&self.reasoning_effort)
            },
            "service_tier": "standard",
            "store": false
        });
        if !self.tools.is_empty() {
            payload["tools"] = json!(
                self.tools
                    .iter()
                    .map(|tool| json!({
                        "type": "function",
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.input_schema
                    }))
                    .collect::<Vec<_>>()
            );
        }
        payload
    }
}

/// One normalized function call returned by Gemini.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AgentToolCall {
    /// Provider call identifier.
    pub call_id: String,
    /// Exact function name.
    pub name: String,
    /// Parsed JSON arguments.
    pub arguments: Value,
}

/// Normalized result of one Gemini text/tool interaction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AgentTurnResponse {
    /// Gemini interaction identifier.
    pub interaction_id: String,
    /// Actual model identifier returned by Gemini.
    pub model: String,
    /// Ordered assistant text.
    pub text: String,
    /// At most one function call.
    pub tool_call: Option<AgentToolCall>,
    /// Provider token usage.
    pub usage: TokenUsage,
}

/// Metadata advertised for one exact Gemini model.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ModelMetadata {
    /// Exact model identifier.
    pub id: String,
    /// Advertised context window, when returned.
    pub context_window_tokens: Option<u64>,
    /// Advertised maximum input, when returned.
    pub max_input_tokens: Option<u64>,
}

pub(crate) fn validate_model(model: &str) -> Result<()> {
    if model.trim().is_empty()
        || model.chars().count() > MAX_MODEL_CHARACTERS
        || !model.bytes().all(|byte| {
            byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'-' | b'_' | b':' | b'/')
        })
    {
        return Err(Error::InvalidInput(
            "model must be an exact safe Gemini model identifier".into(),
        ));
    }
    Ok(())
}

pub(crate) fn parse_tool_calls(value: &Value) -> Result<Option<AgentToolCall>> {
    let mut calls = Vec::new();
    for step in value
        .get("steps")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
    {
        if step.get("type").and_then(Value::as_str) == Some("function_call") {
            calls.push(parse_call(step, calls.len())?);
        }
        for item in step
            .get("content")
            .and_then(Value::as_array)
            .into_iter()
            .flatten()
        {
            if item.get("type").and_then(Value::as_str) == Some("function_call") {
                calls.push(parse_call(item, calls.len())?);
            }
        }
    }
    if calls.len() > 1 {
        return Err(Error::Protocol(
            "Gemini returned multiple function calls; agent tools must run serially".into(),
        ));
    }
    Ok(calls.pop())
}

pub(crate) fn parse_model_metadata(value: &Value, requested: &str) -> ModelMetadata {
    ModelMetadata {
        id: value
            .get("name")
            .or_else(|| value.get("id"))
            .and_then(Value::as_str)
            .and_then(|value| value.strip_prefix("models/").or(Some(value)))
            .unwrap_or(requested)
            .to_owned(),
        context_window_tokens: u64_field(
            value,
            &[
                "contextWindowTokens",
                "context_window_tokens",
                "contextWindow",
            ],
        ),
        max_input_tokens: u64_field(
            value,
            &[
                "inputTokenLimit",
                "input_token_limit",
                "maxInputTokens",
                "max_input_tokens",
            ],
        ),
    }
}

fn parse_call(value: &Value, index: usize) -> Result<AgentToolCall> {
    let arguments = match value.get("arguments") {
        Some(Value::String(arguments)) => serde_json::from_str(arguments).map_err(|_| {
            Error::Protocol("function call contained invalid JSON arguments".into())
        })?,
        Some(Value::Object(_)) => value["arguments"].clone(),
        _ => {
            return Err(Error::Protocol(
                "function call omitted object arguments".into(),
            ));
        }
    };
    let name = value
        .get("name")
        .and_then(Value::as_str)
        .filter(|name| !name.is_empty())
        .ok_or_else(|| Error::Protocol("function call omitted its name".into()))?
        .to_owned();
    Ok(AgentToolCall {
        call_id: value
            .get("call_id")
            .or_else(|| value.get("id"))
            .and_then(Value::as_str)
            .map(str::to_owned)
            .unwrap_or_else(|| format!("gemini-call-{index}")),
        name,
        arguments,
    })
}

fn thinking_level(effort: &str) -> &'static str {
    match effort {
        "none" | "minimal" => "minimal",
        "low" => "low",
        "medium" => "medium",
        "high" | "xhigh" | "max" => "high",
        _ => "medium",
    }
}

fn u64_field(value: &Value, fields: &[&str]) -> Option<u64> {
    fields
        .iter()
        .find_map(|field| value.get(*field).and_then(Value::as_u64))
}

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

    #[test]
    fn request_and_function_response_are_normalized() {
        let mut request = AgentTurnRequest::new("gemini-3.1-pro-preview", "Do the task.");
        request.tools.push(AgentTool {
            name: "call_ktool".into(),
            description: "Call one tool.".into(),
            input_schema: json!({"type": "object"}),
        });
        assert_eq!(
            request.payload()["generation_config"]["thinking_level"],
            "medium"
        );
        let call = parse_tool_calls(&json!({
            "steps": [{
                "type": "model_output",
                "content": [{
                    "type": "function_call",
                    "id": "call-1",
                    "name": "call_ktool",
                    "arguments": {"name": "Read"}
                }]
            }]
        }))
        .unwrap()
        .unwrap();
        assert_eq!(call.arguments["name"], "Read");
    }
}