rig_gemini_grpc/

completion.rs

// ================================================================
//! Google Gemini gRPC Completion Integration
// ================================================================

/// `gemini-2.5-flash` completion model
pub const GEMINI_2_5_FLASH: &str = "gemini-2.5-flash";
/// `gemini-2.0-flash-lite` completion model
pub const GEMINI_2_0_FLASH_LITE: &str = "gemini-2.0-flash-lite";
/// `gemini-2.0-flash` completion model
pub const GEMINI_2_0_FLASH: &str = "gemini-2.0-flash";

use base64::Engine as _;
use rig_core::OneOrMany;
use rig_core::completion::{self, CompletionError, CompletionRequest};
use rig_core::message::{self, MimeType, Reasoning};
use rig_core::providers::gemini::completion::gemini_api_types::{
    Schema as GeminiSchema, tool_parameters_to_schema,
};
use rig_core::telemetry::ProviderResponseExt;
use std::convert::TryFrom;

use super::Client;
use super::proto::{self, GenerateContentRequest, GenerateContentResponse};

// =================================================================
// Rig Implementation Types
// =================================================================

#[derive(Clone, Debug)]
pub struct CompletionModel {
    pub(crate) client: Client,
    pub model: String,
}

impl CompletionModel {
    pub fn new(client: Client, model: impl Into<String>) -> Self {
        Self {
            client,
            model: model.into(),
        }
    }
}

impl completion::CompletionModel for CompletionModel {
    type Response = GenerateContentResponse;
    type StreamingResponse = super::streaming::StreamingCompletionResponse;
    type Client = super::Client;

    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
        Self::new(client.clone(), model)
    }

    async fn completion(
        &self,
        completion_request: CompletionRequest,
    ) -> Result<completion::CompletionResponse<GenerateContentResponse>, CompletionError> {
        let request = create_grpc_request(self.model.clone(), completion_request)?;

        let mut grpc_client = self
            .client
            .grpc_client()
            .map_err(|e| CompletionError::ProviderError(e.to_string()))?;

        let response = grpc_client
            .generate_content(request)
            .await
            .map_err(rpc_error)?
            .into_inner();

        response.try_into()
    }

    async fn stream(
        &self,
        request: CompletionRequest,
    ) -> Result<
        rig_core::streaming::StreamingCompletionResponse<Self::StreamingResponse>,
        CompletionError,
    > {
        super::streaming::stream(self.client.clone(), self.model.clone(), request).await
    }
}

// Map a failed gRPC call into a `CompletionError` that preserves the provider's
// error payload verbatim. gRPC is a non-HTTP transport, so there is no
// `http::StatusCode`; the body is preserved via `from_provider_body` (status:
// None) rather than a Rig-prefixed `ProviderError` diagnostic. Note: tonic does
// not distinguish a server-returned gRPC error from a transport/connection
// failure, so a pure connection error is also preserved here rather than gated
// out as a Rig diagnostic the way Bedrock's typed service errors are.
pub(crate) fn rpc_error(status: tonic::Status) -> CompletionError {
    CompletionError::from_provider_body(status.to_string())
}

