diff --git a/crates/adaptive/src/acg/request_surfaces/mod.rs b/crates/adaptive/src/acg/request_surfaces/mod.rs index 922148fa0..b4834aa56 100644 --- a/crates/adaptive/src/acg/request_surfaces/mod.rs +++ b/crates/adaptive/src/acg/request_surfaces/mod.rs @@ -43,11 +43,13 @@ pub(crate) trait RequestSurfaceApplier: Send + Sync { } impl RequestSurface { - fn from_provider_surface(surface: ProviderSurface) -> Self { + fn from_provider_surface(surface: ProviderSurface) -> Option { match surface { - ProviderSurface::OpenAIChat => Self::OpenAIChat, - ProviderSurface::OpenAIResponses => Self::OpenAIResponses, - ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::OpenAIChat => Some(Self::OpenAIChat), + ProviderSurface::OpenAIResponses => Some(Self::OpenAIResponses), + ProviderSurface::AnthropicMessages => Some(Self::AnthropicMessages), + // Gemini generateContent ACG request editing is intentionally unsupported. + ProviderSurface::GeminiGenerateContent => None, } } @@ -90,13 +92,16 @@ impl RequestSurface { pub(crate) fn resolve_request_surface_from_request( request: &LlmRequest, ) -> crate::acg::Result { - detect_request_surface(&request.content) - .map(RequestSurface::from_provider_surface) - .ok_or_else(|| { - crate::acg::AcgError::Internal( - "unable to resolve request surface from request shape".to_string(), - ) - }) + let Some(surface) = detect_request_surface(&request.content) else { + return Err(crate::acg::AcgError::Internal( + "unable to resolve request surface from request shape".to_string(), + )); + }; + RequestSurface::from_provider_surface(surface).ok_or_else(|| { + crate::acg::AcgError::Internal(format!( + "resolved request surface {surface:?} does not have an ACG applier" + )) + }) } #[cfg_attr(not(test), allow(dead_code))] diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index 36eeca3d1..29f1b6555 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -338,6 +338,28 @@ fn lossy_request_shape(surface: ProviderSurface, content: &Json) -> bool { .is_some_and(|blocks| blocks.iter().any(lossy_system_block)) } ProviderSurface::OpenAIResponses => false, + ProviderSurface::GeminiGenerateContent => { + object + .get("generationConfig") + .is_some_and(|generation_config| { + let Some(generation_config) = generation_config.as_object() else { + return true; + }; + generation_config.keys().any(|key| { + !matches!( + key.as_str(), + "temperature" | "topP" | "maxOutputTokens" | "stopSequences" + ) + }) + }) + || object + .get("systemInstruction") + .is_some_and(lossy_gemini_system_instruction) + || object + .get("contents") + .and_then(Json::as_array) + .is_some_and(|items| items.iter().any(lossy_gemini_content_item)) + } } } @@ -362,6 +384,129 @@ fn lossy_system_block(block: &Json) -> bool { .any(|key| !matches!(key.as_str(), "type" | "text" | "cache_control")) } +/// Whether a Gemini `systemInstruction` would lose detail in the normalized +/// request key. The Gemini codec flattens it to a single system text string, so +/// only one non-empty plain text part without sibling fields is lossless. +fn lossy_gemini_system_instruction(value: &Json) -> bool { + let Some(object) = value.as_object() else { + return true; + }; + if object.keys().any(|key| key != "parts") { + return true; + } + let Some(parts) = object.get("parts").and_then(Json::as_array) else { + return true; + }; + let [part] = parts.as_slice() else { + return true; + }; + let Some(part_object) = part.as_object() else { + return true; + }; + if part_object.len() != 1 { + return true; + } + part_object + .get("text") + .and_then(Json::as_str) + .is_none_or(str::is_empty) +} + +fn lossy_gemini_content_item(item: &Json) -> bool { + let Some(object) = item.as_object() else { + return true; + }; + if object + .keys() + .any(|key| !matches!(key.as_str(), "role" | "parts")) + { + return true; + } + let Some(parts) = object.get("parts").and_then(Json::as_array) else { + return true; + }; + + let mut plain_text_parts = 0usize; + for part in parts { + if lossy_gemini_part(part, &mut plain_text_parts) { + return true; + } + } + plain_text_parts > 1 +} + +fn lossy_gemini_part(part: &Json, plain_text_parts: &mut usize) -> bool { + let Some(object) = part.as_object() else { + return true; + }; + if object.get("thought").and_then(Json::as_bool) == Some(true) { + return true; + } + let data_keys = object + .keys() + .filter(|key| is_gemini_part_data_key(key)) + .collect::>(); + if data_keys.len() > 1 { + return true; + } + let Some(data_key) = data_keys.first().map(|key| key.as_str()) else { + return true; + }; + match data_key { + "text" => { + if !object.get("text").is_some_and(Json::is_string) { + return true; + } + if object.len() == 1 { + *plain_text_parts += 1; + } + false + } + "functionCall" => { + object.keys().any(|key| key != "functionCall") + || match object.get("functionCall").and_then(Json::as_object) { + Some(fc) => { + fc.keys() + .any(|key| !matches!(key.as_str(), "name" | "id" | "args")) + || fc.get("args").is_some_and(|args| !args.is_object()) + } + None => true, + } + } + "functionResponse" => { + object.keys().any(|key| key != "functionResponse") + || match object.get("functionResponse").and_then(Json::as_object) { + Some(fr) => { + fr.keys().any(|key| { + !matches!(key.as_str(), "id" | "name" | "response" | "parts") + }) || match ( + fr.get("id").and_then(Json::as_str), + fr.get("name").and_then(Json::as_str), + ) { + (Some(id), Some(name)) => id != name, + _ => false, + } + } + None => true, + } + } + _ => false, + } +} + +fn is_gemini_part_data_key(key: &str) -> bool { + matches!( + key, + "text" + | "inlineData" + | "fileData" + | "functionCall" + | "functionResponse" + | "executableCode" + | "codeExecutionResult" + ) +} + /// Whether the raw request body carries a non-empty `messages` (chat) or `input` /// (Responses) array — the content a decode is expected to preserve. fn raw_has_messages(request: &LlmRequest) -> bool { diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs index 02a8bf0d4..95c405c78 100644 --- a/crates/adaptive/src/response_cache/replay.rs +++ b/crates/adaptive/src/response_cache/replay.rs @@ -12,8 +12,9 @@ use serde_json::{Map, Value as Json, json}; /// /// A strict streaming client parses only its provider's wire chunks (Anthropic /// `message_start → content_block_delta → … → message_stop`, OpenAI Chat -/// `chat.completion.chunk` deltas, OpenAI Responses lifecycle events) — replaying -/// the aggregate as one frame breaks such clients even though the body is correct. +/// `chat.completion.chunk` deltas, OpenAI Responses lifecycle events, or Gemini +/// `GenerateContentResponse` chunks) — replaying the aggregate as one frame +/// breaks such clients even though the body is correct. /// The surface is detected from the stored aggregate's own shape (the codec /// finalizer's output, or a buffered body of the same shape), so no codec handle /// is needed at the call sites. An unrecognized shape falls back to a @@ -79,6 +80,11 @@ fn synthesize_replay_chunks(aggregate: &Json) -> Option> { ProviderSurface::AnthropicMessages => synthesize_anthropic_chunks(aggregate), ProviderSurface::OpenAIChat => synthesize_chat_chunks(aggregate), ProviderSurface::OpenAIResponses => synthesize_responses_chunks(aggregate), + // Gemini streaming events are GenerateContentResponse objects; a stored + // aggregate is a valid single native chunk. `replay_is_lossy` still + // re-aggregates it and rejects shapes the streaming collector cannot + // preserve exactly. + ProviderSurface::GeminiGenerateContent => vec![aggregate.clone()], }) } diff --git a/crates/adaptive/tests/unit/acg/request_surface_tests.rs b/crates/adaptive/tests/unit/acg/request_surface_tests.rs index b8014f682..6bcc24e1e 100644 --- a/crates/adaptive/tests/unit/acg/request_surface_tests.rs +++ b/crates/adaptive/tests/unit/acg/request_surface_tests.rs @@ -301,6 +301,18 @@ fn test_request_surface_resolution_and_passthrough_support_cover_matrix() { super::resolve_request_surface_from_request(&anthropic_request).unwrap(), RequestSurface::AnthropicMessages ); + let gemini_request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + }), + }; + assert!(matches!( + super::resolve_request_surface_from_request(&gemini_request), + Err(crate::acg::AcgError::Internal(message)) + if message.contains("does not have an ACG applier") + )); assert!(matches!( super::resolve_request_surface_from_request(&invalid_request), Err(crate::acg::AcgError::Internal(message)) diff --git a/crates/adaptive/tests/unit/response_cache/key_tests.rs b/crates/adaptive/tests/unit/response_cache/key_tests.rs index aa03d30f5..54912deb6 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -593,6 +593,155 @@ fn chat_shaped_requests_key_on_the_detected_decode() { ); } +#[test] +fn gemini_shaped_requests_key_on_the_detected_decode() { + let request = request(json!({ + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": {"temperature": 0.0} + })); + let (body, effective_codec) = resolved_body("gemini_generate_content", &request); + assert_eq!(effective_codec, Some("gemini_generate_content")); + assert_ne!( + body, request.content, + "Gemini requests must use the normalized decode when it is lossless" + ); +} + +#[test] +fn gemini_unmodeled_generation_config_fields_do_not_collide() { + let config = cache_all_config(); + let mime_request = |response_mime_type: &str| { + request(json!({ + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": { + "temperature": 0.0, + "responseMimeType": response_mime_type + } + })) + }; + let text = mime_request("text/plain"); + let json = mime_request("application/json"); + for request in [&text, &json] { + let (body, effective_codec) = resolved_body("gemini_generate_content", request); + assert_eq!( + effective_codec, None, + "unmodeled Gemini generationConfig fields must force raw fallback" + ); + assert_eq!( + body, request.content, + "raw fallback must preserve answer-affecting Gemini generationConfig fields" + ); + } + assert_ne!( + key_of("gemini_generate_content", &text, &config), + key_of("gemini_generate_content", &json, &config), + "distinct unmodeled Gemini generationConfig values must not share a cache key" + ); +} + +#[test] +fn gemini_malformed_generation_config_raw_keys() { + let request = request(json!({ + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": "not an object" + })); + let (body, effective_codec) = resolved_body("gemini_generate_content", &request); + assert_eq!( + effective_codec, None, + "malformed Gemini generationConfig must force raw fallback" + ); + assert_eq!(body, request.content); +} + +#[test] +fn gemini_unmodeled_system_instruction_fields_do_not_collide() { + let config = cache_all_config(); + let with_signature = |signature: &str| { + request(json!({ + "model": "gemini-2.5-flash", + "systemInstruction": { + "parts": [{"text": "be concise", "thoughtSignature": signature}] + }, + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })) + }; + let first = with_signature("sig_FIRST=="); + let second = with_signature("sig_SECOND=="); + for request in [&first, &second] { + let (body, effective_codec) = resolved_body("gemini_generate_content", request); + assert_eq!( + effective_codec, None, + "unmodeled Gemini systemInstruction fields must force raw fallback" + ); + assert_eq!( + body, request.content, + "raw fallback must preserve systemInstruction metadata" + ); + } + assert_ne!( + key_of("gemini_generate_content", &first, &config), + key_of("gemini_generate_content", &second, &config), + "distinct Gemini systemInstruction metadata must not share a cache key" + ); +} + +#[test] +fn gemini_function_call_part_metadata_forces_raw_keying() { + let request = request(json!({ + "model": "gemini-2.5-flash", + "contents": [{ + "role": "model", + "parts": [{ + "functionCall": {"id": "call_1", "name": "lookup", "args": {"q": "x"}}, + "thoughtSignature": "sig_CALL==" + }] + }] + })); + let (body, effective_codec) = resolved_body("gemini_generate_content", &request); + assert_eq!( + effective_codec, None, + "functionCall part metadata is provider-native and must raw-key" + ); + assert_eq!(body, request.content); +} + +#[test] +fn gemini_function_response_name_differences_do_not_collide() { + let config = cache_all_config(); + let response_with_name = |name: &str| { + request(json!({ + "model": "gemini-2.5-flash", + "contents": [{ + "role": "user", + "parts": [{ + "functionResponse": { + "id": "call_1", + "name": name, + "response": {"ok": true} + } + }] + }] + })) + }; + let first = response_with_name("lookup"); + let second = response_with_name("search"); + for request in [&first, &second] { + let (body, effective_codec) = resolved_body("gemini_generate_content", request); + assert_eq!( + effective_codec, None, + "Gemini functionResponse.name is native context and must raw-key when it differs from id" + ); + assert_eq!(body, request.content); + } + assert_ne!( + key_of("gemini_generate_content", &first, &config), + key_of("gemini_generate_content", &second, &config) + ); +} + #[test] fn undetectable_shape_falls_back_to_raw_keying() { // No `messages`/`input`/`system` top-level key: no surface detects, so diff --git a/crates/adaptive/tests/unit/response_cache/replay_tests.rs b/crates/adaptive/tests/unit/response_cache/replay_tests.rs index 27e26208c..7d5fa65a8 100644 --- a/crates/adaptive/tests/unit/response_cache/replay_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/replay_tests.rs @@ -78,6 +78,73 @@ fn replay_chunks_roundtrip_through_the_codecs() { } } +#[test] +fn gemini_replay_uses_a_valid_native_stream_event() { + let aggregate = json!({ + "candidates": [{ + "index": 0, + "content": { + "role": "model", + "parts": [ + {"text": "hello", "thoughtSignature": "sig_TEXT=="}, + { + "functionCall": { + "id": "call_1", + "name": "lookup", + "args": {"q": "x"} + }, + "thoughtSignature": "sig_CALL==" + } + ] + }, + "finishReason": "STOP", + "safetyRatings": [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"} + ], + "groundingMetadata": {"webSearchQueries": ["example query"]} + }], + "usageMetadata": { + "promptTokenCount": 9, + "candidatesTokenCount": 3, + "totalTokenCount": 12 + }, + "modelVersion": "gemini-2.5-flash", + "responseId": "resp_1" + }); + let chunks = synthesize_replay_chunks(&aggregate).expect("gemini shape"); + assert_eq!( + chunks, + vec![aggregate.clone()], + "a GenerateContentResponse aggregate is already a native Gemini stream event" + ); + assert!( + !replay_is_lossy(&aggregate), + "the Gemini streaming codec must reassemble the native replay exactly" + ); +} + +#[test] +fn gemini_replay_rejects_multi_candidate_aggregates_as_lossy() { + let aggregate = json!({ + "candidates": [ + { + "index": 0, + "content": {"role": "model", "parts": [{"text": "first"}]}, + "finishReason": "STOP" + }, + { + "index": 1, + "content": {"role": "model", "parts": [{"text": "second"}]}, + "finishReason": "STOP" + } + ] + }); + assert!( + replay_is_lossy(&aggregate), + "Gemini streaming replay must not serve aggregates with candidates the collector cannot preserve" + ); +} + #[test] fn responses_replay_sequence_numbers_are_contiguous() { let aggregate = json!({ diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index f1d55a6cf..5adc048e3 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -160,6 +160,8 @@ pub enum BuiltinLlmCodec { OpenAiResponses, /// Anthropic Messages request and response payloads. AnthropicMessages, + /// Gemini generateContent request and response payloads. + GeminiGenerateContent, } impl BuiltinLlmCodec { @@ -170,6 +172,7 @@ impl BuiltinLlmCodec { Self::OpenAiChat => "openai_chat", Self::OpenAiResponses => "openai_responses", Self::AnthropicMessages => "anthropic_messages", + Self::GeminiGenerateContent => "gemini_generate_content", } } } diff --git a/crates/core/src/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs new file mode 100644 index 000000000..c4137f23d --- /dev/null +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -0,0 +1,2923 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Built-in codec for the Gemini generateContent API. +//! +//! Implements [`LlmCodec`] (request decode/encode) and [`LlmResponseCodec`] +//! (response decode) for the Gemini generateContent API format. + +use std::collections::HashMap; + +use serde::Deserialize; + +use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; +use crate::error::{FlowError, Result}; +use crate::json::Json; + +use super::request::{ + AnnotatedLlmRequest, ContentPart, FunctionCall, FunctionDefinition, GenerationParams, Message, + MessageContent, ToolCall, ToolDefinition, +}; +use super::resolve::{ProviderSurface, ProviderSurfaceDescriptor}; +use super::response::{ + AnnotatedLlmResponse, FinishReason, ResponseToolCall, Usage, estimate_cost_for_provider, + infer_model_provider, +}; +use super::traits::{LlmCodec, LlmResponseCodec}; + +const GEMINI_PROVIDER: &str = "gemini"; + +// --------------------------------------------------------------------------- +// Public codec struct +// --------------------------------------------------------------------------- + +/// Built-in codec for the Gemini generateContent API. +pub struct GeminiGenerateContentCodec; + +pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDescriptor { + surface: ProviderSurface::GeminiGenerateContent, + detect_request: |obj, _hint| obj.contains_key("contents"), + detect_response: detect_gemini_response, + decode_request: |request| GeminiGenerateContentCodec.decode(request), + decode_response: |raw| GeminiGenerateContentCodec.decode_response(raw), + codec_name: "gemini_generate_content", + request_codec: || std::sync::Arc::new(GeminiGenerateContentCodec), + response_codec: || std::sync::Arc::new(GeminiGenerateContentCodec), + streaming_codec: || Box::new(GeminiGenerateContentStreamingCodec::new()), +}; + +// --------------------------------------------------------------------------- +// Private serde intermediates for response decode +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +struct RawGeminiGenerateContentResponse { + candidates: Option>, + #[serde(rename = "usageMetadata")] + usage_metadata: Option, + #[serde(rename = "modelVersion")] + model_version: Option, + #[serde(rename = "responseId")] + response_id: Option, + #[serde(flatten)] + extra: serde_json::Map, +} + +#[derive(Deserialize)] +struct RawCandidate { + content: Option, + #[serde(rename = "finishReason")] + finish_reason: Option, + #[serde(flatten)] + extra: serde_json::Map, +} + +#[derive(Deserialize)] +struct RawContent { + parts: Option>, +} + +#[derive(Deserialize)] +struct RawUsageMetadata { + #[serde(rename = "promptTokenCount")] + prompt_token_count: Option, + #[serde(rename = "candidatesTokenCount")] + candidates_token_count: Option, + #[serde(rename = "totalTokenCount")] + total_token_count: Option, + #[serde(rename = "cachedContentTokenCount")] + cached_content_token_count: Option, + #[serde(rename = "thoughtsTokenCount")] + thoughts_token_count: Option, +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/// Top-level request keys modeled by this codec; all others go into `extra`. +const MODELED_REQUEST_KEYS: &[&str] = &[ + "contents", + "systemInstruction", + "tools", + "generationConfig", + "model", +]; + +/// Map Gemini `finishReason` (and presence of `functionCall` parts) to [`FinishReason`]. +/// +/// Explicit provider reasons are always honored first. The `has_tool_calls` heuristic +/// is only applied when the provider either omits the reason entirely or signals +/// successful generation (`STOP` / `TOOL_CODE` / `FINISH_REASON_UNSPECIFIED` / empty). +/// This prevents `MAX_TOKENS`, safety reasons, or error codes like +/// `MALFORMED_FUNCTION_CALL` from being obscured by the presence of function-call parts. +/// +/// Reference (GenerateContent API): +/// - absent / empty / `FINISH_REASON_UNSPECIFIED` → derive from content (ToolUse or None) +/// - `STOP` → Complete; or ToolUse when function-call parts are present +/// - `TOOL_CODE` → ToolUse +/// - `MAX_TOKENS` → Length +/// - `SAFETY` / `RECITATION` / `BLOCKLIST` / `PROHIBITED_CONTENT` / `SPII` / +/// `LANGUAGE` / `IMAGE_SAFETY` / `IMAGE_PROHIBITED_CONTENT` / `IMAGE_RECITATION` / +/// `ESCALATION` → ContentFilter +/// - anything else → Unknown (e.g. `MALFORMED_FUNCTION_CALL`, `UNEXPECTED_TOOL_CALL`) +fn map_finish_reason(reason: Option<&str>, has_tool_calls: bool) -> Option { + match reason { + // Absent / unspecified: derive from response content. + None | Some("") | Some("FINISH_REASON_UNSPECIFIED") => { + if has_tool_calls { + Some(FinishReason::ToolUse) + } else { + None + } + } + Some(r) => Some(match r { + // Successful stop: ToolUse when tool calls are present, Complete otherwise. + "STOP" => { + if has_tool_calls { + FinishReason::ToolUse + } else { + FinishReason::Complete + } + } + "TOOL_CODE" => FinishReason::ToolUse, + "MAX_TOKENS" => FinishReason::Length, + // All policy / safety terminations map to ContentFilter, regardless of whether + // any function-call parts happened to be present. + "SAFETY" + | "RECITATION" + | "BLOCKLIST" + | "PROHIBITED_CONTENT" + | "SPII" + | "LANGUAGE" + | "IMAGE_SAFETY" + | "IMAGE_PROHIBITED_CONTENT" + | "IMAGE_RECITATION" + | "ESCALATION" => FinishReason::ContentFilter, + // Unknown / future codes (e.g. MALFORMED_FUNCTION_CALL, UNEXPECTED_TOOL_CALL). + other => FinishReason::Unknown(other.to_string()), + }), + } +} + +fn map_prompt_block_reason(reason: Option<&str>) -> Option { + match reason { + None | Some("") | Some("BLOCK_REASON_UNSPECIFIED") => None, + Some("SAFETY" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "IMAGE_SAFETY") => { + Some(FinishReason::ContentFilter) + } + Some(other) => Some(FinishReason::Unknown(other.to_string())), + } +} + +fn detect_gemini_response(obj: &serde_json::Map) -> bool { + obj.get("candidates").is_some_and(Json::is_array) + || obj.get("promptFeedback").is_some_and(|feedback| { + feedback.as_object().is_some_and(|feedback| { + feedback.get("blockReason").is_some() || feedback.get("safetyRatings").is_some() + }) + }) +} + +fn prompt_feedback_block_reason(extra: &serde_json::Map) -> Result> { + let Some(feedback) = extra.get("promptFeedback") else { + return Ok(None); + }; + let Some(feedback) = feedback.as_object() else { + return Err(FlowError::InvalidArgument( + "Gemini response promptFeedback must be an object".into(), + )); + }; + match feedback.get("blockReason") { + Some(Json::String(reason)) => Ok(Some(reason.as_str())), + Some(_) => Err(FlowError::InvalidArgument( + "Gemini response promptFeedback.blockReason must be a string".into(), + )), + None => Ok(None), + } +} + +/// Extract a `tool_call_id` from a serialized `Message::Tool` JSON object. +/// +/// Returns `Err` when the field is absent, empty, or non-string. +fn extract_tool_call_id(obj: &serde_json::Map) -> Result<&str> { + match obj.get("tool_call_id") { + Some(Json::String(s)) if !s.is_empty() => Ok(s.as_str()), + Some(Json::String(_)) => Err(FlowError::InvalidArgument( + "Gemini encoder: Message::Tool has an empty tool_call_id".into(), + )), + Some(_) => Err(FlowError::InvalidArgument( + "Gemini encoder: Message::Tool tool_call_id must be a string".into(), + )), + None => Err(FlowError::Internal( + "Message::Tool has no tool_call_id".into(), + )), + } +} + +/// Parse an optional Gemini `id` field: absent → `Ok(None)`; present non-empty string +/// → `Ok(Some(s))`; present empty string or non-string → `Err`. +fn parse_optional_id(obj: &serde_json::Map, context: &str) -> Result> { + match obj.get("id") { + None => Ok(None), + Some(Json::String(s)) if !s.is_empty() => Ok(Some(s.clone())), + Some(Json::String(_)) => Err(FlowError::InvalidArgument(format!( + "Gemini generateContent {context}.id must be a non-empty string" + ))), + Some(_) => Err(FlowError::InvalidArgument(format!( + "Gemini generateContent {context}.id must be a string" + ))), + } +} + +fn is_gemini_part_data_key(key: &str) -> bool { + matches!( + key, + "text" + | "inlineData" + | "fileData" + | "functionCall" + | "functionResponse" + | "executableCode" + | "codeExecutionResult" + ) +} + +fn gemini_part_data_keys(obj: &serde_json::Map) -> Vec<&str> { + obj.keys() + .map(String::as_str) + .filter(|key| is_gemini_part_data_key(key)) + .collect() +} + +fn validate_single_gemini_part_data_field<'a>( + obj: &'a serde_json::Map, + context: &str, +) -> Result> { + let keys = gemini_part_data_keys(obj); + if keys.len() > 1 { + return Err(FlowError::InvalidArgument(format!( + "Gemini generateContent {context} part must not contain multiple data fields: {}", + keys.join(", ") + ))); + } + Ok(keys.first().copied()) +} + +/// Convert visible Gemini content parts into normalized message content. +/// +/// Thought parts and tool-call/tool-response parts are not user-visible message +/// content. Text-only content stays as `MessageContent::Text` for compatibility; +/// mixed or metadata-bearing Gemini parts are exposed as provider-native content +/// blocks so middleware can inspect and sanitize them. +fn gemini_parts_to_message_content( + parts: &[Json], + context: &str, +) -> Result> { + let mut texts: Vec = Vec::new(); + let mut content_parts: Vec = Vec::new(); + let mut requires_parts = false; + + for part in parts { + let obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument(format!( + "Gemini generateContent {context} parts entry must be an object" + )) + })?; + if obj.get("thought").and_then(Json::as_bool) == Some(true) { + continue; + } + + let Some(data_key) = validate_single_gemini_part_data_field(obj, context)? else { + requires_parts = true; + content_parts.push(ContentPart::ProviderNative { + provider: GEMINI_PROVIDER.into(), + kind: "unknown".into(), + value: part.clone(), + }); + continue; + }; + + match data_key { + "functionCall" | "functionResponse" => continue, + "text" => { + let text = obj.get("text").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument(format!( + "Gemini generateContent {context} parts[].text must be a string" + )) + })?; + let extra: serde_json::Map = obj + .iter() + .filter(|(key, _)| key.as_str() != "text") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + if !extra.is_empty() { + requires_parts = true; + } + texts.push(text.to_string()); + content_parts.push(ContentPart::Text { + text: text.to_string(), + extra, + }); + } + native_key => { + requires_parts = true; + content_parts.push(ContentPart::ProviderNative { + provider: GEMINI_PROVIDER.into(), + kind: native_key.to_string(), + value: part.clone(), + }); + } + } + } + + Ok(if content_parts.is_empty() { + None + } else if requires_parts { + Some(MessageContent::Parts(content_parts)) + } else { + Some(MessageContent::Text(texts.join("\n"))) + }) +} + +fn validate_gemini_nested_function_response_part(part: &Json) -> Result<&str> { + let obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionResponse.parts entry must be an object".into()) + })?; + let data_key = validate_single_gemini_part_data_field(obj, "functionResponse.parts")?; + if matches!(data_key, Some("functionCall" | "functionResponse")) { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse.parts must not contain nested functionCall/functionResponse" + .into(), + )); + } + if data_key == Some("text") && part.get("text").is_some_and(|v| !v.is_string()) { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse.parts[].text must be a string".into(), + )); + } + Ok(data_key.unwrap_or("unknown")) +} + +fn gemini_function_response_to_message_content(fr: &Json) -> Result { + let response = fr + .get("response") + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini functionResponse is missing required 'response'".into(), + ) + })? + .clone(); + let content_str = serde_json::to_string(&response).unwrap_or_else(|_| "{}".into()); + let Some(parts_value) = fr.get("parts") else { + return Ok(MessageContent::Text(content_str)); + }; + let nested_parts = parts_value.as_array().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionResponse.parts must be an array".into()) + })?; + + let mut content_parts = Vec::with_capacity(nested_parts.len() + 1); + content_parts.push(ContentPart::Text { + text: content_str, + extra: Default::default(), + }); + for nested_part in nested_parts { + let kind = validate_gemini_nested_function_response_part(nested_part)?; + content_parts.push(ContentPart::ProviderNative { + provider: GEMINI_PROVIDER.into(), + kind: kind.to_string(), + value: nested_part.clone(), + }); + } + Ok(MessageContent::Parts(content_parts)) +} + +/// Validate and extract content from a Gemini response `parts` array. +/// +/// Skips functionCall and thought parts. The Gemini API spec defines `thought` +/// as a boolean; a non-boolean value is treated as an ordinary field, not a +/// thought part. +fn extract_parts_message_content(parts: &[Json]) -> Result> { + for part in parts { + if part.get("functionResponse").is_some() { + return Err(FlowError::InvalidArgument( + "Gemini response parts must not contain functionResponse".into(), + )); + } + } + gemini_parts_to_message_content(parts, "response") +} + +/// Extract `functionCall` parts from a Gemini `parts` array as [`ResponseToolCall`]s. +/// +/// Uses the `id` field when the model supplies one; falls back to the function name +/// when absent (older models omit it). Returns `Err` for malformed functionCall +/// shapes (non-object, missing name, empty id, non-object args). Non-functionCall +/// parts are skipped; callers must validate text+functionCall conflicts separately. +fn extract_parts_tool_calls(parts: &[Json]) -> Result>> { + let mut calls: Vec = Vec::new(); + for p in parts { + let Some(fc) = p.get("functionCall") else { + continue; + }; + let fc_obj = fc.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini response functionCall must be an object".into()) + })?; + + let name = match fc_obj.get("name") { + Some(Json::String(s)) if !s.is_empty() => s.clone(), + Some(Json::String(_)) => { + return Err(FlowError::InvalidArgument( + "Gemini response functionCall.name must be a non-empty string".into(), + )); + } + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini response functionCall.name must be a string".into(), + )); + } + None => { + return Err(FlowError::InvalidArgument( + "Gemini response functionCall is missing 'name'".into(), + )); + } + }; + + let id = + parse_optional_id(fc_obj, "response functionCall")?.unwrap_or_else(|| name.clone()); + + let arguments = match fc_obj.get("args") { + None => Json::Object(Default::default()), + Some(v) if v.is_object() => v.clone(), + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini response functionCall.args must be an object".into(), + )); + } + }; + + calls.push(ResponseToolCall { + id, + name, + arguments, + }); + } + Ok(if calls.is_empty() { None } else { Some(calls) }) +} + +/// Map Gemini `usageMetadata` to a normalized [`Usage`] and the raw thinking-token count. +/// +/// Returns `(usage, thoughts_token_count)`. The thoughts token count is kept +/// separate because it belongs in `ApiSpecificResponse::GeminiGenerateContent`, not in the +/// provider-neutral `Usage` struct. +fn map_usage(meta: Option, model: Option<&str>) -> (Option, Option) { + let model_provider = infer_model_provider("google", model); + let Some(m) = meta else { + return (None, None); + }; + let thoughts_token_count = m.thoughts_token_count; + let prompt = m.prompt_token_count; + let completion = m.candidates_token_count; + // When totalTokenCount is absent, compute a fallback from every count that is + // available. Thinking tokens count as billable output and must be included even + // when candidatesTokenCount is missing (e.g. thinking-only partial responses). + let total = m.total_token_count.or_else(|| { + let has_any = prompt.is_some() || completion.is_some() || thoughts_token_count.is_some(); + has_any.then(|| { + prompt.unwrap_or(0) + completion.unwrap_or(0) + thoughts_token_count.unwrap_or(0) + }) + }); + // For cost estimation: treat thinking tokens as additional completion tokens + // so the pricing table applies the output-token rate to them even when + // candidatesTokenCount is absent. + let completion_for_cost = match (completion, thoughts_token_count) { + (Some(c), Some(t)) => Some(c + t), + (Some(c), None) => Some(c), + (None, Some(t)) => Some(t), + (None, None) => None, + }; + let usage_for_cost = Usage { + prompt_tokens: prompt, + completion_tokens: completion_for_cost, + total_tokens: total, + cache_read_tokens: m.cached_content_token_count, + cache_write_tokens: None, + cost: None, + }; + let cost = model + .and_then(|m| estimate_cost_for_provider(model_provider.as_deref(), m, &usage_for_cost)); + let usage = Usage { + prompt_tokens: prompt, + completion_tokens: completion, + total_tokens: total, + cache_read_tokens: m.cached_content_token_count, + cache_write_tokens: None, + cost, + }; + (Some(usage), thoughts_token_count) +} + +/// Validate the shape of a Gemini `systemInstruction` value. +fn validate_system_instruction(val: &Json) -> Result<()> { + let obj = val.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini systemInstruction must be an object".into()) + })?; + if obj.get("role").is_some_and(|v| !v.is_string()) { + return Err(FlowError::InvalidArgument( + "Gemini systemInstruction.role must be a string".into(), + )); + } + let parts = obj + .get("parts") + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini systemInstruction must have a 'parts' field".into()) + })? + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini systemInstruction.parts must be an array".into()) + })?; + for part in parts { + let part_obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini systemInstruction.parts entry must be an object".into(), + ) + })?; + let data_key = validate_single_gemini_part_data_field(part_obj, "systemInstruction")?; + if data_key != Some("text") { + return Err(FlowError::InvalidArgument( + "Gemini systemInstruction.parts entries must be text parts".into(), + )); + } + if part.get("text").is_some_and(|v| !v.is_string()) { + return Err(FlowError::InvalidArgument( + "Gemini systemInstruction.parts[].text must be a string".into(), + )); + } + } + Ok(()) +} + +/// Extract the text content from a Gemini `systemInstruction` value. +fn system_instruction_text(val: &Json) -> Option { + let parts = val.get("parts")?.as_array()?; + let text = parts + .iter() + .filter_map(|p| p.get("text")?.as_str()) + .collect::>() + .join("\n"); + if text.is_empty() { None } else { Some(text) } +} + +/// Convert a single Gemini `contents` item to zero or more normalized [`Message`]s. +/// +/// - `functionResponse` parts (user role) → one `Message::Tool` per part, each with +/// `tool_call_id` set to the `id` field when present, falling back to `name`. +/// - `functionCall` parts (model role) → `Message::Assistant { tool_calls: Some([…]) }` +/// - text parts → plain message content +fn gemini_content_to_messages(content: &Json) -> Result> { + // Each contents item must be a JSON object. + let obj = content.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini contents item must be an object".into()) + })?; + + // Role is optional; absent role defaults to "user" per Google's REST spec. + // When present it must be a string; only "user" and "model" are accepted. + let role = match obj.get("role") { + None => "user", + Some(Json::String(s)) if s == "user" => "user", + Some(Json::String(s)) if s == "model" => "model", + Some(Json::String(other)) => { + return Err(FlowError::InvalidArgument(format!( + "Gemini contents item has unsupported role '{other}'; expected 'user' or 'model'" + ))); + } + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini contents item 'role' must be a string".into(), + )); + } + }; + + // `parts` is required and must be an array. + let parts = obj + .get("parts") + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini contents item is missing 'parts'".into()) + })? + .as_array() + .ok_or_else(|| { + FlowError::InvalidArgument("Gemini contents item 'parts' must be an array".into()) + })?; + + // Validate each part; collect functionResponse and functionCall parts with + // strict name checks so invalid items surface as errors rather than silent drops. + let mut fr_parts: Vec<&Json> = Vec::new(); + let mut fn_call_parts: Vec<&Json> = Vec::new(); + + for part in parts { + let part_obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini parts item must be an object".into()) + })?; + let data_key = validate_single_gemini_part_data_field(part_obj, "request")?; + if data_key == Some("functionResponse") { + let fr = part.get("functionResponse").unwrap(); + let fr_obj = fr.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionResponse must be an object".into()) + })?; + let name = fr_obj.get("name").and_then(Json::as_str).unwrap_or(""); + if name.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse is missing a non-empty 'name'".into(), + )); + } + // `response` is required and must be an object (Gemini spec). + match fr_obj.get("response") { + None => { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse is missing required 'response'".into(), + )); + } + Some(r) if !r.is_object() => { + return Err(FlowError::InvalidArgument( + "Gemini functionResponse.response must be an object".into(), + )); + } + _ => {} + } + if let Some(nested_parts) = fr_obj.get("parts") { + let nested_parts = nested_parts.as_array().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini functionResponse.parts must be an array".into(), + ) + })?; + for nested_part in nested_parts { + validate_gemini_nested_function_response_part(nested_part)?; + } + } + parse_optional_id(fr_obj, "functionResponse")?; + fr_parts.push(part); + } else if data_key == Some("functionCall") { + let fc = part.get("functionCall").unwrap(); + let fc_obj = fc.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini functionCall must be an object".into()) + })?; + let name = fc_obj.get("name").and_then(Json::as_str).unwrap_or(""); + if name.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini functionCall is missing a non-empty 'name'".into(), + )); + } + if fc_obj.get("args").is_some_and(|a| !a.is_object()) { + return Err(FlowError::InvalidArgument( + "Gemini functionCall.args must be an object".into(), + )); + } + fn_call_parts.push(part); + } else if data_key == Some("text") && part.get("text").is_some_and(|v| !v.is_string()) { + // A plain text part with a non-string text value has no lossless encoding. + return Err(FlowError::InvalidArgument( + "Gemini parts item 'text' must be a string".into(), + )); + } + } + + // A content item must not mix functionResponse and functionCall parts. + if !fr_parts.is_empty() && !fn_call_parts.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini contents item must not contain both functionResponse and functionCall parts" + .into(), + )); + } + // functionResponse belongs in user-role turns only. + if !fr_parts.is_empty() && role != "user" { + return Err(FlowError::InvalidArgument(format!( + "Gemini functionResponse parts must be in a 'user' role content item, got '{role}'" + ))); + } + // functionCall belongs in model-role turns only. + if !fn_call_parts.is_empty() && role != "model" { + return Err(FlowError::InvalidArgument(format!( + "Gemini functionCall parts must be in a 'model' role content item, got '{role}'" + ))); + } + + // --- functionResponse parts (tool results sent as a user turn) --- + if !fr_parts.is_empty() { + // Reject sibling visible/native parts — they would be silently lost because we + // return only tool messages from this branch. + let has_visible_or_native_content = parts.iter().any(|p| { + p.get("functionResponse").is_none() + && p.get("thought").and_then(Json::as_bool) != Some(true) + }); + if has_visible_or_native_content { + return Err(FlowError::InvalidArgument( + "Gemini contents item must not mix functionResponse with visible/native parts" + .into(), + )); + } + let mut msgs = Vec::with_capacity(fr_parts.len()); + for fr_part in fr_parts { + let fr = fr_part.get("functionResponse").unwrap(); + let name = fr.get("name").and_then(|v| v.as_str()).unwrap().to_string(); + let fr_obj = fr.as_object().unwrap(); // validated above + let id = parse_optional_id(fr_obj, "functionResponse")?.unwrap_or_else(|| name.clone()); + let content = gemini_function_response_to_message_content(fr)?; + msgs.push(Message::Tool { + content, + tool_call_id: id, + }); + } + return Ok(msgs); + } + + // Thought parts carry internal reasoning — exclude from visible content. + let content_opt = gemini_parts_to_message_content(parts, "request")?; + + // --- functionCall parts (model invoking a tool) --- + if !fn_call_parts.is_empty() { + let mut tool_calls: Vec = Vec::with_capacity(fn_call_parts.len()); + for p in &fn_call_parts { + let fc = p.get("functionCall").unwrap(); // guaranteed by validation loop above + let fc_map = fc.as_object().unwrap(); + let name = fc_map + .get("name") + .and_then(Json::as_str) + .unwrap() + .to_string(); + let id = parse_optional_id(fc_map, "functionCall")?.unwrap_or_else(|| name.clone()); + let args = fc_map + .get("args") + .cloned() + .unwrap_or_else(|| Json::Object(Default::default())); + let arguments = serde_json::to_string(&args).unwrap_or_else(|_| "{}".into()); + tool_calls.push(ToolCall { + id, + call_type: "function".into(), + function: FunctionCall { name, arguments }, + }); + } + + return Ok(vec![Message::Assistant { + content: content_opt, + tool_calls: Some(tool_calls), + name: None, + }]); + } + + // --- plain text message (user or model) --- + let content = content_opt.unwrap_or_else(|| MessageContent::Text(String::new())); + let msg = if role == "model" { + Message::Assistant { + content: Some(content), + tool_calls: None, + name: None, + } + } else { + Message::User { + content, + name: None, + } + }; + Ok(vec![msg]) +} + +/// Build the `functionCall` JSON object for a single normalized tool call. +/// +/// Returns `Err` when the tool call has a missing or empty function name, +/// an invalid `id`, or non-object arguments. +/// Used by both the fresh-encode path and the patch path; the patch path +/// additionally merges this into an original part to preserve part-level metadata. +fn tool_call_to_fc_obj(tc: &Json) -> Result> { + let fn_obj = tc.get("function"); + let name = fn_obj + .and_then(|f| f.get("name")) + .and_then(Json::as_str) + .unwrap_or(""); + if name.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: tool call is missing a non-empty function name".into(), + )); + } + let id = match tc.get("id") { + None => None, + Some(Json::String(s)) if !s.is_empty() => Some(s.as_str()), + Some(Json::String(_)) => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: tool call for '{name}' has an empty 'id'" + ))); + } + Some(_) => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: tool call for '{name}' has a non-string 'id'" + ))); + } + }; + let arguments = fn_obj + .and_then(|f| f.get("arguments")) + .and_then(Json::as_str); + // If arguments is present but not valid JSON, surface the error rather than + // silently replacing it with {} which would hide model or interceptor output. + let args: Json = match arguments { + None => Json::Object(Default::default()), + Some(a) => { + let v: Json = serde_json::from_str(a).map_err(|_| { + FlowError::InvalidArgument(format!( + "Gemini encoder: function call '{name}' has arguments that are not valid JSON: {a}" + )) + })?; + if !v.is_object() { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: function call '{name}' arguments must be a JSON object, \ + not {v}" + ))); + } + v + } + }; + let mut fc_obj = serde_json::Map::new(); + fc_obj.insert("name".into(), Json::String(name.to_string())); + if let Some(id) = id { + fc_obj.insert("id".into(), Json::String(id.to_string())); + } + fc_obj.insert("args".into(), args); + Ok(fc_obj) +} + +fn gemini_content_parts_from_normalized(content: &Json) -> Result<(Vec, bool)> { + match content { + Json::Null => Ok((Vec::new(), false)), + Json::String(s) => { + if s.is_empty() { + Ok((Vec::new(), false)) + } else { + Ok((vec![serde_json::json!({"text": s})], false)) + } + } + Json::Array(parts) => { + let mut out = Vec::with_capacity(parts.len()); + for part in parts { + let obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini encoder: content parts must be objects".into(), + ) + })?; + let part_type = match obj.get("type") { + None => "text", + Some(Json::String(s)) => s.as_str(), + Some(other) => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: content part 'type' must be a string, got: {other}" + ))); + } + }; + match part_type { + "text" => { + let text = obj.get("text").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini encoder: text content part must have string 'text'".into(), + ) + })?; + let mut gemini_part: serde_json::Map = obj + .iter() + .filter(|(key, _)| key.as_str() != "type") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + gemini_part.insert("text".into(), Json::String(text.to_string())); + out.push(Json::Object(gemini_part)); + } + "provider_native" => { + out.push(provider_native_gemini_part_value(obj)?); + } + other => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: normalized content part type '{other}' cannot be \ + encoded into Gemini content" + ))); + } + } + } + Ok((out, true)) + } + other => Err(FlowError::InvalidArgument(format!( + "Gemini encoder: message content must be a string, null, or array, got: {other}" + ))), + } +} + +/// Convert a serialized normalized message to a Gemini `contents` item. +/// +/// Returns `Ok(None)` for system messages (handled via `systemInstruction`). +/// Returns `Err(FlowError::InvalidArgument)` for message roles that have no +/// valid Gemini mapping — callers must surface these rather than silently drop +/// the message, which would be data loss. +/// +/// `call_id_to_name` is used to resolve the function name for `Message::Tool` +/// when the caller set `tool_call_id` to the actual call ID. If the ID is not +/// found in the map, `tool_call_id` is used as the name (backward compat with +/// payloads where the name was stored as the ID). +fn normalized_to_gemini_content( + msg_json: &Json, + call_id_to_name: &HashMap, +) -> Result> { + let obj = msg_json + .as_object() + .ok_or_else(|| FlowError::Internal("message is not an object".into()))?; + let role = obj + .get("role") + .and_then(Json::as_str) + .ok_or_else(|| FlowError::Internal("message has no role".into()))?; + + if role == "system" { + return Ok(None); + } + + // Message::Tool → functionResponse in a user turn. + // tool_call_id carries the actual call ID (or function name as fallback). + if role == "tool" { + let call_id = extract_tool_call_id(obj)?; + let fn_name = call_id_to_name + .get(call_id) + .map(String::as_str) + .unwrap_or(call_id); + let content_val = obj.get("content").unwrap_or(&Json::Null); + let payload = function_response_payload_from_tool_content(content_val)?; + let mut fr = serde_json::Map::new(); + fr.insert("id".into(), Json::String(call_id.to_string())); + fr.insert("name".into(), Json::String(fn_name.to_string())); + fr.insert("response".into(), payload.response); + if let Some(parts) = payload.parts { + fr.insert("parts".into(), Json::Array(parts)); + } + return Ok(Some(serde_json::json!({ + "role": "user", + "parts": [{"functionResponse": Json::Object(fr)}] + }))); + } + + // Gemini only accepts "user" and "model". Return an error for anything else + // so callers surface the data loss rather than silently dropping the message. + let gemini_role = match role { + "assistant" => "model", + "user" => "user", + other => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: role '{other}' has no Gemini equivalent \ + (only 'user' and 'assistant' are supported)" + ))); + } + }; + + let content_val = obj.get("content").unwrap_or(&Json::Null); + let (content_parts, _) = gemini_content_parts_from_normalized(content_val)?; + + // Message::Assistant with tool_calls → functionCall parts. + if let Some(tool_calls) = obj.get("tool_calls").and_then(Json::as_array) { + let mut parts = content_parts.clone(); + for tc in tool_calls { + let fc_obj = tool_call_to_fc_obj(tc)?; + parts.push(serde_json::json!({"functionCall": Json::Object(fc_obj)})); + } + if !parts.is_empty() { + return Ok(Some( + serde_json::json!({"role": gemini_role, "parts": parts}), + )); + } + } + + // Plain content message. + let mut parts = content_parts; + if parts.is_empty() { + parts.push(serde_json::json!({"text": ""})); + } + Ok(Some( + serde_json::json!({"role": gemini_role, "parts": parts}), + )) +} + +/// Reject normalized message content that contains non-text parts. +/// +/// Used for normalized surfaces that Gemini can only encode as text +/// (`systemInstruction`). User/assistant and tool content have their own helpers +/// because Gemini-native parts are representable there. +fn reject_non_text_content_parts(content: &Json) -> Result<()> { + let Json::Array(parts) = content else { + return Ok(()); // String or Null — always encodable as text + }; + for part in parts { + let obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini encoder: normalized content parts must be objects".into(), + ) + })?; + let part_type = match obj.get("type") { + None => "text", + Some(Json::String(s)) => s.as_str(), + Some(other) => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: content part 'type' must be a string, got: {other}" + ))); + } + }; + if part_type != "text" { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: normalized content part type '{part_type}' cannot be \ + encoded into a Gemini text part; use a provider-native extra field \ + for non-text content" + ))); + } + match obj.get("text") { + Some(Json::String(_)) => {} + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: text content part must have string 'text'".into(), + )); + } + None => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: text content part is missing 'text'".into(), + )); + } + } + } + Ok(()) +} + +/// Parse a tool-result string into a Gemini `functionResponse.response` object. +/// +/// Gemini requires `response` to be a JSON object. If the string parses as an +/// object it is used directly; any other JSON value (or non-JSON text) is wrapped +/// in `{"output": }` so the response field is always object-shaped. +fn ensure_object_response(content_str: String) -> Json { + match serde_json::from_str::(&content_str) { + Ok(Json::Object(m)) => Json::Object(m), + Ok(other) => serde_json::json!({"output": other}), + Err(_) => serde_json::json!({"output": content_str}), + } +} + +struct GeminiGenerateContentFunctionResponsePayload { + response: Json, + parts: Option>, +} + +fn provider_native_gemini_part_value(obj: &serde_json::Map) -> Result { + let provider = obj.get("provider").and_then(Json::as_str).ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini encoder: provider_native provider must be a string".into(), + ) + })?; + if provider != GEMINI_PROVIDER { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: provider_native content part for provider \ + '{provider}' cannot be encoded by the Gemini generateContent codec" + ))); + } + match obj.get("kind") { + Some(Json::String(s)) if !s.is_empty() => {} + Some(Json::String(_)) => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: provider_native content part kind must be non-empty".into(), + )); + } + _ => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: provider_native kind must be a string".into(), + )); + } + } + let value = obj.get("value").ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini encoder: provider_native content part is missing 'value'".into(), + ) + })?; + let value_obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini encoder: provider_native value must be an object".into()) + })?; + let data_key = validate_single_gemini_part_data_field(value_obj, "provider_native")?; + if matches!(data_key, Some("functionCall" | "functionResponse")) { + return Err(FlowError::InvalidArgument( + "Gemini encoder: provider_native content parts must not encode \ + functionCall or functionResponse; use tool_calls or tool messages" + .into(), + )); + } + if data_key == Some("text") && value.get("text").is_some_and(|v| !v.is_string()) { + return Err(FlowError::InvalidArgument( + "Gemini encoder: provider_native text part must have string 'text'".into(), + )); + } + Ok(value.clone()) +} + +fn function_response_payload_from_tool_content( + content: &Json, +) -> Result { + match content { + Json::Null | Json::String(_) => Ok(GeminiGenerateContentFunctionResponsePayload { + response: ensure_object_response(extract_content_text(content)), + parts: None, + }), + Json::Array(parts) => { + let mut response_texts = Vec::new(); + let mut native_parts = Vec::new(); + for part in parts { + let obj = part.as_object().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini encoder: normalized tool content parts must be objects".into(), + ) + })?; + let part_type = match obj.get("type") { + None => "text", + Some(Json::String(s)) => s.as_str(), + Some(other) => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: content part 'type' must be a string, got: {other}" + ))); + } + }; + match part_type { + "text" => match obj.get("text") { + Some(Json::String(text)) => response_texts.push(text.clone()), + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: text content part must have string 'text'".into(), + )); + } + None => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: text content part is missing 'text'".into(), + )); + } + }, + "provider_native" => { + let value = provider_native_gemini_part_value(obj)?; + validate_gemini_nested_function_response_part(&value)?; + native_parts.push(value); + } + other => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: normalized tool content part type '{other}' cannot be \ + encoded into functionResponse.parts" + ))); + } + } + } + Ok(GeminiGenerateContentFunctionResponsePayload { + response: ensure_object_response(response_texts.join("\n")), + parts: Some(native_parts), + }) + } + other => Err(FlowError::InvalidArgument(format!( + "Gemini encoder: message content must be a string, null, or array, got: {other}" + ))), + } +} + +/// Extract a plain text string from a normalized `content` field. +fn extract_content_text(content: &Json) -> String { + match content { + Json::String(s) => s.clone(), + Json::Array(parts) => parts + .iter() + .filter_map(|p| { + if p.get("type") + .map(|ty| ty.as_str() == Some("text")) + .unwrap_or(true) + { + p.get("text")?.as_str().map(str::to_string) + } else { + p.as_str().map(str::to_string) + } + }) + .collect::>() + .join("\n"), + _ => String::new(), + } +} + +fn insert_serialized( + obj: &mut serde_json::Map, + key: &str, + value: &T, + context: &str, +) -> Result<()> { + let json = serde_json::to_value(value).map_err(|e| { + FlowError::Internal(format!("Gemini generateContent {context} encode: {e}")) + })?; + obj.insert(key.into(), json); + Ok(()) +} + +fn json_f64(v: f64, field: &str) -> Result { + serde_json::Number::from_f64(v) + .map(Json::Number) + .ok_or_else(|| { + FlowError::InvalidArgument(format!( + "Gemini encoder: '{field}' value {v} is not a finite number" + )) + }) +} + +fn gemini_native_tool_fields(group: &serde_json::Map) -> Option { + let native: serde_json::Map = group + .iter() + .filter(|(key, _)| key.as_str() != "functionDeclarations") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + if native.is_empty() { + None + } else { + Some(Json::Object(native)) + } +} + +fn gemini_native_tool_kind(value: &Json) -> String { + value + .as_object() + .and_then(|group| group.keys().next()) + .cloned() + .unwrap_or_else(|| "unknown".to_string()) +} + +fn gemini_native_tool_keys(value: &Json) -> Vec { + let mut keys: Vec = value + .as_object() + .map(|group| { + group + .keys() + .filter(|key| key.as_str() != "functionDeclarations") + .cloned() + .collect() + }) + .unwrap_or_default(); + keys.sort(); + keys +} + +fn take_matching_native_group( + native_groups: &[Json], + native_used: &mut [bool], + expected_keys: &[String], +) -> Option { + for (idx, group) in native_groups.iter().enumerate() { + if native_used[idx] { + continue; + } + if gemini_native_tool_keys(group) == expected_keys { + native_used[idx] = true; + return Some(group.clone()); + } + } + None +} + +// --------------------------------------------------------------------------- +// LlmResponseCodec +// --------------------------------------------------------------------------- + +impl LlmResponseCodec for GeminiGenerateContentCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) + } + + fn decode_response(&self, response: &Json) -> Result { + let raw: RawGeminiGenerateContentResponse = serde_json::from_value(response.clone()) + .map_err(|e| FlowError::Internal(format!("Gemini response decode: {e}")))?; + + let candidate = raw.candidates.as_ref().and_then(|c| c.first()); + + let (message, tool_calls, has_tool_calls) = if let Some(c) = candidate { + let parts = c.content.as_ref().and_then(|ct| ct.parts.as_deref()); + let msg = parts + .map(extract_parts_message_content) + .transpose()? + .flatten(); + let tcs = parts.map(extract_parts_tool_calls).transpose()?.flatten(); + let has = tcs.is_some(); + (msg, tcs, has) + } else { + (None, None, false) + }; + + let prompt_block_reason = prompt_feedback_block_reason(&raw.extra)?; + let finish_reason = candidate + .and_then(|c| map_finish_reason(c.finish_reason.as_deref(), has_tool_calls)) + .or_else(|| map_prompt_block_reason(prompt_block_reason)); + + let model = raw.model_version.clone(); + let (usage, thoughts_tokens) = map_usage(raw.usage_metadata, model.as_deref()); + + // Capture candidate-level metadata (safetyRatings, groundingMetadata, + // citationMetadata, thinking token count, …) that cannot be normalized + // across providers. thoughts_tokens lives here, not in shared Usage. + let api_specific = { + let mut extra = candidate.map(|c| c.extra.clone()).unwrap_or_default(); + extra.remove("api"); + let has_candidate_data = !extra.is_empty() || thoughts_tokens.is_some(); + if !has_candidate_data { + None + } else { + let safety_ratings = extra.remove("safetyRatings"); + let grounding_metadata = extra.remove("groundingMetadata"); + let citation_metadata = extra.remove("citationMetadata"); + Some( + super::response::ApiSpecificResponse::GeminiGenerateContent { + thoughts_tokens, + safety_ratings, + grounding_metadata, + citation_metadata, + extra, + }, + ) + } + }; + + Ok(AnnotatedLlmResponse { + id: raw.response_id, + model, + message, + tool_calls, + finish_reason, + usage, + optimization_summary: None, + api_specific, + extra: raw.extra, + }) + } +} + +// --------------------------------------------------------------------------- +// LlmCodec +// --------------------------------------------------------------------------- + +impl LlmCodec for GeminiGenerateContentCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) + } + + fn decode(&self, request: &LlmRequest) -> Result { + let obj = request + .content + .as_object() + .ok_or_else(|| FlowError::Internal("request content is not an object".into()))?; + + let mut messages: Vec = Vec::new(); + + // Validate and decode systemInstruction when present. + if let Some(sys_val) = obj.get("systemInstruction") { + validate_system_instruction(sys_val)?; + if let Some(text) = system_instruction_text(sys_val) { + let msg = serde_json::from_value::( + serde_json::json!({"role": "system", "content": text}), + ) + .map_err(|e| { + FlowError::Internal(format!("Gemini system instruction decode: {e}")) + })?; + messages.push(msg); + } + } + + // `contents` is required; a present-but-non-array value is malformed. + match obj.get("contents") { + None => { + return Err(FlowError::InvalidArgument( + "Gemini request is missing contents".into(), + )); + } + Some(v) if !v.is_array() => { + return Err(FlowError::InvalidArgument( + "Gemini request contents must be an array".into(), + )); + } + Some(arr) => { + for content in arr.as_array().unwrap() { + messages.extend(gemini_content_to_messages(content)?); + } + } + } + + // generationConfig → GenerationParams + let gen_config = match obj.get("generationConfig") { + Some(v) if !v.is_object() => { + return Err(FlowError::InvalidArgument( + "Gemini generationConfig must be an object".into(), + )); + } + other => other, + }; + let temperature = match gen_config.and_then(|c| c.get("temperature")) { + None => None, + Some(v) => Some(v.as_f64().ok_or_else(|| { + FlowError::InvalidArgument("Gemini temperature must be a number".into()) + })?), + }; + let top_p = match gen_config.and_then(|c| c.get("topP")) { + None => None, + Some(v) => Some(v.as_f64().ok_or_else(|| { + FlowError::InvalidArgument("Gemini topP must be a number".into()) + })?), + }; + let max_tokens = match gen_config.and_then(|c| c.get("maxOutputTokens")) { + None => None, + Some(v) => Some(v.as_u64().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini maxOutputTokens must be a non-negative integer".into(), + ) + })?), + }; + let stop = match gen_config.and_then(|c| c.get("stopSequences")) { + None => None, + Some(v) => { + let parsed = serde_json::from_value::>(v.clone()).ok(); + if parsed.is_none() { + return Err(FlowError::InvalidArgument( + "Gemini stopSequences must be an array of strings".into(), + )); + } + parsed + } + }; + + let params = + if temperature.is_some() || max_tokens.is_some() || top_p.is_some() || stop.is_some() { + Some(GenerationParams { + temperature, + max_tokens, + top_p, + stop, + }) + } else { + None + }; + + // tools[].functionDeclarations → Vec + // Non-modeled fields (parametersJsonSchema, responseJsonSchema, response, behavior, …) + // are captured into FunctionDefinition.extra so they survive encode. + const MODELED_FD_KEYS: &[&str] = &["name", "description", "parameters"]; + + let tools: Option> = match obj.get("tools") { + None => None, + Some(v) if !v.is_array() => { + return Err(FlowError::InvalidArgument( + "Gemini tools must be an array".into(), + )); + } + Some(v) => { + let arr = v.as_array().unwrap(); + let mut defs: Vec = Vec::new(); + for group in arr { + let group_obj = group.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini tools[] entry must be an object".into()) + })?; + let has_function_declarations = group_obj.contains_key("functionDeclarations"); + if let Some(fds_val) = group_obj.get("functionDeclarations") { + let fds = fds_val.as_array().ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini functionDeclarations must be an array".into(), + ) + })?; + for fd in fds { + if !fd.is_object() { + return Err(FlowError::InvalidArgument( + "Gemini functionDeclaration entry must be an object".into(), + )); + } + let name = fd + .get("name") + .and_then(|n| n.as_str()) + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini functionDeclaration must have a non-empty 'name'" + .into(), + ) + })? + .to_string(); + let description = match fd.get("description") { + None => None, + Some(Json::String(s)) => Some(s.clone()), + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini functionDeclaration.description must be a string" + .into(), + )); + } + }; + let parameters = fd.get("parameters").cloned(); + let extra: serde_json::Map = fd + .as_object() + .map(|o| { + o.iter() + .filter(|(k, _)| !MODELED_FD_KEYS.contains(&k.as_str())) + .map(|(k, v)| (k.clone(), v.clone())) + .collect() + }) + .unwrap_or_default(); + defs.push(ToolDefinition::Function { + function: FunctionDefinition { + name, + description, + parameters, + strict: None, + extra, + }, + extra: Default::default(), + }); + } + } + let native_value = if has_function_declarations { + gemini_native_tool_fields(group_obj) + } else { + Some(group.clone()) + }; + if let Some(value) = native_value { + defs.push(ToolDefinition::ProviderNative { + provider: GEMINI_PROVIDER.into(), + kind: gemini_native_tool_kind(&value), + value, + }); + } + } + if defs.is_empty() { None } else { Some(defs) } + } + }; + + // All unrecognized top-level keys go into extra. + let extra: serde_json::Map = obj + .iter() + .filter(|(k, _)| !MODELED_REQUEST_KEYS.contains(&k.as_str())) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + let model = match obj.get("model") { + None => None, + Some(Json::String(s)) => Some(s.clone()), + Some(_) => { + return Err(FlowError::InvalidArgument( + "Gemini request 'model' must be a string".into(), + )); + } + }; + + Ok(AnnotatedLlmRequest { + messages, + instructions: None, + model, + params, + tools, + tool_choice: None, + store: None, + previous_response_id: None, + truncation: None, + reasoning: None, + include: None, + user: None, + metadata: None, + service_tier: None, + parallel_tool_calls: None, + max_output_tokens: None, + max_tool_calls: None, + top_logprobs: None, + stream: None, + api_specific: None, + extra, + }) + } + + fn encode(&self, annotated: &AnnotatedLlmRequest, original: &LlmRequest) -> Result { + let baseline = self.decode(original)?; + validate_gemini_supported_fields(annotated, &baseline)?; + let mut content = original.content.clone(); + let obj = content + .as_object_mut() + .ok_or_else(|| FlowError::Internal("original content is not an object".into()))?; + + if annotated.messages != baseline.messages { + patch_gemini_messages(obj, &annotated.messages, &baseline.messages)?; + } + + if annotated.params != baseline.params { + patch_gemini_params(obj, annotated.params.as_ref())?; + } + + if annotated.tools != baseline.tools { + patch_gemini_tools(obj, annotated.tools.as_ref())?; + } + + if annotated.model != baseline.model { + match &annotated.model { + Some(m) => { + obj.insert("model".into(), Json::String(m.clone())); + } + None => { + obj.remove("model"); + } + } + } + + patch_extra_fields(obj, &baseline.extra, &annotated.extra); + + Ok(LlmRequest { + headers: original.headers.clone(), + content, + }) + } +} + +// --------------------------------------------------------------------------- +// Baseline-aware patch helpers +// --------------------------------------------------------------------------- + +/// Return `InvalidArgument` if the interceptor changed a field that the Gemini +/// encoder cannot represent. These fields are always `None` in the Gemini +/// baseline because the Gemini request format has no equivalent concept. +/// Silently ignoring them would mean data the interceptor intended to act on +/// is lost without any signal. +fn validate_gemini_supported_fields( + annotated: &AnnotatedLlmRequest, + baseline: &AnnotatedLlmRequest, +) -> Result<()> { + for message in &annotated.messages { + match message { + Message::System { name: Some(_), .. } => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: Message::System.name is not representable".into(), + )); + } + Message::User { name: Some(_), .. } => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: Message::User.name is not representable".into(), + )); + } + Message::Assistant { name: Some(_), .. } => { + return Err(FlowError::InvalidArgument( + "Gemini encoder: Message::Assistant.name is not representable".into(), + )); + } + _ => {} + } + } + + macro_rules! reject_if_changed { + ($field:ident) => { + if annotated.$field != baseline.$field { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: field '{}' is not representable in the \ + Gemini generateContent API; use a provider-native extra field instead", + stringify!($field) + ))); + } + }; + } + reject_if_changed!(instructions); + reject_if_changed!(tool_choice); + reject_if_changed!(store); + reject_if_changed!(previous_response_id); + reject_if_changed!(truncation); + reject_if_changed!(reasoning); + reject_if_changed!(include); + reject_if_changed!(user); + reject_if_changed!(metadata); + reject_if_changed!(service_tier); + reject_if_changed!(parallel_tool_calls); + reject_if_changed!(max_output_tokens); + reject_if_changed!(max_tool_calls); + reject_if_changed!(top_logprobs); + reject_if_changed!(stream); + reject_if_changed!(api_specific); + Ok(()) +} + +/// Patch `systemInstruction` and `contents` using prefix/suffix-aligned merging. +/// +/// Computes the longest matching prefix and suffix between the annotated and +/// baseline non-system message lists, then maps each annotated position to its +/// corresponding original `contents` item. Unchanged positions are preserved +/// byte-identically. Pure insertions in the gap between prefix and suffix are +/// encoded fresh. Equal-size gaps (pure edits) pair 1-to-1 with their original +/// item so that `patch_changed_gemini_content` can carry over metadata. +fn patch_gemini_messages( + obj: &mut serde_json::Map, + messages: &[Message], + baseline: &[Message], +) -> Result<()> { + let msgs_json = serde_json::to_value(messages) + .map_err(|e| FlowError::Internal(format!("Gemini messages encode: {e}")))?; + let base_json = serde_json::to_value(baseline) + .map_err(|e| FlowError::Internal(format!("Gemini baseline encode: {e}")))?; + + let msgs_arr = msgs_json.as_array().map(Vec::as_slice).unwrap_or(&[]); + let base_arr = base_json.as_array().map(Vec::as_slice).unwrap_or(&[]); + + // Collect all system messages. Gemini has one systemInstruction so multiple + // system messages are merged by joining their text content with newlines. + let is_sys = |m: &&Json| m.get("role").and_then(Json::as_str) == Some("system"); + let ann_sys_msgs: Vec<&Json> = msgs_arr.iter().filter(is_sys).collect(); + let base_sys_msgs: Vec<&Json> = base_arr.iter().filter(is_sys).collect(); + patch_gemini_system_instruction(obj, &ann_sys_msgs, &base_sys_msgs)?; + + let is_non_sys = |m: &&Json| m.get("role").and_then(Json::as_str) != Some("system"); + let ann_non_sys: Vec<&Json> = msgs_arr.iter().filter(is_non_sys).collect(); + let base_non_sys: Vec<&Json> = base_arr.iter().filter(is_non_sys).collect(); + patch_gemini_non_system_contents(obj, messages, &ann_non_sys, &base_non_sys) +} + +fn patch_gemini_system_instruction( + obj: &mut serde_json::Map, + ann_sys_msgs: &[&Json], + base_sys_msgs: &[&Json], +) -> Result<()> { + if ann_sys_msgs == base_sys_msgs { + return Ok(()); + } + + for m in ann_sys_msgs { + reject_non_text_content_parts(m.get("content").unwrap_or(&Json::Null))?; + } + let text = ann_sys_msgs + .iter() + .filter_map(|m| { + let t = extract_content_text(m.get("content").unwrap_or(&Json::Null)); + if t.is_empty() { None } else { Some(t) } + }) + .collect::>() + .join("\n"); + if text.is_empty() { + obj.remove("systemInstruction"); + return Ok(()); + } + + let mut si: serde_json::Map = obj + .get("systemInstruction") + .and_then(Json::as_object) + .cloned() + .unwrap_or_default(); + if si.get("role").is_some_and(|v| !v.is_string()) { + return Err(FlowError::InvalidArgument( + "Gemini systemInstruction.role must be a string".into(), + )); + } + + let orig_parts = si + .get("parts") + .and_then(Json::as_array) + .cloned() + .unwrap_or_default(); + validate_editable_gemini_system_parts(&orig_parts)?; + si.insert( + "parts".into(), + Json::Array(rebuild_gemini_system_parts(&orig_parts, text)), + ); + obj.insert("systemInstruction".into(), Json::Object(si)); + Ok(()) +} + +fn validate_editable_gemini_system_parts(orig_parts: &[Json]) -> Result<()> { + let text_part_count = orig_parts + .iter() + .filter(|p| { + p.get("thought").and_then(Json::as_bool) != Some(true) && p.get("text").is_some() + }) + .count(); + let has_non_text_non_thought = orig_parts + .iter() + .any(|p| p.get("thought").and_then(Json::as_bool) != Some(true) && p.get("text").is_none()); + if text_part_count > 1 || has_non_text_non_thought { + return Err(FlowError::InvalidArgument( + "Gemini systemInstruction with multiple text parts or non-text parts \ + cannot be edited via the normalized layer; edit the raw provider payload directly" + .into(), + )); + } + Ok(()) +} + +fn rebuild_gemini_system_parts(orig_parts: &[Json], text: String) -> Vec { + let mut new_parts = Vec::with_capacity(orig_parts.len().max(1)); + let mut text_part_placed = false; + for orig_part in orig_parts { + if orig_part.get("thought").and_then(Json::as_bool) == Some(true) { + new_parts.push(orig_part.clone()); + } else if orig_part.get("text").is_some() && !text_part_placed { + let mut p = orig_part.as_object().cloned().unwrap_or_default(); + p.insert("text".into(), Json::String(text.clone())); + new_parts.push(Json::Object(p)); + text_part_placed = true; + } + } + if !text_part_placed { + new_parts.push(serde_json::json!({"text": text})); + } + new_parts +} + +#[derive(Debug)] +struct GeminiGenerateContentMessageAlignment { + prefix_len: usize, + ann_gap_end: usize, + base_gap_end: usize, + ann_gap_len: usize, + base_gap_len: usize, +} + +impl GeminiGenerateContentMessageAlignment { + fn new(ann_non_sys: &[&Json], base_non_sys: &[&Json]) -> Self { + let prefix_len = ann_non_sys + .iter() + .zip(base_non_sys.iter()) + .take_while(|(a, b)| a == b) + .count(); + let max_suffix = ann_non_sys + .len() + .saturating_sub(prefix_len) + .min(base_non_sys.len().saturating_sub(prefix_len)); + let suffix_len = ann_non_sys[ann_non_sys.len().saturating_sub(max_suffix)..] + .iter() + .rev() + .zip( + base_non_sys[base_non_sys.len().saturating_sub(max_suffix)..] + .iter() + .rev(), + ) + .take_while(|(a, b)| a == b) + .count(); + let ann_gap_end = ann_non_sys.len().saturating_sub(suffix_len); + let base_gap_end = base_non_sys.len().saturating_sub(suffix_len); + Self { + prefix_len, + ann_gap_end, + base_gap_end, + ann_gap_len: ann_gap_end.saturating_sub(prefix_len), + base_gap_len: base_gap_end.saturating_sub(prefix_len), + } + } + + fn base_idx_for(&self, i: usize) -> Option { + if i < self.prefix_len { + Some(i) + } else if i >= self.ann_gap_end { + Some(self.base_gap_end + (i - self.ann_gap_end)) + } else if self.ann_gap_len == self.base_gap_len { + Some(self.prefix_len + (i - self.prefix_len)) + } else { + None + } + } +} + +fn gemini_content_idx_of_base_msg( + orig_contents: &[Json], + base_non_sys_len: usize, +) -> Result> { + let mut mapping = Vec::with_capacity(base_non_sys_len); + for (cidx, content) in orig_contents.iter().enumerate() { + let n = gemini_content_to_messages(content)?.len().max(1); + for _ in 0..n { + if mapping.len() < base_non_sys_len { + mapping.push(cidx); + } + } + } + let last = orig_contents.len().saturating_sub(1); + while mapping.len() < base_non_sys_len { + mapping.push(last); + } + Ok(mapping) +} + +fn gemini_content_msg_counts(mapping: &[usize], orig_content_count: usize) -> Vec { + let mut counts = vec![0usize; orig_content_count]; + for &cidx in mapping { + if cidx < counts.len() { + counts[cidx] += 1; + } + } + counts +} + +fn gemini_call_id_to_name(messages: &[Message]) -> HashMap { + messages + .iter() + .filter_map(|m| match m { + Message::Assistant { + tool_calls: Some(tcs), + .. + } => Some( + tcs.iter() + .map(|tc| (tc.id.clone(), tc.function.name.clone())), + ), + _ => None, + }) + .flatten() + .collect() +} + +fn push_fresh_gemini_content( + new_contents: &mut Vec, + message: &Json, + call_id_to_name: &HashMap, +) -> Result<()> { + if let Some(item) = normalized_to_gemini_content(message, call_id_to_name)? { + new_contents.push(item); + } + Ok(()) +} + +fn gemini_run_end( + start: usize, + ann_len: usize, + orig_cidx: usize, + alignment: &GeminiGenerateContentMessageAlignment, + content_idx_of_base_msg: &[usize], +) -> usize { + let mut run_end = start + 1; + while run_end < ann_len { + match alignment.base_idx_for(run_end) { + Some(bidx) if content_idx_of_base_msg[bidx] == orig_cidx => run_end += 1, + _ => break, + } + } + run_end +} + +fn gemini_run_is_unchanged( + start: usize, + run_end: usize, + expected_len: usize, + alignment: &GeminiGenerateContentMessageAlignment, + ann_non_sys: &[&Json], + base_non_sys: &[&Json], +) -> bool { + let run_len = run_end - start; + run_len == expected_len + && (start..run_end).all(|j| { + alignment + .base_idx_for(j) + .and_then(|bidx| base_non_sys.get(bidx)) + .map(|bm| *bm == ann_non_sys[j]) + .unwrap_or(false) + }) +} + +fn patch_gemini_missing_content_run( + run: &[&Json], + call_id_to_name: &HashMap, +) -> Result> { + if run.len() > 1 { + return Err(FlowError::Internal( + "Gemini encode: multiple messages map to a missing content item".into(), + )); + } + match run.first() { + Some(m) => normalized_to_gemini_content(m, call_id_to_name), + None => Ok(None), + } +} + +fn patch_gemini_non_system_contents( + obj: &mut serde_json::Map, + messages: &[Message], + ann_non_sys: &[&Json], + base_non_sys: &[&Json], +) -> Result<()> { + if ann_non_sys == base_non_sys { + return Ok(()); + } + let call_id_to_name = gemini_call_id_to_name(messages); + let orig_contents: Vec = obj + .get("contents") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + let alignment = GeminiGenerateContentMessageAlignment::new(ann_non_sys, base_non_sys); + let content_idx_of_base_msg = + gemini_content_idx_of_base_msg(&orig_contents, base_non_sys.len())?; + let content_msg_count = + gemini_content_msg_counts(&content_idx_of_base_msg, orig_contents.len()); + + let mut new_contents = Vec::new(); + let mut processed_cidxs = std::collections::HashSet::::new(); + let mut i = 0; + while i < ann_non_sys.len() { + let Some(base_idx) = alignment.base_idx_for(i) else { + push_fresh_gemini_content(&mut new_contents, ann_non_sys[i], &call_id_to_name)?; + i += 1; + continue; + }; + let orig_cidx = content_idx_of_base_msg[base_idx]; + if processed_cidxs.contains(&orig_cidx) { + push_fresh_gemini_content(&mut new_contents, ann_non_sys[i], &call_id_to_name)?; + i += 1; + continue; + } + + let run_end = gemini_run_end( + i, + ann_non_sys.len(), + orig_cidx, + &alignment, + &content_idx_of_base_msg, + ); + processed_cidxs.insert(orig_cidx); + let expected_len = content_msg_count.get(orig_cidx).copied().unwrap_or(1); + if gemini_run_is_unchanged( + i, + run_end, + expected_len, + &alignment, + ann_non_sys, + base_non_sys, + ) { + if let Some(orig) = orig_contents.get(orig_cidx) { + new_contents.push(orig.clone()); + } + } else { + let run: Vec<&Json> = ann_non_sys[i..run_end].to_vec(); + let item = orig_contents + .get(orig_cidx) + .map(|orig| patch_changed_gemini_content(orig, &run, &call_id_to_name)) + .unwrap_or_else(|| patch_gemini_missing_content_run(&run, &call_id_to_name))?; + if let Some(item) = item { + new_contents.push(item); + } + } + + i = run_end; + } + obj.insert("contents".into(), Json::Array(new_contents)); + Ok(()) +} + +/// Rebuild a Gemini `contents` item for a position that changed from baseline. +/// +/// Applies one or more normalized-message changes to an original Gemini content +/// item in a single pass over the original parts, preserving native metadata +/// (`thoughtSignature`, `thought`, `inlineData`, unmodeled fields). Surrounding +/// non-call parts keep their original positions; function-call slots follow the +/// annotated call order so reordered calls keep the right signatures. +/// +/// `ann_msgs` is the slice of annotated messages that map to this content item. +/// For most content items it is a single message; for parallel `functionResponse` +/// items it may be multiple consecutive `Message::Tool` entries. +fn patch_changed_gemini_content( + original_item: &Json, + ann_msgs: &[&Json], + call_id_to_name: &HashMap, +) -> Result> { + let orig_obj = original_item + .as_object() + .ok_or_else(|| FlowError::Internal("original contents item is not an object".into()))?; + // Missing role is treated as "user" on decode; mirror that here so valid roleless + // content items are not silently dropped when edited. + let orig_role = orig_obj + .get("role") + .and_then(Json::as_str) + .unwrap_or("user"); + let orig_parts = orig_obj + .get("parts") + .and_then(|v| v.as_array()) + .ok_or_else(|| FlowError::Internal("original contents item has no parts array".into()))?; + + // For multi-message runs (parallel functionResponse), apply changes to one + // representative message first, then chain the rest. + let ann_msg = match ann_msgs.first() { + Some(m) => m, + None => return Ok(None), + }; + let msg_obj = match ann_msg.as_object() { + Some(o) => o, + None => return Ok(None), + }; + + let ann_role = msg_obj + .get("role") + .and_then(Json::as_str) + .unwrap_or(orig_role); + // Apply the same role validation as normalized_to_gemini_content to catch + // unsupported roles injected by interceptors on the edit path. + let gemini_role = match ann_role { + "assistant" => "model", + "tool" => "user", + "user" => "user", + other => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: role '{other}' has no Gemini equivalent \ + (only 'user' and 'assistant' are supported)" + ))); + } + }; + + if ann_role == "tool" { + return patch_gemini_tool_response_content(orig_parts, ann_msgs, call_id_to_name); + } + + patch_gemini_visible_content(orig_parts, msg_obj, gemini_role) +} + +fn gemini_function_response_updates<'a>( + ann_msgs: &'a [&Json], +) -> Result<( + HashMap<&'a str, GeminiGenerateContentFunctionResponsePayload>, + Vec<&'a str>, +)> { + let mut updates = HashMap::new(); + let mut update_order = Vec::new(); + for am in ann_msgs { + let mo = am + .as_object() + .ok_or_else(|| FlowError::Internal("tool message is not an object".into()))?; + let call_id = extract_tool_call_id(mo)?; + let content_val = mo.get("content").unwrap_or(&Json::Null); + let payload = function_response_payload_from_tool_content(content_val)?; + if updates.insert(call_id, payload).is_some() { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: duplicate Message::Tool tool_call_id '{call_id}'" + ))); + } + update_order.push(call_id); + } + Ok((updates, update_order)) +} + +fn function_response_call_id(fr: &Json) -> Option<&str> { + fr.get("id") + .and_then(|v| v.as_str()) + .or_else(|| fr.get("name").and_then(|v| v.as_str())) +} + +fn build_gemini_function_response_part( + call_id: &str, + name: &str, + payload: GeminiGenerateContentFunctionResponsePayload, + original_fr: Option<&Json>, +) -> Json { + let mut new_fr = original_fr + .and_then(Json::as_object) + .cloned() + .unwrap_or_default(); + new_fr.insert("id".into(), Json::String(call_id.to_string())); + new_fr.insert("name".into(), Json::String(name.to_string())); + new_fr.insert("response".into(), payload.response); + if let Some(parts) = payload.parts { + new_fr.insert("parts".into(), Json::Array(parts)); + } + serde_json::json!({"functionResponse": Json::Object(new_fr)}) +} + +fn patch_gemini_tool_response_content( + orig_parts: &[Json], + ann_msgs: &[&Json], + call_id_to_name: &HashMap, +) -> Result> { + let (mut updates, update_order) = gemini_function_response_updates(ann_msgs)?; + let known_ids: std::collections::HashSet<&str> = updates.keys().copied().collect(); + let mut new_parts = Vec::new(); + + for orig_part in orig_parts { + let Some(fr) = orig_part.get("functionResponse") else { + new_parts.push(orig_part.clone()); + continue; + }; + let Some(call_id) = function_response_call_id(fr) else { + new_parts.push(orig_part.clone()); + continue; + }; + if !known_ids.contains(call_id) { + continue; + } + if let Some(payload) = updates.remove(call_id) { + let name = fr.get("name").and_then(Json::as_str).unwrap_or(call_id); + new_parts.push(build_gemini_function_response_part( + call_id, + name, + payload, + Some(fr), + )); + } else { + new_parts.push(orig_part.clone()); + } + } + + for call_id in update_order { + let Some(payload) = updates.remove(call_id) else { + continue; + }; + let fn_name = call_id_to_name + .get(call_id) + .map(String::as_str) + .unwrap_or(call_id); + new_parts.push(build_gemini_function_response_part( + call_id, fn_name, payload, None, + )); + } + + Ok(Some( + serde_json::json!({"role": "user", "parts": new_parts}), + )) +} + +fn original_function_call_entries(orig_parts: &[Json]) -> Vec<(usize, Option<&str>, &str)> { + orig_parts + .iter() + .enumerate() + .filter_map(|(i, p)| { + let fc = p.get("functionCall")?; + let name = fc.get("name")?.as_str()?; + let id = fc.get("id").and_then(Json::as_str); + Some((i, id, name)) + }) + .collect() +} + +fn matching_function_call_entry( + entries: &[(usize, Option<&str>, &str)], + consumed: &std::collections::HashSet, + fn_id: Option<&str>, + fn_name: &str, +) -> Option<(usize, Option)> { + let matched = fn_id + .and_then(|id| { + entries + .iter() + .find(|(idx, orig_id, _)| !consumed.contains(idx) && *orig_id == Some(id)) + }) + .or_else(|| { + entries.iter().find(|(idx, orig_id, name)| { + !consumed.contains(idx) && orig_id.is_none() && *name == fn_name + }) + })?; + Some((matched.0, matched.1.map(str::to_string))) +} + +fn rebuilt_gemini_function_call_parts( + orig_parts: &[Json], + msg_obj: &serde_json::Map, +) -> Result> { + let entries = original_function_call_entries(orig_parts); + let mut consumed = std::collections::HashSet::::new(); + let mut rebuilt = Vec::new(); + let Some(tool_calls) = msg_obj.get("tool_calls").and_then(Json::as_array) else { + return Ok(rebuilt); + }; + + for tc in tool_calls { + let mut fc_obj = tool_call_to_fc_obj(tc)?; + let fn_name = fc_obj + .get("name") + .and_then(Json::as_str) + .unwrap_or("") + .to_string(); + let fn_id = fc_obj.get("id").and_then(Json::as_str).map(str::to_string); + if let Some((orig_idx, orig_id)) = + matching_function_call_entry(&entries, &consumed, fn_id.as_deref(), fn_name.as_str()) + { + consumed.insert(orig_idx); + if orig_id.is_none() && fn_id.as_deref() == Some(fn_name.as_str()) { + fc_obj.remove("id"); + } + let mut part_obj = orig_parts[orig_idx] + .as_object() + .cloned() + .unwrap_or_default(); + part_obj.insert("functionCall".into(), Json::Object(fc_obj)); + rebuilt.push(Json::Object(part_obj)); + } else { + rebuilt.push(serde_json::json!({"functionCall": Json::Object(fc_obj)})); + } + } + Ok(rebuilt) +} + +fn replacement_content_part( + orig_part: &Json, + replacement: Json, + content_is_parts_form: bool, +) -> Json { + if !content_is_parts_form + && orig_part.get("text").is_some() + && replacement.get("text").is_some() + { + let mut obj = orig_part.as_object().cloned().unwrap_or_default(); + obj.insert("text".into(), replacement.get("text").unwrap().clone()); + Json::Object(obj) + } else { + replacement + } +} + +fn merge_gemini_original_parts( + orig_parts: &[Json], + new_content_parts: Vec, + rebuilt_fn_calls: Vec, + content_is_parts_form: bool, +) -> Vec { + let mut parts = Vec::new(); + let mut new_content_parts = new_content_parts.into_iter(); + let mut rebuilt_fn_calls = rebuilt_fn_calls.into_iter(); + let mut content_emitted = false; + + for orig_part in orig_parts { + let is_thought = orig_part.get("thought").and_then(Json::as_bool) == Some(true); + let is_content_part = !is_thought + && orig_part.get("functionCall").is_none() + && orig_part.get("functionResponse").is_none(); + + if orig_part.get("functionCall").is_some() { + if let Some(rebuilt) = rebuilt_fn_calls.next() { + parts.push(rebuilt); + } + } else if is_content_part { + if let Some(replacement) = new_content_parts.next() { + parts.push(replacement_content_part( + orig_part, + replacement, + content_is_parts_form, + )); + content_emitted = true; + } else if !content_is_parts_form + && (orig_part.get("text").is_none() + || (orig_part.get("text").is_some() + && orig_part.get("thoughtSignature").is_some())) + { + parts.push(orig_part.clone()); + } + } else { + parts.push(orig_part.clone()); + } + } + + let remaining_content_parts: Vec = new_content_parts.collect(); + if !content_emitted && !remaining_content_parts.is_empty() { + for part in remaining_content_parts.into_iter().rev() { + parts.insert(0, part); + } + } else { + parts.extend(remaining_content_parts); + } + parts.extend(rebuilt_fn_calls); + if parts.is_empty() { + parts.push(serde_json::json!({"text": ""})); + } + parts +} + +fn patch_gemini_visible_content( + orig_parts: &[Json], + msg_obj: &serde_json::Map, + gemini_role: &str, +) -> Result> { + let content_val = msg_obj.get("content").unwrap_or(&Json::Null); + let (new_content_parts, content_is_parts_form) = + gemini_content_parts_from_normalized(content_val)?; + let rebuilt_fn_calls = rebuilt_gemini_function_call_parts(orig_parts, msg_obj)?; + let parts = merge_gemini_original_parts( + orig_parts, + new_content_parts, + rebuilt_fn_calls, + content_is_parts_form, + ); + + Ok(Some( + serde_json::json!({"role": gemini_role, "parts": parts}), + )) +} + +/// Patch `generationConfig` with the modeled params, preserving unmodeled keys. +/// +/// Modeled keys: `temperature`, `topP`, `maxOutputTokens`, `stopSequences`. +/// All other keys (e.g. `responseMimeType`, `responseSchema`, `thinkingConfig`) +/// are preserved regardless of whether `params` is `Some` or `None`. +const MODELED_GEN_CONFIG_KEYS: &[&str] = + &["temperature", "topP", "maxOutputTokens", "stopSequences"]; + +fn patch_gemini_params( + obj: &mut serde_json::Map, + params: Option<&GenerationParams>, +) -> Result<()> { + let Some(params) = params else { + // Params cleared: remove only modeled keys; keep provider-native fields. + if let Some(gc) = obj + .get_mut("generationConfig") + .and_then(|v| v.as_object_mut()) + { + for key in MODELED_GEN_CONFIG_KEYS { + gc.remove(*key); + } + } + // If generationConfig is now empty (or was absent), drop the key entirely. + if obj + .get("generationConfig") + .and_then(|v| v.as_object()) + .map(|m| m.is_empty()) + .unwrap_or(false) + { + obj.remove("generationConfig"); + } + return Ok(()); + }; + + let mut gen_config: serde_json::Map = obj + .get("generationConfig") + .and_then(|v| v.as_object()) + .cloned() + .unwrap_or_default(); + + match params.temperature { + Some(t) => { + gen_config.insert("temperature".into(), json_f64(t, "temperature")?); + } + None => { + gen_config.remove("temperature"); + } + } + match params.top_p { + Some(p) => { + gen_config.insert("topP".into(), json_f64(p, "topP")?); + } + None => { + gen_config.remove("topP"); + } + } + match params.max_tokens { + Some(n) => { + gen_config.insert("maxOutputTokens".into(), Json::from(n)); + } + None => { + gen_config.remove("maxOutputTokens"); + } + } + match ¶ms.stop { + Some(stop) => { + insert_serialized(&mut gen_config, "stopSequences", stop, "stopSequences")?; + } + None => { + gen_config.remove("stopSequences"); + } + } + + if gen_config.is_empty() { + obj.remove("generationConfig"); + } else { + obj.insert("generationConfig".into(), Json::Object(gen_config)); + } + Ok(()) +} + +/// Patch the `functionDeclarations` group inside `tools`, preserving original +/// group order, group-level sibling fields, and all other tool groups +/// (googleSearch, codeExecution, …) in their original positions. +/// +/// Provider-native fields captured in `FunctionDefinition.extra` +/// (parametersJsonSchema, responseJsonSchema, response, behavior, …) are merged +/// back into the encoded functionDeclaration, with modeled fields (name, +/// description, parameters) taking precedence. +fn patch_gemini_tools( + obj: &mut serde_json::Map, + tools: Option<&Vec>, +) -> Result<()> { + // Build functionDeclaration objects and provider-native Gemini tool groups. + let fn_declarations: Vec = { + let mut out = Vec::new(); + if let Some(ts) = tools { + for td in ts { + match td { + ToolDefinition::Function { function: fd, .. } => { + if fd.name.is_empty() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: FunctionDefinition.name must be non-empty".into(), + )); + } + if fd.strict.is_some() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: FunctionDefinition.strict is not supported; \ + remove it or use a provider-native extra field" + .into(), + )); + } + // Start with extra (provider-native fields), then overlay modeled fields. + let mut fdobj: serde_json::Map = fd.extra.clone(); + fdobj.insert("name".into(), Json::String(fd.name.clone())); + if let Some(ref desc) = fd.description { + fdobj.insert("description".into(), Json::String(desc.clone())); + } + if let Some(ref params) = fd.parameters { + fdobj.insert("parameters".into(), params.clone()); + } + out.push(Json::Object(fdobj)); + } + ToolDefinition::ProviderNative { + provider, kind: _, .. + } if provider == GEMINI_PROVIDER => {} + ToolDefinition::ProviderNative { provider, kind, .. } => { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: ProviderNative tool '{kind}' (provider '{provider}') \ + cannot be represented on the Gemini surface" + ))); + } + } + } + } + out + }; + let native_groups: Vec = tools + .into_iter() + .flatten() + .filter_map(|td| match td { + ToolDefinition::ProviderNative { + provider, value, .. + } if provider == GEMINI_PROVIDER => Some(value), + _ => None, + }) + .map(|value| { + if !value.is_object() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: ProviderNative tool value must be an object".into(), + )); + } + if value.get("functionDeclarations").is_some() { + return Err(FlowError::InvalidArgument( + "Gemini encoder: ProviderNative tool value must not contain \ + functionDeclarations; use ToolDefinition::Function instead" + .into(), + )); + } + Ok(value.clone()) + }) + .collect::>>()?; + + // Walk the original tools array in order: replace the FIRST functionDeclarations + // group with the rebuilt list, merge any native sibling fields through the + // normalized ProviderNative item, replace native-only groups from the normalized + // list, and append any newly added native groups. If there was no original + // functionDeclarations group, append the new one at the end. + let orig_tools = obj + .get("tools") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + + // If the original request had multiple functionDeclarations groups, the decode path + // already flattened them into a single normalized list so there is no way to know + // which function belongs in which group after an edit. Return an error rather than + // silently collapsing all functions into the first group and losing the rest. + let fn_decl_group_count = orig_tools + .iter() + .filter(|g| g.get("functionDeclarations").is_some()) + .count(); + if fn_decl_group_count > 1 { + return Err(FlowError::InvalidArgument(format!( + "Gemini encoder: the original request has {fn_decl_group_count} \ + functionDeclarations groups; editing tools with multiple groups is not \ + supported because the decode path flattens them and grouping cannot be \ + recovered. Use provider-native extra fields to manage multi-group tools." + ))); + } + let mut new_groups: Vec = Vec::with_capacity(orig_tools.len()); + let mut fn_group_placed = false; + let mut native_used = vec![false; native_groups.len()]; + for orig_group in &orig_tools { + if orig_group.get("functionDeclarations").is_some() { + if !fn_group_placed { + let native_sibling_keys = gemini_native_tool_keys(orig_group); + let mut group = serde_json::Map::new(); + if !fn_declarations.is_empty() { + group.insert( + "functionDeclarations".into(), + Json::Array(fn_declarations.clone()), + ); + } + if !native_sibling_keys.is_empty() + && let Some(native_group) = take_matching_native_group( + &native_groups, + &mut native_used, + &native_sibling_keys, + ) + && let Some(native_obj) = native_group.as_object() + { + group.extend(native_obj.clone()); + } + if !group.is_empty() { + new_groups.push(Json::Object(group)); + } + fn_group_placed = true; + } + // Unreachable: the fn_decl_group_count > 1 guard above returns Err before + // this loop when tools changed, and this function is only called when tools + // changed. A second functionDeclarations group can never be reached here. + } else { + let native_keys = gemini_native_tool_keys(orig_group); + if let Some(native_group) = + take_matching_native_group(&native_groups, &mut native_used, &native_keys) + { + new_groups.push(native_group); + } + } + } + if !fn_group_placed && !fn_declarations.is_empty() { + new_groups.push(serde_json::json!({"functionDeclarations": fn_declarations})); + } + new_groups.extend( + native_groups + .iter() + .enumerate() + .filter(|(idx, _)| !native_used[*idx]) + .map(|(_, group)| group.clone()), + ); + + if new_groups.is_empty() { + obj.remove("tools"); + } else { + obj.insert("tools".into(), Json::Array(new_groups)); + } + Ok(()) +} + +/// Overlay extra-field changes from `annotated` onto `obj`, guided by `baseline`. +fn patch_extra_fields( + obj: &mut serde_json::Map, + baseline: &serde_json::Map, + annotated: &serde_json::Map, +) { + for key in baseline.keys().filter(|k| !annotated.contains_key(*k)) { + obj.remove(key); + } + for (key, value) in annotated { + if baseline.get(key) != Some(value) { + obj.insert(key.clone(), value.clone()); + } + } +} + +// --------------------------------------------------------------------------- +// Streaming codec +// --------------------------------------------------------------------------- + +/// Streaming counterpart to [`GeminiGenerateContentCodec`]. +/// +/// Accumulates Gemini server-sent event chunks and assembles a complete response +/// that [`GeminiGenerateContentCodec::decode_response`] can consume. +pub struct GeminiGenerateContentStreamingCodec { + state: std::sync::Arc>, +} + +impl GeminiGenerateContentStreamingCodec { + /// Creates a fresh streaming codec with empty accumulator state. + pub fn new() -> Self { + Self { + state: std::sync::Arc::new(std::sync::Mutex::new( + GeminiGenerateContentStreamingState::default(), + )), + } + } +} + +impl Default for GeminiGenerateContentStreamingCodec { + fn default() -> Self { + Self::new() + } +} + +impl super::streaming::StreamingCodec for GeminiGenerateContentStreamingCodec { + fn collector(&self) -> crate::api::runtime::LlmCollectorFn { + let state = std::sync::Arc::clone(&self.state); + Box::new(move |event: Json| -> Result<()> { + let mut guard = state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + guard.observe(&event)?; + Ok(()) + }) + } + + fn finalizer(&self) -> crate::api::runtime::LlmFinalizerFn { + let state = std::sync::Arc::clone(&self.state); + Box::new(move || -> Json { + let mut guard = state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + std::mem::take(&mut *guard).finalize() + }) + } +} + +#[derive(Debug, Default)] +struct GeminiGenerateContentStreamingState { + parts: Vec, + candidate_index: Option, + finish_reason: Option, + usage_metadata: Option, + model_version: Option, + response_id: Option, + /// Merged candidate-level extra fields (safetyRatings, groundingMetadata, + /// citationMetadata, avgLogprobs, etc.) accumulated across SSE chunks. + candidate_extra: serde_json::Map, +} + +impl GeminiGenerateContentStreamingState { + fn push_text_part(&mut self, text: &str, part_obj: &serde_json::Map) { + let mut extra: serde_json::Map = part_obj + .iter() + .filter(|(k, _)| k.as_str() != "text") + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + if let Some(last_obj) = self.parts.last_mut().and_then(Json::as_object_mut) { + let current_has_metadata = !extra.is_empty(); + let previous_has_metadata = last_obj.keys().any(|k| k != "text"); + let can_merge = last_obj.get("text").is_some_and(Json::is_string) + && !current_has_metadata + && !previous_has_metadata; + if can_merge { + let Some(Json::String(existing)) = last_obj.get_mut("text") else { + return; + }; + existing.push_str(text); + for (k, v) in extra { + last_obj.insert(k, v); + } + return; + } + } + + extra.insert("text".into(), Json::String(text.to_string())); + self.parts.push(Json::Object(extra)); + } + + fn observe(&mut self, event: &Json) -> Result<()> { + if let Some(candidates) = event.get("candidates").and_then(Json::as_array) + && let Some(candidate) = candidates.first() + { + if candidates.len() > 1 { + return Err(FlowError::InvalidArgument( + "Gemini streaming chunks with multiple candidates are not supported".into(), + )); + } + let candidate_obj = candidate.as_object().ok_or_else(|| { + FlowError::InvalidArgument("Gemini streaming candidate must be an object".into()) + })?; + let index = candidate_obj + .get("index") + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini streaming candidate index is required".into(), + ) + })? + .as_u64() + .ok_or_else(|| { + FlowError::InvalidArgument( + "Gemini streaming candidate index must be an unsigned integer".into(), + ) + })?; + if let Some(previous_index) = self.candidate_index { + if previous_index != index { + return Err(FlowError::InvalidArgument( + "Gemini streaming candidate index changed across chunks".into(), + )); + } + } else { + if index != 0 { + return Err(FlowError::InvalidArgument( + "Gemini streaming only supports candidate index 0".into(), + )); + } + self.candidate_index = Some(index); + } + + if let Some(parts) = candidate_obj + .get("content") + .and_then(|c| c.get("parts")) + .and_then(Json::as_array) + { + for part in parts { + if !part.is_object() { + return Err(FlowError::InvalidArgument( + "Gemini streaming parts entry must be an object".into(), + )); + } + let part_obj = part.as_object().unwrap(); + let data_key = validate_single_gemini_part_data_field(part_obj, "streaming")?; + // Preserve thought parts in the provider-native aggregate; the response + // decoder filters them out of the normalized message. + if part.get("thought").and_then(Json::as_bool) == Some(true) { + self.parts.push(part.clone()); + continue; + } + + match data_key { + Some("text") => { + let text_val = part.get("text").unwrap(); + match text_val.as_str() { + Some(s) => { + self.push_text_part(s, part_obj); + } + None => { + return Err(FlowError::InvalidArgument( + "Gemini streaming parts[].text must be a string".into(), + )); + } + } + } + Some("functionCall") => { + self.parts.push(part.clone()); + } + Some("functionResponse") => { + return Err(FlowError::InvalidArgument( + "Gemini streaming response parts must not contain functionResponse" + .into(), + )); + } + Some(_) | None => { + self.parts.push(part.clone()); + } + } + } + } + if let Some(reason_val) = candidate.get("finishReason") { + match reason_val.as_str() { + Some(s) => self.finish_reason = Some(s.to_string()), + None => { + return Err(FlowError::InvalidArgument( + "Gemini streaming candidate finishReason must be a string".into(), + )); + } + } + } + // Collect candidate-level metadata fields (safetyRatings, groundingMetadata, + // citationMetadata, avgLogprobs, etc.) that non-streaming decode preserves + // in ApiSpecificResponse::GeminiGenerateContent. Later chunks overwrite earlier ones for + // the same key (last-wins), matching the non-streaming behaviour. + for (k, v) in candidate_obj { + if !matches!(k.as_str(), "content" | "finishReason" | "index") { + self.candidate_extra.insert(k.clone(), v.clone()); + } + } + } + if let Some(usage) = event.get("usageMetadata") { + self.usage_metadata = Some(usage.clone()); + } + if let Some(mv_val) = event.get("modelVersion") { + match mv_val.as_str() { + Some(s) => self.model_version = Some(s.to_string()), + None => { + return Err(FlowError::InvalidArgument( + "Gemini streaming event modelVersion must be a string".into(), + )); + } + } + } + if let Some(rid_val) = event.get("responseId") { + match rid_val.as_str() { + Some(s) => self.response_id = Some(s.to_string()), + None => { + return Err(FlowError::InvalidArgument( + "Gemini streaming event responseId must be a string".into(), + )); + } + } + } + Ok(()) + } + + fn finalize(self) -> Json { + let mut candidate_obj = serde_json::Map::new(); + candidate_obj.insert( + "content".into(), + serde_json::json!({"role": "model", "parts": self.parts}), + ); + if let Some(reason) = self.finish_reason { + candidate_obj.insert("finishReason".into(), Json::String(reason)); + } + candidate_obj.insert( + "index".into(), + Json::from(self.candidate_index.unwrap_or(0)), + ); + // Merge accumulated candidate-level metadata so that decode_response can + // populate ApiSpecificResponse::GeminiGenerateContent with the same fields as non-streaming. + for (k, v) in self.candidate_extra { + candidate_obj.entry(k).or_insert(v); + } + let candidate = Json::Object(candidate_obj); + + let mut output = serde_json::Map::new(); + output.insert("candidates".to_string(), Json::Array(vec![candidate])); + if let Some(usage) = self.usage_metadata { + output.insert("usageMetadata".to_string(), usage); + } + if let Some(mv) = self.model_version { + output.insert("modelVersion".to_string(), Json::String(mv)); + } + if let Some(rid) = self.response_id { + output.insert("responseId".to_string(), Json::String(rid)); + } + Json::Object(output) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +#[path = "../../tests/unit/codec/gemini_generate_content_tests.rs"] +mod tests; diff --git a/crates/core/src/codec/mod.rs b/crates/core/src/codec/mod.rs index a92f00826..4a587533e 100644 --- a/crates/core/src/codec/mod.rs +++ b/crates/core/src/codec/mod.rs @@ -15,6 +15,7 @@ //! provider codec from a raw payload when no codec annotation is present. pub mod anthropic; +pub mod gemini_generate_content; pub mod model_pricing; pub mod openai_chat; pub mod openai_responses; diff --git a/crates/core/src/codec/resolve.rs b/crates/core/src/codec/resolve.rs index 4dc7cbf0b..f54652291 100644 --- a/crates/core/src/codec/resolve.rs +++ b/crates/core/src/codec/resolve.rs @@ -14,7 +14,7 @@ use super::request::AnnotatedLlmRequest; use super::response::AnnotatedLlmResponse; use super::streaming::StreamingCodec; use super::traits::{LlmCodec, LlmResponseCodec}; -use super::{anthropic, openai_chat, openai_responses}; +use super::{anthropic, gemini_generate_content, openai_chat, openai_responses}; /// A built-in provider request/response surface. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -25,6 +25,8 @@ pub enum ProviderSurface { OpenAIResponses, /// Anthropic Messages. AnthropicMessages, + /// Gemini generateContent. + GeminiGenerateContent, } /// Request shape detector; the optional `&str` is a provider hint a codec may use @@ -67,15 +69,16 @@ pub(crate) static BUILTIN_PROVIDER_SURFACES: &[ProviderSurfaceDescriptor] = &[ openai_responses::PROVIDER_SURFACE, anthropic::PROVIDER_SURFACE, openai_chat::PROVIDER_SURFACE, + gemini_generate_content::PROVIDER_SURFACE, ]; /// Detect the request surface from a raw request body by top-level key. /// /// Priority: OpenAI Responses (`input`/`instructions`) > Anthropic Messages -/// (`system`) > OpenAI Chat (`messages`). `None` when no key matches or `body` -/// is not an object. This is a best-effort heuristic: an Anthropic request that -/// omits the optional top-level `system` is indistinguishable from OpenAI Chat -/// and classifies as `OpenAIChat`. +/// (`system`) > OpenAI Chat (`messages`) > Gemini generateContent (`contents`). +/// `None` when no key matches or `body` is not an object. This is a best-effort heuristic: an +/// Anthropic request that omits the optional top-level `system` is +/// indistinguishable from OpenAI Chat and classifies as `OpenAIChat`. #[must_use] pub fn detect_request_surface(body: &Json) -> Option { detect_request_surface_with_hint(body, None) @@ -152,19 +155,19 @@ fn descriptor_for(surface: ProviderSurface) -> &'static ProviderSurfaceDescripto ProviderSurface::OpenAIChat => &openai_chat::PROVIDER_SURFACE, ProviderSurface::OpenAIResponses => &openai_responses::PROVIDER_SURFACE, ProviderSurface::AnthropicMessages => &anthropic::PROVIDER_SURFACE, + ProviderSurface::GeminiGenerateContent => &gemini_generate_content::PROVIDER_SURFACE, } } impl ProviderSurface { - /// The canonical codec name for this surface (e.g. `"openai_chat"`), the - /// inverse of [`Self::from_codec_name`]. + /// The canonical codec name for this surface (e.g. `"openai_chat"`). #[must_use] pub fn codec_name(self) -> &'static str { descriptor_for(self).codec_name } - /// Resolves a canonical codec name to its surface, or `None` when `name` is - /// not a built-in provider codec. + /// Resolves a canonical codec name to its surface, or `None` when `name` + /// is not a built-in provider codec. #[must_use] pub fn from_codec_name(name: &str) -> Option { BUILTIN_PROVIDER_SURFACES diff --git a/crates/core/src/plugins/nemo_guardrails/component.rs b/crates/core/src/plugins/nemo_guardrails/component.rs index 82a5f0e63..9689c5bab 100644 --- a/crates/core/src/plugins/nemo_guardrails/component.rs +++ b/crates/core/src/plugins/nemo_guardrails/component.rs @@ -282,7 +282,7 @@ crate::editor_config! { codec => { label: "codec", kind: Enum, - values: ["openai_chat", "openai_responses", "anthropic_messages"], + values: ["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"], optional: true, }, input => { label: "input", kind: Boolean }, @@ -430,7 +430,12 @@ fn mode_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::sc fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema { string_enum_schema( generator, - &["openai_chat", "openai_responses", "anthropic_messages"], + &[ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate_content", + ], None, ) } @@ -921,7 +926,10 @@ fn validate_codec_requirements( "nemo_guardrails.unsupported_value", Some(NEMO_GUARDRAILS_PLUGIN_KIND.to_string()), Some("codec".to_string()), - "codec must be 'openai_chat', 'openai_responses', or 'anthropic_messages'".to_string(), + format!( + "codec must be one of: {}", + supported_codec_names().join(", ") + ), ); } } diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs index 504a798cb..78ebf92a2 100644 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ b/crates/core/src/plugins/nemo_guardrails/python.rs @@ -891,6 +891,7 @@ enum LocalGuardrailsCodec { OpenAIChat, OpenAIResponses, AnthropicMessages, + GeminiGenerateContent, } impl LocalGuardrailsCodec { @@ -899,6 +900,7 @@ impl LocalGuardrailsCodec { Self::OpenAIChat => ProviderSurface::OpenAIChat, Self::OpenAIResponses => ProviderSurface::OpenAIResponses, Self::AnthropicMessages => ProviderSurface::AnthropicMessages, + Self::GeminiGenerateContent => ProviderSurface::GeminiGenerateContent, } } @@ -907,6 +909,7 @@ impl LocalGuardrailsCodec { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, } } @@ -1347,6 +1350,26 @@ fn extract_stream_text(codec: LocalGuardrailsCodec, chunk: &Json) -> Option { + let candidates = chunk.get("candidates")?.as_array()?; + let parts = candidates + .first()? + .get("content")? + .get("parts")? + .as_array()?; + let mut texts = vec![]; + for part in parts { + if part.get("thought").and_then(Json::as_bool) == Some(true) { + continue; + } + if let Some(text) = part.get("text").and_then(Json::as_str) + && !text.is_empty() + { + texts.push(text); + } + } + (!texts.is_empty()).then(|| texts.join("")) + } } } diff --git a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs new file mode 100644 index 000000000..f1216562f --- /dev/null +++ b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs @@ -0,0 +1,4760 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Unit tests for GeminiGenerateContentCodec in the NeMo Relay core crate. + +use super::*; +use serde_json::json; + +use super::super::request::{ContentPart, Message, MessageContent, ToolDefinition}; +use super::super::response::FinishReason; +use super::super::streaming::StreamingCodec; + +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; + +// ------------------------------------------------------------------- +// Helpers +// ------------------------------------------------------------------- + +fn make_request(content: Json) -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content, + } +} + +// =================================================================== +// codec_identity +// =================================================================== + +#[test] +fn test_codec_identity_is_gemini_builtin() { + let codec = GeminiGenerateContentCodec; + assert_eq!( + LlmCodec::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent), + "GeminiGenerateContentCodec must not return Opaque; PII sanitization depends on a known identity" + ); +} + +#[test] +fn test_response_codec_identity_is_gemini_builtin() { + let codec = GeminiGenerateContentCodec; + assert_eq!( + ::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent), + "GeminiGenerateContentCodec response codec must not return Opaque" + ); +} + +// =================================================================== +// Response decode tests +// =================================================================== + +#[test] +fn test_decode_response_text() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"text": "Hello, world!"}] + }, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15 + }, + "modelVersion": "gemini-2.0-flash" + }); + + let resp = codec.decode_response(&response).unwrap(); + assert_eq!( + resp.message, + Some(MessageContent::Text("Hello, world!".into())) + ); + assert_eq!(resp.finish_reason, Some(FinishReason::Complete)); + assert_eq!(resp.model, Some("gemini-2.0-flash".into())); + + let usage = resp.usage.unwrap(); + assert_eq!(usage.prompt_tokens, Some(10)); + assert_eq!(usage.completion_tokens, Some(5)); + assert_eq!(usage.total_tokens, Some(15)); +} + +#[test] +fn test_decode_response_native_part_as_provider_native_content() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "ran code"}, + {"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "sk-code-secret"}} + ] + }, + "finishReason": "STOP" + }] + }); + + let resp = codec.decode_response(&response).unwrap(); + let Some(MessageContent::Parts(parts)) = resp.message else { + panic!("expected mixed Gemini response parts to decode as MessageContent::Parts"); + }; + + assert!(matches!( + &parts[0], + ContentPart::Text { text, .. } if text == "ran code" + )); + match &parts[1] { + ContentPart::ProviderNative { + provider, + kind, + value, + } => { + assert_eq!(provider, "gemini"); + assert_eq!(kind, "codeExecutionResult"); + assert_eq!( + value["codeExecutionResult"]["output"], + json!("sk-code-secret") + ); + } + other => panic!("expected Gemini ProviderNative response part, got {other:?}"), + } +} + +#[test] +fn test_decode_response_response_id() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}, "finishReason": "STOP"}], + "responseId": "resp-abc-123", + "usageMetadata": {"promptTokenCount": 1} + }); + let resp = codec.decode_response(&response).unwrap(); + assert_eq!( + resp.id.as_deref(), + Some("resp-abc-123"), + "responseId must be mapped to AnnotatedLlmResponse.id" + ); +} + +#[test] +fn test_decode_response_function_call() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{ + "functionCall": { + "name": "get_weather", + "args": {"location": "NYC"} + } + }] + }, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": { + "promptTokenCount": 20, + "candidatesTokenCount": 10, + "totalTokenCount": 30 + } + }); + + let resp = codec.decode_response(&response).unwrap(); + assert_eq!(resp.finish_reason, Some(FinishReason::ToolUse)); + assert!(resp.message.is_none()); + + let tool_calls = resp.tool_calls.unwrap(); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].arguments, json!({"location": "NYC"})); + // id is derived from name when Gemini omits one + assert_eq!(tool_calls[0].id, "get_weather"); +} + +#[test] +fn test_decode_response_finish_reason_table() { + let codec = GeminiGenerateContentCodec; + + // (finishReason string, expected FinishReason) + // NOTE: these cases have no functionCall parts in the response, so has_tool_calls = false. + let cases: &[(&str, Option)] = &[ + ("STOP", Some(FinishReason::Complete)), + ("MAX_TOKENS", Some(FinishReason::Length)), + ("TOOL_CODE", Some(FinishReason::ToolUse)), + // Error / malfunction codes must not be overridden by tool-call heuristic. + ( + "MALFORMED_FUNCTION_CALL", + Some(FinishReason::Unknown("MALFORMED_FUNCTION_CALL".into())), + ), + ( + "UNEXPECTED_TOOL_CALL", + Some(FinishReason::Unknown("UNEXPECTED_TOOL_CALL".into())), + ), + // Text safety / policy reasons — all must map to ContentFilter, not Unknown. + ("SAFETY", Some(FinishReason::ContentFilter)), + ("RECITATION", Some(FinishReason::ContentFilter)), + ("BLOCKLIST", Some(FinishReason::ContentFilter)), + ("PROHIBITED_CONTENT", Some(FinishReason::ContentFilter)), + ("SPII", Some(FinishReason::ContentFilter)), + // Image safety and other policy reasons — also ContentFilter. + ("LANGUAGE", Some(FinishReason::ContentFilter)), + ("IMAGE_SAFETY", Some(FinishReason::ContentFilter)), + ( + "IMAGE_PROHIBITED_CONTENT", + Some(FinishReason::ContentFilter), + ), + ("IMAGE_RECITATION", Some(FinishReason::ContentFilter)), + ("ESCALATION", Some(FinishReason::ContentFilter)), + // Unspecified / not-yet-finished must map to None. + ("FINISH_REASON_UNSPECIFIED", None), + // Unknown future values map to Unknown. + ( + "FUTURE_REASON", + Some(FinishReason::Unknown("FUTURE_REASON".into())), + ), + ]; + + for (reason_str, expected) in cases { + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": []}, + "finishReason": reason_str, + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 0} + }); + let resp = codec.decode_response(&response).unwrap(); + assert_eq!( + &resp.finish_reason, expected, + "finishReason={reason_str} must decode to {expected:?}" + ); + } +} + +#[test] +fn test_decode_response_prompt_feedback_block_reason_content_filter() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "promptFeedback": { + "blockReason": "SAFETY", + "safetyRatings": [{"category": "HARM_CATEGORY_DANGEROUS_CONTENT"}] + }, + "usageMetadata": {"promptTokenCount": 12}, + "modelVersion": "gemini-2.0-flash" + }); + + let resp = codec.decode_response(&response).unwrap(); + assert_eq!(resp.finish_reason, Some(FinishReason::ContentFilter)); + assert!(resp.message.is_none()); + assert!(resp.tool_calls.is_none()); + assert_eq!( + resp.extra + .get("promptFeedback") + .and_then(|feedback| feedback.get("blockReason")) + .and_then(Json::as_str), + Some("SAFETY"), + "promptFeedback must remain available as top-level response extra" + ); +} + +#[test] +fn test_decode_response_prompt_feedback_block_reason_must_be_string() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "promptFeedback": {"blockReason": 123} + }); + + let err = codec.decode_response(&response).unwrap_err(); + assert!( + err.to_string().contains("blockReason must be a string"), + "unexpected error: {err}" + ); +} + +#[test] +fn test_decode_response_cached_content_token_count() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "cached response"}]}, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": { + "promptTokenCount": 100, + "candidatesTokenCount": 10, + "totalTokenCount": 110, + "cachedContentTokenCount": 80 + } + }); + + let resp = codec.decode_response(&response).unwrap(); + let usage = resp.usage.unwrap(); + assert_eq!(usage.prompt_tokens, Some(100)); + assert_eq!(usage.completion_tokens, Some(10)); + assert_eq!(usage.cache_read_tokens, Some(80)); + assert_eq!(usage.cache_write_tokens, None); +} + +#[test] +fn test_decode_response_no_candidates() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 0, "totalTokenCount": 5} + }); + + let resp = codec.decode_response(&response).unwrap(); + assert!(resp.message.is_none()); + assert!(resp.tool_calls.is_none()); + assert!(resp.finish_reason.is_none()); +} + +#[test] +fn test_decode_response_extra_fields_preserved() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hi"}]}, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + "unknownFutureField": "value" + }); + + let resp = codec.decode_response(&response).unwrap(); + assert!(resp.extra.contains_key("unknownFutureField")); +} + +#[test] +fn test_decode_response_candidate_extra_removes_reserved_api_key() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hi"}]}, + "index": 0, + "api": "bad_discriminator", + "futureField": true + }] + }); + + let resp = codec.decode_response(&response).unwrap(); + let Some(super::super::response::ApiSpecificResponse::GeminiGenerateContent { extra, .. }) = + resp.api_specific + else { + panic!("expected Gemini api_specific metadata"); + }; + assert!(extra.get("api").is_none()); + assert_eq!(extra.get("futureField"), Some(&json!(true))); +} + +#[test] +fn test_decode_response_candidate_only_reserved_api_key_has_no_api_specific() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hi"}]}, + "api": "bad_discriminator" + }] + }); + + let resp = codec.decode_response(&response).unwrap(); + assert!( + resp.api_specific.is_none(), + "reserved api discriminator alone must not create empty Gemini api_specific metadata" + ); +} + +// =================================================================== +// Request decode tests +// =================================================================== + +#[test] +fn test_decode_contents_with_system_instruction() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "Hello"}]}, + {"role": "model", "parts": [{"text": "Hi there"}]}, + {"role": "user", "parts": [{"text": "What's the weather?"}]} + ], + "systemInstruction": { + "parts": [{"text": "You are a helpful assistant."}] + } + })); + + let annotated = codec.decode(&request).unwrap(); + + assert!( + matches!(&annotated.messages[0], Message::System { content: MessageContent::Text(t), .. } if t == "You are a helpful assistant.") + ); + assert!( + matches!(&annotated.messages[1], Message::User { content: MessageContent::Text(t), .. } if t == "Hello") + ); + assert!( + matches!(&annotated.messages[2], Message::Assistant { content: Some(MessageContent::Text(t)), .. } if t == "Hi there") + ); + assert!( + matches!(&annotated.messages[3], Message::User { content: MessageContent::Text(t), .. } if t == "What's the weather?") + ); + assert_eq!(annotated.messages.len(), 4); +} + +#[test] +fn test_decode_function_declarations() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{ + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"] + } + }] + }] + })); + + let annotated = codec.decode(&request).unwrap(); + let tools = annotated.tools.unwrap(); + assert_eq!(tools.len(), 1); + let nemo_relay_types::codec::request::ToolDefinition::Function { ref function, .. } = tools[0] + else { + panic!("expected Function variant"); + }; + assert_eq!(function.name, "get_weather"); + assert_eq!( + function.description.as_deref(), + Some("Get the weather for a location") + ); + assert!(function.parameters.is_some()); +} + +#[test] +fn test_decode_function_declaration_preserves_provider_fields() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{ + "name": "my_tool", + "description": "A tool", + "parameters": {"type": "object"}, + "parametersJsonSchema": {"$schema": "draft-2020-12"}, + "responseJsonSchema": {"type": "object"}, + "behavior": "BLOCKING" + }] + }] + })); + + let annotated = codec.decode(&request).unwrap(); + let tools = annotated.tools.unwrap(); + let nemo_relay_types::codec::request::ToolDefinition::Function { ref function, .. } = tools[0] + else { + panic!("expected Function variant"); + }; + assert!( + function.extra.contains_key("parametersJsonSchema"), + "parametersJsonSchema must be in extra" + ); + assert!( + function.extra.contains_key("responseJsonSchema"), + "responseJsonSchema must be in extra" + ); + assert!( + function.extra.contains_key("behavior"), + "behavior must be in extra" + ); + assert!( + !function.extra.contains_key("name"), + "modeled fields must not appear in extra" + ); +} + +#[test] +fn test_decode_generation_config() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": { + "temperature": 0.7, + "topP": 0.9, + "maxOutputTokens": 1024, + "stopSequences": ["stop1", "stop2"] + } + })); + + let annotated = codec.decode(&request).unwrap(); + let params = annotated.params.unwrap(); + assert!((params.temperature.unwrap() - 0.7).abs() < 1e-9); + assert!((params.top_p.unwrap() - 0.9).abs() < 1e-9); + assert_eq!(params.max_tokens, Some(1024)); + assert_eq!( + params.stop.as_deref(), + Some(&["stop1".to_string(), "stop2".to_string()][..]) + ); +} + +#[test] +fn test_decode_extra_fields_captured() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "safetySettings": [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}], + "cachedContent": "cachedContent/abc" + })); + + let annotated = codec.decode(&request).unwrap(); + assert!(annotated.extra.contains_key("safetySettings")); + assert!(annotated.extra.contains_key("cachedContent")); + assert!(!annotated.extra.contains_key("contents")); +} + +// =================================================================== +// Request encode tests +// =================================================================== + +#[test] +fn test_encode_round_trip_preserves_extra_fields() { + let codec = GeminiGenerateContentCodec; + let original_json = json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "safetySettings": [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}] + }); + let original = make_request(original_json); + + let annotated = codec.decode(&original).unwrap(); + let re_encoded = codec.encode(&annotated, &original).unwrap(); + + assert!(re_encoded.content.get("safetySettings").is_some()); +} + +/// Verify that "assistant" → "model" mapping works when a fresh Message::Assistant +/// is encoded (not decoded from an existing Gemini request). +#[test] +fn test_encode_role_assistant_becomes_model() { + let codec = GeminiGenerateContentCodec; + // Decode a single-user-turn request to get a valid baseline. + let original_json = json!({ + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] + }); + let original = make_request(original_json); + let mut annotated = codec.decode(&original).unwrap(); + + // Intercept adds a fresh assistant reply (not decoded from Gemini — this + // exercises the normalized_to_gemini_content path directly). + annotated.messages.push(Message::Assistant { + content: Some(MessageContent::Text("Hi there".into())), + tool_calls: None, + name: None, + }); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded.content.get("contents").unwrap().as_array().unwrap(); + + assert_eq!(contents.len(), 2); + let second_role = contents[1].get("role").unwrap().as_str().unwrap(); + assert_eq!( + second_role, "model", + "Message::Assistant must encode as role 'model'" + ); +} + +/// Decode-then-encode preserves systemInstruction when unchanged. +#[test] +fn test_encode_preserves_system_instruction_when_unchanged() { + let codec = GeminiGenerateContentCodec; + let original_json = json!({ + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "You are an assistant."}]} + }); + let original = make_request(original_json); + + let annotated = codec.decode(&original).unwrap(); + let re_encoded = codec.encode(&annotated, &original).unwrap(); + + let sys = re_encoded.content.get("systemInstruction").unwrap(); + let parts = sys.get("parts").unwrap().as_array().unwrap(); + assert_eq!( + parts[0].get("text").unwrap().as_str().unwrap(), + "You are an assistant." + ); +} + +/// Decode-then-encode preserves tools when unchanged (round-trip). +#[test] +fn test_encode_preserves_tools_when_unchanged() { + let codec = GeminiGenerateContentCodec; + let original_json = json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{ + "name": "search", + "description": "Search the web", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}} + }] + }] + }); + let original = make_request(original_json); + + let annotated = codec.decode(&original).unwrap(); + let re_encoded = codec.encode(&annotated, &original).unwrap(); + + let tools = re_encoded.content.get("tools").unwrap().as_array().unwrap(); + assert_eq!(tools.len(), 1); + let fds = tools[0] + .get("functionDeclarations") + .unwrap() + .as_array() + .unwrap(); + assert_eq!(fds.len(), 1); + assert_eq!(fds[0].get("name").unwrap().as_str().unwrap(), "search"); + assert!(fds[0].get("parameters").is_some()); +} + +/// Decode-then-encode preserves generationConfig when unchanged (round-trip). +#[test] +fn test_encode_preserves_generation_config_when_unchanged() { + let codec = GeminiGenerateContentCodec; + let original_json = json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": { + "temperature": 0.5, + "maxOutputTokens": 512 + } + }); + let original = make_request(original_json); + + let annotated = codec.decode(&original).unwrap(); + let re_encoded = codec.encode(&annotated, &original).unwrap(); + + let gc = re_encoded.content.get("generationConfig").unwrap(); + assert!((gc.get("temperature").unwrap().as_f64().unwrap() - 0.5).abs() < 1e-9); + assert_eq!(gc.get("maxOutputTokens").unwrap().as_u64().unwrap(), 512); +} + +/// Interceptor can edit the system message; the new text appears in the output. +#[test] +fn test_encode_system_instruction_edit() { + let codec = GeminiGenerateContentCodec; + let original_json = json!({ + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "old prompt"}]} + }); + let original = make_request(original_json); + + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::System { content, .. } = msg { + *content = MessageContent::Text("new prompt".into()); + } + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let sys_text = encoded + .content + .get("systemInstruction") + .and_then(|s| s.get("parts")) + .and_then(|p| p.as_array()) + .and_then(|a| a.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str) + .unwrap(); + assert_eq!(sys_text, "new prompt"); +} + +// =================================================================== +// Streaming tests +// =================================================================== + +#[test] +fn test_streaming_two_chunks_accumulated() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "Hello "}]}, + "index": 0 + }], + "modelVersion": "gemini-2.0-flash" + })) + .unwrap(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "world!"}]}, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 3, + "totalTokenCount": 8 + } + })) + .unwrap(); + + let assembled = finalizer(); + assert_eq!(assembled["candidates"][0]["index"].as_u64(), Some(0)); + + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + assert_eq!( + resp.message, + Some(MessageContent::Text("Hello world!".into())) + ); + assert_eq!(resp.finish_reason, Some(FinishReason::Complete)); + + let usage = resp.usage.unwrap(); + assert_eq!(usage.prompt_tokens, Some(5)); + assert_eq!(usage.completion_tokens, Some(3)); + assert_eq!(usage.total_tokens, Some(8)); +} + +#[test] +fn test_streaming_rejects_missing_candidate_index() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + + let err = collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "orphan"}]} + }] + })) + .expect_err("candidate index is required to avoid corrupt streaming aggregates"); + + assert!( + err.to_string().contains("candidate index is required"), + "unexpected error: {err}" + ); +} + +#[test] +fn test_streaming_rejects_nonzero_candidate_index() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + + let err = collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "second"}]}, + "index": 1 + }] + })) + .expect_err("nonzero candidate indexes cannot be reassembled losslessly"); + + assert!( + err.to_string().contains("only supports candidate index 0"), + "unexpected error: {err}" + ); +} + +#[test] +fn test_streaming_rejects_candidate_index_change() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "first"}]}, + "index": 0 + }] + })) + .unwrap(); + + let err = collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "second"}]}, + "index": 1 + }] + })) + .expect_err("candidate index changes must not be merged into candidate 0"); + + assert!( + err.to_string() + .contains("candidate index changed across chunks"), + "unexpected error: {err}" + ); +} + +#[test] +fn test_streaming_rejects_multiple_candidates_in_one_chunk() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + + let err = collector(json!({ + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "first"}]}, + "index": 0 + }, + { + "content": {"role": "model", "parts": [{"text": "second"}]}, + "index": 1 + } + ] + })) + .expect_err("multi-candidate chunks cannot be reassembled by the single-candidate state"); + + assert!( + err.to_string().contains("multiple candidates"), + "unexpected error: {err}" + ); +} + +#[test] +fn test_streaming_preserves_text_part_metadata_from_empty_final_chunk() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "answer"}]}, + "index": 0 + }] + })) + .unwrap(); + + collector(json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"text": "", "thoughtSignature": "sig_STREAM=="}] + }, + "finishReason": "STOP", + "index": 0 + }] + })) + .unwrap(); + + let assembled = finalizer(); + let parts = assembled["candidates"][0]["content"]["parts"] + .as_array() + .unwrap(); + assert_eq!(parts.len(), 2); + assert_eq!(parts[0]["text"].as_str(), Some("answer")); + assert!( + parts[0].get("thoughtSignature").is_none(), + "signature from the empty final chunk must not be merged onto the previous part" + ); + assert_eq!(parts[1]["text"].as_str(), Some("")); + assert_eq!( + parts[1]["thoughtSignature"].as_str(), + Some("sig_STREAM=="), + "streamed text metadata must survive finalization" + ); + + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + let Some(MessageContent::Parts(parts)) = resp.message else { + panic!("metadata-bearing streamed text must decode as content parts"); + }; + assert!(matches!( + &parts[0], + ContentPart::Text { text, extra } if text == "answer" && extra.is_empty() + )); + assert!(matches!( + &parts[1], + ContentPart::Text { text, extra } + if text.is_empty() + && extra.get("thoughtSignature").and_then(Json::as_str) == Some("sig_STREAM==") + )); +} + +#[test] +fn test_streaming_keeps_non_empty_signed_text_part_separate_from_plain_text() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "first"}]}, + "index": 0 + }] + })) + .unwrap(); + + collector(json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"text": "second", "thoughtSignature": "sig_SECOND=="}] + }, + "finishReason": "STOP", + "index": 0 + }] + })) + .unwrap(); + + let assembled = finalizer(); + let parts = assembled["candidates"][0]["content"]["parts"] + .as_array() + .unwrap(); + assert_eq!(parts.len(), 2); + assert_eq!(parts[0]["text"].as_str(), Some("first")); + assert!( + parts[0].get("thoughtSignature").is_none(), + "signature from the second text part must not be moved onto the first" + ); + assert_eq!(parts[1]["text"].as_str(), Some("second")); + assert_eq!(parts[1]["thoughtSignature"].as_str(), Some("sig_SECOND==")); +} + +#[test] +fn test_streaming_finalize_valid_response_shape() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "chunk"}]}, + "finishReason": "MAX_TOKENS", + "index": 0 + }], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 2, + "totalTokenCount": 12 + }, + "modelVersion": "gemini-1.5-pro" + })) + .unwrap(); + + let assembled = finalizer(); + + assert!(assembled.get("candidates").is_some_and(Json::is_array)); + assert!(assembled.get("usageMetadata").is_some()); + assert_eq!( + assembled.get("modelVersion").and_then(Json::as_str), + Some("gemini-1.5-pro") + ); + + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + assert_eq!(resp.finish_reason, Some(FinishReason::Length)); + assert_eq!(resp.model, Some("gemini-1.5-pro".into())); +} + +#[test] +fn test_streaming_last_usage_metadata_wins() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{"content": {"role": "model", "parts": [{"text": "a"}]}, "index": 0}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2} + })) + .unwrap(); + + collector(json!({ + "candidates": [{"content": {"role": "model", "parts": [{"text": "b"}]}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15} + })) + .unwrap(); + + let assembled = finalizer(); + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + + let usage = resp.usage.unwrap(); + assert_eq!(usage.prompt_tokens, Some(10)); + assert_eq!(usage.completion_tokens, Some(5)); + assert_eq!(resp.message, Some(MessageContent::Text("ab".into()))); +} + +#[test] +fn test_streaming_preserves_native_response_parts() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "ran code"}, + {"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "sk-stream-code"}} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + })) + .unwrap(); + + let assembled = finalizer(); + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + let Some(MessageContent::Parts(parts)) = resp.message else { + panic!("expected streamed native Gemini part to survive finalization"); + }; + + assert!(matches!( + &parts[0], + ContentPart::Text { text, .. } if text == "ran code" + )); + match &parts[1] { + ContentPart::ProviderNative { + provider, + kind, + value, + } => { + assert_eq!(provider, "gemini"); + assert_eq!(kind, "codeExecutionResult"); + assert_eq!( + value["codeExecutionResult"]["output"], + json!("sk-stream-code") + ); + } + other => panic!("expected Gemini ProviderNative streaming part, got {other:?}"), + } +} + +#[test] +fn test_streaming_preserves_native_before_text_order() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"inlineData": {"mimeType": "image/png", "data": "abc123=="}}, + {"text": "caption", "thoughtSignature": "sig_TEXT=="} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + })) + .unwrap(); + + let assembled = finalizer(); + let parts = assembled["candidates"][0]["content"]["parts"] + .as_array() + .unwrap(); + assert!( + parts[0].get("inlineData").is_some(), + "streaming finalizer must preserve native part position before text" + ); + assert_eq!(parts[1]["text"].as_str(), Some("caption")); + assert_eq!(parts[1]["thoughtSignature"].as_str(), Some("sig_TEXT==")); +} + +#[test] +fn test_streaming_does_not_merge_text_across_native_part() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "before"}, + {"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "42"}}, + {"text": "after", "thoughtSignature": "sig_AFTER=="} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + })) + .unwrap(); + + let assembled = finalizer(); + let parts = assembled["candidates"][0]["content"]["parts"] + .as_array() + .unwrap(); + assert_eq!(parts[0]["text"].as_str(), Some("before")); + assert!( + parts[1].get("codeExecutionResult").is_some(), + "native part must stay between the two streamed text parts" + ); + assert_eq!(parts[2]["text"].as_str(), Some("after")); + assert_eq!(parts[2]["thoughtSignature"].as_str(), Some("sig_AFTER==")); +} + +#[test] +fn test_streaming_preserves_adjacent_signed_text_parts() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "first", "thoughtSignature": "sig_FIRST=="}, + {"text": "second", "thoughtSignature": "sig_SECOND=="} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + })) + .unwrap(); + + let assembled = finalizer(); + let parts = assembled["candidates"][0]["content"]["parts"] + .as_array() + .unwrap(); + assert_eq!(parts[0]["text"].as_str(), Some("first")); + assert_eq!(parts[0]["thoughtSignature"].as_str(), Some("sig_FIRST==")); + assert_eq!(parts[1]["text"].as_str(), Some("second")); + assert_eq!(parts[1]["thoughtSignature"].as_str(), Some("sig_SECOND==")); +} + +// =================================================================== +// Thought-signature and native-part preservation +// =================================================================== + +#[test] +fn test_encode_lossless_when_messages_unchanged() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + { + "role": "user", + "parts": [ + {"text": "what is 2+2?"}, + {"inlineData": {"mimeType": "image/png", "data": "abc123=="}} + ] + }, + { + "role": "model", + "parts": [ + {"thought": true, "text": "let me think..."}, + {"thought": true, "thoughtSignature": "sig_XYZ_abc=="}, + {"functionCall": {"name": "calculator", "id": "call_1", "args": {"op": "add", "a": 2, "b": 2}}} + ] + } + ], + "systemInstruction": {"parts": [{"text": "Be helpful."}]}, + "generationConfig": {"temperature": 0.5}, + "safetySettings": [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}] + })); + + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + + assert_eq!( + encoded.content, original.content, + "encode(decode(req), req) must be byte-identical to req" + ); +} + +#[test] +fn test_encode_thought_signature_preserved_when_system_message_changes() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "call something"}]}, + { + "role": "model", + "parts": [ + {"thought": true, "thoughtSignature": "sig_CRITICAL=="}, + {"functionCall": {"name": "my_fn", "id": "call_99", "args": {"x": 1}}} + ] + } + ], + "systemInstruction": {"parts": [{"text": "old system prompt"}]} + })); + + let mut annotated = codec.decode(&original).unwrap(); + + for msg in annotated.messages.iter_mut() { + if let Message::System { content, .. } = msg { + *content = MessageContent::Text("new system prompt".into()); + } + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + + let sys_text = encoded + .content + .get("systemInstruction") + .and_then(|s| s.get("parts")) + .and_then(|p| p.as_array()) + .and_then(|a| a.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str) + .unwrap(); + assert_eq!(sys_text, "new system prompt"); + + assert_eq!( + encoded.content.get("contents"), + original.content.get("contents"), + "thoughtSignature must survive when only the system message changed" + ); +} + +/// Multi-turn continuation: interceptor appends a tool result and a new user turn. +/// Earlier turns must be preserved exactly; the new tool-result turn must carry +/// the correct id, name, and response fields. +#[test] +fn test_encode_thought_signature_preserved_in_multi_turn_continuation() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "call the tool"}]}, + { + "role": "model", + "parts": [ + {"thought": true, "thoughtSignature": "sig_MUST_SURVIVE=="}, + {"functionCall": {"name": "my_fn", "id": "call_1", "args": {}}} + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + annotated.messages.push(Message::Tool { + content: MessageContent::Text(r#"{"output": "done"}"#.into()), + tool_call_id: "call_1".into(), + }); + annotated.messages.push(Message::User { + content: MessageContent::Text("thanks, continue".into()), + name: None, + }); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + + // Unchanged turns are preserved byte-identically. + let orig_contents = original + .content + .get("contents") + .unwrap() + .as_array() + .unwrap(); + assert_eq!( + &contents[0], &orig_contents[0], + "unchanged user turn at position 0 must be preserved byte-identically" + ); + assert_eq!( + &contents[1], &orig_contents[1], + "model turn with thoughtSignature at position 1 must be preserved byte-identically" + ); + assert_eq!(contents.len(), 4); + + // New tool-result turn at position 2 must have correct id, name, role, and response. + let tool_turn = &contents[2]; + assert_eq!( + tool_turn.get("role").and_then(Json::as_str), + Some("user"), + "tool result must use role 'user'" + ); + let fr = tool_turn + .get("parts") + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("functionResponse")) + .expect("tool result must contain a functionResponse part"); + assert_eq!( + fr.get("id").and_then(Json::as_str), + Some("call_1"), + "functionResponse.id must be the actual call ID, not the function name" + ); + assert_eq!( + fr.get("name").and_then(Json::as_str), + Some("my_fn"), + "functionResponse.name must be the function name looked up from the assistant turn" + ); + assert!( + fr.get("response").is_some(), + "functionResponse.response must be present" + ); + + // New user turn at position 3. + assert_eq!(contents[3].get("role").and_then(Json::as_str), Some("user")); +} + +/// thoughtSignature on a functionCall part itself must survive when an interceptor +/// edits the function call arguments (triggers the rebuild path). +#[test] +fn test_encode_thought_signature_on_function_call_part_survives_edit() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "run the tool"}]}, + { + "role": "model", + "parts": [ + { + "functionCall": {"name": "my_fn", "id": "call_1", "args": {"x": 1}}, + "thoughtSignature": "ABCDEF==" + } + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Intercept changes the function call arguments. + if let Message::Assistant { + tool_calls: Some(tcs), + .. + } = &mut annotated.messages[1] + && let Some(tc) = tcs.first_mut() + { + tc.function.arguments = r#"{"x": 42}"#.to_string(); + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + let model_parts = contents[1].get("parts").and_then(Json::as_array).unwrap(); + + let fc_part = model_parts + .iter() + .find(|p| p.get("functionCall").is_some()) + .expect("encoded model turn must contain a functionCall part"); + assert_eq!( + fc_part.get("thoughtSignature").and_then(Json::as_str), + Some("ABCDEF=="), + "thoughtSignature on the functionCall part must survive when args are edited" + ); + // Verify the args were actually updated. + let args = fc_part + .get("functionCall") + .and_then(|fc| fc.get("args")) + .unwrap(); + assert_eq!(args.get("x").and_then(Json::as_i64), Some(42)); +} + +// =================================================================== +// inlineData round-trip +// =================================================================== + +#[test] +fn test_encode_inline_data_preserved_when_messages_unchanged() { + let codec = GeminiGenerateContentCodec; + let image_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQ=="; + let original = make_request(json!({ + "contents": [ + { + "role": "user", + "parts": [ + {"text": "what is in this image?"}, + {"inlineData": {"mimeType": "image/png", "data": image_data}} + ] + } + ] + })); + + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + + let parts = encoded.content.get("contents").unwrap().as_array().unwrap()[0] + .get("parts") + .unwrap() + .as_array() + .unwrap(); + + let inline = parts + .iter() + .find(|p| p.get("inlineData").is_some()) + .unwrap(); + assert_eq!( + inline + .get("inlineData") + .unwrap() + .get("data") + .and_then(Json::as_str), + Some(image_data), + "inlineData must be byte-identical after round-trip" + ); +} + +#[test] +fn test_encode_inline_data_preserved_when_text_changes() { + let codec = GeminiGenerateContentCodec; + let image_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQ=="; + let original = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"text": "what is in this image?"}, + {"inlineData": {"mimeType": "image/png", "data": image_data}} + ] + }] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + if let Message::User { content, .. } = &mut annotated.messages[0] { + *content = MessageContent::Text("Please describe this image in detail.".into()); + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let parts = encoded.content.get("contents").unwrap().as_array().unwrap()[0] + .get("parts") + .unwrap() + .as_array() + .unwrap(); + + assert_eq!( + parts[0].get("text").and_then(Json::as_str), + Some("Please describe this image in detail."), + "text part must be updated to the interceptor's new text" + ); + let inline = parts + .iter() + .find(|p| p.get("inlineData").is_some()) + .expect("inlineData must survive when interceptor only changed the text"); + assert_eq!( + inline + .get("inlineData") + .unwrap() + .get("data") + .and_then(Json::as_str), + Some(image_data) + ); +} + +#[test] +fn test_decode_inline_data_as_provider_native_content_part() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"text": "what is in this file?"}, + {"inlineData": {"mimeType": "text/plain", "data": "sk-file-secret"}} + ] + }] + })); + + let annotated = codec.decode(&request).unwrap(); + let Message::User { + content: MessageContent::Parts(parts), + .. + } = &annotated.messages[0] + else { + panic!("expected mixed Gemini content to decode as MessageContent::Parts"); + }; + + assert!(matches!( + &parts[0], + ContentPart::Text { text, .. } if text == "what is in this file?" + )); + match &parts[1] { + ContentPart::ProviderNative { + provider, + kind, + value, + } => { + assert_eq!(provider, "gemini"); + assert_eq!(kind, "inlineData"); + assert_eq!( + value["inlineData"]["data"], + json!("sk-file-secret"), + "native inlineData must be visible to normalized middleware" + ); + } + other => panic!("expected Gemini ProviderNative content part, got {other:?}"), + } +} + +#[test] +fn test_encode_patches_provider_native_inline_data_content_part() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"text": "inspect this"}, + {"inlineData": {"mimeType": "text/plain", "data": "sk-file-secret"}} + ] + }] + })); + + let mut annotated = codec.decode(&original).unwrap(); + let Message::User { + content: MessageContent::Parts(parts), + .. + } = &mut annotated.messages[0] + else { + panic!("expected native content parts"); + }; + match &mut parts[1] { + ContentPart::ProviderNative { value, .. } => { + value["inlineData"]["data"] = json!("[REDACTED]"); + } + other => panic!("expected Gemini ProviderNative content part, got {other:?}"), + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content["contents"][0]["parts"][1]["inlineData"]["data"], + json!("[REDACTED]"), + "editing the normalized provider-native content part must update the raw Gemini part" + ); +} + +// =================================================================== +// Thought part filtering +// =================================================================== + +#[test] +fn test_decode_response_filters_thought_parts() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"thought": true, "text": "internal reasoning the user should not see"}, + {"text": "this is the actual answer"} + ] + }, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5} + }); + let ann = codec.decode_response(&response).unwrap(); + assert_eq!( + ann.message, + Some(MessageContent::Text("this is the actual answer".into())), + "thought parts must not leak into the normalized message" + ); +} + +#[test] +fn test_streaming_filters_thought_parts() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"thought": true, "text": "reasoning"}]}, + "index": 0 + }] + })) + .unwrap(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "final answer"}]}, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 3} + })) + .unwrap(); + + let assembled = finalizer(); + let assembled_parts = assembled["candidates"][0]["content"]["parts"] + .as_array() + .unwrap(); + assert_eq!( + assembled_parts[0]["thought"].as_bool(), + Some(true), + "thought chunks must survive in the provider-native streaming aggregate" + ); + assert_eq!(assembled_parts[0]["text"].as_str(), Some("reasoning")); + assert_eq!(assembled_parts[1]["text"].as_str(), Some("final answer")); + + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + assert_eq!( + resp.message, + Some(MessageContent::Text("final answer".into())), + "thought chunks must not appear in the streamed message" + ); +} + +/// When a model turn in the request history has both a thought part (thought: true) +/// and a functionCall part, the thought text must NOT appear in the decoded +/// Message::Assistant.content. Only visible (non-thought) text should appear. +#[test] +fn test_decode_request_thought_text_does_not_leak_into_assistant_content() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "search for it"}]}, + { + "role": "model", + "parts": [ + {"thought": true, "text": "I should call the search function"}, + {"functionCall": {"id": "call_1", "name": "search", "args": {"q": "test"}}} + ] + } + ] + })); + + let annotated = codec.decode(&request).unwrap(); + let asst = annotated + .messages + .iter() + .find(|m| matches!(m, Message::Assistant { .. })) + .expect("must have assistant message"); + if let Message::Assistant { + content, + tool_calls, + .. + } = asst + { + assert!( + content.is_none(), + "thought text must not appear in Message::Assistant.content; got: {:?}", + content + ); + assert!(tool_calls.is_some(), "functionCall must still be decoded"); + } +} + +/// A model turn with both visible text and a functionCall: visible text appears in +/// content, thought text does not. +#[test] +fn test_decode_request_visible_text_survives_when_thought_present() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{ + "role": "model", + "parts": [ + {"thought": true, "text": "let me think"}, + {"text": "here is my answer"}, + {"functionCall": {"id": "c1", "name": "fn", "args": {}}} + ] + }] + })); + + let annotated = codec.decode(&request).unwrap(); + let asst = annotated + .messages + .iter() + .find(|m| matches!(m, Message::Assistant { .. })) + .expect("must have assistant message"); + if let Message::Assistant { content, .. } = asst { + assert_eq!( + content.as_ref().map(|c| match c { + MessageContent::Text(t) => t.as_str(), + _ => "", + }), + Some("here is my answer"), + "only non-thought visible text must appear in content" + ); + } +} + +// =================================================================== +// Mixed tool groups +// =================================================================== + +#[test] +fn test_encode_preserves_native_tool_groups_when_tools_unchanged() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [ + {"functionDeclarations": [{"name": "get_weather", "description": "Get weather"}]}, + {"googleSearch": {}}, + {"codeExecution": {}} + ] + })); + + let annotated = codec.decode(&original).unwrap(); + let normalized_tools = annotated.tools.as_ref().unwrap(); + assert!( + normalized_tools.iter().any( + |td| matches!(td, ToolDefinition::ProviderNative { provider, kind, .. } + if provider == "gemini" && kind == "googleSearch") + ), + "native googleSearch group must be visible as a Gemini ProviderNative tool" + ); + assert!( + normalized_tools.iter().any( + |td| matches!(td, ToolDefinition::ProviderNative { provider, kind, .. } + if provider == "gemini" && kind == "codeExecution") + ), + "native codeExecution group must be visible as a Gemini ProviderNative tool" + ); + let encoded = codec.encode(&annotated, &original).unwrap(); + + let tools = encoded + .content + .get("tools") + .and_then(Json::as_array) + .unwrap(); + assert!( + tools.iter().any(|g| g.get("googleSearch").is_some()), + "googleSearch group must survive unchanged tools encode" + ); + assert!( + tools.iter().any(|g| g.get("codeExecution").is_some()), + "codeExecution group must survive unchanged tools encode" + ); + assert!( + tools + .iter() + .any(|g| g.get("functionDeclarations").is_some()), + "functionDeclarations group must survive" + ); +} + +#[test] +fn test_decode_native_only_tool_group_as_provider_native() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"googleSearch": {"apiKey": "sk-tool-secret"}}] + })); + + let annotated = codec.decode(&original).unwrap(); + let tools = annotated.tools.as_ref().unwrap(); + assert_eq!(tools.len(), 1); + match &tools[0] { + ToolDefinition::ProviderNative { + provider, + kind, + value, + } => { + assert_eq!(provider, "gemini"); + assert_eq!(kind, "googleSearch"); + assert_eq!( + value, + &json!({"googleSearch": {"apiKey": "sk-tool-secret"}}) + ); + } + other => panic!("expected Gemini ProviderNative tool, got {other:?}"), + } +} + +#[test] +fn test_decode_mixed_tool_group_exposes_native_siblings() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{"name": "lookup"}], + "googleSearch": {"apiKey": "sk-tool-secret"} + }] + })); + + let annotated = codec.decode(&original).unwrap(); + let tools = annotated.tools.as_ref().unwrap(); + assert_eq!(tools.len(), 2); + assert!(matches!(tools[0], ToolDefinition::Function { .. })); + match &tools[1] { + ToolDefinition::ProviderNative { + provider, + kind, + value, + } => { + assert_eq!(provider, "gemini"); + assert_eq!(kind, "googleSearch"); + assert_eq!( + value, + &json!({"googleSearch": {"apiKey": "sk-tool-secret"}}) + ); + } + other => panic!("expected Gemini ProviderNative sibling fields, got {other:?}"), + } +} + +#[test] +fn test_encode_preserves_native_tool_groups_when_functions_change() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [ + {"functionDeclarations": [{"name": "old_fn"}]}, + {"googleSearch": {}} + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + if let Some(tools) = annotated.tools.as_mut() { + for td in tools.iter_mut() { + if let nemo_relay_types::codec::request::ToolDefinition::Function { function, .. } = td + { + function.name = "new_fn".into(); + } + } + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let tools = encoded + .content + .get("tools") + .and_then(Json::as_array) + .unwrap(); + + assert!( + tools.iter().any(|g| g.get("googleSearch").is_some()), + "googleSearch must survive when only function declarations changed" + ); + + let fn_group = tools + .iter() + .find(|g| g.get("functionDeclarations").is_some()) + .unwrap(); + let fns = fn_group + .get("functionDeclarations") + .and_then(Json::as_array) + .unwrap(); + assert_eq!(fns[0].get("name").and_then(Json::as_str), Some("new_fn")); +} + +/// Provider-native functionDeclaration fields (parametersJsonSchema, responseJsonSchema, +/// response, behavior) must survive a decode → edit description → encode round-trip. +#[test] +fn test_encode_preserves_provider_tool_fields_when_description_changes() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{ + "name": "my_tool", + "description": "old description", + "parameters": {"type": "object"}, + "parametersJsonSchema": {"$schema": "draft-2020-12"}, + "responseJsonSchema": {"type": "object"}, + "behavior": "BLOCKING" + }] + }] + })); + + let mut annotated = codec.decode(&original).unwrap(); + if let Some(tools) = annotated.tools.as_mut() { + for td in tools.iter_mut() { + if let nemo_relay_types::codec::request::ToolDefinition::Function { function, .. } = td + { + function.description = Some("new description".into()); + } + } + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let tools = encoded + .content + .get("tools") + .and_then(Json::as_array) + .unwrap(); + let fn_group = tools + .iter() + .find(|g| g.get("functionDeclarations").is_some()) + .unwrap(); + let fd = &fn_group + .get("functionDeclarations") + .and_then(Json::as_array) + .unwrap()[0]; + + assert_eq!( + fd.get("description").and_then(Json::as_str), + Some("new description") + ); + assert!( + fd.get("parametersJsonSchema").is_some(), + "parametersJsonSchema must survive" + ); + assert!( + fd.get("responseJsonSchema").is_some(), + "responseJsonSchema must survive" + ); + assert_eq!( + fd.get("behavior").and_then(Json::as_str), + Some("BLOCKING"), + "behavior must survive" + ); +} + +#[test] +fn test_encode_patches_gemini_provider_native_tool_group() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"googleSearch": {"apiKey": "sk-tool-secret"}}] + })); + + let mut annotated = codec.decode(&original).unwrap(); + let tools = annotated.tools.as_mut().unwrap(); + match &mut tools[0] { + ToolDefinition::ProviderNative { value, .. } => { + *value = json!({"googleSearch": {"apiKey": "[REDACTED]"}}); + } + other => panic!("expected Gemini ProviderNative tool, got {other:?}"), + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content.get("tools").unwrap(), + &json!([{"googleSearch": {"apiKey": "[REDACTED]"}}]) + ); +} + +#[test] +fn test_encode_patches_mixed_gemini_provider_native_tool_group() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{"name": "lookup"}], + "googleSearch": {"apiKey": "sk-tool-secret"} + }] + })); + + let mut annotated = codec.decode(&original).unwrap(); + let tools = annotated.tools.as_mut().unwrap(); + match &mut tools[1] { + ToolDefinition::ProviderNative { value, .. } => { + *value = json!({"googleSearch": {"apiKey": "[REDACTED]"}}); + } + other => panic!("expected Gemini ProviderNative sibling fields, got {other:?}"), + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content.get("tools").unwrap(), + &json!([{ + "functionDeclarations": [{"name": "lookup"}], + "googleSearch": {"apiKey": "[REDACTED]"} + }]) + ); +} + +#[test] +fn test_encode_deleting_mixed_native_sibling_does_not_rehome_later_native_group() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [ + { + "functionDeclarations": [{"name": "lookup"}], + "googleSearch": {} + }, + {"codeExecution": {}} + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + annotated.tools.as_mut().unwrap().retain(|tool| { + !matches!( + tool, + ToolDefinition::ProviderNative { provider, kind, .. } + if provider == "gemini" && kind == "googleSearch" + ) + }); + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content.get("tools").unwrap(), + &json!([ + {"functionDeclarations": [{"name": "lookup"}]}, + {"codeExecution": {}} + ]), + "deleting a native sibling must not merge a later native-only group \ + into the functionDeclarations group" + ); +} + +// =================================================================== +// Function-call correlation IDs +// =================================================================== + +#[test] +fn test_decode_response_function_call_uses_id_field() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "my_fn", "id": "call_abc123", "args": {"x": 1}}} + ] + }, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 5} + }); + + let ann = codec.decode_response(&response).unwrap(); + let tc = ann.tool_calls.unwrap(); + assert_eq!(tc.len(), 1); + assert_eq!( + tc[0].id, "call_abc123", + "must use Gemini-provided id, not function name" + ); + assert_eq!(tc[0].name, "my_fn"); +} + +#[test] +fn test_decode_response_function_call_fallback_to_name_when_no_id() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "fallback_fn", "args": {}}}] + }, + "finishReason": "STOP" + }], + "usageMetadata": {} + }); + + let ann = codec.decode_response(&response).unwrap(); + let tc = ann.tool_calls.unwrap(); + assert_eq!( + tc[0].id, "fallback_fn", + "must fall back to function name when id is absent" + ); +} + +/// Two simultaneous calls to the same function must have distinct IDs. +#[test] +fn test_decode_response_multi_call_same_function_distinct_ids() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "search", "id": "call_001", "args": {"q": "a"}}}, + {"functionCall": {"name": "search", "id": "call_002", "args": {"q": "b"}}} + ] + }, + "finishReason": "STOP" + }], + "usageMetadata": {} + }); + + let ann = codec.decode_response(&response).unwrap(); + let tc = ann.tool_calls.unwrap(); + assert_eq!(tc.len(), 2); + assert_eq!(tc[0].id, "call_001"); + assert_eq!(tc[1].id, "call_002"); + assert_ne!( + tc[0].id, tc[1].id, + "parallel calls to same function must have distinct IDs" + ); + assert_eq!(tc[0].name, "search"); + assert_eq!(tc[1].name, "search"); +} + +/// After decode → encode, the functionResponse must contain the actual call ID (not +/// the function name) in the `id` field, and the correct function name in `name`. +#[test] +fn test_encode_function_response_id_not_name() { + let codec = GeminiGenerateContentCodec; + // A multi-turn request where Gemini provided a functionCall with an explicit id. + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "call my_fn"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"id": "call_abc123", "name": "my_fn", "args": {"x": 1}}} + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Interceptor appends the tool result, using the ID from the decoded ToolCall. + annotated.messages.push(Message::Tool { + content: MessageContent::Text(r#"{"output": 42}"#.into()), + tool_call_id: "call_abc123".into(), + }); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + assert_eq!(contents.len(), 3); + + let tool_turn = &contents[2]; + assert_eq!(tool_turn.get("role").and_then(Json::as_str), Some("user")); + let fr = tool_turn + .get("parts") + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("functionResponse")) + .expect("must have functionResponse"); + assert_eq!( + fr.get("id").and_then(Json::as_str), + Some("call_abc123"), + "functionResponse.id must be the actual call ID" + ); + assert_eq!( + fr.get("name").and_then(Json::as_str), + Some("my_fn"), + "functionResponse.name must be the function name, not the call ID" + ); + // Sanity check: "call_abc123" must NOT appear as the name. + assert_ne!( + fr.get("name").and_then(Json::as_str), + Some("call_abc123"), + "functionResponse.name must not be the call ID" + ); +} + +/// When a functionResponse is decoded from request history, the `id` field +/// (not the name) is stored in tool_call_id. +#[test] +fn test_decode_function_response_uses_id_not_name() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "hi"}]}, + {"role": "model", "parts": [ + {"functionCall": {"id": "call_xyz", "name": "my_fn", "args": {}}} + ]}, + {"role": "user", "parts": [ + {"functionResponse": {"id": "call_xyz", "name": "my_fn", "response": {"val": 1}}} + ]} + ] + })); + + let annotated = codec.decode(&request).unwrap(); + // messages: [assistant_with_toolcall, tool_result] + // (no system message, so indices are direct) + let tool_msg = annotated + .messages + .iter() + .find(|m| matches!(m, Message::Tool { .. })) + .expect("must decode a Message::Tool"); + if let Message::Tool { tool_call_id, .. } = tool_msg { + assert_eq!( + tool_call_id, "call_xyz", + "tool_call_id must be the actual id field, not the function name" + ); + } +} + +/// When a system message is added (system → systemInstruction, not contents), the +/// original contents items are left untouched. +#[test] +fn test_encode_system_message_added_leaves_contents_unchanged() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "original user"}]}, + {"role": "model", "parts": [{"text": "original model"}]} + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + annotated.messages.insert( + 0, + Message::System { + content: MessageContent::Text("system context".into()), + name: None, + }, + ); + + let encoded = codec.encode(&annotated, &original).unwrap(); + + assert!(encoded.content.get("systemInstruction").is_some()); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + assert_eq!(contents.len(), 2); + assert_eq!(contents[0].get("role").and_then(Json::as_str), Some("user")); + assert_eq!( + contents[1].get("role").and_then(Json::as_str), + Some("model") + ); +} + +/// Inserting a new non-system message increases the contents length; the new +/// message must be encoded correctly and the appended position must use +/// normalized_to_gemini_content (not a wrong original as base). +#[test] +fn test_encode_new_user_message_appended_encodes_correctly() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]} + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Append a new user message (contents grows from 2 → 3). + annotated.messages.push(Message::User { + content: MessageContent::Text("follow up question".into()), + name: None, + }); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + + assert_eq!(contents.len(), 3); + // Original positions unchanged. + assert_eq!(contents[0].get("role").and_then(Json::as_str), Some("user")); + assert_eq!( + contents[1].get("role").and_then(Json::as_str), + Some("model") + ); + // New position encoded fresh — correct role and text. + assert_eq!(contents[2].get("role").and_then(Json::as_str), Some("user")); + let new_text = contents[2] + .get("parts") + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str); + assert_eq!( + new_text, + Some("follow up question"), + "new message text must be encoded correctly" + ); +} + +// =================================================================== +// Streaming: function call parts accumulation +// =================================================================== + +#[test] +fn test_streaming_accumulates_function_call_parts() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"thought": true, "text": "deciding..."}]}, + "index": 0 + }] + })) + .unwrap(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [ + {"functionCall": {"name": "tool_a", "id": "c1", "args": {"x": 1}}} + ]}, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 4, "totalTokenCount": 14} + })) + .unwrap(); + + let assembled = finalizer(); + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + + assert_eq!(resp.finish_reason, Some(FinishReason::ToolUse)); + let tc = resp.tool_calls.unwrap(); + assert_eq!(tc.len(), 1); + assert_eq!(tc[0].id, "c1"); + assert_eq!(tc[0].name, "tool_a"); +} + +/// When two parallel calls to the SAME function each carry a distinct +/// thoughtSignature, both signatures must survive a decode → edit-args → encode +/// round-trip. A HashMap lookup would drop the first entry (same key), leaving +/// both encoded parts with the second part's signature. +#[test] +fn test_encode_thought_signature_preserved_for_parallel_same_function_calls() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "run two searches"}]}, + { + "role": "model", + "parts": [ + { + "functionCall": {"name": "search", "id": "c1", "args": {"q": "first"}}, + "thoughtSignature": "sig_FIRST==" + }, + { + "functionCall": {"name": "search", "id": "c2", "args": {"q": "second"}}, + "thoughtSignature": "sig_SECOND==" + } + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Interceptor edits the first call's args; second is unchanged. + if let Message::Assistant { + tool_calls: Some(tcs), + .. + } = &mut annotated.messages[1] + && let Some(tc) = tcs.first_mut() + { + tc.function.arguments = r#"{"q": "first-edited"}"#.to_string(); + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + let model_parts = contents[1].get("parts").and_then(Json::as_array).unwrap(); + + // There must be exactly two functionCall parts. + let fc_parts: Vec<&Json> = model_parts + .iter() + .filter(|p| p.get("functionCall").is_some()) + .collect(); + assert_eq!(fc_parts.len(), 2, "both functionCall parts must be encoded"); + + // The first encoded part must carry sig_FIRST (not sig_SECOND). + assert_eq!( + fc_parts[0].get("thoughtSignature").and_then(Json::as_str), + Some("sig_FIRST=="), + "first functionCall part must retain its own thoughtSignature" + ); + // The second encoded part must carry sig_SECOND. + assert_eq!( + fc_parts[1].get("thoughtSignature").and_then(Json::as_str), + Some("sig_SECOND=="), + "second functionCall part must retain its own thoughtSignature" + ); + // Verify the first call's args were actually updated. + let first_args = fc_parts[0] + .get("functionCall") + .and_then(|fc| fc.get("args")) + .unwrap(); + assert_eq!( + first_args.get("q").and_then(Json::as_str), + Some("first-edited") + ); +} + +/// Reordering same-name parallel function calls must not swap their +/// thoughtSignature fields. Match original parts by Gemini's stable call ID, +/// falling back to name only for provider payloads that omitted an ID. +#[test] +fn test_encode_reordered_same_name_function_calls_match_signatures_by_id() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "run two searches"}]}, + { + "role": "model", + "parts": [ + { + "functionCall": {"name": "search", "id": "c1", "args": {"q": "first"}}, + "thoughtSignature": "sig_FIRST==" + }, + { + "functionCall": {"name": "search", "id": "c2", "args": {"q": "second"}}, + "thoughtSignature": "sig_SECOND==" + } + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + if let Message::Assistant { + tool_calls: Some(tcs), + .. + } = &mut annotated.messages[1] + { + tcs.swap(0, 1); + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + let model_parts = contents[1].get("parts").and_then(Json::as_array).unwrap(); + let fc_parts: Vec<&Json> = model_parts + .iter() + .filter(|p| p.get("functionCall").is_some()) + .collect(); + + assert_eq!(fc_parts.len(), 2); + assert_eq!(fc_parts[0]["functionCall"]["id"].as_str(), Some("c2")); + assert_eq!( + fc_parts[0].get("thoughtSignature").and_then(Json::as_str), + Some("sig_SECOND==") + ); + assert_eq!(fc_parts[1]["functionCall"]["id"].as_str(), Some("c1")); + assert_eq!( + fc_parts[1].get("thoughtSignature").and_then(Json::as_str), + Some("sig_FIRST==") + ); +} + +/// An interceptor that inserts a new system message at position 0 shifts the message-list +/// indices but MUST NOT cause thoughtSignature bleed onto the new position or cause the +/// model turn's thoughtSignature to be lost. +/// +/// System messages map to `systemInstruction`, not to `contents`, so the contents array +/// indices are unaffected by system-message insertion. The model turn stays at contents +/// index 1; only the message-list index changes from 1 to 2. +#[test] +fn test_encode_insertion_does_not_bleed_thought_signature() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "user message"}]}, + { + "role": "model", + "parts": [{ + "functionCall": {"id": "c1", "name": "fn", "args": {}}, + "thoughtSignature": "sig_abc" + }] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Interceptor inserts a new system message at position 0. + // Message list becomes: [new_system, user, model] + annotated.messages.insert( + 0, + Message::System { + content: MessageContent::Text("new system prompt".into()), + name: None, + }, + ); + + let encoded = codec.encode(&annotated, &original).unwrap(); + + // systemInstruction must carry the new text. + let sys = encoded + .content + .get("systemInstruction") + .expect("systemInstruction must be set"); + let sys_text = sys + .get("parts") + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str) + .unwrap(); + assert_eq!(sys_text, "new system prompt"); + + // contents array must still have exactly 2 items (system → systemInstruction only). + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + assert_eq!( + contents.len(), + 2, + "system insertion must not change the contents array length" + ); + + // Position 0 (user): must NOT have thoughtSignature. + let user_turn = &contents[0]; + assert_eq!(user_turn.get("role").and_then(Json::as_str), Some("user")); + let user_parts = user_turn.get("parts").and_then(Json::as_array).unwrap(); + assert!( + user_parts + .iter() + .all(|p| p.get("thoughtSignature").is_none()), + "user turn must not have thoughtSignature" + ); + + // Position 1 (model): must retain thoughtSignature from the original. + let model_turn = &contents[1]; + assert_eq!(model_turn.get("role").and_then(Json::as_str), Some("model")); + let model_parts = model_turn.get("parts").and_then(Json::as_array).unwrap(); + let fc_part = model_parts + .iter() + .find(|p| p.get("functionCall").is_some()) + .expect("model turn must still have functionCall part"); + assert_eq!( + fc_part.get("thoughtSignature").and_then(Json::as_str), + Some("sig_abc"), + "thoughtSignature must survive system-message insertion at position 0" + ); + assert_eq!( + fc_part + .get("functionCall") + .and_then(|fc| fc.get("id")) + .and_then(Json::as_str), + Some("c1"), + "functionCall id must survive system-message insertion" + ); + assert_eq!( + fc_part + .get("functionCall") + .and_then(|fc| fc.get("name")) + .and_then(Json::as_str), + Some("fn"), + "functionCall name must survive system-message insertion" + ); +} + +// =================================================================== +// Prepend non-system user message — content-index insertion regression +// =================================================================== + +/// Prepending a new non-system user message must not corrupt the model turn's +/// thoughtSignature. The new message should be encoded fresh; the original user +/// and model turns must be preserved byte-identically. +#[test] +fn test_encode_prepend_user_message_preserves_thought_signature() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "original question"}]}, + { + "role": "model", + "parts": [{ + "functionCall": {"id": "c1", "name": "search", "args": {}}, + "thoughtSignature": "sig_PREPEND_TEST==" + }] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Prepend a brand-new user message before the existing user message. + // Message list becomes: [new_user, old_user, model_with_signature] + annotated.messages.insert( + 0, + Message::User { + content: MessageContent::Text("context preamble".into()), + name: None, + }, + ); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + + assert_eq!(contents.len(), 3, "prepend must produce 3 content items"); + + // Position 0: freshly encoded new user message. + assert_eq!(contents[0].get("role").and_then(Json::as_str), Some("user")); + let new_text = contents[0] + .get("parts") + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str); + assert_eq!( + new_text, + Some("context preamble"), + "prepended message must have the correct text" + ); + + // Positions 1 and 2: original items preserved byte-identically. + let orig_contents = original + .content + .get("contents") + .unwrap() + .as_array() + .unwrap(); + assert_eq!( + &contents[1], &orig_contents[0], + "original user turn at position 1 must be preserved byte-identically" + ); + assert_eq!( + &contents[2], &orig_contents[1], + "model turn with thoughtSignature at position 2 must be preserved byte-identically" + ); + + // Explicit check that the thoughtSignature survived. + let model_parts = contents[2].get("parts").and_then(Json::as_array).unwrap(); + let fc_part = model_parts + .iter() + .find(|p| p.get("functionCall").is_some()) + .expect("model turn must have functionCall part"); + assert_eq!( + fc_part.get("thoughtSignature").and_then(Json::as_str), + Some("sig_PREPEND_TEST=="), + "thoughtSignature must survive prepend of a non-system user message" + ); +} + +// =================================================================== +// thoughtSignature on a regular text part +// =================================================================== + +/// A model turn where the response text part itself carries a thoughtSignature +/// must have that signature preserved when an interceptor edits the text. +#[test] +fn test_encode_thought_signature_on_text_part_survives_edit() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "say hello"}]}, + { + "role": "model", + "parts": [{"text": "Hello!", "thoughtSignature": "sig_TEXT=="}] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Interceptor changes the assistant's reply text. + if let Message::Assistant { content, .. } = &mut annotated.messages[1] { + *content = Some(MessageContent::Text("Goodbye!".into())); + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + let model_parts = contents[1].get("parts").and_then(Json::as_array).unwrap(); + + let text_part = model_parts + .first() + .expect("model turn must have a text part"); + assert_eq!( + text_part.get("text").and_then(Json::as_str), + Some("Goodbye!"), + "text must be updated to the interceptor's new value" + ); + assert_eq!( + text_part.get("thoughtSignature").and_then(Json::as_str), + Some("sig_TEXT=="), + "thoughtSignature on a regular text part must survive when text is edited" + ); +} + +// =================================================================== +// Parallel functionResponse parts in a single content item +// =================================================================== + +/// A single user-role content item with multiple functionResponse parts +/// (Gemini's parallel-call response format) must decode to one Message::Tool +/// per part — not collapse to just the first. +#[test] +fn test_decode_multiple_function_responses_in_one_turn() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "run two tools"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"id": "c1", "name": "tool_a", "args": {}}}, + {"functionCall": {"id": "c2", "name": "tool_b", "args": {}}} + ] + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "c1", "name": "tool_a", "response": {"r": 1}}}, + {"functionResponse": {"id": "c2", "name": "tool_b", "response": {"r": 2}}} + ] + } + ] + })); + + let annotated = codec.decode(&request).unwrap(); + + // Expect two separate Message::Tool entries from the multi-functionResponse turn. + let tool_msgs: Vec<&Message> = annotated + .messages + .iter() + .filter(|m| matches!(m, Message::Tool { .. })) + .collect(); + assert_eq!( + tool_msgs.len(), + 2, + "each functionResponse part must produce a separate Message::Tool" + ); + + let ids: Vec<&str> = tool_msgs + .iter() + .filter_map(|m| { + if let Message::Tool { tool_call_id, .. } = m { + Some(tool_call_id.as_str()) + } else { + None + } + }) + .collect(); + assert!( + ids.contains(&"c1"), + "first functionResponse id must be decoded" + ); + assert!( + ids.contains(&"c2"), + "second functionResponse id must be decoded" + ); +} + +#[test] +fn test_decode_function_response_with_native_sibling_errors() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"functionResponse": {"id": "c1", "name": "lookup", "response": {"ok": true}}}, + {"inlineData": {"mimeType": "text/plain", "data": "sk-hidden-secret"}} + ] + }] + })); + + assert!( + codec.decode(&request).is_err(), + "native sibling content beside functionResponse would be hidden by Message::Tool decode" + ); +} + +#[test] +fn test_function_response_nested_parts_are_exposed_and_patchable() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "show my ordered instrument"}]}, + {"role": "model", "parts": [{ + "functionCall": {"id": "call_img", "name": "get_image", "args": {}}, + "thoughtSignature": "sig_CALL==" + }]}, + {"role": "user", "parts": [{ + "functionResponse": { + "id": "call_img", + "name": "get_image", + "response": {"image_ref": {"$ref": "instrument.jpg"}}, + "parts": [{ + "inlineData": { + "displayName": "instrument.jpg", + "mimeType": "image/jpeg", + "data": "sk-image-secret" + } + }] + } + }]} + ] + })); + + let mut annotated = codec.decode(&request).unwrap(); + let Message::Tool { content, .. } = &mut annotated.messages[2] else { + panic!("expected functionResponse to decode as Message::Tool"); + }; + let MessageContent::Parts(parts) = content else { + panic!("functionResponse.parts must be exposed as normalized content parts"); + }; + assert!(matches!( + &parts[0], + ContentPart::Text { text, .. } + if text.contains("instrument.jpg") && text.contains("image_ref") + )); + match &mut parts[1] { + ContentPart::ProviderNative { + provider, + kind, + value, + } => { + assert_eq!(provider.as_str(), "gemini"); + assert_eq!(kind.as_str(), "inlineData"); + assert_eq!(value["inlineData"]["data"], json!("sk-image-secret")); + value["inlineData"]["data"] = json!("[REDACTED]"); + } + other => panic!("expected nested Gemini functionResponse part, got {other:?}"), + } + + let encoded = codec.encode(&annotated, &request).unwrap(); + let fr = &encoded.content["contents"][2]["parts"][0]["functionResponse"]; + assert_eq!( + fr["response"]["image_ref"]["$ref"], + json!("instrument.jpg"), + "functionResponse.response must remain object-shaped" + ); + assert_eq!( + fr["parts"][0]["inlineData"]["data"], + json!("[REDACTED]"), + "nested functionResponse.parts must be editable through normalized ProviderNative content" + ); +} + +// =================================================================== +// Request decode validation +// =================================================================== + +#[test] +fn test_decode_rejects_missing_contents() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({"model": "gemini-2.0-flash"})); + assert!( + codec.decode(&request).is_err(), + "missing contents must return an error" + ); +} + +#[test] +fn test_decode_rejects_non_array_contents() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({"contents": "not an array"})); + assert!( + codec.decode(&request).is_err(), + "non-array contents must return an error" + ); +} + +#[test] +fn test_decode_rejects_non_object_generation_config() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": "not an object" + })); + assert!( + codec.decode(&request).is_err(), + "non-object generationConfig must return an error" + ); +} + +#[test] +fn test_decode_rejects_invalid_stop_sequences() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": {"stopSequences": "not an array"} + })); + assert!( + codec.decode(&request).is_err(), + "non-array stopSequences must return an error" + ); +} + +#[test] +fn test_decode_rejects_non_array_tools() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": "not an array" + })); + assert!( + codec.decode(&request).is_err(), + "non-array tools must return an error" + ); +} + +#[test] +fn test_decode_rejects_function_declaration_without_name() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"functionDeclarations": [{"description": "no name"}]}] + })); + assert!( + codec.decode(&request).is_err(), + "functionDeclaration without name must return an error" + ); +} + +// =================================================================== +// Encoder: invalid role rejection +// =================================================================== + +/// Encoding a message with a role unknown to Gemini (e.g. "developer") +/// must return an error — silently dropping it would be data loss. +#[test] +fn test_encode_unsupported_role_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hello"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Inject a message with an unsupported Gemini role. + annotated.messages.push( + serde_json::from_value(json!({"role": "developer", "content": "system note"})).unwrap(), + ); + let result = codec.encode(&annotated, &original); + assert!( + result.is_err(), + "encode must return an error for an unsupported role rather than silently dropping the message" + ); +} + +// =================================================================== +// Parallel functionResponse: deletion and insertion +// =================================================================== + +/// Deleting one of two parallel tool results must produce a single-part +/// functionResponse content item — not resurrect the deleted result. +#[test] +fn test_encode_delete_one_parallel_tool_response() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "run two tools"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"id": "c1", "name": "tool_a", "args": {}}}, + {"functionCall": {"id": "c2", "name": "tool_b", "args": {}}} + ] + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "c1", "name": "tool_a", "response": {"r": 1}}}, + {"functionResponse": {"id": "c2", "name": "tool_b", "response": {"r": 2}}} + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Delete the second tool result (tool_b / c2) from the annotated list. + annotated + .messages + .retain(|m| !matches!(m, Message::Tool { tool_call_id, .. } if tool_call_id == "c2")); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + + // Find the functionResponse content item(s). + let fr_items: Vec<&Json> = contents + .iter() + .filter(|c| { + c.get("parts") + .and_then(Json::as_array) + .map(|p| p.iter().any(|part| part.get("functionResponse").is_some())) + .unwrap_or(false) + }) + .collect(); + + // The deleted tool result must not appear in any functionResponse. + for item in &fr_items { + let parts = item.get("parts").and_then(Json::as_array).unwrap(); + for part in parts { + if let Some(fr) = part.get("functionResponse") { + assert_ne!( + fr.get("id").and_then(Json::as_str), + Some("c2"), + "deleted tool result c2 must not appear in encoded contents" + ); + assert_ne!( + fr.get("name").and_then(Json::as_str), + Some("tool_b"), + "deleted tool result tool_b must not appear in encoded contents" + ); + } + } + } + // Exactly one functionResponse part (c1) must survive. + let fr_part_count: usize = fr_items + .iter() + .map(|item| { + item.get("parts") + .and_then(Json::as_array) + .map(|p| { + p.iter() + .filter(|part| part.get("functionResponse").is_some()) + .count() + }) + .unwrap_or(0) + }) + .sum(); + assert_eq!( + fr_part_count, 1, + "exactly one functionResponse (c1) must survive deletion of c2" + ); +} + +/// Inserting a new tool response between two existing parallel results must +/// not duplicate or resurrect the original item. +#[test] +fn test_encode_insert_between_parallel_tool_responses() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "run two tools"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"id": "c1", "name": "tool_a", "args": {}}}, + {"functionCall": {"id": "c2", "name": "tool_b", "args": {}}} + ] + }, + { + "role": "user", + "parts": [ + {"functionResponse": {"id": "c1", "name": "tool_a", "response": {"r": 1}}}, + {"functionResponse": {"id": "c2", "name": "tool_b", "response": {"r": 2}}} + ] + } + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + + // Find positions of c1 and c2 tool messages. + let c1_pos = annotated + .messages + .iter() + .position(|m| matches!(m, Message::Tool { tool_call_id, .. } if tool_call_id == "c1")) + .unwrap(); + + // Insert a new tool result (c3) between c1 and c2. + annotated.messages.insert( + c1_pos + 1, + Message::Tool { + content: MessageContent::Text(r#"{"r": 99}"#.into()), + tool_call_id: "c3".into(), + }, + ); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + + // Collect all functionResponse parts across all content items. + let fr_ids: Vec<&str> = contents + .iter() + .flat_map(|c| { + c.get("parts") + .and_then(Json::as_array) + .into_iter() + .flatten() + }) + .filter_map(|part| { + part.get("functionResponse") + .and_then(|fr| fr.get("id")) + .and_then(Json::as_str) + }) + .collect(); + + // c1, c2, and the new c3 must all appear exactly once — no duplicates. + assert_eq!( + fr_ids.iter().filter(|&&id| id == "c1").count(), + 1, + "c1 must appear exactly once" + ); + assert_eq!( + fr_ids.iter().filter(|&&id| id == "c2").count(), + 1, + "c2 must appear exactly once" + ); + assert_eq!( + fr_ids.iter().filter(|&&id| id == "c3").count(), + 1, + "new c3 must be encoded" + ); +} + +// =================================================================== +// thoughtsTokenCount accounting +// =================================================================== + +/// thoughtsTokenCount must be stored in ApiSpecificResponse::GeminiGenerateContent and +/// must be included in the fallback total_tokens calculation. +#[test] +fn test_decode_response_thoughts_token_count() { + use super::super::response::ApiSpecificResponse; + + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "let me think..."}]}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20 + // totalTokenCount intentionally absent to test fallback + } + }); + + let resp = codec.decode_response(&response).unwrap(); + let usage = resp.usage.unwrap(); + + // completion_tokens reflects only candidatesTokenCount (not thinking tokens). + assert_eq!(usage.completion_tokens, Some(5)); + // Fallback total must include thinking tokens: 10 + 5 + 20 = 35. + assert_eq!( + usage.total_tokens, + Some(35), + "fallback total_tokens must include thoughtsTokenCount" + ); + + // thoughts_tokens must be in ApiSpecificResponse::GeminiGenerateContent. + match resp.api_specific { + Some(ApiSpecificResponse::GeminiGenerateContent { + thoughts_tokens, .. + }) => { + assert_eq!( + thoughts_tokens, + Some(20), + "thoughtsTokenCount must be in api_specific" + ); + } + other => panic!("expected ApiSpecificResponse::GeminiGenerateContent, got: {other:?}"), + } +} + +// Roleless contents decode to user messages (Google allows omitting role). +#[test] +fn test_decode_roleless_content_treated_as_user() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [ + {"parts": [{"text": "hello without a role"}]}, + {"role": "model", "parts": [{"text": "hi"}]} + ] + })); + let annotated = codec.decode(&request).unwrap(); + let non_sys: Vec<&Message> = annotated + .messages + .iter() + .filter(|m| !matches!(m, Message::System { .. })) + .collect(); + assert_eq!(non_sys.len(), 2); + assert!( + matches!(non_sys[0], Message::User { .. }), + "roleless content must decode as a user message" + ); + if let Message::User { content, .. } = non_sys[0] { + assert_eq!( + content, + &MessageContent::Text("hello without a role".into()), + "roleless content text must be preserved" + ); + } +} + +// Wrong numeric types in generationConfig must return an error. +#[test] +fn test_decode_rejects_non_numeric_temperature() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": {"temperature": "hot"} + })); + assert!( + codec.decode(&request).is_err(), + "non-numeric temperature must return an error" + ); +} + +#[test] +fn test_decode_rejects_non_numeric_max_output_tokens() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": {"maxOutputTokens": "a lot"} + })); + assert!( + codec.decode(&request).is_err(), + "non-integer maxOutputTokens must return an error" + ); +} + +// Streaming finalizer must propagate responseId when present. +#[test] +fn test_streaming_propagates_response_id() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hello"}]}, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2}, + "responseId": "stream-resp-xyz", + "modelVersion": "gemini-2.0-flash" + })) + .unwrap(); + + let assembled = finalizer(); + + assert_eq!( + assembled.get("responseId").and_then(Json::as_str), + Some("stream-resp-xyz"), + "streaming finalizer must propagate responseId from SSE events" + ); + + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + assert_eq!( + resp.id.as_deref(), + Some("stream-resp-xyz"), + "responseId from streaming must survive decode_response" + ); +} + +// MAX_TOKENS + functionCall parts → Length, not ToolUse. +// Explicit reasons must not be overridden by the tool-call presence heuristic. +#[test] +fn test_decode_response_max_tokens_with_function_call_is_length_not_tool_use() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "fn", "args": {}}}] + }, + "finishReason": "MAX_TOKENS", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 5} + }); + let resp = codec.decode_response(&response).unwrap(); + assert_eq!( + resp.finish_reason, + Some(FinishReason::Length), + "MAX_TOKENS must map to Length even when functionCall parts are present" + ); +} + +// SAFETY + functionCall parts → ContentFilter, not ToolUse. +#[test] +fn test_decode_response_safety_with_function_call_is_content_filter_not_tool_use() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "fn", "args": {}}}] + }, + "finishReason": "SAFETY", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 5} + }); + let resp = codec.decode_response(&response).unwrap(); + assert_eq!( + resp.finish_reason, + Some(FinishReason::ContentFilter), + "SAFETY must map to ContentFilter even when functionCall parts are present" + ); +} + +// STOP + functionCall parts → ToolUse (unchanged behaviour). +#[test] +fn test_decode_response_stop_with_function_call_is_tool_use() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "fn", "args": {}}}] + }, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": {"promptTokenCount": 5} + }); + let resp = codec.decode_response(&response).unwrap(); + assert_eq!( + resp.finish_reason, + Some(FinishReason::ToolUse), + "STOP + functionCall must map to ToolUse" + ); +} + +// Roleless content items (no explicit "role") that are later edited must not +// be silently dropped by patch_changed_gemini_content. +#[test] +fn test_encode_roleless_original_item_survives_edit() { + let codec = GeminiGenerateContentCodec; + // A Gemini request whose user content item has no "role" field. + let original = make_request(json!({ + "contents": [{"parts": [{"text": "original text"}]}] + })); + + let mut annotated = codec.decode(&original).unwrap(); + // Change the text (triggers the patch path, not the unchanged path). + if let Message::User { content, .. } = &mut annotated.messages[0] { + *content = MessageContent::Text("edited text".into()); + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let contents = encoded + .content + .get("contents") + .and_then(Json::as_array) + .unwrap(); + assert_eq!( + contents.len(), + 1, + "roleless item must not be dropped on edit" + ); + let text = contents[0] + .get("parts") + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str); + assert_eq!( + text, + Some("edited text"), + "edited text must appear in output" + ); +} + +// Multiple system messages must be merged into one systemInstruction. +#[test] +fn test_encode_multiple_system_messages_merged() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Inject two system messages. + annotated.messages.insert( + 0, + Message::System { + content: MessageContent::Text("part one".into()), + name: None, + }, + ); + annotated.messages.insert( + 1, + Message::System { + content: MessageContent::Text("part two".into()), + name: None, + }, + ); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let sys = encoded + .content + .get("systemInstruction") + .and_then(|s| s.get("parts")) + .and_then(Json::as_array) + .and_then(|p| p.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str) + .unwrap(); + assert!( + sys.contains("part one") && sys.contains("part two"), + "multiple system messages must be merged into one systemInstruction" + ); +} + +// Encoding unsupported fields must error, not silently ignore. +#[test] +fn test_encode_unsupported_field_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Set a field that has no Gemini equivalent. + annotated.max_output_tokens = Some(512); + let result = codec.encode(&annotated, &original); + assert!( + result.is_err(), + "setting max_output_tokens must return an error" + ); +} + +// Unparsable tool-call arguments must return an error rather than becoming {}. +#[test] +fn test_encode_invalid_tool_call_arguments_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.messages.push(Message::Assistant { + content: None, + name: None, + tool_calls: Some(vec![super::super::request::ToolCall { + id: "c1".into(), + call_type: "function".into(), + function: super::super::request::FunctionCall { + name: "fn".into(), + arguments: "not valid json {{".into(), + }, + }]), + }); + let result = codec.encode(&annotated, &original); + assert!( + result.is_err(), + "invalid tool-call arguments must return an error" + ); +} + +// Thinking-only partial usage (no candidatesTokenCount). +#[test] +fn test_decode_response_thoughts_only_usage_computes_fallback_total() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": []}, + "finishReason": "STOP" + }], + "usageMetadata": { + "promptTokenCount": 10, + // candidatesTokenCount intentionally absent + "thoughtsTokenCount": 30 + } + }); + let resp = codec.decode_response(&response).unwrap(); + let usage = resp.usage.unwrap(); + assert_eq!( + usage.completion_tokens, None, + "candidatesTokenCount absent → completion_tokens None" + ); + assert_eq!( + usage.total_tokens, + Some(40), + "fallback total must be prompt(10) + thoughts(30) = 40 even without candidatesTokenCount" + ); +} + +// Decode validation: table-driven negative tests for malformed contents items. +#[test] +fn test_decode_contents_item_validation() { + let codec = GeminiGenerateContentCodec; + + let cases: &[(&str, Json)] = &[ + ( + "contents item not an object", + json!({"contents": ["not an object"]}), + ), + ("parts missing", json!({"contents": [{"role": "user"}]})), + ( + "parts not an array", + json!({"contents": [{"role": "user", "parts": "bad"}]}), + ), + ( + "explicit unknown role", + json!({"contents": [{"role": "system", "parts": []}]}), + ), + ( + "part not an object", + json!({"contents": [{"role": "user", "parts": ["not an object"]}]}), + ), + ( + "functionCall not an object", + json!({"contents": [{"role": "model", "parts": [{"functionCall": "bad"}]}]}), + ), + ( + "functionCall missing name", + json!({"contents": [{"role": "model", "parts": [{"functionCall": {"args": {}}}]}]}), + ), + ( + "functionCall empty name", + json!({"contents": [{"role": "model", "parts": [{"functionCall": {"name": "", "args": {}}}]}]}), + ), + ( + "functionResponse not an object", + json!({"contents": [{"role": "user", "parts": [{"functionResponse": "bad"}]}]}), + ), + ( + "functionResponse missing name", + json!({"contents": [{"role": "user", "parts": [{"functionResponse": {"response": {}}}]}]}), + ), + ( + "functionResponse.response not an object", + json!({"contents": [{"role": "user", "parts": [{"functionResponse": {"name": "fn", "response": "bad"}}]}]}), + ), + ]; + + for (label, body) in cases { + let req = make_request(body.clone()); + assert!( + codec.decode(&req).is_err(), + "case '{label}' must return an error" + ); + } +} + +// Valid roleless contents item must decode successfully as a user message. +#[test] +fn test_decode_roleless_contents_item_is_valid() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"parts": [{"text": "hello"}]}] + })); + let ann = codec.decode(&req).unwrap(); + assert!( + matches!(&ann.messages[0], Message::User { content: MessageContent::Text(t), .. } if t == "hello"), + "roleless contents item must decode as a user message" + ); +} + +// Encoding a normalized message whose content has non-text parts must error. +#[test] +fn test_encode_non_text_content_part_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Replace the user message with one that has an image-url part in its content. + annotated.messages[0] = serde_json::from_value(json!({ + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}} + ] + })) + .unwrap(); + let result = codec.encode(&annotated, &original); + assert!( + result.is_err(), + "non-text content part (image_url) must return an error, not be silently dropped" + ); +} + +// Encoding a newly-inserted user message with non-text content must also error. +#[test] +fn test_encode_inserted_non_text_content_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "original"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.messages.push( + serde_json::from_value(json!({ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}} + ] + })) + .unwrap(), + ); + let result = codec.encode(&annotated, &original); + assert!( + result.is_err(), + "inserting a message with non-text content must error, not produce empty text" + ); +} + +// Editing tools when the original request has multiple functionDeclarations +// groups must return an error rather than silently collapsing them. +#[test] +fn test_encode_multiple_fn_decl_groups_error_on_edit() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [ + {"functionDeclarations": [{"name": "fn_a", "description": "group 1"}]}, + {"functionDeclarations": [{"name": "fn_b", "description": "group 2"}]} + ] + })); + + let mut annotated = codec.decode(&original).unwrap(); + // Change a tool to trigger the tools-changed path. + if let Some(tools) = annotated.tools.as_mut() { + for td in tools.iter_mut() { + if let nemo_relay_types::codec::request::ToolDefinition::Function { function, .. } = td + { + function.description = Some("edited".into()); + } + } + } + + let result = codec.encode(&annotated, &original); + assert!( + result.is_err(), + "editing tools when original has multiple functionDeclarations groups must error" + ); +} + +// Streaming: candidate metadata (safetyRatings, groundingMetadata) must survive finalization. +#[test] +fn test_streaming_candidate_extras_survive_finalize() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let finalizer = streaming_codec.finalizer(); + + collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hello"}]}, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"} + ], + "groundingMetadata": {"webSearchQueries": ["example query"]} + }], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2} + })) + .unwrap(); + + let assembled = finalizer(); + + let codec = GeminiGenerateContentCodec; + let resp = codec.decode_response(&assembled).unwrap(); + + use super::super::response::ApiSpecificResponse; + match &resp.api_specific { + Some(ApiSpecificResponse::GeminiGenerateContent { + safety_ratings, + grounding_metadata, + .. + }) => { + assert!( + safety_ratings.is_some(), + "safetyRatings from streaming must survive in ApiSpecificResponse::GeminiGenerateContent" + ); + assert!( + grounding_metadata.is_some(), + "groundingMetadata from streaming must survive in ApiSpecificResponse::GeminiGenerateContent" + ); + } + other => panic!("expected ApiSpecificResponse::GeminiGenerateContent, got: {other:?}"), + } +} + +// Present-but-non-string role must error. +#[test] +fn test_decode_non_string_role_errors() { + let codec = GeminiGenerateContentCodec; + for bad_role in [json!(123), json!(null), json!(true), json!([])] { + let req = make_request(json!({ + "contents": [{"role": bad_role, "parts": [{"text": "hi"}]}] + })); + assert!( + codec.decode(&req).is_err(), + "role={bad_role} must error; only string 'user'/'model' are accepted" + ); + } +} + +// functionResponse.response is required; missing must error. +#[test] +fn test_decode_function_response_missing_response_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [ + {"functionResponse": {"id": "c1", "name": "fn"}} + ]}] + })); + assert!( + codec.decode(&req).is_err(), + "missing functionResponse.response must error" + ); +} + +#[test] +fn test_function_response_content_helper_missing_response_errors() { + assert!(gemini_function_response_to_message_content(&json!({"name": "fn"})).is_err()); +} + +// functionResponse.response must be an object. +#[test] +fn test_decode_function_response_non_object_response_errors() { + let codec = GeminiGenerateContentCodec; + for bad in [json!("string"), json!([1, 2]), json!(42)] { + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [ + {"functionResponse": {"name": "fn", "response": bad}} + ]}] + })); + assert!( + codec.decode(&req).is_err(), + "response={bad} must error; response must be an object" + ); + } +} + +// Tool encode: non-object content is wrapped in {"output": ...}. +#[test] +fn test_encode_tool_content_non_object_is_wrapped() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "go"}]}, + {"role": "model", "parts": [{"functionCall": {"id": "c1", "name": "fn", "args": {}}}]}, + {"role": "user", "parts": [{"functionResponse": {"id": "c1", "name": "fn", "response": {"val": 1}}}]} + ] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Interceptor sets tool content to a bare number (parses as JSON but not object). + if let Message::Tool { content, .. } = &mut annotated.messages[2] { + *content = MessageContent::Text("42".into()); + } + let encoded = codec.encode(&annotated, &original).unwrap(); + let fr = encoded.content["contents"][2]["parts"][0]["functionResponse"]["response"].clone(); + assert!( + fr.is_object(), + "non-object tool content must be wrapped as an object" + ); + assert_eq!(fr["output"], json!(42)); +} + +// System message with non-text content must error. +#[test] +fn test_encode_system_message_non_text_content_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.messages.insert( + 0, + serde_json::from_value(json!({ + "role": "system", + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}] + })) + .unwrap(), + ); + assert!( + codec.encode(&annotated, &original).is_err(), + "system message with image content must error" + ); +} + +// Tool message with non-text normalized content must error. +#[test] +fn test_encode_tool_message_non_text_content_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "go"}]}, + {"role": "model", "parts": [{"functionCall": {"id": "c1", "name": "fn", "args": {}}}]}, + {"role": "user", "parts": [{"functionResponse": {"id": "c1", "name": "fn", "response": {}}}]} + ] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Use image_url — a ContentPart variant that round-trips through serde but has + // no Gemini encoding, so the guard must reject it. + if let Message::Tool { content, .. } = &mut annotated.messages[2] { + *content = serde_json::from_value(json!([ + {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}} + ])) + .unwrap(); + } + assert!( + codec.encode(&annotated, &original).is_err(), + "tool message with image content must error" + ); +} + +// Refusal content must error (not silently drop). +#[test] +fn test_encode_refusal_content_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.messages.push( + serde_json::from_value(json!({ + "role": "assistant", + "content": [{"type": "refusal", "refusal": "I cannot help"}] + })) + .unwrap(), + ); + assert!( + codec.encode(&annotated, &original).is_err(), + "refusal content must error, not silently vanish" + ); +} + +// tools[] validation negative cases. +#[test] +fn test_decode_tools_validation_negative() { + let codec = GeminiGenerateContentCodec; + let cases: &[(&str, Json)] = &[ + ( + "tools entry not an object", + json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": ["bad"]}), + ), + ( + "functionDeclarations not an array", + json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": "bad"}]}), + ), + ( + "functionDeclaration not an object", + json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": ["bad"]}]}), + ), + ( + "functionDeclaration empty name", + json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": [{"name": ""}]}]}), + ), + ( + "functionDeclaration missing name", + json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": [{"description": "no name"}]}]}), + ), + ]; + for (label, body) in cases { + assert!( + codec.decode(&make_request(body.clone())).is_err(), + "case '{label}' must error" + ); + } +} + +// functionCall.args must be an object on request decode. +#[test] +fn test_decode_function_call_non_object_args_errors() { + let codec = GeminiGenerateContentCodec; + for bad in [json!([1, 2]), json!("string"), json!(42)] { + let req = make_request(json!({ + "contents": [{ + "role": "model", + "parts": [{"functionCall": {"name": "fn", "args": bad}}] + }] + })); + assert!( + codec.decode(&req).is_err(), + "functionCall.args={bad} must error on request decode" + ); + } +} + +// functionCall arguments must be a JSON object on encode. +#[test] +fn test_encode_function_call_non_object_args_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({"contents": [{"role": "user", "parts": [{"text": "go"}]}]})); + let mut annotated = codec.decode(&original).unwrap(); + annotated.messages.push(Message::Assistant { + content: None, + name: None, + tool_calls: Some(vec![super::super::request::ToolCall { + id: "c1".into(), + call_type: "function".into(), + function: super::super::request::FunctionCall { + name: "fn".into(), + arguments: r#"["not","object"]"#.into(), + }, + }]), + }); + assert!( + codec.encode(&annotated, &original).is_err(), + "array arguments must error; Gemini requires object args" + ); +} + +// Unchanged payload with native inlineData round-trips byte-identically. +#[test] +fn test_encode_native_inline_data_round_trips_unchanged() { + let codec = GeminiGenerateContentCodec; + let image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQ=="; + let original = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"text": "what is this?"}, + {"inlineData": {"mimeType": "image/png", "data": image}} + ] + }], + "generationConfig": {"temperature": 0.5} + })); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content, original.content, + "encode(decode(req), req) must be identical to req" + ); +} + +// generationConfig unmodeled keys are preserved when params are cleared. +#[test] +fn test_encode_clearing_params_preserves_unmodeled_gen_config_fields() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": { + "temperature": 0.7, + "responseMimeType": "application/json", + "responseSchema": {"type": "object"} + } + })); + let mut annotated = codec.decode(&original).unwrap(); + // Clear params so temperature is removed. + annotated.params = None; + let encoded = codec.encode(&annotated, &original).unwrap(); + let gc = encoded + .content + .get("generationConfig") + .expect("generationConfig must remain"); + assert!( + gc.get("temperature").is_none(), + "temperature must be removed" + ); + assert_eq!( + gc["responseMimeType"], + json!("application/json"), + "responseMimeType must survive" + ); + assert!( + gc.get("responseSchema").is_some(), + "responseSchema must survive" + ); +} + +// All modeled keys cleared, empty generationConfig is removed. +#[test] +fn test_encode_clearing_params_removes_empty_gen_config() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": {"temperature": 0.5, "maxOutputTokens": 256} + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.params = None; + let encoded = codec.encode(&annotated, &original).unwrap(); + assert!( + encoded.content.get("generationConfig").is_none(), + "generationConfig must be removed when all modeled keys are cleared and no unmodeled keys remain" + ); +} + +// Response decode: malformed functionCall parts error rather than silently drop. +#[test] +fn test_decode_response_malformed_function_call_errors() { + let codec = GeminiGenerateContentCodec; + let cases: &[(&str, Json)] = &[ + ( + "functionCall is non-object", + json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": "not-an-object"}]}, "finishReason": "STOP"}] + }), + ), + ( + "functionCall missing name", + json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": {"args": {}}}]}, "finishReason": "STOP"}] + }), + ), + ( + "functionCall empty name", + json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": {"name": "", "args": {}}}]}, "finishReason": "STOP"}] + }), + ), + ( + "functionCall.id non-string", + json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": {"name": "fn", "id": 123, "args": {}}}]}, "finishReason": "STOP"}] + }), + ), + ( + "functionCall.id empty", + json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": {"name": "fn", "id": "", "args": {}}}]}, "finishReason": "STOP"}] + }), + ), + ( + "functionCall.args non-object", + json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": {"name": "fn", "args": [1, 2]}}]}, "finishReason": "STOP"}] + }), + ), + ]; + for (label, body) in cases { + assert!( + codec.decode_response(body).is_err(), + "case '{label}' must error" + ); + } +} + +// STOP + malformed functionCall must error, not return Complete. +#[test] +fn test_decode_response_stop_with_malformed_function_call_errors() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{"content": {"role": "model", + "parts": [{"functionCall": {"name": "", "args": {}}}]}, + "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 5} + }); + assert!( + codec.decode_response(&response).is_err(), + "STOP + malformed functionCall must error, not return Complete" + ); +} + +// Request decode: present functionCall.id must be a non-empty string. +#[test] +fn test_decode_request_function_call_id_validation() { + let codec = GeminiGenerateContentCodec; + let cases: &[(&str, Json)] = &[ + ( + "functionCall.id non-string", + json!({ + "contents": [{"role": "model", + "parts": [{"functionCall": {"name": "fn", "id": 99, "args": {}}}]}] + }), + ), + ( + "functionCall.id empty", + json!({ + "contents": [{"role": "model", + "parts": [{"functionCall": {"name": "fn", "id": "", "args": {}}}]}] + }), + ), + ]; + for (label, body) in cases { + assert!( + codec.decode(&make_request(body.clone())).is_err(), + "case '{label}' must error" + ); + } +} + +// Request decode: present functionResponse.id must be a non-empty string. +#[test] +fn test_decode_request_function_response_id_validation() { + let codec = GeminiGenerateContentCodec; + let cases: &[(&str, Json)] = &[ + ( + "functionResponse.id non-string", + json!({ + "contents": [{"role": "user", + "parts": [{"functionResponse": {"name": "fn", "id": 99, "response": {}}}]}] + }), + ), + ( + "functionResponse.id empty", + json!({ + "contents": [{"role": "user", + "parts": [{"functionResponse": {"name": "fn", "id": "", "response": {}}}]}] + }), + ), + ]; + for (label, body) in cases { + assert!( + codec.decode(&make_request(body.clone())).is_err(), + "case '{label}' must error" + ); + } +} + +// systemInstruction validation. +#[test] +fn test_decode_system_instruction_validation() { + let codec = GeminiGenerateContentCodec; + let base_contents = json!([{"role": "user", "parts": [{"text": "hi"}]}]); + let cases: &[(&str, Json)] = &[ + ( + "systemInstruction not an object", + json!({"contents": base_contents, "systemInstruction": "bad"}), + ), + ( + "systemInstruction missing parts", + json!({"contents": base_contents, "systemInstruction": {}}), + ), + ( + "systemInstruction.parts not an array", + json!({"contents": base_contents, "systemInstruction": {"parts": "bad"}}), + ), + ( + "systemInstruction.parts entry not an object", + json!({"contents": base_contents, "systemInstruction": {"parts": ["bad"]}}), + ), + ( + "systemInstruction.parts[].text not a string", + json!({"contents": base_contents, "systemInstruction": {"parts": [{"text": 123}]}}), + ), + ( + "systemInstruction.parts native part", + json!({ + "contents": base_contents, + "systemInstruction": { + "parts": [ + {"text": "ok"}, + { + "inlineData": { + "mimeType": "text/plain", + "data": "sk-system-secret" + } + } + ] + } + }), + ), + ]; + for (label, body) in cases { + assert!( + codec.decode(&make_request(body.clone())).is_err(), + "case '{label}' must error" + ); + } +} + +// model must be a string when present. +#[test] +fn test_decode_non_string_model_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "model": 42 + })); + assert!(codec.decode(&req).is_err(), "non-string model must error"); +} + +// Encode: empty tool-call name must error. +#[test] +fn test_encode_empty_tool_call_name_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]})); + let mut annotated = codec.decode(&original).unwrap(); + annotated.messages.push(Message::Assistant { + content: None, + name: None, + tool_calls: Some(vec![super::super::request::ToolCall { + id: "c1".into(), + call_type: "function".into(), + function: super::super::request::FunctionCall { + name: "".into(), // empty name + arguments: "{}".into(), + }, + }]), + }); + assert!( + codec.encode(&annotated, &original).is_err(), + "empty tool-call function name must error" + ); +} + +// Encode: empty tool_call_id must error. +#[test] +fn test_encode_empty_tool_call_id_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "go"}]}, + {"role": "model", "parts": [{"functionCall": {"id": "c1", "name": "fn", "args": {}}}]}, + {"role": "user", "parts": [{"functionResponse": {"id": "c1", "name": "fn", "response": {}}}]} + ] + })); + let mut annotated = codec.decode(&original).unwrap(); + // Give the tool message an empty tool_call_id (force it via JSON). + annotated.messages[2] = serde_json::from_value(json!({ + "role": "tool", + "tool_call_id": "", + "content": "{}" + })) + .unwrap(); + assert!( + codec.encode(&annotated, &original).is_err(), + "empty tool_call_id must error" + ); +} + +// Encode: empty FunctionDefinition.name must error. +#[test] +fn test_encode_empty_function_definition_name_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"functionDeclarations": [{"name": "good_fn"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + if let Some(tools) = annotated.tools.as_mut() { + for td in tools.iter_mut() { + if let nemo_relay_types::codec::request::ToolDefinition::Function { function, .. } = td + { + function.name = "".into(); + } + } + } + assert!( + codec.encode(&annotated, &original).is_err(), + "empty FunctionDefinition.name must error" + ); +} + +// Encode: FunctionDefinition.strict must error when set. +#[test] +fn test_encode_function_definition_strict_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"functionDeclarations": [{"name": "fn"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + if let Some(tools) = annotated.tools.as_mut() { + for td in tools.iter_mut() { + if let nemo_relay_types::codec::request::ToolDefinition::Function { function, .. } = td + { + function.strict = Some(true); + } + } + } + assert!( + codec.encode(&annotated, &original).is_err(), + "FunctionDefinition.strict is not representable in Gemini and must error" + ); +} + +// Decode: non-string description in functionDeclaration must error. +#[test] +fn test_decode_function_declaration_non_string_description_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"functionDeclarations": [{"name": "fn", "description": 42}]}] + })); + assert!( + codec.decode(&req).is_err(), + "non-string functionDeclaration.description must error" + ); +} + +// Encode: systemInstruction sibling fields are preserved when text changes. +#[test] +fn test_encode_system_instruction_preserves_sibling_fields() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": { + "role": "user", + "parts": [{"text": "old prompt"}], + "nativeField": "keep me" + } + })); + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::System { content, .. } = msg { + *content = MessageContent::Text("new prompt".into()); + } + } + let encoded = codec.encode(&annotated, &original).unwrap(); + let si = encoded.content.get("systemInstruction").unwrap(); + assert_eq!( + si.get("parts") + .and_then(|p| p.as_array()) + .and_then(|a| a.first()) + .and_then(|p| p.get("text")) + .and_then(Json::as_str), + Some("new prompt"), + "new system text must appear in parts" + ); + assert_eq!( + si.get("nativeField").and_then(Json::as_str), + Some("keep me"), + "systemInstruction native sibling fields must be preserved when text changes" + ); + assert_eq!( + si.get("role").and_then(Json::as_str), + Some("user"), + "systemInstruction.role must be preserved" + ); +} + +// Helper: text parts without an explicit normalized type are still text. +#[test] +fn test_extract_content_text_treats_missing_part_type_as_text() { + assert_eq!( + extract_content_text(&json!([{"text": "new prompt"}])), + "new prompt", + "missing normalized part type should be interpreted as text, matching encoder validation" + ); +} + +// systemInstruction.role non-string must error on decode (validated in validate_system_instruction). +#[test] +fn test_decode_system_instruction_non_string_role_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": {"role": 123, "parts": [{"text": "system"}]} + })); + assert!( + codec.decode(&req).is_err(), + "systemInstruction.role non-string must error on decode" + ); +} + +// Mixed functionResponse + functionCall in one content item must error. +#[test] +fn test_decode_mixed_fr_and_fc_parts_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"functionResponse": {"name": "fn", "response": {}}}, + {"functionCall": {"name": "fn", "args": {}}} + ] + }] + })); + assert!( + codec.decode(&req).is_err(), + "content item with both functionResponse and functionCall parts must error" + ); +} + +// Non-numeric topP must error. +#[test] +fn test_decode_rejects_non_numeric_top_p() { + let codec = GeminiGenerateContentCodec; + let request = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "generationConfig": {"topP": "high"} + })); + assert!( + codec.decode(&request).is_err(), + "non-numeric topP must error" + ); +} + +// ProviderNative tool variants owned by a different provider must error on encode. +#[test] +fn test_encode_mismatched_provider_native_tool_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.tools = Some(vec![ToolDefinition::ProviderNative { + provider: "openai_chat".into(), + kind: "web_search".into(), + value: serde_json::json!({"type": "web_search_preview"}), + }]); + assert!( + codec.encode(&annotated, &original).is_err(), + "ProviderNative tool from another provider must error" + ); +} + +#[test] +fn test_encode_provider_native_tool_with_function_declarations_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + annotated.tools = Some(vec![ToolDefinition::ProviderNative { + provider: "gemini".into(), + kind: "functionDeclarations".into(), + value: serde_json::json!({"functionDeclarations": [{"name": "fn"}]}), + }]); + assert!( + codec.encode(&annotated, &original).is_err(), + "functionDeclarations must use ToolDefinition::Function" + ); +} + +#[test] +fn test_encode_system_message_name_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "systemInstruction": {"parts": [{"text": "be helpful"}]}, + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::System { name, .. } = msg { + *name = Some("system-name".into()); + } + } + assert!( + codec.encode(&annotated, &original).is_err(), + "Gemini cannot represent system message names" + ); +} + +#[test] +fn test_encode_user_message_name_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + let mut annotated = codec.decode(&original).unwrap(); + if let Message::User { name, .. } = &mut annotated.messages[0] { + *name = Some("user-name".into()); + } + assert!( + codec.encode(&annotated, &original).is_err(), + "Gemini cannot represent user message names" + ); +} + +#[test] +fn test_encode_assistant_message_name_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "hi"}]}, + {"role": "model", "parts": [{"text": "hello"}]} + ] + })); + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::Assistant { name, .. } = msg { + *name = Some("assistant-name".into()); + } + } + assert!( + codec.encode(&annotated, &original).is_err(), + "Gemini cannot represent assistant message names" + ); +} + +#[test] +fn test_encode_assistant_tool_call_message_name_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [ + {"role": "user", "parts": [{"text": "use a tool"}]}, + { + "role": "model", + "parts": [{"functionCall": {"id": "call_1", "name": "lookup", "args": {}}}] + } + ] + })); + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::Assistant { name, .. } = msg { + *name = Some("assistant-name".into()); + } + } + assert!( + codec.encode(&annotated, &original).is_err(), + "Gemini cannot represent assistant names on tool-call messages" + ); +} + +#[test] +fn test_encode_previous_response_id_returns_error() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]})); + let mut annotated = codec.decode(&original).unwrap(); + annotated.previous_response_id = Some("prev-123".into()); + assert!( + codec.encode(&annotated, &original).is_err(), + "previous_response_id is not representable in Gemini and must error" + ); +} + +#[test] +fn test_decode_non_string_text_part_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": 42}]}] + })); + assert!( + codec.decode(&req).is_err(), + "parts[].text with non-string value must error" + ); +} + +// functionResponse on model-role content must error. +#[test] +fn test_decode_function_response_on_model_role_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{ + "role": "model", + "parts": [{"functionResponse": {"name": "fn", "response": {}}}] + }] + })); + assert!( + codec.decode(&req).is_err(), + "functionResponse in a 'model' role content item must error" + ); +} + +// functionCall on user-role content must error. +#[test] +fn test_decode_function_call_on_user_role_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [{"functionCall": {"name": "fn", "args": {}}}] + }] + })); + assert!( + codec.decode(&req).is_err(), + "functionCall in a 'user' role content item must error" + ); +} + +// functionResponse mixed with visible text parts must error. +#[test] +fn test_decode_function_response_mixed_with_text_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [ + {"functionResponse": {"name": "fn", "response": {}}}, + {"text": "visible text"} + ] + }] + })); + assert!( + codec.decode(&req).is_err(), + "functionResponse mixed with visible text parts must error" + ); +} + +// Response decode: non-string text value must error. +#[test] +fn test_decode_response_non_string_text_part_errors() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": 42}]}, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 1} + }); + assert!( + codec.decode_response(&response).is_err(), + "response parts[].text with non-string value must error" + ); +} + +// systemInstruction part-level metadata is preserved when text changes. +#[test] +fn test_encode_system_instruction_part_metadata_preserved() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": { + "parts": [{"text": "old prompt", "nativePartField": "keep-me"}] + } + })); + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::System { content, .. } = msg { + *content = MessageContent::Text("new prompt".into()); + } + } + let encoded = codec.encode(&annotated, &original).unwrap(); + let part = &encoded.content["systemInstruction"]["parts"][0]; + assert_eq!(part["text"].as_str(), Some("new prompt")); + assert_eq!( + part["nativePartField"].as_str(), + Some("keep-me"), + "part-level native fields must be preserved when system text changes" + ); +} + +// Unknown sibling fields on a known text part are metadata, not additional data-union fields. +#[test] +fn test_encode_text_part_unknown_metadata_preserved() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{ + "role": "user", + "parts": [{"text": "old text", "nativePartField": "keep-me"}] + }] + })); + let mut annotated = codec.decode(&original).unwrap(); + if let Message::User { content, .. } = &mut annotated.messages[0] { + match content { + MessageContent::Parts(parts) => { + if let ContentPart::Text { text, .. } = &mut parts[0] { + *text = "new text".into(); + } + } + other => panic!("expected metadata-bearing text part, got {other:?}"), + } + } + + let encoded = codec.encode(&annotated, &original).unwrap(); + let part = &encoded.content["contents"][0]["parts"][0]; + assert_eq!(part["text"].as_str(), Some("new text")); + assert_eq!(part["nativePartField"].as_str(), Some("keep-me")); +} + +// Request decode: part with both text and functionResponse must error. +#[test] +fn test_decode_part_with_text_and_function_response_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "user", "parts": [ + {"text": "visible", "functionResponse": {"name": "fn", "response": {}}} + ]}] + })); + assert!( + codec.decode(&req).is_err(), + "part with both 'text' and 'functionResponse' must error" + ); +} + +// Request decode: part with both text and functionCall must error. +#[test] +fn test_decode_part_with_text_and_function_call_errors() { + let codec = GeminiGenerateContentCodec; + let req = make_request(json!({ + "contents": [{"role": "model", "parts": [ + {"text": "visible", "functionCall": {"name": "fn", "args": {}}} + ]}] + })); + assert!( + codec.decode(&req).is_err(), + "part with both 'text' and 'functionCall' must error" + ); +} + +// Response decode: non-object part must error. +#[test] +fn test_decode_response_non_object_part_errors() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{"content": {"role": "model", "parts": ["not an object"]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1} + }); + assert!( + codec.decode_response(&response).is_err(), + "response part that is not an object must error" + ); +} + +// Response decode: part with both text and functionCall must error. +#[test] +fn test_decode_response_part_text_and_function_call_errors() { + let codec = GeminiGenerateContentCodec; + let response = json!({ + "candidates": [{"content": {"role": "model", "parts": [ + {"text": "hello", "functionCall": {"name": "fn", "args": {}}} + ]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1} + }); + assert!( + codec.decode_response(&response).is_err(), + "response part with both 'text' and 'functionCall' must error" + ); +} + +// Streaming: part with both text and functionCall must error via collector. +#[test] +fn test_streaming_part_text_and_function_call_errors() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + + let result = collector(json!({ + "candidates": [{ + "content": {"role": "model", "parts": [ + {"text": "hello", "functionCall": {"name": "fn", "args": {}}} + ]}, + "index": 0 + }] + })); + assert!( + result.is_err(), + "streaming part with both 'text' and 'functionCall' must error via collector" + ); +} + +// Streaming: non-object part must error. +#[test] +fn test_streaming_non_object_part_errors() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let result = collector(json!({ + "candidates": [{"content": {"role": "model", "parts": ["not an object"]}, "index": 0}] + })); + assert!( + result.is_err(), + "streaming non-object part must error via collector" + ); +} + +// Streaming: non-string text value must error. +#[test] +fn test_streaming_non_string_text_errors() { + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); + let mut collector = streaming_codec.collector(); + let result = collector(json!({ + "candidates": [{"content": {"role": "model", "parts": [{"text": 42}]}, "index": 0}] + })); + assert!( + result.is_err(), + "streaming part with non-string text must error via collector" + ); +} + +// systemInstruction with multiple text parts must error on edit. +#[test] +fn test_encode_system_instruction_multiple_text_parts_edit_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": { + "parts": [{"text": "part one"}, {"text": "part two"}] + } + })); + let mut annotated = codec.decode(&original).unwrap(); + for msg in annotated.messages.iter_mut() { + if let Message::System { content, .. } = msg { + *content = MessageContent::Text("new system".into()); + } + } + assert!( + codec.encode(&annotated, &original).is_err(), + "editing systemInstruction with multiple text parts must error" + ); +} + +// systemInstruction with a native non-text part must error on decode. +#[test] +fn test_decode_system_instruction_non_text_part_errors() { + let codec = GeminiGenerateContentCodec; + let original = make_request(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": { + "parts": [{"text": "ok"}, {"nativePart": "value"}] + } + })); + assert!( + codec.decode(&original).is_err(), + "systemInstruction with non-text native parts must error" + ); +} diff --git a/crates/core/tests/unit/codec/parity_tests.rs b/crates/core/tests/unit/codec/parity_tests.rs index 955b38e1e..64ab62077 100644 --- a/crates/core/tests/unit/codec/parity_tests.rs +++ b/crates/core/tests/unit/codec/parity_tests.rs @@ -3,10 +3,13 @@ //! Cross-provider codec parity tests for the NeMo Relay core crate. //! -//! Each test builds the same logical scenario in all three built-in provider -//! schemas (OpenAI Chat Completions, Anthropic Messages, OpenAI Responses) and -//! asserts that detection plus normalization produce agreeing output. Where -//! the schemas legitimately diverge, the divergence is asserted explicitly: +//! Each test builds the same logical scenario in the three OpenAI/Anthropic +//! built-in provider schemas (OpenAI Chat Completions, Anthropic Messages, +//! OpenAI Responses) and asserts that detection plus normalization produce +//! agreeing output. Gemini generateContent-specific semantics (contents array, +//! functionResponse role split, thinking tokens, etc.) are covered in +//! gemini_generate_content_tests.rs. Where the schemas legitimately diverge, +//! the divergence is asserted explicitly: //! the asymmetry is part of the parity contract, and a change here means one //! codec drifted from the others. @@ -604,7 +607,11 @@ fn test_request_unknown_hint_matches_hintless_normalization() { for body in &bodies { let request = req(body.clone()); let baseline = normalize_request(&request).expect("canonical body decodes"); - for hint in ["gemini", "not-a-provider", "anthropic.count_tokens"] { + for hint in [ + "gemini_generate_content", + "not-a-provider", + "anthropic.count_tokens", + ] { assert_eq!( normalize_request_with_hint(&request, Some(hint)).as_ref(), Some(&baseline), diff --git a/crates/core/tests/unit/codec/resolve_tests.rs b/crates/core/tests/unit/codec/resolve_tests.rs index 70907a933..58fce0f97 100644 --- a/crates/core/tests/unit/codec/resolve_tests.rs +++ b/crates/core/tests/unit/codec/resolve_tests.rs @@ -26,6 +26,7 @@ fn builtin_provider_surface_registry_keeps_request_priority() { ProviderSurface::OpenAIResponses, ProviderSurface::AnthropicMessages, ProviderSurface::OpenAIChat, + ProviderSurface::GeminiGenerateContent, ] ); } @@ -84,6 +85,57 @@ fn detect_request_none_for_unknown_or_non_object() { assert_eq!(detect_request_surface(&json!("string")), None); } +#[test] +fn detect_request_gemini_by_contents() { + assert_eq!( + detect_request_surface(&json!({"contents": []})), + Some(ProviderSurface::GeminiGenerateContent) + ); + assert_eq!( + detect_request_surface(&json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]})), + Some(ProviderSurface::GeminiGenerateContent) + ); + // Higher-priority surfaces still win when their discriminators are present. + assert_eq!( + detect_request_surface(&json!({"contents": [], "messages": []})), + Some(ProviderSurface::OpenAIChat), + "messages key wins over contents in priority order" + ); +} + +#[test] +fn detect_response_gemini_by_candidates() { + assert_eq!( + detect_response_surface(&json!({"candidates": []})), + Some(ProviderSurface::GeminiGenerateContent) + ); + assert_eq!( + detect_response_surface(&json!({ + "candidates": [{"content": {"parts": [{"text": "hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {} + })), + Some(ProviderSurface::GeminiGenerateContent) + ); + // A scalar candidates key does not match (must be an array). + assert_eq!( + detect_response_surface(&json!({"candidates": "not-array"})), + None + ); +} + +#[test] +fn detect_response_gemini_by_prompt_feedback_block() { + assert_eq!( + detect_response_surface(&json!({ + "promptFeedback": { + "blockReason": "SAFETY", + "safetyRatings": [] + } + })), + Some(ProviderSurface::GeminiGenerateContent) + ); +} + // --------------------------------------------------------------------------- // detect_response_surface (strict; ambiguity -> None) // --------------------------------------------------------------------------- @@ -346,7 +398,7 @@ fn hint_other_or_unknown_provider_stays_chat() { Some("anthropic.count_tokens"), Some("anthropic.preview"), Some("passthrough"), - Some("gemini"), + Some("gemini_generate_content"), None, ] { assert_eq!( @@ -395,10 +447,11 @@ fn hint_does_not_classify_non_object_or_keyless() { // Provider-codec factory (name<->surface mapping + codec construction) // --------------------------------------------------------------------------- -const ALL_SURFACES: [ProviderSurface; 3] = [ +const ALL_SURFACES: [ProviderSurface; 4] = [ ProviderSurface::OpenAIChat, ProviderSurface::OpenAIResponses, ProviderSurface::AnthropicMessages, + ProviderSurface::GeminiGenerateContent, ]; #[test] @@ -423,11 +476,29 @@ fn codec_name_uses_canonical_spellings() { ProviderSurface::AnthropicMessages.codec_name(), "anthropic_messages" ); + assert_eq!( + ProviderSurface::GeminiGenerateContent.codec_name(), + "gemini_generate_content" + ); +} + +#[test] +fn from_codec_name_rejects_ambiguous_gemini_name() { + assert_eq!( + ProviderSurface::from_codec_name("gemini"), + None, + "Gemini codec names must name the concrete API surface" + ); + assert_eq!( + ProviderSurface::GeminiGenerateContent.codec_name(), + "gemini_generate_content", + "the canonical Gemini codec spelling names generateContent explicitly" + ); } #[test] fn from_codec_name_is_none_for_unknown_names() { - assert_eq!(ProviderSurface::from_codec_name("gemini"), None); + assert_eq!(ProviderSurface::from_codec_name("generate_content"), None); assert_eq!(ProviderSurface::from_codec_name(""), None); assert_eq!(ProviderSurface::from_codec_name("OpenAIChat"), None); } @@ -436,7 +507,12 @@ fn from_codec_name_is_none_for_unknown_names() { fn supported_codec_names_track_the_builtin_registry() { assert_eq!( supported_codec_names(), - vec!["openai_responses", "anthropic_messages", "openai_chat"] + vec![ + "openai_responses", + "anthropic_messages", + "openai_chat", + "gemini_generate_content" + ] ); let from_registry: Vec<_> = BUILTIN_PROVIDER_SURFACES .iter() @@ -480,6 +556,17 @@ fn request_codec_decodes_each_surface() { .messages .is_empty() ); + + let gemini = req(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + assert!( + !request_codec(ProviderSurface::GeminiGenerateContent) + .decode(&gemini) + .expect("gemini request decodes") + .messages + .is_empty() + ); } #[test] @@ -519,6 +606,18 @@ fn response_codec_decodes_each_surface() { .response_text(), Some("yo") ); + + let gemini = json!({ + "candidates": [{"content": {"parts": [{"text": "hi gemini"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1} + }); + assert_eq!( + response_codec(ProviderSurface::GeminiGenerateContent) + .decode_response(&gemini) + .expect("gemini response decodes") + .response_text(), + Some("hi gemini") + ); } #[test] diff --git a/crates/core/tests/unit/codec/response_tests.rs b/crates/core/tests/unit/codec/response_tests.rs index 1dfe9047f..62e85a5ab 100644 --- a/crates/core/tests/unit/codec/response_tests.rs +++ b/crates/core/tests/unit/codec/response_tests.rs @@ -1595,6 +1595,64 @@ fn test_api_specific_anthropic_messages_round_trip() { assert_eq!(api, deserialized); } +#[test] +fn test_api_specific_gemini_generate_content_round_trip_empty_extra() { + let api = ApiSpecificResponse::GeminiGenerateContent { + thoughts_tokens: Some(3), + safety_ratings: None, + grounding_metadata: None, + citation_metadata: None, + extra: serde_json::Map::new(), + }; + let json_val = serde_json::to_value(&api).unwrap(); + assert_eq!(json_val["api"], json!("gemini_generate_content")); + let deserialized: ApiSpecificResponse = serde_json::from_value(json_val).unwrap(); + assert_eq!(api, deserialized); +} + +#[test] +fn test_api_specific_gemini_generate_content_round_trip_extra() { + let api = ApiSpecificResponse::GeminiGenerateContent { + thoughts_tokens: None, + safety_ratings: Some(json!([{"category": "HARM_CATEGORY_HATE_SPEECH"}])), + grounding_metadata: None, + citation_metadata: None, + extra: serde_json::Map::from_iter([("urlContextMetadata".into(), json!({"url": "u"}))]), + }; + let json_val = serde_json::to_value(&api).unwrap(); + assert_eq!(json_val["api"], json!("gemini_generate_content")); + assert_eq!(json_val["urlContextMetadata"], json!({"url": "u"})); + let deserialized: ApiSpecificResponse = serde_json::from_value(json_val).unwrap(); + assert_eq!(api, deserialized); +} + +#[test] +fn test_api_specific_gemini_generate_content_extra_cannot_override_api_tag() { + let api = ApiSpecificResponse::GeminiGenerateContent { + thoughts_tokens: None, + safety_ratings: None, + grounding_metadata: None, + citation_metadata: None, + extra: serde_json::Map::from_iter([ + ("api".into(), json!("not_the_tag")), + ("futureField".into(), json!(true)), + ]), + }; + let json_val = serde_json::to_value(&api).unwrap(); + assert_eq!( + json_val["api"], + json!("gemini_generate_content"), + "Gemini extra.api must not overwrite the enum discriminator" + ); + assert_eq!(json_val["futureField"], json!(true)); + let deserialized: ApiSpecificResponse = serde_json::from_value(json_val).unwrap(); + let ApiSpecificResponse::GeminiGenerateContent { extra, .. } = deserialized else { + panic!("expected Gemini generateContent metadata"); + }; + assert!(extra.get("api").is_none()); + assert_eq!(extra.get("futureField"), Some(&json!(true))); +} + #[test] fn test_api_specific_custom_round_trip() { let api = ApiSpecificResponse::Custom { diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs index 153d88b4d..f838f1da2 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs @@ -399,7 +399,12 @@ fn schema_contains_every_supported_nemo_guardrails_option() { assert!(schema_property_has_enum( &schema, "codec", - &["openai_chat", "openai_responses", "anthropic_messages"] + &[ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate_content" + ] )); assert!(schema_property_has_default( &schema, @@ -640,8 +645,15 @@ fn assert_invalid_remote_identity_and_codec() { }))); assert!(bad_codec.has_errors()); assert!(bad_codec.diagnostics.iter().any(|diag| { - diag.message - .contains("codec must be 'openai_chat', 'openai_responses', or 'anthropic_messages'") + diag.message.contains("codec must be one of:") + && [ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate_content", + ] + .iter() + .all(|name| diag.message.contains(name)) })); let unsupported_remote_codec = validate_plugin_config(&plugin_config(json!({ @@ -676,6 +688,25 @@ fn assert_invalid_remote_identity_and_codec() { .contains("remote mode currently supports only codec = 'openai_chat'") }) ); + + let unsupported_remote_gemini_codec = validate_plugin_config(&plugin_config(json!({ + "mode": "remote", + "codec": "gemini_generate_content", + "remote": { + "endpoint": "http://localhost:8000", + "config_id": "default" + } + }))); + assert!(unsupported_remote_gemini_codec.has_errors()); + assert!( + unsupported_remote_gemini_codec + .diagnostics + .iter() + .any(|diag| { + diag.message + .contains("remote mode currently supports only codec = 'openai_chat'") + }) + ); } fn assert_remote_tool_surface_validation() { diff --git a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs b/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs index 3353112bd..ce173c97a 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/local_python_tests.rs @@ -689,6 +689,10 @@ fn local_codec_and_rewrite_helpers_cover_all_provider_surfaces() { LocalGuardrailsCodec::AnthropicMessages, ProviderSurface::AnthropicMessages, ), + ( + LocalGuardrailsCodec::GeminiGenerateContent, + ProviderSurface::GeminiGenerateContent, + ), ] { assert_eq!(codec.provider_surface(), surface); assert_eq!( @@ -1074,6 +1078,14 @@ fn stream_text_extraction_handles_supported_codecs() { ), Some("hello".to_string()) ); + // Gemini: visible text parts reach the guardrail worker. + assert_eq!( + extract_stream_text( + LocalGuardrailsCodec::GeminiGenerateContent, + &json!({"candidates": [{"content": {"parts": [{"text": "visible"}]}, "index": 0}]}) + ), + Some("visible".to_string()) + ); for (codec, chunk) in [ (LocalGuardrailsCodec::OpenAIChat, Json::Null), ( @@ -1088,11 +1100,37 @@ fn stream_text_extraction_handles_supported_codecs() { LocalGuardrailsCodec::AnthropicMessages, json!({"type": "content_block_delta", "delta": {"type": "input_json_delta"}}), ), + (LocalGuardrailsCodec::GeminiGenerateContent, Json::Null), ] { assert_eq!(extract_stream_text(codec, &chunk), None); } } +#[test] +fn stream_text_extraction_gemini_skips_thought_parts() { + // A thought chunk (thought: true) must NOT reach the guardrail worker. + assert_eq!( + extract_stream_text( + LocalGuardrailsCodec::GeminiGenerateContent, + &json!({"candidates": [{"content": {"parts": [{"thought": true, "text": "internal reasoning"}]}, "index": 0}]}) + ), + None, + "thought parts must not be forwarded to the guardrail worker" + ); + // A chunk with both a thought part and a visible part: only the visible text is forwarded. + assert_eq!( + extract_stream_text( + LocalGuardrailsCodec::GeminiGenerateContent, + &json!({"candidates": [{"content": {"parts": [ + {"thought": true, "text": "reasoning"}, + {"text": "answer"} + ]}, "index": 0}]}) + ), + Some("answer".to_string()), + "only non-thought text must reach the guardrail worker" + ); +} + #[cfg(unix)] async fn install_local_plugin(config: &NeMoGuardrailsConfig) { let component_config = serde_json::to_value(config) diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 112286a8b..0675080f2 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -962,6 +962,17 @@ struct FfiCodecHandle *nemo_relay_openai_responses_codec_new(void); */ struct FfiCodecHandle *nemo_relay_anthropic_messages_codec_new(void); +/** + * Create a new Gemini generateContent API codec handle. + * + * The returned handle implements both request codec (decode/encode) and + * response codec (decode_response). Free with `nemo_relay_codec_free`. + * + * # Safety + * Caller must free the returned handle via `nemo_relay_codec_free`. + */ +struct FfiCodecHandle *nemo_relay_gemini_generate_content_codec_new(void); + /** * Execute an LLM call end-to-end: run conditional-execution guardrails (on raw * request), then request intercepts, sanitize-request guardrails, execution diff --git a/crates/ffi/src/api/llm.rs b/crates/ffi/src/api/llm.rs index 62060d846..31ba4f8e9 100644 --- a/crates/ffi/src/api/llm.rs +++ b/crates/ffi/src/api/llm.rs @@ -372,6 +372,23 @@ pub extern "C" fn nemo_relay_anthropic_messages_codec_new() -> *mut FfiCodecHand })) } +/// Create a new Gemini generateContent API codec handle. +/// +/// The returned handle implements both request codec (decode/encode) and +/// response codec (decode_response). Free with `nemo_relay_codec_free`. +/// +/// # Safety +/// Caller must free the returned handle via `nemo_relay_codec_free`. +#[unsafe(no_mangle)] +pub extern "C" fn nemo_relay_gemini_generate_content_codec_new() -> *mut FfiCodecHandle { + Box::into_raw(Box::new(FfiCodecHandle { + codec: Arc::new(nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec), + response_codec: Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), + })) +} + struct ParsedExecuteInputs { name: String, request: LlmRequest, diff --git a/crates/ffi/tests/unit/types_tests.rs b/crates/ffi/tests/unit/types_tests.rs index 074ddf60d..15f0aa0cf 100644 --- a/crates/ffi/tests/unit/types_tests.rs +++ b/crates/ffi/tests/unit/types_tests.rs @@ -760,14 +760,20 @@ fn test_annotated_event_accessors_and_codec_handles() { let openai_chat = api::nemo_relay_openai_chat_codec_new(); let openai_responses = api::nemo_relay_openai_responses_codec_new(); let anthropic = api::nemo_relay_anthropic_messages_codec_new(); + let gemini = api::nemo_relay_gemini_generate_content_codec_new(); assert!(!openai_chat.is_null()); assert!(!openai_responses.is_null()); assert!(!anthropic.is_null()); + assert!( + !gemini.is_null(), + "GeminiGenerateContentCodec FFI constructor must return a non-null handle" + ); unsafe { nemo_relay_codec_free(openai_chat); nemo_relay_codec_free(openai_responses); nemo_relay_codec_free(anthropic); + nemo_relay_codec_free(gemini); nemo_relay_codec_free(std::ptr::null_mut()); } } diff --git a/crates/node/pii_redaction.d.ts b/crates/node/pii_redaction.d.ts index 5a5f20e76..969eae122 100644 --- a/crates/node/pii_redaction.d.ts +++ b/crates/node/pii_redaction.d.ts @@ -39,7 +39,7 @@ export interface Config { tool_output?: boolean; mark?: boolean; priority?: number; - codec?: 'openai_chat' | 'openai_responses' | 'anthropic_messages' | string; + codec?: 'openai_chat' | 'openai_responses' | 'anthropic_messages' | 'gemini_generate_content' | string; builtin?: BuiltinConfig; local?: LocalModelConfig; policy?: ConfigPolicy; diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index d4b77b92e..1ee433eee 100644 --- a/crates/node/plugin.d.ts +++ b/crates/node/plugin.d.ts @@ -9,7 +9,7 @@ import type { LlmCodec, LlmResponseCodec } from './typed'; /** Codec identity available while a managed LLM event is sanitized. */ export type LlmCodecIdentity = | { kind: 'none' } - | { kind: 'builtin'; id: 'openai_chat' | 'openai_responses' | 'anthropic_messages' } + | { kind: 'builtin'; id: 'openai_chat' | 'openai_responses' | 'anthropic_messages' | 'gemini_generate_content' } | { kind: 'runtime'; id: string } | { kind: 'opaque' }; diff --git a/crates/node/src/types/mod.rs b/crates/node/src/types/mod.rs index 2a36c4dc9..4c8095991 100644 --- a/crates/node/src/types/mod.rs +++ b/crates/node/src/types/mod.rs @@ -480,6 +480,67 @@ impl OpenAIResponsesCodec { } } +/// Built-in codec for the Gemini generateContent API. +/// +/// Implements both request codec (decode/encode) and response codec +/// (decodeResponse). Construct with `new GeminiGenerateContentCodec()`. +#[napi(js_name = "GeminiGenerateContentCodec")] +pub struct GeminiGenerateContentCodec { + pub(crate) inner_codec: std::sync::Arc, + pub(crate) inner_response_codec: std::sync::Arc, +} + +#[napi] +impl GeminiGenerateContentCodec { + #[napi(constructor)] + pub fn new() -> Self { + Self { + inner_codec: std::sync::Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), + inner_response_codec: std::sync::Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), + } + } + + /// Decode an opaque LLM request into structured form. + #[napi] + pub fn decode(&self, request: Json) -> napi::Result { + let llm_req: CoreLlmRequest = serde_json::from_value(request) + .map_err(|e| napi::Error::from_reason(format!("invalid LlmRequest: {e}")))?; + let annotated = self + .inner_codec + .decode(&llm_req) + .map_err(|e| napi::Error::from_reason(e.to_string()))?; + serde_json::to_value(&annotated).map_err(|e| napi::Error::from_reason(e.to_string())) + } + + /// Encode structured changes back into an opaque LLM request. + #[napi] + pub fn encode(&self, annotated: Json, original: Json) -> napi::Result { + let ann: AnnotatedLlmRequest = serde_json::from_value(annotated) + .map_err(|e| napi::Error::from_reason(format!("invalid AnnotatedLlmRequest: {e}")))?; + let orig: CoreLlmRequest = serde_json::from_value(original) + .map_err(|e| napi::Error::from_reason(format!("invalid LlmRequest: {e}")))?; + let result = self + .inner_codec + .encode(&ann, &orig) + .map_err(|e| napi::Error::from_reason(e.to_string()))?; + serde_json::to_value(&result).map_err(|e| napi::Error::from_reason(e.to_string())) + } + + /// Decode a raw LLM response into structured form. + #[napi(js_name = "decodeResponse")] + pub fn decode_response(&self, response: Json) -> napi::Result { + let annotated = self + .inner_response_codec + .decode_response(&response) + .map_err(|e| napi::Error::from_reason(e.to_string()))?; + serde_json::to_value(&annotated).map_err(|e| napi::Error::from_reason(e.to_string())) + } +} + /// Built-in codec for the Anthropic Messages API. /// /// Implements both request codec (decode/encode) and response codec diff --git a/crates/node/tests/types_tests.mjs b/crates/node/tests/types_tests.mjs index 45db9ad25..f4812c16d 100644 --- a/crates/node/tests/types_tests.mjs +++ b/crates/node/tests/types_tests.mjs @@ -23,6 +23,7 @@ describe('Type constants', () => { assert.equal(typeof lib.OpenAIChatCodec, 'function'); assert.equal(typeof lib.OpenAIResponsesCodec, 'function'); assert.equal(typeof lib.AnthropicMessagesCodec, 'function'); + assert.equal(typeof lib.GeminiGenerateContentCodec, 'function'); }); it('scope type enum values', () => { @@ -79,3 +80,154 @@ describe('ScopeStack', () => { assert.ok(stack instanceof ScopeStack); }); }); + +// =========================================================================== +// GeminiGenerateContentCodec +// =========================================================================== + +describe('GeminiGenerateContentCodec', () => { + const { GeminiGenerateContentCodec } = lib; + + it('instantiates', () => { + const codec = new GeminiGenerateContentCodec(); + assert.ok(codec instanceof GeminiGenerateContentCodec); + }); + + it('decode returns an AnnotatedLLMRequest with messages and no params', () => { + const codec = new GeminiGenerateContentCodec(); + const req = { + headers: {}, + content: { + contents: [ + { role: 'user', parts: [{ text: 'hello' }] }, + { role: 'model', parts: [{ text: 'hi' }] }, + ], + systemInstruction: { parts: [{ text: 'Be helpful.' }] }, + }, + }; + const annotated = codec.decode(req); + const msgs = annotated.messages; + assert.equal(msgs.length, 3, 'system + user + model = 3 messages'); + assert.equal(msgs[0].role, 'system'); + assert.equal(msgs[1].role, 'user'); + assert.equal(msgs[2].role, 'assistant', 'model role must normalize to assistant'); + }); + + it('decode captures generationConfig into params', () => { + const codec = new GeminiGenerateContentCodec(); + const req = { + headers: {}, + content: { + contents: [{ role: 'user', parts: [{ text: 'hi' }] }], + generationConfig: { temperature: 0.5, maxOutputTokens: 256 }, + }, + }; + const annotated = codec.decode(req); + assert.ok(annotated.params !== null && annotated.params !== undefined, 'params must be set'); + assert.ok(Math.abs(annotated.params.temperature - 0.5) < 1e-6); + // Rust serializes GenerationParams fields in snake_case (max_tokens, not maxTokens) + assert.equal(annotated.params.max_tokens, 256); + }); + + it('encode round-trips extra fields', () => { + const codec = new GeminiGenerateContentCodec(); + const req = { + headers: {}, + content: { + contents: [{ role: 'user', parts: [{ text: 'hi' }] }], + safetySettings: [{ category: 'HARM_CATEGORY_HATE_SPEECH', threshold: 'BLOCK_NONE' }], + }, + }; + const annotated = codec.decode(req); + const reEncoded = codec.encode(annotated, req); + assert.ok( + Array.isArray(reEncoded.content.safetySettings), + 'safetySettings must survive encode round-trip', + ); + }); + + it('decodeResponse extracts text and finish reason', () => { + const codec = new GeminiGenerateContentCodec(); + const raw = { + candidates: [{ + content: { role: 'model', parts: [{ text: 'Hello!' }] }, + finishReason: 'STOP', + index: 0, + }], + usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 2, totalTokenCount: 7 }, + modelVersion: 'gemini-2.0-flash', + }; + const resp = codec.decodeResponse(raw); + // message is a plain string (MessageContent::Text serializes to a string, not {text: ...}) + assert.equal(resp.message, 'Hello!'); + // finish_reason is snake_case; value matches FinishReason::Complete serialized as "complete" + assert.equal(resp.finish_reason, 'complete'); + assert.equal(resp.model, 'gemini-2.0-flash'); + // usage fields are also snake_case + assert.equal(resp.usage?.prompt_tokens, 5); + }); + + it('decodeResponse maps functionCall id correctly', () => { + const codec = new GeminiGenerateContentCodec(); + const raw = { + candidates: [{ + content: { + role: 'model', + parts: [{ functionCall: { id: 'call_xyz', name: 'my_fn', args: { x: 1 } } }], + }, + finishReason: 'STOP', + index: 0, + }], + usageMetadata: {}, + }; + const resp = codec.decodeResponse(raw); + // tool_calls is snake_case + assert.ok(Array.isArray(resp.tool_calls), 'tool_calls must be an array'); + assert.equal(resp.tool_calls[0].id, 'call_xyz', 'id must come from functionCall.id, not the function name'); + assert.equal(resp.tool_calls[0].name, 'my_fn'); + // Sanity: confirm the id is NOT the function name + assert.notEqual(resp.tool_calls[0].id, 'my_fn', 'id must not be the function name'); + }); + + it('decode throws for malformed Gemini requests', () => { + const codec = new GeminiGenerateContentCodec(); + assert.throws( + () => codec.decode({ headers: {}, content: { contents: 'not an array' } }), + /contents must be an array/, + ); + }); + + it('encode throws for malformed annotated requests and codec failures', () => { + const codec = new GeminiGenerateContentCodec(); + const original = { + headers: {}, + content: { + contents: [{ role: 'user', parts: [{ text: 'hi' }] }], + }, + }; + assert.throws( + () => codec.encode({ messages: 'not an array' }, original), + /invalid AnnotatedLlmRequest/, + ); + + const annotated = codec.decode(original); + annotated.messages.push({ role: 'developer', content: 'unsupported role' }); + assert.throws( + () => codec.encode(annotated, original), + /no Gemini equivalent/, + ); + }); + + it('decodeResponse throws for malformed Gemini responses', () => { + const codec = new GeminiGenerateContentCodec(); + assert.throws( + () => codec.decodeResponse({ + candidates: [{ + content: { role: 'model', parts: [{ text: 42 }] }, + finishReason: 'STOP', + }], + }), + /parts.*text must be a string/, + ); + }); +}); diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 4c8528582..4460e107b 100644 --- a/crates/pii-redaction/README.md +++ b/crates/pii-redaction/README.md @@ -33,7 +33,7 @@ NeMo Relay PII Redaction allows you to: - Use built-in detector presets as first-party detectors for common PII, structured secrets, and cloud credentials. - Handle codec-aware LLMs with overlay support for `openai_chat`, - `openai_responses`, and `anthropic_messages`. + `openai_responses`, `anthropic_messages`, and `gemini_generate_content`. - Remove conversational trajectory content while preserving event structure, tool-call identity, model attribution, routing, usage, and cost analytics. - Use the `local_model` config contract and provider registration surface for diff --git a/crates/pii-redaction/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index 8852f8543..bc52cf8d9 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -304,6 +304,9 @@ impl CompiledBuiltinBackend { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) => { Some(ProviderSurface::AnthropicMessages) } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) => { + Some(ProviderSurface::GeminiGenerateContent) + } LlmCodecIdentity::Runtime(_) | LlmCodecIdentity::Opaque => None, } } @@ -396,7 +399,19 @@ impl CompiledBuiltinBackend { .get("choices") .and_then(Json::as_array) .is_some_and(|choices| choices.len() > 1) - && self.targets_normalized_openai_chat_choice() + && self.targets_normalized_single_projected_response() + { + return None; + } + // Gemini responses with multiple candidates: the normalized layer only projects + // candidate[0], so candidate[1+] would survive in the raw payload unredacted. + // Fail closed identically to the OpenAI Chat multi-choice guard. + if surface == ProviderSurface::GeminiGenerateContent + && payload + .get("candidates") + .and_then(Json::as_array) + .is_some_and(|candidates| candidates.len() > 1) + && self.targets_normalized_single_projected_response() { return None; } @@ -446,7 +461,7 @@ impl CompiledBuiltinBackend { }) } - fn targets_normalized_openai_chat_choice(&self) -> bool { + fn targets_normalized_single_projected_response(&self) -> bool { self.target_paths.iter().any(|path| { json_pointer_segments(path) .and_then(|segments| segments.into_iter().next()) diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 1e21dfe43..38e581d15 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -7,7 +7,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use nemo_relay::codec::resolve::supported_codec_names; +use nemo_relay::codec::resolve::{ProviderSurface, supported_codec_names}; use nemo_relay::plugin::{ ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError, PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior, @@ -274,7 +274,7 @@ nemo_relay::editor_config! { codec => { label: "codec", kind: Enum, - values: ["openai_chat", "openai_responses", "anthropic_messages"], + values: ["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"], optional: true, }, profiles => { label: "profiles", kind: List, list: &PII_REDACTION_PROFILE_LIST_ITEM }, @@ -498,11 +498,8 @@ fn custom_mark_payload_policy_schema( #[cfg(feature = "schema")] fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema { - string_enum_schema( - generator, - &["openai_chat", "openai_responses", "anthropic_messages"], - None, - ) + let codec_names = supported_codec_names(); + string_enum_schema(generator, &codec_names, None) } #[cfg(feature = "schema")] @@ -1130,14 +1127,15 @@ fn validate_codec_requirements( return; }; - if !supported_codec_names().contains(&codec) { + if ProviderSurface::from_codec_name(codec).is_none() { + let supported = supported_codec_names(); push_policy_diag( diagnostics, policy.unsupported_value, "pii_redaction.unsupported_value", Some(PII_REDACTION_PLUGIN_KIND.to_string()), Some("codec".to_string()), - "codec must be 'openai_chat', 'openai_responses', or 'anthropic_messages'".to_string(), + format!("codec must be one of: {}", supported.join(", ")), ); } } diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index e74da6e26..dc896053d 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -12,6 +12,7 @@ pub(crate) enum BuiltinCodecName { OpenAIChat, OpenAIResponses, AnthropicMessages, + GeminiGenerateContent, } impl BuiltinCodecName { @@ -20,6 +21,7 @@ impl BuiltinCodecName { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, } } @@ -32,10 +34,145 @@ impl BuiltinCodecName { Self::OpenAIChat => overlay_openai_chat_response(payload, annotated), Self::OpenAIResponses => overlay_openai_responses_response(payload, annotated), Self::AnthropicMessages => overlay_anthropic_response(payload, annotated), + Self::GeminiGenerateContent => overlay_gemini_response(payload, annotated), } } } +fn gemini_message_parts_for_overlay(message: Option<&MessageContent>) -> Option> { + let MessageContent::Parts(parts) = message? else { + return None; + }; + Some( + parts + .iter() + .filter_map(|part| match part { + ContentPart::Text { text, extra } => { + let mut obj = extra.clone(); + obj.insert("text".into(), Json::String(text.clone())); + Some(Json::Object(obj)) + } + ContentPart::ProviderNative { + provider, value, .. + } if provider == "gemini" => Some(value.clone()), + ContentPart::ImageUrl { .. } + | ContentPart::Image { .. } + | ContentPart::Audio { .. } + | ContentPart::File { .. } + | ContentPart::Refusal { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } + | ContentPart::ProviderNative { .. } => None, + }) + .collect(), + ) +} + +fn overlay_gemini_response(mut payload: Json, annotated: &AnnotatedLlmResponse) -> Json { + let Some(root) = payload.as_object_mut() else { + return payload; + }; + + set_optional_string_field(root, "responseId", annotated.id.as_deref()); + set_optional_string_field(root, "modelVersion", annotated.model.as_deref()); + + let Some(candidate) = root + .get_mut("candidates") + .and_then(Json::as_array_mut) + .and_then(|arr| arr.first_mut()) + .and_then(Json::as_object_mut) + else { + return payload; + }; + + // finishReason is NOT overlaid here. The normalized FinishReason::ToolUse can + // originate from STOP + functionCall parts, so re-emitting it as TOOL_CODE would + // corrupt a response that the API legitimately sent as STOP. The raw provider + // value is already correct and must be preserved. + + let Some(parts) = candidate + .get_mut("content") + .and_then(Json::as_object_mut) + .and_then(|c| c.get_mut("parts")) + .and_then(Json::as_array_mut) + else { + return payload; + }; + + if let Some(message_parts) = gemini_message_parts_for_overlay(annotated.message.as_ref()) { + let mut sanitized = message_parts.into_iter(); + parts.retain_mut(|part| { + let is_thought = part.get("thought").and_then(Json::as_bool) == Some(true); + let is_tool_call = part.get("functionCall").is_some(); + if is_thought || is_tool_call { + return true; + } + let Some(next) = sanitized.next() else { + return false; + }; + *part = next; + true + }); + parts.extend(sanitized); + } else { + // Overlay fallback text into the first visible text part. The normalized + // text may contain embedded newlines, so splitting would confuse text + // content with Gemini part boundaries. + let message_text = annotated_message_text(annotated.message.as_ref()); + let mut wrote_text = false; + parts.retain_mut(|part| { + let is_thought = part.get("thought").and_then(Json::as_bool) == Some(true); + if part.get("text").is_none() || is_thought { + return true; + } + let Some(p) = part.as_object_mut() else { + return false; + }; + if wrote_text { + return false; + } + let Some(text) = message_text.as_deref() else { + return false; + }; + set_optional_string_field(p, "text", Some(text)); + wrote_text = true; + true + }); + } + + // Overlay functionCall parts. + if let Some(tool_calls) = annotated.tool_calls.as_deref() { + let mut sanitized = tool_calls.iter(); + parts.retain_mut(|part| { + let Some(fc) = part + .as_object_mut() + .and_then(|p| p.get_mut("functionCall")) + .and_then(Json::as_object_mut) + else { + return true; // not a functionCall part — keep + }; + let Some(sc) = sanitized.next() else { + return false; // no sanitized call left — remove this part + }; + if fc.contains_key("id") { + set_optional_string_field(fc, "id", Some(sc.id.as_str())); + } + set_optional_string_field(fc, "name", Some(sc.name.as_str())); + fc.insert("args".into(), sc.arguments.clone()); + true + }); + } else { + // No tool calls in sanitized response: remove all functionCall parts. + parts.retain(|part| { + part.as_object() + .map(|p| !p.contains_key("functionCall")) + .unwrap_or(true) + }); + } + + payload +} + fn overlay_openai_chat_response(mut payload: Json, annotated: &AnnotatedLlmResponse) -> Json { let Some(root) = payload.as_object_mut() else { return payload; diff --git a/crates/pii-redaction/tests/coverage/overlay_tests.rs b/crates/pii-redaction/tests/coverage/overlay_tests.rs index 00e877c21..621eaa5f7 100644 --- a/crates/pii-redaction/tests/coverage/overlay_tests.rs +++ b/crates/pii-redaction/tests/coverage/overlay_tests.rs @@ -133,3 +133,310 @@ fn anthropic_overlay_preserves_full_multiline_text_in_single_text_block() { assert_eq!(blocks[0]["text"], json!("line one\nline two")); } + +fn gemini_annotated( + message: Option<&str>, + tool_calls: Option>, + id: Option<&str>, + model: Option<&str>, +) -> AnnotatedLlmResponse { + AnnotatedLlmResponse { + id: id.map(String::from), + model: model.map(String::from), + message: message.map(|t| nemo_relay::codec::request::MessageContent::Text(t.into())), + tool_calls, + finish_reason: None, + usage: None, + optimization_summary: None, + api_specific: None, + extra: Default::default(), + } +} + +#[test] +fn gemini_overlay_redacts_candidate_text() { + let payload = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "raw secret text"}]}, + "finishReason": "STOP", + "index": 0 + }], + "modelVersion": "gemini-2.0-flash" + }); + + let annotated = gemini_annotated(Some("[REDACTED]"), None, None, None); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + + let text = result["candidates"][0]["content"]["parts"][0]["text"] + .as_str() + .unwrap(); + assert_eq!( + text, "[REDACTED]", + "Gemini overlay must redact candidate text" + ); +} + +#[test] +fn gemini_overlay_preserves_embedded_newline_text_and_thought_parts() { + let payload = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "raw line one\nraw line two"}, + {"text": "raw second part"}, + {"text": "", "thought": true, "thoughtSignature": "sig-THOUGHT"} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = gemini_annotated(Some("[REDACTED]\nkept together"), None, None, None); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + let parts = result["candidates"][0]["content"]["parts"] + .as_array() + .expect("Gemini parts array"); + + assert_eq!(parts.len(), 2); + assert_eq!(parts[0]["text"], json!("[REDACTED]\nkept together")); + assert_eq!(parts[1]["thought"], json!(true)); + assert_eq!(parts[1]["thoughtSignature"], json!("sig-THOUGHT")); +} + +#[test] +fn gemini_overlay_does_not_add_absent_response_id_or_model_version() { + let payload = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hi"}]}, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = gemini_annotated(Some("hi"), None, None, None); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + assert!(result.get("responseId").is_none()); + assert!(result.get("modelVersion").is_none()); +} + +#[test] +fn gemini_overlay_redacts_provider_native_candidate_part() { + let payload = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "ran code"}, + {"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "sk-code-secret"}} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Parts(vec![ + ContentPart::Text { + text: "ran code".into(), + extra: Default::default(), + }, + ContentPart::ProviderNative { + provider: "gemini".into(), + kind: "codeExecutionResult".into(), + value: json!({ + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "[REDACTED]" + } + }), + }, + ])), + ..Default::default() + }; + + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + assert_eq!( + result["candidates"][0]["content"]["parts"][1]["codeExecutionResult"]["output"], + json!("[REDACTED]"), + "Gemini overlay must write sanitized provider-native response parts back to raw payload" + ); +} + +#[test] +fn gemini_overlay_updates_tool_call_args() { + let payload = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "search", "id": "c1", "args": {"secret": "raw"}}} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = gemini_annotated( + None, + Some(vec![tool_call( + "c1", + "search", + json!({"secret": "[REDACTED]"}), + )]), + None, + None, + ); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + + let args = &result["candidates"][0]["content"]["parts"][0]["functionCall"]["args"]; + assert_eq!( + args["secret"], + json!("[REDACTED]"), + "Gemini overlay must redact functionCall args" + ); +} + +#[test] +fn gemini_overlay_does_not_synthesize_missing_function_call_id() { + let payload = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "search", "args": {"secret": "raw"}}} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = gemini_annotated( + None, + Some(vec![tool_call( + "search", + "search", + json!({"secret": "[REDACTED]"}), + )]), + None, + None, + ); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + let fc = &result["candidates"][0]["content"]["parts"][0]["functionCall"]; + + assert!(fc.get("id").is_none()); + assert_eq!(fc["name"], json!("search")); + assert_eq!(fc["args"]["secret"], json!("[REDACTED]")); +} + +#[test] +fn gemini_overlay_removes_extra_function_call_parts() { + let payload = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "one", "id": "c1", "args": {"secret": "raw-1"}}}, + {"functionCall": {"name": "two", "id": "c2", "args": {"secret": "raw-2"}}}, + {"text": "", "thought": true, "thoughtSignature": "sig-KEEP"} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = gemini_annotated( + None, + Some(vec![tool_call( + "c1", + "one", + json!({"secret": "[REDACTED]"}), + )]), + None, + None, + ); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + let parts = result["candidates"][0]["content"]["parts"] + .as_array() + .expect("Gemini parts array"); + + assert_eq!(parts.len(), 2); + assert_eq!(parts[0]["functionCall"]["id"], json!("c1")); + assert_eq!( + parts[0]["functionCall"]["args"]["secret"], + json!("[REDACTED]") + ); + assert_eq!(parts[1]["thoughtSignature"], json!("sig-KEEP")); +} + +#[test] +fn gemini_overlay_updates_response_id_and_model_version() { + let payload = json!({ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "hi"}]}, + "finishReason": "STOP", + "index": 0 + }], + "responseId": "resp-old", + "modelVersion": "gemini-old" + }); + + // Annotated view carries the sanitizer-approved id/model. + let annotated = gemini_annotated(Some("hi"), None, Some("resp-abc"), Some("gemini-2.0-flash")); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + + assert_eq!( + result["responseId"], + json!("resp-abc"), + "overlay must write annotated.id to responseId" + ); + assert_eq!( + result["modelVersion"], + json!("gemini-2.0-flash"), + "overlay must write annotated.model to modelVersion" + ); +} + +#[test] +fn gemini_overlay_does_not_overwrite_finish_reason() { + // A STOP response with a functionCall part: normalized finish_reason is ToolUse, + // but the raw finishReason in the payload is STOP and must not be overwritten. + let payload = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "fn", "id": "c1", "args": {}}}] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let annotated = AnnotatedLlmResponse { + finish_reason: Some(nemo_relay::codec::response::FinishReason::ToolUse), + tool_calls: Some(vec![tool_call("c1", "fn", json!({}))]), + ..Default::default() + }; + + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); + + assert_eq!( + result["candidates"][0]["finishReason"].as_str(), + Some("STOP"), + "Gemini overlay must not overwrite native finishReason with the derived ToolUse value" + ); +} diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index 2e80f1172..81276be95 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -1794,6 +1794,32 @@ fn validate_allows_llm_surfaces_without_codec() { assert!(report.diagnostics.is_empty(), "{report:?}"); } +#[test] +fn validate_rejects_ambiguous_gemini_codec_name() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + + let report = validate_plugin_config(&plugin_config(json!({ + "mode": "builtin", + "codec": "gemini", + "input": true, + "output": true, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/messages/0/content"] + } + }))); + + assert!(report.diagnostics.iter().any(|diag| { + diag.field.as_deref() == Some("codec") + && diag.message.contains("gemini_generate_content") + && !diag.message.contains("gemini,") + && !diag.message.ends_with("gemini") + })); +} + #[test] fn validate_rejects_regex_replace_without_pattern() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -4347,6 +4373,338 @@ async fn builtin_backend_removes_targeted_message_names_and_ignores_missing_norm clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn builtin_backend_sanitizes_gemini_generate_content_via_codec() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/messages/0/content"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-generate-content"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "sk-gemini-secret"}]}], + }), + }; + + let _handle = llm_call( + LlmCallParams::builder() + .name("gemini-generate-content") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input(), + Some(&json!({ + "headers": {}, + "content": { + "contents": [{"role": "user", "parts": [{"text": "[REDACTED]"}]}] + } + })) + ); + + deregister_subscriber("pii-redaction-gemini-generate-content").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_gemini_provider_native_tools_via_codec() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/tools/1/value/googleSearch/apiKey"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-native-tools"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{"name": "lookup"}], + "googleSearch": {"apiKey": "sk-tool-secret"} + }] + }), + }; + + let _handle = llm_call( + LlmCallParams::builder() + .name("gemini_generate_content") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input(), + Some(&json!({ + "headers": {}, + "content": { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{ + "functionDeclarations": [{"name": "lookup"}], + "googleSearch": {"apiKey": "[REDACTED]"} + }] + } + })) + ); + + deregister_subscriber("pii-redaction-gemini-native-tools").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_gemini_provider_native_request_content_via_codec() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/messages/0/content/1/value/fileData/fileUri"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-native-request-content"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{ + "role": "user", + "parts": [ + {"text": "summarize this file"}, + { + "fileData": { + "mimeType": "text/plain", + "fileUri": "https://example.test/sk-file-secret" + } + } + ] + }] + }), + }; + + let _handle = llm_call( + LlmCallParams::builder() + .name("gemini_generate_content") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input().unwrap()["content"]["contents"][0]["parts"][1]["fileData"]["fileUri"], + json!("https://example.test/[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("sk-file-secret") + ); + + deregister_subscriber("pii-redaction-gemini-native-request-content").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_gemini_function_response_nested_parts_via_codec() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/messages/2/content/1/value/inlineData/data"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-function-response-parts"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [ + {"role": "user", "parts": [{"text": "show my ordered instrument"}]}, + {"role": "model", "parts": [{ + "functionCall": {"id": "call_img", "name": "get_image", "args": {}}, + "thoughtSignature": "sig_CALL==" + }]}, + {"role": "user", "parts": [{ + "functionResponse": { + "id": "call_img", + "name": "get_image", + "response": {"image_ref": {"$ref": "instrument.jpg"}}, + "parts": [{ + "inlineData": { + "displayName": "instrument.jpg", + "mimeType": "image/jpeg", + "data": "sk-image-secret" + } + }] + } + }]} + ] + }), + }; + + let _handle = llm_call( + LlmCallParams::builder() + .name("gemini_generate_content") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input().unwrap()["content"]["contents"][2]["parts"][0]["functionResponse"] + ["parts"][0]["inlineData"]["data"], + json!("[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("sk-image-secret") + ); + + deregister_subscriber("pii-redaction-gemini-function-response-parts").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_gemini_provider_native_response_content_via_codec() { + use crate::codec::gemini_generate_content::GeminiGenerateContentCodec; + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message/1/value/codeExecutionResult/output"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-native-response-content"); + let response = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [ + {"text": "ran code"}, + {"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "sk-code-secret"}} + ] + }, + "finishReason": "STOP", + "index": 0 + }] + }); + + let _ = llm_call_execute( + LlmCallExecuteParams::builder() + .name("gemini_generate_content") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "run code"}]}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(GeminiGenerateContentCodec)) + .build(), + ) + .await + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert_eq!( + captured_events[1].output().unwrap()["candidates"][0]["content"]["parts"][1]["codeExecutionResult"] + ["output"], + json!("[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[1]) + .unwrap() + .contains("sk-code-secret") + ); + + deregister_subscriber("pii-redaction-gemini-native-response-content").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[tokio::test] async fn builtin_backend_omits_request_and_annotation_for_unsafe_normalized_array_removal() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -5182,3 +5540,156 @@ async fn builtin_backend_sanitizes_openai_responses_output_text_alias_on_stream_ deregister_subscriber("pii-redaction-openai-responses-output-text-stream").unwrap(); clear_plugin_configuration().unwrap(); } + +#[tokio::test] +async fn builtin_backend_omits_multi_candidate_gemini_normalized_response() { + use crate::codec::gemini_generate_content::GeminiGenerateContentCodec; + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-multi-candidate-response"); + let response = json!({ + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "sk-first-secret"}]}, + "finishReason": "STOP", + "index": 0 + }, + { + "content": {"role": "model", "parts": [{"text": "sk-second-secret"}]}, + "finishReason": "STOP", + "index": 1 + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 10} + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("gemini-multi-candidate") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "hello"}]}] + }), + }) + // noop_openai_chat_exec_fn is codec-agnostic: it returns whatever JSON is passed. + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(GeminiGenerateContentCodec)) + .build(), + ) + .await + .unwrap(); + + // The raw response is returned unchanged (fail-closed: no partial redaction). + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + // The end event must carry no annotated output (omitted due to multi-candidate). + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + // Neither secret must appear in the serialized event. + let serialized_end = serde_json::to_string(&captured_events[1]).unwrap(); + assert!(!serialized_end.contains("sk-first-secret")); + assert!(!serialized_end.contains("sk-second-secret")); + + deregister_subscriber("pii-redaction-gemini-multi-candidate-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_raw_multi_candidate_gemini_response() { + use crate::codec::gemini_generate_content::GeminiGenerateContentCodec; + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini_generate_content", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/candidates/1/content/parts/0/text"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-gemini-raw-multi-candidate-response"); + let response = json!({ + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "sk-first-secret"}]}, + "finishReason": "STOP", + "index": 0 + }, + { + "content": {"role": "model", "parts": [{"text": "sk-second-secret"}]}, + "finishReason": "STOP", + "index": 1 + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 10} + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("gemini-raw-multi-candidate") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "hello"}]}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(GeminiGenerateContentCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!( + result, response, + "PII sanitization must not mutate the caller result" + ); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + let output = captured_events[1].output().expect("sanitized output event"); + assert_eq!( + output["candidates"][0]["content"]["parts"][0]["text"], + json!("sk-first-secret"), + "raw targeting must not trigger normalized multi-candidate omission" + ); + assert_eq!( + output["candidates"][1]["content"]["parts"][0]["text"], + json!("[REDACTED]") + ); + assert!(captured_events[1].annotated_response().is_some()); + + deregister_subscriber("pii-redaction-gemini-raw-multi-candidate-response").unwrap(); + clear_plugin_configuration().unwrap(); +} diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index b6bf3da7d..6302e3c14 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -59,6 +59,9 @@ pub enum BuiltinLlmCodec { /// Anthropic Messages. #[serde(rename = "anthropic_messages")] AnthropicMessages, + /// Gemini generateContent. + #[serde(rename = "gemini_generate_content")] + GeminiGenerateContent, } /// Per-call LLM codec identity delivered to native plugins. @@ -3445,6 +3448,7 @@ fn llm_codec_identity_from_native( "openai_chat" => BuiltinLlmCodec::OpenAiChat, "openai_responses" => BuiltinLlmCodec::OpenAiResponses, "anthropic_messages" => BuiltinLlmCodec::AnthropicMessages, + "gemini_generate_content" => BuiltinLlmCodec::GeminiGenerateContent, _ => { set_last_error(host, &format!("unknown built-in LLM codec ID: {id}")); return Err(NemoRelayStatus::InvalidArg); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 54f801702..6fbc17803 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -49,10 +49,10 @@ use uuid::Uuid; use crate::convert::{json_to_py, opt_py_to_json, opt_py_to_timestamp, py_to_json}; use crate::py_callable; use crate::py_types::{ - PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyLLMAttributes, PyLLMHandle, PyLLMRequest, - PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, - PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyThreadScopeStackBinding, - PyToolAttributes, PyToolHandle, + PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyGeminiGenerateContentCodec, + PyLLMAttributes, PyLLMHandle, PyLLMRequest, PyLlmStream, PyOpenAIChatCodec, + PyOpenAIResponsesCodec, PyPropagationContext, PyScopeAttributes, PyScopeHandle, PyScopeStack, + PyScopeType, PyThreadScopeStackBinding, PyToolAttributes, PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -258,6 +258,9 @@ fn py_llm_response_codec( if let Ok(builtin) = c.extract::>() { return Some(builtin.inner_response_codec.clone()); } + if let Ok(builtin) = c.extract::>() { + return Some(builtin.inner_response_codec.clone()); + } // Fall back to wrapping the Python object as a custom response codec Some(Arc::new(py_callable::PyLlmResponseCodecWrapper { py_codec: c.clone().unbind(), @@ -279,6 +282,9 @@ fn py_llm_codec(codec: Option<&Bound<'_, PyAny>>) -> Option> { if let Ok(builtin) = codec.extract::>() { return Some(builtin.inner_codec.clone()); } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } Some(Arc::new(py_callable::PyLlmCodecWrapper { py_codec: codec.clone().unbind(), })) diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index 1b3a22875..a7a56238d 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -1015,3 +1015,72 @@ impl PyAnthropicMessagesCodec { "" } } + +/// Built-in codec for the Gemini generateContent API. +/// +/// Implements both ``LlmCodec`` (decode/encode for requests) and +/// ``LlmResponseCodec`` (decode_response for responses). +/// +/// Example: +/// ```python +/// from nemo_relay.codecs import GeminiGenerateContentCodec +/// codec = GeminiGenerateContentCodec() +/// annotated_req = codec.decode(request) +/// annotated_resp = codec.decode_response(response) +/// ``` +#[pyclass(name = "GeminiGenerateContentCodec")] +pub struct PyGeminiGenerateContentCodec { + pub(crate) inner_codec: Arc, + pub(crate) inner_response_codec: Arc, +} + +#[pymethods] +impl PyGeminiGenerateContentCodec { + #[new] + pub(crate) fn new() -> Self { + Self { + inner_codec: Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), + inner_response_codec: Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), + } + } + + /// Parse an opaque ``LlmRequest`` into a structured ``AnnotatedLLMRequest``. + pub(crate) fn decode(&self, request: &PyLLMRequest) -> PyResult { + self.inner_codec + .decode(&request.inner) + .map(|r| PyAnnotatedLLMRequest { inner: r }) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } + + /// Merge structured changes back into the opaque request. + pub(crate) fn encode( + &self, + annotated: &PyAnnotatedLLMRequest, + original: &PyLLMRequest, + ) -> PyResult { + self.inner_codec + .encode(&annotated.inner, &original.inner) + .map(|r| PyLLMRequest { inner: r }) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } + + /// Parse a raw JSON response into a structured ``AnnotatedLLMResponse``. + pub(crate) fn decode_response( + &self, + response: &Bound<'_, PyAny>, + ) -> PyResult { + let json = py_to_json(response)?; + self.inner_response_codec + .decode_response(&json) + .map(|r| PyAnnotatedLLMResponse { inner: r }) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } + + pub(crate) fn __repr__(&self) -> &'static str { + "" + } +} diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index c1679f4ef..f559c39c8 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -186,6 +186,7 @@ fn register_codec_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/crates/types/src/codec/response.rs b/crates/types/src/codec/response.rs index 731617966..8b3a6abf8 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -6,7 +6,8 @@ //! This module defines [`AnnotatedLlmResponse`] and its supporting types //! for structured, API-agnostic access to LLM response data. -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde_json::Map; use crate::Json; @@ -190,11 +191,16 @@ fn default_cost_currency() -> String { /// Normalized reason why the model stopped generating. /// /// Maps from provider-specific stop reasons: -/// - **Complete**: OpenAI Chat `"stop"`, Anthropic `"end_turn"`, Responses `"completed"` -/// - **Length**: OpenAI Chat `"length"`, Anthropic `"max_tokens"`, Responses incomplete+max_output_tokens -/// - **ToolUse**: OpenAI Chat `"tool_calls"`, Anthropic `"tool_use"` -/// - **ContentFilter**: OpenAI Chat `"content_filter"`, Responses incomplete+content_filter +/// - **Complete**: OpenAI Chat `"stop"`, Anthropic `"end_turn"`, Responses `"completed"`, +/// Gemini generateContent `"STOP"` (without function-call parts) +/// - **Length**: OpenAI Chat `"length"`, Anthropic `"max_tokens"`, Responses incomplete+max_output_tokens, +/// Gemini generateContent `"MAX_TOKENS"` +/// - **ToolUse**: OpenAI Chat `"tool_calls"`, Anthropic `"tool_use"`, +/// Gemini generateContent `"TOOL_CODE"` (unconditional) or `"STOP"` when function-call parts are present +/// - **ContentFilter**: OpenAI Chat `"content_filter"`, Responses incomplete+content_filter, +/// Gemini generateContent `"SAFETY"` / `"RECITATION"` / `"BLOCKLIST"` / `"PROHIBITED_CONTENT"` and other policy codes /// - **Unknown**: Forward-compatible catch-all for unrecognized reasons +/// (e.g. Gemini generateContent `"MALFORMED_FUNCTION_CALL"`, `"UNEXPECTED_TOOL_CALL"`) #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FinishReason { @@ -229,7 +235,7 @@ impl FinishReason { /// Unlike the request-side `ToolCall` (which stores arguments as a JSON /// string per OpenAI convention), response tool calls store arguments as /// parsed [`Json`]. Codecs parse OpenAI's string arguments during decode; -/// Anthropic's `input` is already parsed JSON. +/// Anthropic's `input` and Gemini generateContent `args` are already parsed JSON objects. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ResponseToolCall { /// Unique identifier for this tool call. @@ -326,6 +332,32 @@ pub enum ApiSpecificResponse { content_blocks: Option>, }, + /// Gemini generateContent API-specific fields. + #[serde(rename = "gemini_generate_content")] + GeminiGenerateContent { + /// Tokens consumed by the model's internal reasoning (Gemini generateContent thinking). + #[serde(skip_serializing_if = "Option::is_none")] + thoughts_tokens: Option, + /// Candidate-level safety ratings from the Gemini generateContent response. + #[serde(skip_serializing_if = "Option::is_none")] + safety_ratings: Option, + /// Grounding metadata (web search attribution, etc.). + #[serde(skip_serializing_if = "Option::is_none")] + grounding_metadata: Option, + /// Citation metadata for grounded responses. + #[serde(skip_serializing_if = "Option::is_none")] + citation_metadata: Option, + /// Any remaining candidate-level fields not modeled above. + #[serde( + flatten, + default, + skip_serializing_if = "gemini_extra_is_empty", + serialize_with = "serialize_gemini_extra", + deserialize_with = "deserialize_gemini_extra" + )] + extra: serde_json::Map, + }, + /// Custom/unknown API -- catch-all for user-implemented codecs. #[serde(rename = "custom")] Custom { @@ -336,6 +368,31 @@ pub enum ApiSpecificResponse { }, } +fn gemini_extra_is_empty(extra: &Map) -> bool { + extra.keys().all(|key| key == "api") +} + +fn serialize_gemini_extra(extra: &Map, serializer: S) -> Result +where + S: Serializer, +{ + let filtered = extra + .iter() + .filter(|(key, _)| key.as_str() != "api") + .map(|(key, value)| (key.clone(), value.clone())) + .collect::>(); + filtered.serialize(serializer) +} + +fn deserialize_gemini_extra<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let mut extra = Map::::deserialize(deserializer)?; + extra.remove("api"); + Ok(extra) +} + // --------------------------------------------------------------------------- // Helper methods // --------------------------------------------------------------------------- diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index fb15f886f..92133ae4a 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -154,6 +154,8 @@ pub enum BuiltinLlmCodec { OpenAiResponses, /// Anthropic Messages. AnthropicMessages, + /// Gemini generateContent request and response payloads. + GeminiGenerateContent, } /// Per-call LLM codec identity supplied to worker sanitizers. @@ -2157,6 +2159,9 @@ fn codec_identity_from_proto( Some("anthropic_messages") => { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) } + Some("gemini_generate_content") => { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) + } _ => LlmCodecIdentity::Opaque, }, Some(LlmCodecKind::Runtime) => codec_id @@ -2675,3 +2680,7 @@ fn rustc_version_runtime() -> String { .unwrap_or("unknown") .to_string() } + +#[cfg(test)] +#[path = "../tests/unit/codec_identity_tests.rs"] +mod codec_identity_tests; diff --git a/crates/worker/tests/unit/codec_identity_tests.rs b/crates/worker/tests/unit/codec_identity_tests.rs new file mode 100644 index 000000000..022e6dceb --- /dev/null +++ b/crates/worker/tests/unit/codec_identity_tests.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn test_gemini_codec_identity_decoded_as_builtin_not_opaque() { + use nemo_relay_worker_proto::v1::LlmCodecIdentity as ProtoIdentity; + use nemo_relay_worker_proto::v1::LlmCodecKind; + + let proto = ProtoIdentity { + kind: LlmCodecKind::Builtin as i32, + id: Some("gemini_generate_content".to_string()), + }; + let identity = codec_identity_from_proto(Some(&proto)); + assert_eq!( + identity, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent), + "Gemini generateContent codec id must decode to BuiltIn(GeminiGenerateContent), not Opaque" + ); +} diff --git a/docs/about-nemo-relay/concepts/codecs.mdx b/docs/about-nemo-relay/concepts/codecs.mdx index d68b70aed..cbcc97e89 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -104,7 +104,9 @@ changing a key overlays it. - `messages`, content parts, function calls and results, tools, and tool choice expose portable components when the provider formats have equivalent meaning. - `api_specific` is a tagged, mutable field for modeled Anthropic Messages, - OpenAI Chat Completions, or OpenAI Responses controls. + OpenAI Chat Completions, OpenAI Responses, or Gemini `generateContent` + controls. Response annotations use the same pattern for provider-only + metadata, including Gemini `generateContent` thinking tokens. - Provider-only messages, content blocks, input items, tools, and tool choices use `{ provider, kind, value }`, where `value` is the exact native JSON. - Top-level `extra` preserves unknown or unmodeled fields. @@ -155,9 +157,9 @@ reused without changing public binding APIs. ### Provider Schema Extraction Provider codecs extract provider schema fields. The built-in codecs for OpenAI -Chat Completions, OpenAI Responses, and Anthropic Messages recognize their -request and response payloads and map them into `AnnotatedLlmRequest` or -`AnnotatedLlmResponse`. +Chat Completions, OpenAI Responses, Anthropic Messages, and Gemini +`generateContent` recognize their request and response payloads and map them +into `AnnotatedLlmRequest` or `AnnotatedLlmResponse`. When a managed LLM event already has an annotation, subscribers and exporters consume that annotation. When an event has only raw provider JSON, best-effort @@ -201,12 +203,15 @@ hints let Anthropic Messages requests without a top-level `system` field decode through the Anthropic codec instead of being inferred as OpenAI Chat payloads from their fields alone. -The `nemo-relay` gateway always enables matching request codecs for -`/v1/messages`, `/v1/chat/completions`, and `/v1/responses`, for both buffered -and streaming execution. Count-token, model, probe, and non-LLM passthrough -routes do not enable request codecs. Gateway request intercepts must therefore -edit generation bodies through `annotated_request`; raw `request.content` -remains writable on routes without a request codec. +The `nemo-relay` gateway always enables matching OpenAI Chat, OpenAI Responses, +and Anthropic Messages request codecs for the provider generation routes it +proxies: `/v1/messages`, `/v1/chat/completions`, and `/v1/responses`, for both +buffered and streaming execution. Gemini `generateContent` does not yet have a +gateway route, so `GeminiGenerateContentCodec` is available for direct framework +use rather than gateway auto-selection. Count-token, model, probe, and non-LLM +passthrough routes do not enable request codecs automatically. Gateway request +intercepts must therefore edit generation bodies through `annotated_request`; +raw `request.content` remains writable on routes without a request codec. ### Agent Payload Extraction diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 9b0e82ff8..b06c58ab1 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -346,7 +346,7 @@ In-process Rust and the typed native Rust SDK expose enum variants. The raw native ABI exposes the same information through `codec_kind` and `codec_id`. `codec.kind` is `none` for a call with no codec, `builtin` for -Relay's built-in `openai_chat`, `openai_responses`, and `anthropic_messages` +Relay's built-in `openai_chat`, `openai_responses`, `anthropic_messages`, and `gemini_generate_content` codecs, `runtime` for a named runtime-registered codec, and `opaque` for an active codec without a registered identity. `codec.id` is present only for `builtin` and `runtime`. Do not infer a provider from an opaque request shape. diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index 82167c28e..5a940e8df 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -223,7 +223,7 @@ by `NemoRelayNativeLlmSanitizeRequestContext` or `NemoRelayNativeLlmSanitizeResponseContext`. Each context contains structured codec identity and a borrowed, callback-lifetime codec handle. `codec_kind` is `None`, `BuiltIn`, `Runtime`, or `Opaque`. `codec_id` is present for `BuiltIn` -(one of `openai_chat`, `openai_responses`, or `anthropic_messages`) and +(one of `openai_chat`, `openai_responses`, `anthropic_messages`, or `gemini_generate_content`) and `Runtime`, and null for `None` and `Opaque`. The request handle supports host operations to decode an `LlmRequest` into an diff --git a/docs/configure-plugins/adaptive/response-cache.mdx b/docs/configure-plugins/adaptive/response-cache.mdx index f1cb692cb..9a7faa1f6 100644 --- a/docs/configure-plugins/adaptive/response-cache.mdx +++ b/docs/configure-plugins/adaptive/response-cache.mdx @@ -308,14 +308,15 @@ Only complete, replayable answers are stored: the consumer while being assembled into one aggregate response — the same shape a buffered call stores, so buffered and streaming calls share one keyspace. On a hit the stored answer is replayed as provider-native chunks - (OpenAI Chat deltas, OpenAI Responses lifecycle events, or Anthropic - Messages events), so strict streaming clients parse it like a live stream. + (OpenAI Chat deltas, OpenAI Responses lifecycle events, Anthropic Messages + events, or Gemini `generateContent` response chunks), so strict streaming + clients parse it like a live stream. Provider-terminal token-limited Chat (`finish_reason = "length"`) and Anthropic (`stop_reason = "max_tokens"`) answers can be stored. OpenAI Responses answers with `status = "incomplete"`, streams without a terminal - event, and streams whose content cannot be replayed faithfully (for example - thinking blocks) are never stored. A streaming request whose surface cannot - be inferred runs live, uncached. + event, and streams whose content cannot be replayed faithfully are never + stored. A streaming request whose surface cannot be inferred runs live, + uncached. Streaming publication is write-behind: Relay can report end-of-stream before the backend write finishes. `wait_for_idle()` does not wait for cache writes. diff --git a/docs/configure-plugins/nemo-guardrails/configuration.mdx b/docs/configure-plugins/nemo-guardrails/configuration.mdx index d34013a0d..f4915c252 100644 --- a/docs/configure-plugins/nemo-guardrails/configuration.mdx +++ b/docs/configure-plugins/nemo-guardrails/configuration.mdx @@ -57,7 +57,7 @@ The following table compares remote and local backend support: | Managed `tool_input` | Not supported against the stock Guardrails remote contract | Supported | | Managed `tool_output` | Supported | Supported | | `request_defaults` pass-through | Supported | Not supported | -| Codec support | `openai_chat` | `openai_chat`, `openai_responses`, `anthropic_messages` | +| Codec support | `openai_chat` | `openai_chat`, `openai_responses`, `anthropic_messages`, `gemini_generate_content` | | Runtime availability | Any runtime that includes the remote backend | Runtimes that can start `python3 >= 3.11` with `nemoguardrails==0.22.0` installed | ## Remote Mode @@ -340,6 +340,7 @@ The current built-in local mode supports managed LLM execution with: - `openai_chat` - `openai_responses` - `anthropic_messages` +- `gemini_generate_content` ### Managed Tool Boundary diff --git a/docs/configure-plugins/pii-redaction/configuration.mdx b/docs/configure-plugins/pii-redaction/configuration.mdx index 392999639..e1a77cf12 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -94,8 +94,8 @@ higher-precedence profiles placed before lower-precedence profiles. For a gateway that routes requests to more than one supported provider, omit `codec`. Relay selects the active codec for each managed call, so the same -policy safely applies to OpenAI Chat, OpenAI Responses, and Anthropic Messages -traffic: +policy safely applies to OpenAI Chat, OpenAI Responses, Anthropic Messages, and +Gemini `generateContent` traffic: ```toml [[components]] @@ -134,7 +134,7 @@ The following table compares the available PII redaction backends: | Managed `tool_input` | Supported | Not implemented | | Managed `tool_output` | Supported | Not implemented | | Built-in actions | `remove`, `redact`, `regex_replace`, `hash`, `mask` | N/A | -| Codec support | `openai_chat`, `openai_responses`, `anthropic_messages` | Runtime-specific future implementation | +| Codec support | `openai_chat`, `openai_responses`, `anthropic_messages`, `gemini_generate_content` | Runtime-specific future implementation | | Runtime availability | Any runtime that includes the `nemo-relay-pii-redaction` plugin crate | Runtimes that install a local backend provider | ## Built-in Mode diff --git a/docs/integrate-into-frameworks/provider-codecs.mdx b/docs/integrate-into-frameworks/provider-codecs.mdx index 4012984a2..5df68707e 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -91,6 +91,7 @@ Use the built-in provider codecs when the framework payload already matches a su - `OpenAIChatCodec`: OpenAI Chat Completions-compatible requests and responses. - `OpenAIResponsesCodec`: OpenAI Responses-compatible requests and responses. - `AnthropicMessagesCodec`: Anthropic Messages-compatible requests and responses. +- `GeminiGenerateContentCodec`: Gemini `generateContent`-compatible requests and responses. ## Provider Codec Roles @@ -106,13 +107,17 @@ The built-in provider codecs expose the same core methods: | OpenAI Chat | `nemo_relay.codecs.OpenAIChatCodec` | `OpenAIChatCodec` from `nemo-relay-node` | `decode`, `encode`, `decode_response` / `decodeResponse` | | OpenAI Responses | `nemo_relay.codecs.OpenAIResponsesCodec` | `OpenAIResponsesCodec` from `nemo-relay-node` | `decode`, `encode`, `decode_response` / `decodeResponse` | | Anthropic Messages | `nemo_relay.codecs.AnthropicMessagesCodec` | `AnthropicMessagesCodec` from `nemo-relay-node` | `decode`, `encode`, `decode_response` / `decodeResponse` | +| Gemini `generateContent` | `nemo_relay.codecs.GeminiGenerateContentCodec` | `GeminiGenerateContentCodec` from `nemo-relay-node` | `decode`, `encode`, `decode_response` / `decodeResponse` | Choose the provider codec that matches the payload shape the framework already sends to the provider. Do not translate to a different provider shape only to make the codec fit. -The `nemo-relay` gateway selects these request codecs automatically for -`/v1/messages`, `/v1/chat/completions`, and `/v1/responses`, for both buffered -and streaming calls. Count-token, model, probe, and non-LLM passthrough routes -do not use request codecs. +The `nemo-relay` gateway selects the matching request codec automatically on the +provider generation routes it proxies: `/v1/messages`, `/v1/chat/completions`, +and `/v1/responses`, for both buffered and streaming calls. Gemini +`generateContent` does not yet have a gateway route, so use +`GeminiGenerateContentCodec` directly when a framework sends +`generateContent`-shaped payloads. Count-token, model, probe, and non-LLM +passthrough routes do not use request codecs. ## Example: Add a System Message with a Provider Codec diff --git a/docs/integrate-into-frameworks/provider-response-codecs.mdx b/docs/integrate-into-frameworks/provider-response-codecs.mdx index 205fc9533..8b6a7dd2a 100644 --- a/docs/integrate-into-frameworks/provider-response-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-response-codecs.mdx @@ -361,13 +361,15 @@ Relay does not model are dropped. Built-in codecs normalize provider field names as follows: -| Normalized field | OpenAI Chat | OpenAI Responses | Anthropic Messages | -|---|---|---|---| -| `prompt_tokens` | `prompt_tokens` | `input_tokens` | `input_tokens` | -| `completion_tokens` | `completion_tokens` | `output_tokens` | `output_tokens` | -| `total_tokens` | `total_tokens` | `total_tokens` | computed | -| `cache_read_tokens` | `prompt_tokens_details.cached_tokens` | `input_tokens_details.cached_tokens` | `cache_read_input_tokens` | -| `cache_write_tokens` | — | — | `cache_creation_input_tokens` | +| Normalized Field | OpenAI Chat | OpenAI Responses | Anthropic Messages | Gemini `generateContent` | +|---|---|---|---|---| +| `prompt_tokens` | `prompt_tokens` | `input_tokens` | `input_tokens` | `promptTokenCount` | +| `completion_tokens` | `completion_tokens` | `output_tokens` | `output_tokens` | `candidatesTokenCount` | +| `total_tokens` | `total_tokens` | `total_tokens` | computed | `totalTokenCount` (or computed as prompt + candidates + thinking) | +| `cache_read_tokens` | `prompt_tokens_details.cached_tokens` | `input_tokens_details.cached_tokens` | `cache_read_input_tokens` | `cachedContentTokenCount` | +| `cache_write_tokens` | — | — | `cache_creation_input_tokens` | — | + +Gemini `generateContent` thinking tokens (`thoughtsTokenCount`) are stored in `api_specific.thoughts_tokens` rather than `completion_tokens`, because Google bills them as output tokens but reports them separately. The cost estimate folds them into the effective output-token count so that the pricing table reflects the real billing cost. Built-in codecs preserve only modeled provider-specific usage details under `api_specific`; other usage fields are dropped. For example, OpenAI Responses @@ -440,9 +442,10 @@ The following behaviors are intentional in this release but can change later: The built-in provider codecs also implement response decoding: -- `OpenAIChatCodec` -- `OpenAIResponsesCodec` -- `AnthropicMessagesCodec` +- `OpenAIChatCodec` — OpenAI Chat Completions API +- `OpenAIResponsesCodec` — OpenAI Responses API +- `AnthropicMessagesCodec` — Anthropic Messages API +- `GeminiGenerateContentCodec` — Google Gemini `generateContent` API Choose the codec that matches the actual provider response shape. For example, do not use `OpenAIChatCodec` for an OpenAI Responses API payload only because both came from an OpenAI-compatible provider. diff --git a/docs/integrate-into-frameworks/using-codecs.mdx b/docs/integrate-into-frameworks/using-codecs.mdx index c79eea0b6..5a529a32c 100644 --- a/docs/integrate-into-frameworks/using-codecs.mdx +++ b/docs/integrate-into-frameworks/using-codecs.mdx @@ -38,7 +38,7 @@ Typed value codecs are different from provider codecs: | Codec Type | Purpose | Common Use | |---|---|---| | Typed value codec | Converts application values to and from JSON. | Dataclasses, Pydantic models, TypeScript object shapes, custom framework types. | -| Provider codec | Converts provider-specific LLM requests and responses to annotated NeMo Relay request or response data. | OpenAI Chat, OpenAI Responses, Anthropic Messages, custom provider payloads. | +| Provider codec | Converts provider-specific LLM requests and responses to annotated NeMo Relay request or response data. | OpenAI Chat, OpenAI Responses, Anthropic Messages, Gemini `generateContent`, custom provider payloads. | Use this page for typed value codecs. Use [Provider Codecs](/integrate-into-frameworks/provider-codecs) when request intercepts or request-side middleware need normalized LLM messages, tools, model names, and generation parameters, or when subscribers and exporters need provider response annotations. diff --git a/go/nemo_relay/callbacks_test.go b/go/nemo_relay/callbacks_test.go index 5394b5618..8905e7b7e 100644 --- a/go/nemo_relay/callbacks_test.go +++ b/go/nemo_relay/callbacks_test.go @@ -96,6 +96,7 @@ func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) openAIChat := "openai_chat" openAIResponses := "openai_responses" anthropicMessages := "anthropic_messages" + gemini := "gemini_generate_content" runtimeCodec := "com.example.chat.v1" cases := []codecIdentityTestCase{ @@ -103,6 +104,7 @@ func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) {"openai chat", 1, &openAIChat, LLMCodecBuiltin}, {"openai responses", 1, &openAIResponses, LLMCodecBuiltin}, {"anthropic messages", 1, &anthropicMessages, LLMCodecBuiltin}, + {"gemini_generate_content", 1, &gemini, LLMCodecBuiltin}, {"runtime", 2, &runtimeCodec, LLMCodecRuntime}, {"opaque", 3, nil, LLMCodecOpaque}, {"unknown", 99, nil, LLMCodecOpaque}, diff --git a/go/nemo_relay/coverage_gap_test.go b/go/nemo_relay/coverage_gap_test.go index f53a9d82e..da56733e2 100644 --- a/go/nemo_relay/coverage_gap_test.go +++ b/go/nemo_relay/coverage_gap_test.go @@ -473,7 +473,8 @@ func testWrapperAndCodecFinalizersRun(t *testing.T) { chatCodec := NewOpenAIChatCodec() responsesCodec := NewOpenAIResponsesCodec() anthropicCodec := NewAnthropicMessagesCodec() - if chatCodec == nil || responsesCodec == nil || anthropicCodec == nil { + geminiCodec := NewGeminiGenerateContentCodec() + if chatCodec == nil || responsesCodec == nil || anthropicCodec == nil || geminiCodec == nil { t.Fatal("expected non-nil codec handles") } @@ -484,6 +485,7 @@ func testWrapperAndCodecFinalizersRun(t *testing.T) { chatCodec = nil responsesCodec = nil anthropicCodec = nil + geminiCodec = nil for i := 0; i < 8; i++ { runtime.GC() @@ -491,3 +493,81 @@ func testWrapperAndCodecFinalizersRun(t *testing.T) { time.Sleep(10 * time.Millisecond) } } + +func TestGeminiGenerateContentCodecFunctionCallID(t *testing.T) { + runTestWithScopeStack(t, testGeminiGenerateContentCodecFunctionCallID) +} + +func testGeminiGenerateContentCodecFunctionCallID(t *testing.T) { + // Verify that NewGeminiGenerateContentCodec returns a usable handle and that, + // when processing a Gemini response with an explicit functionCall.id, + // the annotated response carries the actual id (not the function name). + geminiResp := json.RawMessage(`{ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"functionCall": {"id": "call_abc123", "name": "my_fn", "args": {"x": 1}}}] + }, + "finishReason": "STOP", + "index": 0 + }], + "usageMetadata": {} + }`) + + executor := func(_ json.RawMessage) (json.RawMessage, error) { + return geminiResp, nil + } + + capturedEvents, cleanup := registerLlmCodecEventCollector(t) + defer cleanup() + + _, err := LlmCallExecute( + "gemini_fn_id_test", + map[string]interface{}{ + "headers": map[string]interface{}{}, + "content": map[string]interface{}{ + "contents": []interface{}{ + map[string]interface{}{ + "role": "user", + "parts": []interface{}{ + map[string]interface{}{"text": "call my_fn"}, + }, + }, + }, + }, + }, + executor, + WithLLMResponseCodec(NewGeminiGenerateContentCodec()), + ) + if err != nil { + t.Fatalf("LlmCallExecute with GeminiGenerateContentCodec failed: %v", err) + } + + events := capturedEvents() + _, endEvent := requireLlmScopeEvents(t, events) + + var annotated map[string]interface{} + if err := json.Unmarshal(endEvent.AnnotatedResponse(), &annotated); err != nil { + t.Fatalf("AnnotatedResponse not valid JSON: %v", err) + } + + toolCalls, ok := annotated["tool_calls"].([]interface{}) + if !ok || len(toolCalls) == 0 { + t.Fatalf("expected tool_calls in annotated response, got %#v", annotated) + } + tc, ok := toolCalls[0].(map[string]interface{}) + if !ok { + t.Fatalf("expected tool_calls[0] to be a map, got %T", toolCalls[0]) + } + id, _ := tc["id"].(string) + if id != "call_abc123" { + t.Errorf("functionCall id must be 'call_abc123', got %q", id) + } + name, _ := tc["name"].(string) + if name != "my_fn" { + t.Errorf("functionCall name must be 'my_fn', got %q", name) + } + if id == name { + t.Error("id must differ from name — id must not be the function name") + } +} diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 8646c122a..85ddb45fd 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -116,6 +116,7 @@ extern int32_t nemo_relay_llm_stream_call_execute( extern FfiCodecHandle* nemo_relay_openai_chat_codec_new(void); extern FfiCodecHandle* nemo_relay_openai_responses_codec_new(void); extern FfiCodecHandle* nemo_relay_anthropic_messages_codec_new(void); +extern FfiCodecHandle* nemo_relay_gemini_generate_content_codec_new(void); extern void nemo_relay_codec_free(FfiCodecHandle* handle); extern void nemo_relay_set_last_error_message(const char* msg); @@ -922,9 +923,9 @@ func WithLLMCodec(codec CodecFunc) LLMCallOption { // CodecHandle wraps an opaque FFI codec handle that carries both request // codec (decode/encode) and response codec (decode_response) implementations. -// Create via [NewOpenAIChatCodec], [NewOpenAIResponsesCodec], or -// [NewAnthropicMessagesCodec]. The handle is automatically freed when -// garbage collected. +// Create via [NewOpenAIChatCodec], [NewOpenAIResponsesCodec], +// [NewAnthropicMessagesCodec], or [NewGeminiGenerateContentCodec]. The handle is +// automatically freed when garbage collected. type CodecHandle struct { ptr *C.FfiCodecHandle } @@ -977,9 +978,25 @@ func NewAnthropicMessagesCodec() *CodecHandle { return h } +// NewGeminiGenerateContentCodec creates a codec for the Gemini generateContent API. +// +// The returned handle can be passed to [WithLLMCodec] or +// [WithLLMResponseCodec] to enable structured request and response handling for +// Gemini generateContent payloads. +func NewGeminiGenerateContentCodec() *CodecHandle { + h := &CodecHandle{ptr: C.nemo_relay_gemini_generate_content_codec_new()} + runtime.SetFinalizer(h, func(h *CodecHandle) { + if h.ptr != nil { + C.nemo_relay_codec_free(h.ptr) + h.ptr = nil + } + }) + return h +} + // WithLLMResponseCodec sets the response codec for this LLM call. -// Pass a CodecHandle created by [NewOpenAIChatCodec], -// [NewOpenAIResponsesCodec], or [NewAnthropicMessagesCodec]. +// Pass a CodecHandle created by [NewOpenAIChatCodec], [NewOpenAIResponsesCodec], +// [NewAnthropicMessagesCodec], or [NewGeminiGenerateContentCodec]. // The codec handle is kept alive for the duration of the FFI call via // runtime.KeepAlive, so it is safe to pass an inline-constructed handle. func WithLLMResponseCodec(codec *CodecHandle) LLMCallOption { diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index cc025751c..6481de01e 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1157,6 +1157,26 @@ class AnthropicMessagesCodec: """Decode an Anthropic response into a normalized response view.""" ... +class GeminiGenerateContentCodec: + """Built-in codec for Gemini generateContent requests and responses. + + Summary: + Native codec bridge for Gemini generateContent payloads. + """ + + def __init__(self) -> None: + """Create a Gemini generateContent codec.""" + ... + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: + """Decode a Gemini generateContent request into a normalized request view.""" + ... + def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMRequest: + """Encode a normalized request back into Gemini generateContent shape.""" + ... + def decode_response(self, response: _Json) -> AnnotatedLLMResponse: + """Decode a Gemini response into a normalized response view.""" + ... + class AdaptiveRuntime: """Hosted adaptive runtime bridge implemented by the native extension. diff --git a/python/nemo_relay/codecs.py b/python/nemo_relay/codecs.py index 7e9877fb6..98cdb2e32 100644 --- a/python/nemo_relay/codecs.py +++ b/python/nemo_relay/codecs.py @@ -47,6 +47,7 @@ async def impl(request: LLMRequest): from nemo_relay._native import ( AnnotatedLLMRequest, AnthropicMessagesCodec, + GeminiGenerateContentCodec, LLMRequest, OpenAIChatCodec, OpenAIResponsesCodec, @@ -161,6 +162,7 @@ def decode_response(self, response: Json) -> "AnnotatedLLMResponse": __all__ = [ "AnnotatedLLMRequest", "AnthropicMessagesCodec", + "GeminiGenerateContentCodec", "LlmCodec", "LlmResponseCodec", "OpenAIChatCodec", diff --git a/python/nemo_relay/codecs.pyi b/python/nemo_relay/codecs.pyi index 5bb540295..d77ed6168 100644 --- a/python/nemo_relay/codecs.pyi +++ b/python/nemo_relay/codecs.pyi @@ -163,9 +163,48 @@ class AnthropicMessagesCodec: """ ... +class GeminiGenerateContentCodec: + """Built-in codec for Gemini generateContent requests and responses.""" + + def __init__(self) -> None: ... + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: + """Decode a Gemini generateContent request. + + Args: + request: Raw generateContent request payload. + + Returns: + AnnotatedLLMRequest: Normalized request representation. + """ + ... + + def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMRequest: + """Encode a normalized request back into generateContent format. + + Args: + annotated: Normalized request after intercept edits. + original: Original generateContent request. + + Returns: + LLMRequest: Updated generateContent request payload. + """ + ... + + def decode_response(self, response: Json) -> AnnotatedLLMResponse: + """Decode a Gemini generateContent response. + + Args: + response: Raw generateContent response payload. + + Returns: + AnnotatedLLMResponse: Normalized response representation. + """ + ... + __all__ = [ "AnnotatedLLMRequest", "AnthropicMessagesCodec", + "GeminiGenerateContentCodec", "LlmCodec", "LlmResponseCodec", "OpenAIChatCodec", diff --git a/python/nemo_relay/pii_redaction.py b/python/nemo_relay/pii_redaction.py index 53c2b77f8..99dea9047 100644 --- a/python/nemo_relay/pii_redaction.py +++ b/python/nemo_relay/pii_redaction.py @@ -136,7 +136,9 @@ class PiiRedactionConfig: tool_output: bool = True mark: bool = True priority: int = 100 - codec: Literal["openai_chat", "openai_responses", "anthropic_messages"] | str | None = None + codec: Literal["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"] | str | None = ( + None + ) builtin: BuiltinConfig | None = None local: LocalModelConfig | None = None policy: ConfigPolicy = field(default_factory=ConfigPolicy) diff --git a/python/nemo_relay/pii_redaction.pyi b/python/nemo_relay/pii_redaction.pyi index 244f6a3ef..d7719d5e0 100644 --- a/python/nemo_relay/pii_redaction.pyi +++ b/python/nemo_relay/pii_redaction.pyi @@ -58,7 +58,9 @@ class PiiRedactionConfig: tool_output: bool = ... mark: bool = ... priority: int = ... - codec: Literal["openai_chat", "openai_responses", "anthropic_messages"] | str | None = ... + codec: ( + Literal["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"] | str | None + ) = ... builtin: BuiltinConfig | None = ... local: LocalModelConfig | None = ... policy: ConfigPolicy = field(default_factory=ConfigPolicy) diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 7ab918d20..0094ed8e3 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -779,6 +779,7 @@ def response_sanitizer(response: Json, codec_context: LlmSanitizeResponseContext (pb.LLM_CODEC_KIND_BUILTIN, "openai_chat"), (pb.LLM_CODEC_KIND_BUILTIN, "openai_responses"), (pb.LLM_CODEC_KIND_BUILTIN, "anthropic_messages"), + (pb.LLM_CODEC_KIND_BUILTIN, "gemini_generate_content"), (pb.LLM_CODEC_KIND_RUNTIME, "com.example.chat.v1"), (pb.LLM_CODEC_KIND_OPAQUE, None), ]: @@ -830,6 +831,8 @@ def response_sanitizer(response: Json, codec_context: LlmSanitizeResponseContext "response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "anthropic_messages")), ), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "gemini_generate_content"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "gemini_generate_content"))), ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("runtime", "com.example.chat.v1"))), ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("runtime", "com.example.chat.v1"))), ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("opaque"))), diff --git a/python/tests/test_builtin_codecs.py b/python/tests/test_builtin_codecs.py index 2e794a0a5..00458be45 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -4,8 +4,9 @@ """Tests for built-in codec Python classes and LlmResponseCodec protocol. Covers: -- Built-in codec construction (OpenAIChatCodec, OpenAIResponsesCodec, AnthropicMessagesCodec) -- Built-in codec decode/encode/decode_response methods +- Built-in codec construction for OpenAIChatCodec, OpenAIResponsesCodec, + AnthropicMessagesCodec, and GeminiGenerateContentCodec +- Built-in codec decode/encode/decode_response methods for all four providers - LlmResponseCodec protocol - response_codec parameter accepts object (not string) """ @@ -22,7 +23,7 @@ llm, subscribers, ) -from nemo_relay.codecs import AnthropicMessagesCodec, OpenAIChatCodec, OpenAIResponsesCodec +from nemo_relay.codecs import AnthropicMessagesCodec, GeminiGenerateContentCodec, OpenAIChatCodec, OpenAIResponsesCodec # --------------------------------------------------------------------------- # 1. Built-in codec construction @@ -66,6 +67,18 @@ def test_anthropic_messages_codec_has_methods(self): assert hasattr(codec, "encode") assert hasattr(codec, "decode_response") + def test_gemini_codec_constructable(self): + """GeminiGenerateContentCodec() is constructable.""" + codec = GeminiGenerateContentCodec() + assert codec is not None + + def test_gemini_codec_has_methods(self): + """GeminiGenerateContentCodec has decode, encode, decode_response methods.""" + codec = GeminiGenerateContentCodec() + assert hasattr(codec, "decode") + assert hasattr(codec, "encode") + assert hasattr(codec, "decode_response") + # --------------------------------------------------------------------------- # 2. Built-in codec decode/encode round-trip @@ -272,6 +285,119 @@ def test_anthropic_messages_decode_response(self): assert annotated.model == "claude-3-sonnet-20240229" assert annotated.response_text() == "Hello!" + def test_gemini_codec_decode(self): + """GeminiGenerateContentCodec.decode() returns AnnotatedLLMRequest with messages and params.""" + codec = GeminiGenerateContentCodec() + request = LLMRequest( + {}, + { + "model": "gemini-2.0-flash", + "contents": [{"role": "user", "parts": [{"text": "Hi"}]}], + "generationConfig": {"temperature": 0.5, "maxOutputTokens": 256}, + "systemInstruction": {"parts": [{"text": "Be concise."}]}, + }, + ) + annotated = codec.decode(request) + assert isinstance(annotated, AnnotatedLLMRequest) + assert annotated.model == "gemini-2.0-flash" + # System message comes first, then the user message. + assert len(annotated.messages) == 2 + assert annotated.messages[0]["role"] == "system" + assert annotated.messages[1]["role"] == "user" + assert annotated.params is not None + + def test_gemini_codec_encode_round_trip(self): + """GeminiGenerateContentCodec.encode(decode(req), req) is idempotent when nothing changes.""" + codec = GeminiGenerateContentCodec() + original = LLMRequest( + {}, + { + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "generationConfig": {"temperature": 0.7}, + "safetySettings": [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}], + }, + ) + annotated = codec.decode(original) + re_encoded = codec.encode(annotated, original) + assert re_encoded.content == original.content, ( + "encode(decode(req), req) must be idempotent when nothing changes" + ) + + def test_gemini_codec_decode_response_text(self): + """GeminiGenerateContentCodec.decode_response() extracts text and usage from a generateContent response.""" + codec = GeminiGenerateContentCodec() + response = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "Hello from Gemini!"}]}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + "totalTokenCount": 12, + }, + "modelVersion": "gemini-2.0-flash", + } + annotated = codec.decode_response(response) + assert isinstance(annotated, AnnotatedLLMResponse) + assert annotated.response_text() == "Hello from Gemini!" + assert annotated.finish_reason == "complete" + assert annotated.model == "gemini-2.0-flash" + assert annotated.usage is not None + assert annotated.usage["prompt_tokens"] == 8 + + def test_gemini_codec_decode_response_safety_finish_reason(self): + """GeminiGenerateContentCodec maps SAFETY finish reason to 'content_filter', not 'unknown'.""" + codec = GeminiGenerateContentCodec() + response = { + "candidates": [ + { + "content": {"role": "model", "parts": []}, + "finishReason": "SAFETY", + "index": 0, + } + ], + "usageMetadata": {"promptTokenCount": 5}, + } + annotated = codec.decode_response(response) + assert annotated.finish_reason == "content_filter", ( + "SAFETY finish reason must map to content_filter, not unknown" + ) + + def test_gemini_codec_decode_response_function_call(self): + """GeminiGenerateContentCodec.decode_response() extracts functionCall parts as tool_calls.""" + codec = GeminiGenerateContentCodec() + response = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "get_weather", + "id": "call_abc", + "args": {"location": "NYC"}, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 10}, + } + annotated = codec.decode_response(response) + assert annotated.has_tool_calls() is True + tool_calls = annotated.tool_calls + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0]["id"] == "call_abc" + assert tool_calls[0]["name"] == "get_weather" + # --------------------------------------------------------------------------- # 4. LlmResponseCodec protocol @@ -292,6 +418,7 @@ def test_builtin_codecs_satisfy_protocol(self): assert isinstance(OpenAIChatCodec(), LlmResponseCodec) assert isinstance(OpenAIResponsesCodec(), LlmResponseCodec) assert isinstance(AnthropicMessagesCodec(), LlmResponseCodec) + assert isinstance(GeminiGenerateContentCodec(), LlmResponseCodec) # --------------------------------------------------------------------------- @@ -564,14 +691,21 @@ def test_no_builtin_codecs_tuple(self): class TestBuiltinCodecImports: def test_importable_from_codecs_module(self): """Built-in codecs are importable from nemo_relay.codecs.""" - from nemo_relay.codecs import AnthropicMessagesCodec, OpenAIChatCodec, OpenAIResponsesCodec + from nemo_relay.codecs import ( + AnthropicMessagesCodec, + GeminiGenerateContentCodec, + OpenAIChatCodec, + OpenAIResponsesCodec, + ) assert OpenAIChatCodec is not None assert OpenAIResponsesCodec is not None assert AnthropicMessagesCodec is not None + assert GeminiGenerateContentCodec is not None def test_not_reexported_from_top_level(self): """Built-in codecs are not re-exported from nemo_relay.""" assert not hasattr(nemo_relay, "OpenAIChatCodec") assert not hasattr(nemo_relay, "OpenAIResponsesCodec") assert not hasattr(nemo_relay, "AnthropicMessagesCodec") + assert not hasattr(nemo_relay, "GeminiGenerateContentCodec")