// Helper function to create gRPC request from Rig's CompletionRequest
pub(crate) fn create_grpc_request(
    model: String,
    completion_request: CompletionRequest,
) -> Result<GenerateContentRequest, CompletionError> {
    let CompletionRequest {
        model: _,
        preamble,
        chat_history,
        documents: _,
        tools,
        temperature,
        max_tokens,
        tool_choice: _,
        additional_params: _,
        output_schema: _,
        record_telemetry_content: _,
    } = completion_request;

    let (history_system, chat_history) = split_system_messages_from_history(chat_history);
    let mut contents = Vec::new();

    // Convert chat history to gRPC Content messages
    for msg in chat_history {
        contents.push(rig_message_to_grpc_content(msg)?);
    }

    // Handle system instruction (preamble)
    let mut system_parts = Vec::new();
    if let Some(preamble) = preamble
        && !preamble.is_empty()
    {
        system_parts.push(proto::Part {
            data: Some(proto::part::Data::Text(preamble)),
            thought: false,
            thought_signature: Vec::new(),
            part_metadata: None,
        });
    }
    for content in history_system {
        if !content.is_empty() {
            system_parts.push(proto::Part {
                data: Some(proto::part::Data::Text(content)),
                thought: false,
                thought_signature: Vec::new(),
                part_metadata: None,
            });
        }
    }
    let system_instruction = if system_parts.is_empty() {
        None
    } else {
        Some(proto::Content {
            parts: system_parts,
            role: "model".to_string(),
        })
    };

    // Handle generation config
    let generation_config = if temperature.is_some() || max_tokens.is_some() {
        Some(proto::GenerationConfig {
            temperature: temperature.map(|t| t as f32),
            max_output_tokens: max_tokens.map(|t| t as i32),
            ..Default::default()
        })
    } else {
        None
    };

    // Handle tools (functions)
    let tools = if !tools.is_empty() {
        let function_declarations = tools
            .into_iter()
            .map(|tool| {
                Ok(proto::FunctionDeclaration {
                    name: tool.name,
                    description: tool.description,
                    parameters: tool_parameters_to_proto_schema(&tool.parameters)?,
                    ..Default::default()
                })
            })
            .collect::<Result<Vec<_>, CompletionError>>()?;

        vec![proto::Tool {
            function_declarations,
            code_execution: None,
        }]
    } else {
        vec![]
    };

    Ok(GenerateContentRequest {
        model: format!("models/{}", model),
        contents,
        tools,
        safety_settings: vec![],
        generation_config,
        tool_config: None,
        system_instruction,
        cached_content: String::new(),
    })
}

// Convert Rig message to gRPC Content
fn rig_message_to_grpc_content(msg: message::Message) -> Result<proto::Content, CompletionError> {
    match msg {
        message::Message::System { .. } => Err(CompletionError::RequestError(
            "System messages must be sent via Gemini gRPC system_instruction".into(),
        )),
        message::Message::User { content } => {
            let parts = content
                .into_iter()
                .map(rig_user_content_to_grpc_part)
                .collect::<Result<Vec<_>, _>>()?;

            Ok(proto::Content {
                parts,
                role: "user".to_string(),
            })
        }
        message::Message::Assistant { content, .. } => {
            let parts = content
                .into_iter()
                .map(rig_assistant_content_to_grpc_part)
                .collect::<Result<Vec<_>, _>>()?;

            Ok(proto::Content {
                parts,
                role: "model".to_string(),
            })
        }
    }
}

fn split_system_messages_from_history(
    history: OneOrMany<message::Message>,
) -> (Vec<String>, Vec<message::Message>) {
    let mut system = Vec::new();
    let mut remaining = Vec::new();

    for message in history {
        match message {
            message::Message::System { content } => system.push(content),
            other => remaining.push(other),
        }
    }

    (system, remaining)
}

// Convert Rig UserContent to gRPC Part
fn rig_user_content_to_grpc_part(
    content: message::UserContent,
) -> Result<proto::Part, CompletionError> {
    match content {
        message::UserContent::Text(message::Text { text, .. }) => Ok(proto::Part {
            data: Some(proto::part::Data::Text(text)),
            thought: false,
            thought_signature: Vec::new(),
            part_metadata: None,
        }),
        message::UserContent::ToolResult(result) => {
            let mut values = result
                .content
                .into_iter()
                .map(|content| match content {
                    message::ToolResultContent::Text(t) => Ok(serde_json::Value::String(t.text)),
                    message::ToolResultContent::Json { value } => Ok(value),
                    message::ToolResultContent::Image(_) => Err(CompletionError::RequestError(
                        "Gemini gRPC does not support images in tool results".into(),
                    )),
                })
                .collect::<Result<Vec<_>, _>>()?;
            let result_value = if values.len() == 1 {
                values.remove(0)
            } else {
                serde_json::Value::Array(values)
            };

            let response_struct =
                json_to_prost_struct(serde_json::json!({ "result": result_value }))?;

            Ok(proto::Part {
                data: Some(proto::part::Data::FunctionResponse(
                    proto::FunctionResponse {
                        name: result.id,
                        response: Some(response_struct),
                        id: result.call_id.unwrap_or_default(),
                    },
                )),
                thought: false,
                thought_signature: Vec::new(),
                part_metadata: None,
            })
        }
        message::UserContent::Image(img) => {
            let Some(media_type) = img.media_type else {
                return Err(CompletionError::RequestError(
                    "Media type for image is required for Gemini".into(),
                ));
            };

            match media_type {
                message::ImageMediaType::JPEG
                | message::ImageMediaType::PNG
                | message::ImageMediaType::WEBP
                | message::ImageMediaType::HEIC
                | message::ImageMediaType::HEIF => {}
                _ => {
                    return Err(CompletionError::RequestError(
                        format!("Unsupported image media type {media_type:?}").into(),
                    ));
                }
            }

            let mime_type = media_type.to_mime_type().to_string();

            let data = match img.data {
                message::DocumentSourceKind::Url(file_uri) => {
                    return Ok(proto::Part {
                        data: Some(proto::part::Data::FileData(proto::FileData {
                            mime_type,
                            file_uri,
                        })),
                        thought: false,
                        thought_signature: Vec::new(),
                        part_metadata: None,
                    });
                }
                message::DocumentSourceKind::Raw(bytes) => bytes,
                message::DocumentSourceKind::Base64(data)
                | message::DocumentSourceKind::String(data) => decode_base64_bytes(&data)?,
                message::DocumentSourceKind::Unknown => {
                    return Err(CompletionError::RequestError(
                        "Image content has no body".into(),
                    ));
                }
                _ => {
                    return Err(CompletionError::RequestError(
                        "Unsupported document source kind".into(),
                    ));
                }
            };

            Ok(proto::Part {
                data: Some(proto::part::Data::InlineData(proto::Blob {
                    mime_type,
                    data,
                })),
                thought: false,
                thought_signature: Vec::new(),
                part_metadata: None,
            })
        }
        _ => Err(CompletionError::RequestError(
            "Unsupported user content type".into(),
        )),
    }
}

// Convert Rig AssistantContent to gRPC Part
fn rig_assistant_content_to_grpc_part(
    content: message::AssistantContent,
) -> Result<proto::Part, CompletionError> {
    match content {
        message::AssistantContent::Text(message::Text { text, .. }) => Ok(proto::Part {
            data: Some(proto::part::Data::Text(text)),
            thought: false,
            thought_signature: Vec::new(),
            part_metadata: None,
        }),
        message::AssistantContent::ToolCall(tool_call) => {
            let args = json_to_prost_struct(tool_call.function.arguments)?;

            Ok(proto::Part {
                data: Some(proto::part::Data::FunctionCall(proto::FunctionCall {
                    name: tool_call.function.name,
                    args: Some(args),
                    id: tool_call.call_id.unwrap_or(tool_call.id),
                })),
                thought: false,
                thought_signature: decode_optional_base64(tool_call.signature)?,
                part_metadata: None,
            })
        }
        message::AssistantContent::Reasoning(reasoning) => Ok(proto::Part {
            data: Some(proto::part::Data::Text(reasoning.display_text())),
            thought: true,
            thought_signature: decode_optional_base64(
                reasoning.first_signature().map(|s| s.to_string()),
            )?,
            part_metadata: None,
        }),
        _ => Err(CompletionError::RequestError(
            "Unsupported assistant content type".into(),
        )),
    }
}

// Convert gRPC GenerateContentResponse to Rig CompletionResponse
impl TryFrom<GenerateContentResponse> for completion::CompletionResponse<GenerateContentResponse> {
    type Error = CompletionError;

    fn try_from(response: GenerateContentResponse) -> Result<Self, Self::Error> {
        let candidate = response.candidates.first().ok_or_else(|| {
            CompletionError::ResponseError("No response candidates in response".into())
        })?;

        let content_ref = candidate.content.as_ref().ok_or_else(|| {
            CompletionError::ResponseError(format!(
                "Gemini candidate missing content (finish_reason={})",
                candidate.finish_reason
            ))
        })?;

        let mut assistant_contents = Vec::new();

        for part in &content_ref.parts {
            let assistant_content = match &part.data {
                Some(proto::part::Data::Text(text)) => {
                    if part.thought {
                        completion::AssistantContent::Reasoning(Reasoning::new_with_signature(
                            text,
                            encode_optional_base64(&part.thought_signature),
                        ))
                    } else {
                        completion::AssistantContent::text(text)
                    }
                }
                Some(proto::part::Data::InlineData(inline_data)) => {
                    let mime_type = message::MediaType::from_mime_type(&inline_data.mime_type);
                    match mime_type {
                        Some(message::MediaType::Image(media_type)) => {
                            let b64 =
                                base64::engine::general_purpose::STANDARD.encode(&inline_data.data);
                            completion::AssistantContent::image_base64(
                                b64,
                                Some(media_type),
                                Some(message::ImageDetail::default()),
                            )
                        }
                        _ => {
                            return Err(CompletionError::ResponseError(format!(
                                "Unsupported media type {mime_type:?}"
                            )));
                        }
                    }
                }
                Some(proto::part::Data::FunctionCall(function_call)) => {
                    let args = function_call
                        .args
                        .as_ref()
                        .map(prost_struct_to_json)
                        .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));

                    let mut tool_call = message::ToolCall::new(
                        if function_call.id.is_empty() {
                            function_call.name.clone()
                        } else {
                            function_call.id.clone()
                        },
                        message::ToolFunction::new(function_call.name.clone(), args),
                    );

                    if !function_call.id.is_empty() {
                        tool_call = tool_call.with_call_id(function_call.id.clone());
                    }

                    tool_call =
                        tool_call.with_signature(encode_optional_base64(&part.thought_signature));

                    completion::AssistantContent::ToolCall(tool_call)
                }
                _ => {
                    return Err(CompletionError::ResponseError(
                        "Response did not contain a message or tool call".into(),
                    ));
                }
            };

            assistant_contents.push(assistant_content);
        }

        let choice = OneOrMany::many(assistant_contents).map_err(|_| {
            CompletionError::ResponseError(
                "Response contained no message or tool call (empty)".to_owned(),
            )
        })?;

        let usage = response
            .usage_metadata
            .as_ref()
            .map(|usage| completion::Usage {
                input_tokens: usage.prompt_token_count as u64,
                output_tokens: usage.candidates_token_count as u64,
                total_tokens: usage.total_token_count as u64,
                cached_input_tokens: usage.cached_content_token_count as u64,
                cache_creation_input_tokens: 0,
                tool_use_prompt_tokens: 0,
                reasoning_tokens: 0,
            })
            .unwrap_or_default();

        Ok(completion::CompletionResponse {
            choice,
            usage,
            raw_response: response,
            message_id: None,
        })
    }
}

// Implement ProviderResponseExt for telemetry
impl ProviderResponseExt for GenerateContentResponse {
    type OutputMessage = proto::Candidate;
    type Usage = proto::UsageMetadata;

    fn get_response_id(&self) -> Option<String> {
        if self.response_id.is_empty() {
            None
        } else {
            Some(self.response_id.clone())
        }
    }

    fn get_response_model_name(&self) -> Option<String> {
        if self.model_version.is_empty() {
            None
        } else {
            Some(self.model_version.clone())
        }
    }

    fn get_output_messages(&self) -> Vec<Self::OutputMessage> {
        self.candidates.clone()
    }

    fn get_text_response(&self) -> Option<String> {
        self.candidates.first().and_then(|c| {
            c.content.as_ref().and_then(|content| {
                let text: Vec<String> = content
                    .parts
                    .iter()
                    .filter_map(|part| {
                        if let Some(proto::part::Data::Text(text)) = &part.data {
                            Some(text.clone())
                        } else {
                            None
                        }
                    })
                    .collect();

                if text.is_empty() {
                    None
                } else {
                    Some(text.join("\n"))
                }
            })
        })
    }

    fn get_usage(&self) -> Option<Self::Usage> {
        self.usage_metadata
    }
}

fn decode_base64_bytes(input: &str) -> Result<Vec<u8>, CompletionError> {
    let data = input.trim();

    // Allow `data:<mime>;base64,<data>` inputs.
    let data = if let Some(rest) = data.strip_prefix("data:") {
        rest.split_once(',').map(|(_, b64)| b64).unwrap_or(data)
    } else {
        data
    };

    let mut last_err: Option<String> = None;

    for engine in [
        &base64::engine::general_purpose::STANDARD,
        &base64::engine::general_purpose::URL_SAFE,
        &base64::engine::general_purpose::STANDARD_NO_PAD,
        &base64::engine::general_purpose::URL_SAFE_NO_PAD,
    ] {
        match engine.decode(data) {
            Ok(bytes) => return Ok(bytes),
            Err(err) => last_err = Some(err.to_string()),
        }
    }

    let err = last_err.unwrap_or_else(|| "unknown base64 decode error".to_string());
    Err(CompletionError::RequestError(
        format!("Invalid base64 data: {err}").into(),
    ))
}

fn decode_optional_base64(sig: Option<String>) -> Result<Vec<u8>, CompletionError> {
    let Some(sig) = sig else {
        return Ok(Vec::new());
    };
    decode_base64_bytes(&sig)
}

fn encode_optional_base64(bytes: &[u8]) -> Option<String> {
    if bytes.is_empty() {
        None
    } else {
        Some(base64::engine::general_purpose::STANDARD.encode(bytes))
    }
}

fn json_to_prost_struct(value: serde_json::Value) -> Result<proto::Struct, CompletionError> {
    match value {
        serde_json::Value::Object(map) => Ok(proto::Struct {
            fields: map
                .into_iter()
                .map(|(k, v)| (k, json_to_prost_value(v)))
                .collect(),
        }),
        _ => Err(CompletionError::RequestError(
            "Expected a JSON object for google.protobuf.Struct".into(),
        )),
    }
}

fn json_to_prost_value(value: serde_json::Value) -> proto::Value {
    match value {
        serde_json::Value::Null => proto::Value {
            kind: Some(proto::value::Kind::NullValue(
                proto::NullValue::NullValue as i32,
            )),
        },
        serde_json::Value::Bool(b) => proto::Value {
            kind: Some(proto::value::Kind::BoolValue(b)),
        },
        serde_json::Value::Number(n) => proto::Value {
            kind: Some(proto::value::Kind::NumberValue(
                n.as_f64().unwrap_or_default(),
            )),
        },
        serde_json::Value::String(s) => proto::Value {
            kind: Some(proto::value::Kind::StringValue(s)),
        },
        serde_json::Value::Array(items) => proto::Value {
            kind: Some(proto::value::Kind::ListValue(proto::ListValue {
                values: items.into_iter().map(json_to_prost_value).collect(),
            })),
        },
        serde_json::Value::Object(map) => proto::Value {
            kind: Some(proto::value::Kind::StructValue(proto::Struct {
                fields: map
                    .into_iter()
                    .map(|(k, v)| (k, json_to_prost_value(v)))
                    .collect(),
            })),
        },
    }
}

fn prost_struct_to_json(st: &proto::Struct) -> serde_json::Value {
    let mut out = serde_json::Map::with_capacity(st.fields.len());
    for (k, v) in &st.fields {
        out.insert(k.clone(), prost_value_to_json(v));
    }
    serde_json::Value::Object(out)
}

fn prost_value_to_json(v: &proto::Value) -> serde_json::Value {
    match &v.kind {
        None | Some(proto::value::Kind::NullValue(_)) => serde_json::Value::Null,
        Some(proto::value::Kind::BoolValue(b)) => serde_json::Value::Bool(*b),
        Some(proto::value::Kind::NumberValue(n)) => serde_json::Number::from_f64(*n)
            .map(serde_json::Value::Number)
            .unwrap_or(serde_json::Value::Null),
        Some(proto::value::Kind::StringValue(s)) => serde_json::Value::String(s.clone()),
        Some(proto::value::Kind::StructValue(st)) => prost_struct_to_json(st),
        Some(proto::value::Kind::ListValue(list)) => {
            serde_json::Value::Array(list.values.iter().map(prost_value_to_json).collect())
        }
    }
}

// Convert the JSON Schema carried by `ToolDefinition.parameters` into the typed
// `proto::Schema` expected by `FunctionDeclaration.parameters`.
//
// Without this, every tool was sent to Gemini with `parameters = None`, which
// caused the model to invoke tools with no argument shape (issue #1710).
//
// An empty object schema (`{"type": "object", "properties": {}}`, the default
// when a tool takes no arguments) is mapped to `None` rather than a vacuous
// schema, matching the convention used by `rig-core::providers::gemini`.
fn tool_parameters_to_proto_schema(
    value: &serde_json::Value,
) -> Result<Option<proto::Schema>, CompletionError> {
    tool_parameters_to_schema(value.clone()).map(|schema| schema.map(gemini_schema_to_proto_schema))
}

fn gemini_schema_to_proto_schema(schema: GeminiSchema) -> proto::Schema {
    proto::Schema {
        r#type: json_type_to_proto_type(&schema.r#type) as i32,
        format: schema.format.unwrap_or_default(),
        description: schema.description.unwrap_or_default(),
        nullable: schema.nullable.unwrap_or(false),
        r#enum: schema.r#enum.unwrap_or_default(),
        items: schema
            .items
            .map(|items| Box::new(gemini_schema_to_proto_schema(*items))),
        properties: schema
            .properties
            .unwrap_or_default()
            .into_iter()
            .map(|(name, schema)| (name, gemini_schema_to_proto_schema(schema)))
            .collect(),
        required: schema.required.unwrap_or_default(),
    }
}

fn json_type_to_proto_type(t: &str) -> proto::Type {
    match t {
        "string" => proto::Type::String,
        "number" => proto::Type::Number,
        "integer" => proto::Type::Integer,
        "boolean" => proto::Type::Boolean,
        "array" => proto::Type::Array,
        "object" => proto::Type::Object,
        "null" => proto::Type::Null,
        _ => proto::Type::Unspecified,
    }
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod tests {
    use super::*;

    // ============================================================
    // rpc_error — pins the from_provider_body usage on the RPC error path
    // ============================================================

    #[test]
    fn rpc_error_preserves_status_text_without_http_status() {
        let status = tonic::Status::unavailable("boom");
        let expected = status.to_string();

        let err = rpc_error(status);

        // The raw provider error text is preserved verbatim, and there is no
        // HTTP status because gRPC is a non-HTTP transport.
        assert_eq!(err.provider_response_body(), Some(expected.as_str()));
        assert_eq!(err.provider_response_status(), None);
    }

    #[test]
    fn test_decode_base64_bytes_accepts_url_safe_with_padding() {
        assert!(matches!(
            decode_base64_bytes("_-wgVQA="),
            Ok(bytes) if bytes == vec![0xFF, 0xEC, 0x20, 0x55, 0x00]
        ));
    }

    #[test]
    fn test_decode_base64_bytes_accepts_url_safe_no_pad() {
        assert!(matches!(
            decode_base64_bytes("_-wgVQA"),
            Ok(bytes) if bytes == vec![0xFF, 0xEC, 0x20, 0x55, 0x00]
        ));
    }

    #[test]
    fn test_decode_base64_bytes_accepts_standard_no_pad() {
        assert!(matches!(
            decode_base64_bytes("Zg"),
            Ok(bytes) if bytes == b"f".to_vec()
        ));
    }

    #[test]
    fn test_decode_base64_bytes_accepts_data_uri_prefix() {
        assert!(matches!(
            decode_base64_bytes("data:text/plain;base64,Zm9v"),
            Ok(bytes) if bytes == b"foo".to_vec()
        ));
    }

    // ============================================================
    // tool_parameters_to_proto_schema — regression coverage for #1710
    // ============================================================

    #[test]
    fn tool_params_empty_object_maps_to_none() {
        let v = serde_json::json!({"type": "object", "properties": {}});
        assert!(tool_parameters_to_proto_schema(&v).unwrap().is_none());
    }

    #[test]
    fn tool_params_null_maps_to_none() {
        assert!(
            tool_parameters_to_proto_schema(&serde_json::Value::Null)
                .unwrap()
                .is_none()
        );
    }

    #[test]
    fn tool_params_object_with_scalar_properties_round_trips() {
        let v = serde_json::json!({
            "type": "object",
            "properties": {
                "city":      { "type": "string",  "description": "City name" },
                "max_price": { "type": "integer", "description": "Cap, USD"  }
            },
            "required": ["city"]
        });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");
        assert_eq!(schema.r#type, proto::Type::Object as i32);
        assert_eq!(schema.required, vec!["city".to_string()]);
        assert_eq!(schema.properties.len(), 2);

        let city = schema.properties.get("city").expect("city prop");
        assert_eq!(city.r#type, proto::Type::String as i32);
        assert_eq!(city.description, "City name");

        let max_price = schema.properties.get("max_price").expect("max_price prop");
        assert_eq!(max_price.r#type, proto::Type::Integer as i32);
    }

    #[test]
    fn tool_params_array_with_typed_items() {
        let v = serde_json::json!({
            "type": "array",
            "items": { "type": "string" }
        });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");
        assert_eq!(schema.r#type, proto::Type::Array as i32);
        let items = schema.items.expect("items");
        assert_eq!(items.r#type, proto::Type::String as i32);
    }

    #[test]
    fn tool_params_enum_strings_preserved() {
        let v = serde_json::json!({
            "type": "string",
            "enum": ["celsius", "fahrenheit"]
        });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");
        assert_eq!(schema.r#type, proto::Type::String as i32);
        assert_eq!(
            schema.r#enum,
            vec!["celsius".to_string(), "fahrenheit".to_string()]
        );
    }

    #[test]
    fn tool_params_resolves_defs_ref_properties() {
        let v = serde_json::json!({
            "type": "object",
            "properties": {
                "destination": { "$ref": "#/$defs/Destination" }
            },
            "required": ["destination"],
            "$defs": {
                "Destination": {
                    "type": "object",
                    "properties": {
                        "city": { "type": "string" },
                        "country_code": { "type": "string" }
                    },
                    "required": ["city"]
                }
            }
        });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");
        let destination = schema
            .properties
            .get("destination")
            .expect("destination prop");

        assert_eq!(destination.r#type, proto::Type::Object as i32);
        assert_eq!(destination.required, vec!["city".to_string()]);
        assert_eq!(
            destination
                .properties
                .get("city")
                .expect("city prop")
                .r#type,
            proto::Type::String as i32
        );
    }

    #[test]
    fn tool_params_nullable_type_array_preserves_non_null_type() {
        let v = serde_json::json!({
            "type": "object",
            "properties": {
                "nickname": { "type": ["null", "string"] }
            }
        });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");
        let nickname = schema.properties.get("nickname").expect("nickname prop");

        assert_eq!(nickname.r#type, proto::Type::String as i32);
        assert!(nickname.nullable);
    }

    #[test]
    fn tool_params_any_of_uses_non_null_schema() {
        let v = serde_json::json!({
            "anyOf": [
                { "type": "null" },
                {
                    "type": "object",
                    "properties": {
                        "query": { "type": "string" }
                    },
                    "required": ["query"]
                }
            ]
        });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");

        assert_eq!(schema.r#type, proto::Type::Object as i32);
        assert!(schema.nullable);
        assert_eq!(schema.required, vec!["query".to_string()]);
        assert_eq!(
            schema.properties.get("query").expect("query prop").r#type,
            proto::Type::String as i32
        );
    }

    #[test]
    fn tool_params_array_without_items_defaults_to_string_items() {
        let v = serde_json::json!({ "type": "array" });

        let schema = tool_parameters_to_proto_schema(&v)
            .expect("schema conversion")
            .expect("schema");

        assert_eq!(schema.r#type, proto::Type::Array as i32);
        assert_eq!(
            schema.items.expect("items").r#type,
            proto::Type::String as i32
        );
    }

    #[test]
    fn create_grpc_request_populates_tool_parameters() {
        use rig_core::completion::ToolDefinition;

        let tool = ToolDefinition {
            name: "get_weather".to_string(),
            description: "Look up the current weather for a city.".to_string(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "city": { "type": "string", "description": "City name" }
                },
                "required": ["city"]
            }),
        };

        let req = create_grpc_request(
            "gemini-2.5-flash".to_string(),
            CompletionRequest {
                model: None,
                preamble: None,
                chat_history: OneOrMany::one(message::Message::user("forecast in Berlin?")),
                documents: Vec::new(),
                tools: vec![tool],
                temperature: None,
                max_tokens: None,
                tool_choice: None,
                additional_params: None,
                output_schema: None,
                record_telemetry_content: false,
            },
        )
        .expect("request build");

        assert_eq!(req.tools.len(), 1);
        let tool = req.tools.first().expect("tool entry");
        let decl = tool
            .function_declarations
            .first()
            .expect("function declaration");
        assert_eq!(decl.name, "get_weather");

        // The regression in #1710 was `parameters: None` here.
        let params = decl.parameters.as_ref().expect("parameters populated");
        assert_eq!(params.r#type, proto::Type::Object as i32);
        assert_eq!(params.required, vec!["city".to_string()]);
        assert!(params.properties.contains_key("city"));
    }
}