From 2c7dda6f6171478140118637667a1ffaf2023839 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 15:16:50 -0700 Subject: [PATCH 01/11] feat(codec): add Gemini generateContent codec Add Gemini generateContent request, response, and streaming codecs across the runtime and language bindings. Wire Gemini into built-in codec resolution, PII redaction overlays, adaptive response cache paths, NeMo Guardrails codec configuration, FFI, Python, Node.js, Go, and docs. Add focused Gemini codec, PII, adaptive cache, binding, and documentation coverage. Signed-off-by: Sahiti Bulusu --- .../adaptive/src/acg/request_surfaces/mod.rs | 11 +- crates/adaptive/src/response_cache/key.rs | 2 +- crates/adaptive/src/response_cache/replay.rs | 5 + .../tests/unit/response_cache/key_tests.rs | 15 + .../tests/unit/response_cache/replay_tests.rs | 28 + crates/core/src/api/runtime/callbacks.rs | 3 + crates/core/src/codec/gemini.rs | 2868 ++++++++++ crates/core/src/codec/mod.rs | 1 + crates/core/src/codec/resolve.rs | 14 +- .../src/plugins/nemo_guardrails/component.rs | 14 +- .../src/plugins/nemo_guardrails/python.rs | 23 + .../tests/fixtures/native_plugin/.gitignore | 5 + crates/core/tests/unit/codec/gemini_tests.rs | 4603 +++++++++++++++++ crates/core/tests/unit/codec/parity_tests.rs | 8 +- crates/core/tests/unit/codec/resolve_tests.rs | 88 +- .../nemo_guardrails/component_tests.rs | 17 +- .../nemo_guardrails/local_python_tests.rs | 35 + crates/ffi/nemo_relay.h | 11 + crates/ffi/src/api/llm.rs | 15 + crates/ffi/tests/unit/types_tests.rs | 6 + crates/node/pii_redaction.d.ts | 2 +- crates/node/plugin.d.ts | 2 +- crates/node/src/types/mod.rs | 57 + crates/node/tests/types_tests.mjs | 101 + crates/pii-redaction/README.md | 2 +- crates/pii-redaction/src/builtin.rs | 26 + crates/pii-redaction/src/component.rs | 14 +- crates/pii-redaction/src/overlay.rs | 148 + .../tests/coverage/overlay_tests.rs | 179 + .../tests/unit/component_tests.rs | 351 ++ crates/plugin/src/lib.rs | 4 + crates/python/src/py_api/mod.rs | 10 +- crates/python/src/py_types/codecs.rs | 65 + crates/python/src/py_types/mod.rs | 1 + crates/types/src/codec/response.rs | 35 +- crates/worker/src/lib.rs | 27 + docs/about-nemo-relay/concepts/codecs.mdx | 10 +- docs/about-nemo-relay/concepts/middleware.mdx | 2 +- .../dynamic-plugins/native-dynamic/about.mdx | 2 +- .../nemo-guardrails/configuration.mdx | 3 +- .../pii-redaction/configuration.mdx | 4 +- .../provider-codecs.mdx | 2 + .../provider-response-codecs.mdx | 23 +- .../using-codecs.mdx | 2 +- go/nemo_relay/callbacks_test.go | 3 + go/nemo_relay/coverage_gap_test.go | 82 +- go/nemo_relay/nemo_relay.go | 27 +- python/nemo_relay/_native.pyi | 20 + python/nemo_relay/codecs.py | 2 + python/nemo_relay/codecs.pyi | 39 + python/nemo_relay/pii_redaction.py | 2 +- python/nemo_relay/pii_redaction.pyi | 2 +- python/tests/plugin/test_worker_sdk.py | 3 + python/tests/test_builtin_codecs.py | 141 +- 54 files changed, 9095 insertions(+), 70 deletions(-) create mode 100644 crates/core/src/codec/gemini.rs create mode 100644 crates/core/tests/fixtures/native_plugin/.gitignore create mode 100644 crates/core/tests/unit/codec/gemini_tests.rs diff --git a/crates/adaptive/src/acg/request_surfaces/mod.rs b/crates/adaptive/src/acg/request_surfaces/mod.rs index 922148fa0..e590dd4a6 100644 --- a/crates/adaptive/src/acg/request_surfaces/mod.rs +++ b/crates/adaptive/src/acg/request_surfaces/mod.rs @@ -43,11 +43,12 @@ 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), + _ => None, } } @@ -91,7 +92,7 @@ pub(crate) fn resolve_request_surface_from_request( request: &LlmRequest, ) -> crate::acg::Result { detect_request_surface(&request.content) - .map(RequestSurface::from_provider_surface) + .and_then(RequestSurface::from_provider_surface) .ok_or_else(|| { crate::acg::AcgError::Internal( "unable to resolve request surface from request shape".to_string(), diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index 36eeca3d1..bc0dc6a65 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -337,7 +337,7 @@ fn lossy_request_shape(surface: ProviderSurface, content: &Json) -> bool { .and_then(Json::as_array) .is_some_and(|blocks| blocks.iter().any(lossy_system_block)) } - ProviderSurface::OpenAIResponses => false, + ProviderSurface::OpenAIResponses | ProviderSurface::Gemini => false, } } diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs index 02a8bf0d4..7f191246f 100644 --- a/crates/adaptive/src/response_cache/replay.rs +++ b/crates/adaptive/src/response_cache/replay.rs @@ -79,6 +79,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::Gemini => vec![aggregate.clone()], }) } diff --git a/crates/adaptive/tests/unit/response_cache/key_tests.rs b/crates/adaptive/tests/unit/response_cache/key_tests.rs index aa03d30f5..781545db5 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -593,6 +593,21 @@ 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", &request); + assert_eq!(effective_codec, Some("gemini")); + assert_ne!( + body, request.content, + "Gemini requests must use the normalized decode when it is lossless" + ); +} + #[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..a87a26c56 100644 --- a/crates/adaptive/tests/unit/response_cache/replay_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/replay_tests.rs @@ -78,6 +78,34 @@ 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"}]}, + "finishReason": "STOP" + }], + "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 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..0c2e6a7fb 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. + Gemini, } impl BuiltinLlmCodec { @@ -170,6 +172,7 @@ impl BuiltinLlmCodec { Self::OpenAiChat => "openai_chat", Self::OpenAiResponses => "openai_responses", Self::AnthropicMessages => "anthropic_messages", + Self::Gemini => "gemini", } } } diff --git a/crates/core/src/codec/gemini.rs b/crates/core/src/codec/gemini.rs new file mode 100644 index 000000000..40da37ce5 --- /dev/null +++ b/crates/core/src/codec/gemini.rs @@ -0,0 +1,2868 @@ +// 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 GeminiCodec; + +pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDescriptor { + surface: ProviderSurface::Gemini, + detect_request: |obj, _hint| obj.contains_key("contents"), + detect_response: detect_gemini_response, + decode_request: |request| GeminiCodec.decode(request), + decode_response: |raw| GeminiCodec.decode_response(raw), + codec_name: "gemini", + request_codec: || std::sync::Arc::new(GeminiCodec), + response_codec: || std::sync::Arc::new(GeminiCodec), + streaming_codec: || Box::new(GeminiStreamingCodec::new()), +}; + +// --------------------------------------------------------------------------- +// Private serde intermediates for response decode +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +struct RawGeminiResponse { + 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 {context}.id must be a non-empty string" + ))), + Some(_) => Err(FlowError::InvalidArgument(format!( + "Gemini {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 {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 {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 {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").unwrap().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::Gemini`, 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 GeminiFunctionResponsePayload { + 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 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(GeminiFunctionResponsePayload { + 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(GeminiFunctionResponsePayload { + 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 {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 GeminiCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini) + } + + fn decode_response(&self, response: &Json) -> Result { + let raw: RawGeminiResponse = 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 cand_extra = candidate.map(|c| c.extra.clone()).unwrap_or_default(); + let has_candidate_data = !cand_extra.is_empty() || thoughts_tokens.is_some(); + if !has_candidate_data { + None + } else { + let mut extra = cand_extra; + let safety_ratings = extra.remove("safetyRatings"); + let grounding_metadata = extra.remove("groundingMetadata"); + let citation_metadata = extra.remove("citationMetadata"); + Some(super::response::ApiSpecificResponse::Gemini { + 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 GeminiCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini) + } + + 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 GeminiMessageAlignment { + prefix_len: usize, + ann_gap_end: usize, + base_gap_end: usize, + ann_gap_len: usize, + base_gap_len: usize, +} + +impl GeminiMessageAlignment { + 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: &GeminiMessageAlignment, + 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: &GeminiMessageAlignment, + 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 = GeminiMessageAlignment::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, GeminiFunctionResponsePayload>, + 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: GeminiFunctionResponsePayload, + 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 [`GeminiCodec`]. +/// +/// Accumulates Gemini server-sent event chunks and assembles a complete response +/// that [`GeminiCodec::decode_response`] can consume. +pub struct GeminiStreamingCodec { + state: std::sync::Arc>, +} + +impl GeminiStreamingCodec { + /// Creates a fresh streaming codec with empty accumulator state. + pub fn new() -> Self { + Self { + state: std::sync::Arc::new(std::sync::Mutex::new(GeminiStreamingState::default())), + } + } +} + +impl Default for GeminiStreamingCodec { + fn default() -> Self { + Self::new() + } +} + +impl super::streaming::StreamingCodec for GeminiStreamingCodec { + 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 GeminiStreamingState { + parts: Vec, + 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 GeminiStreamingState { + 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) + && (text.is_empty() || (!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 let Some(parts) = candidate + .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")?; + // Skip thought parts before inspecting text, matching the non-streaming path. + if part.get("thought").and_then(Json::as_bool) == Some(true) { + 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::Gemini. Later chunks overwrite earlier ones for + // the same key (last-wins), matching the non-streaming behaviour. + if let Some(cand_obj) = candidate.as_object() { + for (k, v) in cand_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(0u32)); + // Merge accumulated candidate-level metadata so that decode_response can + // populate ApiSpecificResponse::Gemini 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_tests.rs"] +mod tests; diff --git a/crates/core/src/codec/mod.rs b/crates/core/src/codec/mod.rs index a92f00826..eb84a0240 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; 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..6bf03af90 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, 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. + Gemini, } /// 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::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 (`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,6 +155,7 @@ 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::Gemini => &gemini::PROVIDER_SURFACE, } } diff --git a/crates/core/src/plugins/nemo_guardrails/component.rs b/crates/core/src/plugins/nemo_guardrails/component.rs index 82a5f0e63..fe58df310 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"], 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", + ], 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..9aaffdf30 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, + Gemini, } impl LocalGuardrailsCodec { @@ -899,6 +900,7 @@ impl LocalGuardrailsCodec { Self::OpenAIChat => ProviderSurface::OpenAIChat, Self::OpenAIResponses => ProviderSurface::OpenAIResponses, Self::AnthropicMessages => ProviderSurface::AnthropicMessages, + Self::Gemini => ProviderSurface::Gemini, } } @@ -907,6 +909,7 @@ impl LocalGuardrailsCodec { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::Gemini => Self::Gemini, } } @@ -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/fixtures/native_plugin/.gitignore b/crates/core/tests/fixtures/native_plugin/.gitignore new file mode 100644 index 000000000..c160acc34 --- /dev/null +++ b/crates/core/tests/fixtures/native_plugin/.gitignore @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Library fixture — Cargo.lock is not committed for library crates. +Cargo.lock diff --git a/crates/core/tests/unit/codec/gemini_tests.rs b/crates/core/tests/unit/codec/gemini_tests.rs new file mode 100644 index 000000000..cef096ed7 --- /dev/null +++ b/crates/core/tests/unit/codec/gemini_tests.rs @@ -0,0 +1,4603 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Unit tests for GeminiCodec 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 = GeminiCodec; + assert_eq!( + LlmCodec::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), + "GeminiCodec must not return Opaque; PII sanitization depends on a known identity" + ); +} + +#[test] +fn test_response_codec_identity_is_gemini_builtin() { + let codec = GeminiCodec; + assert_eq!( + ::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), + "GeminiCodec response codec must not return Opaque" + ); +} + +// =================================================================== +// Response decode tests +// =================================================================== + +#[test] +fn test_decode_response_text() { + let codec = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + + // (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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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")); +} + +// =================================================================== +// Request decode tests +// =================================================================== + +#[test] +fn test_decode_contents_with_system_instruction() { + let codec = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + // 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiStreamingCodec::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(); + + let codec = GeminiCodec; + 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_preserves_text_part_metadata_from_empty_final_chunk() { + let streaming_codec = GeminiStreamingCodec::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 part = &assembled["candidates"][0]["content"]["parts"][0]; + assert_eq!(part["text"].as_str(), Some("answer")); + assert_eq!( + part["thoughtSignature"].as_str(), + Some("sig_STREAM=="), + "streamed text metadata must survive finalization" + ); + + let codec = GeminiCodec; + 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.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 = GeminiStreamingCodec::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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiStreamingCodec::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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 codec = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + // 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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::Gemini 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 = GeminiCodec; + 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::Gemini. + match resp.api_specific { + Some(ApiSpecificResponse::Gemini { + thoughts_tokens, .. + }) => { + assert_eq!( + thoughts_tokens, + Some(20), + "thoughtsTokenCount must be in api_specific" + ); + } + other => panic!("expected ApiSpecificResponse::Gemini, got: {other:?}"), + } +} + +// Roleless contents decode to user messages (Google allows omitting role). +#[test] +fn test_decode_roleless_content_treated_as_user() { + let codec = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + // 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiCodec; + let resp = codec.decode_response(&assembled).unwrap(); + + use super::super::response::ApiSpecificResponse; + match &resp.api_specific { + Some(ApiSpecificResponse::Gemini { + safety_ratings, + grounding_metadata, + .. + }) => { + assert!( + safety_ratings.is_some(), + "safetyRatings from streaming must survive in ApiSpecificResponse::Gemini" + ); + assert!( + grounding_metadata.is_some(), + "groundingMetadata from streaming must survive in ApiSpecificResponse::Gemini" + ); + } + other => panic!("expected ApiSpecificResponse::Gemini, got: {other:?}"), + } +} + +// Present-but-non-string role must error. +#[test] +fn test_decode_non_string_role_errors() { + let codec = GeminiCodec; + 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 = GeminiCodec; + 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" + ); +} + +// functionResponse.response must be an object. +#[test] +fn test_decode_function_response_non_object_response_errors() { + let codec = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiCodec; + 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 = GeminiStreamingCodec::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 = GeminiStreamingCodec::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 = GeminiStreamingCodec::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 = GeminiCodec; + 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 = GeminiCodec; + 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..51675d01e 100644 --- a/crates/core/tests/unit/codec/parity_tests.rs +++ b/crates/core/tests/unit/codec/parity_tests.rs @@ -3,9 +3,11 @@ //! 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 +//! 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-specific semantics (contents array, functionResponse +//! role split, thinking tokens, etc.) are covered in gemini_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. diff --git a/crates/core/tests/unit/codec/resolve_tests.rs b/crates/core/tests/unit/codec/resolve_tests.rs index 70907a933..454c11902 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::Gemini, ] ); } @@ -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::Gemini) + ); + assert_eq!( + detect_request_surface(&json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]})), + Some(ProviderSurface::Gemini) + ); + // 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::Gemini) + ); + assert_eq!( + detect_response_surface(&json!({ + "candidates": [{"content": {"parts": [{"text": "hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {} + })), + Some(ProviderSurface::Gemini) + ); + // 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::Gemini) + ); +} + // --------------------------------------------------------------------------- // detect_response_surface (strict; ambiguity -> None) // --------------------------------------------------------------------------- @@ -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::Gemini, ]; #[test] @@ -423,11 +476,12 @@ fn codec_name_uses_canonical_spellings() { ProviderSurface::AnthropicMessages.codec_name(), "anthropic_messages" ); + assert_eq!(ProviderSurface::Gemini.codec_name(), "gemini"); } #[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 +490,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" + ] ); let from_registry: Vec<_> = BUILTIN_PROVIDER_SURFACES .iter() @@ -480,6 +539,17 @@ fn request_codec_decodes_each_surface() { .messages .is_empty() ); + + let gemini = req(json!({ + "contents": [{"role": "user", "parts": [{"text": "hi"}]}] + })); + assert!( + !request_codec(ProviderSurface::Gemini) + .decode(&gemini) + .expect("gemini request decodes") + .messages + .is_empty() + ); } #[test] @@ -519,6 +589,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::Gemini) + .decode_response(&gemini) + .expect("gemini response decodes") + .response_text(), + Some("hi gemini") + ); } #[test] 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..c5d18d523 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" + ] )); assert!(schema_property_has_default( &schema, @@ -639,10 +644,12 @@ 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'") - })); + assert!( + bad_codec + .diagnostics + .iter() + .any(|diag| { diag.message.contains("codec must be one of:") }) + ); let unsupported_remote_codec = validate_plugin_config(&plugin_config(json!({ "mode": "remote", 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..a82d5e32e 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,7 @@ fn local_codec_and_rewrite_helpers_cover_all_provider_surfaces() { LocalGuardrailsCodec::AnthropicMessages, ProviderSurface::AnthropicMessages, ), + (LocalGuardrailsCodec::Gemini, ProviderSurface::Gemini), ] { assert_eq!(codec.provider_surface(), surface); assert_eq!( @@ -1074,6 +1075,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::Gemini, + &json!({"candidates": [{"content": {"parts": [{"text": "visible"}]}, "index": 0}]}) + ), + Some("visible".to_string()) + ); for (codec, chunk) in [ (LocalGuardrailsCodec::OpenAIChat, Json::Null), ( @@ -1088,11 +1097,37 @@ fn stream_text_extraction_handles_supported_codecs() { LocalGuardrailsCodec::AnthropicMessages, json!({"type": "content_block_delta", "delta": {"type": "input_json_delta"}}), ), + (LocalGuardrailsCodec::Gemini, 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::Gemini, + &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::Gemini, + &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..61a7fe26a 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_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..43890baf3 100644 --- a/crates/ffi/src/api/llm.rs +++ b/crates/ffi/src/api/llm.rs @@ -372,6 +372,21 @@ 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_codec_new() -> *mut FfiCodecHandle { + Box::into_raw(Box::new(FfiCodecHandle { + codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), + response_codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), + })) +} + 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..a6a09ecc0 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_codec_new(); assert!(!openai_chat.is_null()); assert!(!openai_responses.is_null()); assert!(!anthropic.is_null()); + assert!( + !gemini.is_null(), + "GeminiCodec 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..f546a99dc 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' | string; builtin?: BuiltinConfig; local?: LocalModelConfig; policy?: ConfigPolicy; diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index d4b77b92e..1afb84457 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' } | { kind: 'runtime'; id: string } | { kind: 'opaque' }; diff --git a/crates/node/src/types/mod.rs b/crates/node/src/types/mod.rs index 2a36c4dc9..e9b58eb7f 100644 --- a/crates/node/src/types/mod.rs +++ b/crates/node/src/types/mod.rs @@ -480,6 +480,63 @@ impl OpenAIResponsesCodec { } } +/// Built-in codec for the Gemini generateContent API. +/// +/// Implements both request codec (decode/encode) and response codec +/// (decodeResponse). Construct with `new GeminiCodec()`. +#[napi(js_name = "GeminiCodec")] +pub struct GeminiCodec { + pub(crate) inner_codec: std::sync::Arc, + pub(crate) inner_response_codec: std::sync::Arc, +} + +#[napi] +impl GeminiCodec { + #[napi(constructor)] + pub fn new() -> Self { + Self { + inner_codec: std::sync::Arc::new(nemo_relay::codec::gemini::GeminiCodec), + inner_response_codec: std::sync::Arc::new(nemo_relay::codec::gemini::GeminiCodec), + } + } + + /// 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..1108ed27d 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.GeminiCodec, 'function'); }); it('scope type enum values', () => { @@ -79,3 +80,103 @@ describe('ScopeStack', () => { assert.ok(stack instanceof ScopeStack); }); }); + +// =========================================================================== +// GeminiCodec +// =========================================================================== + +describe('GeminiCodec', () => { + const { GeminiCodec } = lib; + + it('instantiates', () => { + const codec = new GeminiCodec(); + assert.ok(codec instanceof GeminiCodec); + }); + + it('decode returns an AnnotatedLLMRequest with messages and no params', () => { + const codec = new GeminiCodec(); + const req = new LlmRequest({}, { + 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 GeminiCodec(); + const req = new LlmRequest({}, { + 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 GeminiCodec(); + const req = new LlmRequest({}, { + 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 GeminiCodec(); + 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 GeminiCodec(); + 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'); + }); +}); diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 4c8528582..4d8904137 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`. - 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..c0d4db900 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -304,6 +304,7 @@ impl CompiledBuiltinBackend { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) => { Some(ProviderSurface::AnthropicMessages) } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini) => Some(ProviderSurface::Gemini), LlmCodecIdentity::Runtime(_) | LlmCodecIdentity::Opaque => None, } } @@ -400,6 +401,18 @@ impl CompiledBuiltinBackend { { 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::Gemini + && payload + .get("candidates") + .and_then(Json::as_array) + .is_some_and(|candidates| candidates.len() > 1) + && self.targets_normalized_gemini_candidate() + { + return None; + } let codec_name = BuiltinCodecName::from_provider_surface(surface); let annotated = codec.decode_response(&payload).ok()?; let annotated_json = serde_json::to_value(&annotated).ok()?; @@ -458,6 +471,19 @@ impl CompiledBuiltinBackend { }) }) } + + fn targets_normalized_gemini_candidate(&self) -> bool { + self.target_paths.iter().any(|path| { + json_pointer_segments(path) + .and_then(|segments| segments.into_iter().next()) + .is_some_and(|segment| { + matches!( + segment.as_str(), + "message" | "tool_calls" | "finish_reason" | "api_specific" + ) + }) + }) + } } pub(super) fn tool_sanitize_callback(backend: CompiledBuiltinBackend) -> ToolSanitizeFn { diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 1e21dfe43..72fa57af7 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -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"], optional: true, }, profiles => { label: "profiles", kind: List, list: &PII_REDACTION_PROFILE_LIST_ITEM }, @@ -500,7 +500,12 @@ fn custom_mark_payload_policy_schema( 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", + ], None, ) } @@ -1137,7 +1142,10 @@ fn validate_codec_requirements( "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_codec_names().join(", ") + ), ); } } diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index e74da6e26..071335cef 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, + Gemini, } impl BuiltinCodecName { @@ -20,6 +21,7 @@ impl BuiltinCodecName { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::Gemini => Self::Gemini, } } @@ -32,10 +34,156 @@ 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::Gemini => 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 text: update all non-thought text parts, splitting multi-line normalized + // text across them when there are multiple. + let message_text = annotated_message_text(annotated.message.as_ref()); + let text_part_count = parts + .iter() + .filter(|p| { + p.get("text").is_some() && p.get("thought").and_then(Json::as_bool) != Some(true) + }) + .count(); + let mut non_thought_text_count = 0usize; + for part in parts.iter_mut() { + let is_thought = part.get("thought").and_then(Json::as_bool) == Some(true); + if part.get("text").is_none() || is_thought { + continue; + } + let Some(p) = part.as_object_mut() else { + continue; + }; + let value = if text_part_count <= 1 { + message_text.as_deref() + } else { + // Multiple text parts: split on newline and assign one chunk each. + message_text + .as_deref() + .map(|t| t.split('\n').collect::>()) + .as_ref() + .and_then(|chunks| chunks.get(non_thought_text_count).copied()) + .or_else(|| { + (non_thought_text_count == 0) + .then_some(message_text.as_deref()) + .flatten() + }) + }; + set_optional_string_field(p, "text", value); + non_thought_text_count += 1; + } + } + + // 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 + }; + 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..13d675de1 100644 --- a/crates/pii-redaction/tests/coverage/overlay_tests.rs +++ b/crates/pii-redaction/tests/coverage/overlay_tests.rs @@ -133,3 +133,182 @@ 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::Gemini.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_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::Gemini.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::Gemini.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_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::Gemini.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::Gemini.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..410aa6c98 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -4347,6 +4347,283 @@ async fn builtin_backend_removes_targeted_message_names_and_ignores_missing_norm 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", + "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") + .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", + "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") + .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", + "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") + .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::GeminiCodec; + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini", + "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") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "run code"}]}] + }), + }) + .func(noop_openai_chat_exec_fn(response)) + .response_codec(Arc::new(GeminiCodec)) + .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 +5459,77 @@ 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::GeminiCodec; + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini", + "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(GeminiCodec)) + .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(); +} diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index b6bf3da7d..18d0fff37 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")] + Gemini, } /// 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" => BuiltinLlmCodec::Gemini, _ => { 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..2dcd292d6 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -49,8 +49,8 @@ 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, + PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyGeminiCodec, PyLLMAttributes, PyLLMHandle, + PyLLMRequest, PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyThreadScopeStackBinding, PyToolAttributes, PyToolHandle, }; @@ -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..8ee050174 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -1015,3 +1015,68 @@ 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 GeminiCodec +/// codec = GeminiCodec() +/// annotated_req = codec.decode(request) +/// annotated_resp = codec.decode_response(response) +/// ``` +#[pyclass(name = "GeminiCodec")] +pub struct PyGeminiCodec { + pub(crate) inner_codec: Arc, + pub(crate) inner_response_codec: Arc, +} + +#[pymethods] +impl PyGeminiCodec { + #[new] + pub(crate) fn new() -> Self { + Self { + inner_codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), + inner_response_codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), + } + } + + /// 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..ac716fca8 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..d15c5936d 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -190,11 +190,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 `"STOP"` (without function-call parts) +/// - **Length**: OpenAI Chat `"length"`, Anthropic `"max_tokens"`, Responses incomplete+max_output_tokens, +/// Gemini `"MAX_TOKENS"` +/// - **ToolUse**: OpenAI Chat `"tool_calls"`, Anthropic `"tool_use"`, +/// Gemini `"TOOL_CODE"` (unconditional) or `"STOP"` when function-call parts are present +/// - **ContentFilter**: OpenAI Chat `"content_filter"`, Responses incomplete+content_filter, +/// Gemini `"SAFETY"` / `"RECITATION"` / `"BLOCKLIST"` / `"PROHIBITED_CONTENT"` and other policy codes /// - **Unknown**: Forward-compatible catch-all for unrecognized reasons +/// (e.g. Gemini `"MALFORMED_FUNCTION_CALL"`, `"UNEXPECTED_TOOL_CALL"`) #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FinishReason { @@ -229,7 +234,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's `args` are already parsed JSON objects. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ResponseToolCall { /// Unique identifier for this tool call. @@ -326,6 +331,26 @@ pub enum ApiSpecificResponse { content_blocks: Option>, }, + /// Gemini generateContent API-specific fields. + #[serde(rename = "gemini")] + Gemini { + /// Tokens consumed by the model's internal reasoning (Gemini thinking). + #[serde(skip_serializing_if = "Option::is_none")] + thoughts_tokens: Option, + /// Candidate-level safety ratings from the Gemini 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, skip_serializing_if = "serde_json::Map::is_empty")] + extra: serde_json::Map, + }, + /// Custom/unknown API -- catch-all for user-implemented codecs. #[serde(rename = "custom")] Custom { diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index fb15f886f..0029e3ee9 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. + Gemini, } /// Per-call LLM codec identity supplied to worker sanitizers. @@ -2157,6 +2159,7 @@ fn codec_identity_from_proto( Some("anthropic_messages") => { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) } + Some("gemini") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), _ => LlmCodecIdentity::Opaque, }, Some(LlmCodecKind::Runtime) => codec_id @@ -2675,3 +2678,27 @@ fn rustc_version_runtime() -> String { .unwrap_or("unknown") .to_string() } + +#[cfg(test)] +mod tests { + use super::*; + + /// Verify that BuiltinLlmCodec has a Gemini variant and that codec_identity_from_proto + /// decodes the "gemini" id to the correct built-in identity (not Opaque). + #[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".to_string()), + }; + let identity = codec_identity_from_proto(Some(&proto)); + assert_eq!( + identity, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), + "Gemini codec id must decode to BuiltIn(Gemini), not Opaque" + ); + } +} diff --git a/docs/about-nemo-relay/concepts/codecs.mdx b/docs/about-nemo-relay/concepts/codecs.mdx index 0049d840e..8bed91b9d 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -88,7 +88,9 @@ changing a key overlays it. - `messages`, content parts, function calls and results, tools, and tool choice expose portable components when the shapes have shared semantics. - `api_specific` is a tagged, mutable surface for modeled Anthropic Messages, - OpenAI Chat Completions, or OpenAI Responses controls. + OpenAI Chat Completions, or OpenAI Responses request controls. Response + annotations use the same pattern for provider-only metadata, including Gemini + 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` is reserved for unknown future fields. @@ -118,9 +120,9 @@ normalization logic without changing runtime ownership or public binding APIs. ### Provider Schema Extraction Provider schema extraction is codec-owned. Built-in provider surfaces, such as -OpenAI Chat Completions, OpenAI Responses, and Anthropic Messages, each own the -logic that recognizes their request and response shapes and maps them into -`AnnotatedLlmRequest` or `AnnotatedLlmResponse`. +OpenAI Chat Completions, OpenAI Responses, Anthropic Messages, and Gemini +generateContent, each own the logic that recognizes their request and response +shapes and maps 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 diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 19d583755..b46637ca8 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -299,7 +299,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` 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..ebd09f2c7 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`) 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/nemo-guardrails/configuration.mdx b/docs/configure-plugins/nemo-guardrails/configuration.mdx index d34013a0d..b9141495b 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` | | 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` ### Managed Tool Boundary diff --git a/docs/configure-plugins/pii-redaction/configuration.mdx b/docs/configure-plugins/pii-redaction/configuration.mdx index 392999639..faa88702a 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -94,7 +94,7 @@ 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 +policy safely applies to OpenAI Chat, OpenAI Responses, Anthropic Messages, and Gemini traffic: ```toml @@ -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` | 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..708b7f1d6 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. +- `GeminiCodec`: Gemini generateContent-compatible requests and responses. ## Provider Codec Roles @@ -106,6 +107,7 @@ 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 | `nemo_relay.codecs.GeminiCodec` | `GeminiCodec` 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. diff --git a/docs/integrate-into-frameworks/provider-response-codecs.mdx b/docs/integrate-into-frameworks/provider-response-codecs.mdx index 205fc9533..c4ea2e3cc 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 | +|---|---|---|---|---| +| `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 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 +- `GeminiCodec` — 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..6e4887094 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..ddf71a51c 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" runtimeCodec := "com.example.chat.v1" cases := []codecIdentityTestCase{ @@ -103,6 +104,8 @@ func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) {"openai chat", 1, &openAIChat, LLMCodecBuiltin}, {"openai responses", 1, &openAIResponses, LLMCodecBuiltin}, {"anthropic messages", 1, &anthropicMessages, LLMCodecBuiltin}, + // Gemini must be LLMCodecBuiltin, not Opaque — PII sanitization depends on this. + {"gemini", 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..ef3c1165b 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 := NewGeminiCodec() + 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 TestGeminiCodecFunctionCallID(t *testing.T) { + runTestWithScopeStack(t, testGeminiCodecFunctionCallID) +} + +func testGeminiCodecFunctionCallID(t *testing.T) { + // Verify that NewGeminiCodec 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(NewGeminiCodec()), + ) + if err != nil { + t.Fatalf("LlmCallExecute with GeminiCodec 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..f26524c95 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_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 [NewGeminiCodec]. The handle is +// automatically freed when garbage collected. type CodecHandle struct { ptr *C.FfiCodecHandle } @@ -977,9 +978,25 @@ func NewAnthropicMessagesCodec() *CodecHandle { return h } +// NewGeminiCodec 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 NewGeminiCodec() *CodecHandle { + h := &CodecHandle{ptr: C.nemo_relay_gemini_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 [NewGeminiCodec]. // 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..cba8889a2 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 GeminiCodec: + """Built-in codec for Gemini generateContent requests and responses. + + Summary: + Native codec bridge for Gemini generateContent payloads. + """ + + def __init__(self) -> None: + """Create a Gemini 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..22ac68b7b 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, + GeminiCodec, LLMRequest, OpenAIChatCodec, OpenAIResponsesCodec, @@ -161,6 +162,7 @@ def decode_response(self, response: Json) -> "AnnotatedLLMResponse": __all__ = [ "AnnotatedLLMRequest", "AnthropicMessagesCodec", + "GeminiCodec", "LlmCodec", "LlmResponseCodec", "OpenAIChatCodec", diff --git a/python/nemo_relay/codecs.pyi b/python/nemo_relay/codecs.pyi index 5bb540295..c76d15c0f 100644 --- a/python/nemo_relay/codecs.pyi +++ b/python/nemo_relay/codecs.pyi @@ -163,9 +163,48 @@ class AnthropicMessagesCodec: """ ... +class GeminiCodec: + """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", + "GeminiCodec", "LlmCodec", "LlmResponseCodec", "OpenAIChatCodec", diff --git a/python/nemo_relay/pii_redaction.py b/python/nemo_relay/pii_redaction.py index 53c2b77f8..4c48d87e2 100644 --- a/python/nemo_relay/pii_redaction.py +++ b/python/nemo_relay/pii_redaction.py @@ -136,7 +136,7 @@ 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"] | 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..1815f054c 100644 --- a/python/nemo_relay/pii_redaction.pyi +++ b/python/nemo_relay/pii_redaction.pyi @@ -58,7 +58,7 @@ 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"] | 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..29d49d7e5 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"), (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"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "gemini"))), ("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..b32df1418 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -4,8 +4,8 @@ """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 (OpenAIChatCodec, OpenAIResponsesCodec, AnthropicMessagesCodec, GeminiCodec) +- Built-in codec decode/encode/decode_response methods for all four providers - LlmResponseCodec protocol - response_codec parameter accepts object (not string) """ @@ -22,7 +22,7 @@ llm, subscribers, ) -from nemo_relay.codecs import AnthropicMessagesCodec, OpenAIChatCodec, OpenAIResponsesCodec +from nemo_relay.codecs import AnthropicMessagesCodec, GeminiCodec, OpenAIChatCodec, OpenAIResponsesCodec # --------------------------------------------------------------------------- # 1. Built-in codec construction @@ -66,6 +66,18 @@ def test_anthropic_messages_codec_has_methods(self): assert hasattr(codec, "encode") assert hasattr(codec, "decode_response") + def test_gemini_codec_constructable(self): + """GeminiCodec() is constructable.""" + codec = GeminiCodec() + assert codec is not None + + def test_gemini_codec_has_methods(self): + """GeminiCodec has decode, encode, decode_response methods.""" + codec = GeminiCodec() + assert hasattr(codec, "decode") + assert hasattr(codec, "encode") + assert hasattr(codec, "decode_response") + # --------------------------------------------------------------------------- # 2. Built-in codec decode/encode round-trip @@ -272,6 +284,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): + """GeminiCodec.decode() returns AnnotatedLLMRequest with messages and params.""" + codec = GeminiCodec() + 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): + """GeminiCodec.encode(decode(req), req) is idempotent when nothing changes.""" + codec = GeminiCodec() + 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): + """GeminiCodec.decode_response() extracts text and usage from a generateContent response.""" + codec = GeminiCodec() + 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): + """GeminiCodec maps SAFETY finish reason to 'content_filter', not 'unknown'.""" + codec = GeminiCodec() + 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): + """GeminiCodec.decode_response() extracts functionCall parts as tool_calls.""" + codec = GeminiCodec() + 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 +417,7 @@ def test_builtin_codecs_satisfy_protocol(self): assert isinstance(OpenAIChatCodec(), LlmResponseCodec) assert isinstance(OpenAIResponsesCodec(), LlmResponseCodec) assert isinstance(AnthropicMessagesCodec(), LlmResponseCodec) + assert isinstance(GeminiCodec(), LlmResponseCodec) # --------------------------------------------------------------------------- @@ -564,14 +690,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, + GeminiCodec, + OpenAIChatCodec, + OpenAIResponsesCodec, + ) assert OpenAIChatCodec is not None assert OpenAIResponsesCodec is not None assert AnthropicMessagesCodec is not None + assert GeminiCodec 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, "GeminiCodec") From d9faaafdaca5158a722c8540d7f158108535b26e Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 21:25:06 -0700 Subject: [PATCH 02/11] fix(codec): clarify Gemini generateContent naming Signed-off-by: Sahiti Bulusu --- crates/adaptive/src/response_cache/key.rs | 2 +- crates/adaptive/src/response_cache/replay.rs | 2 +- .../tests/unit/response_cache/key_tests.rs | 4 +- crates/core/src/api/runtime/callbacks.rs | 4 +- .../{gemini.rs => gemini_generate_content.rs} | 115 +++--- crates/core/src/codec/mod.rs | 2 +- crates/core/src/codec/resolve.rs | 12 +- .../src/plugins/nemo_guardrails/component.rs | 4 +- .../src/plugins/nemo_guardrails/python.rs | 8 +- ...ts.rs => gemini_generate_content_tests.rs} | 340 +++++++++--------- crates/core/tests/unit/codec/parity_tests.rs | 13 +- crates/core/tests/unit/codec/resolve_tests.rs | 27 +- .../nemo_guardrails/component_tests.rs | 2 +- .../nemo_guardrails/local_python_tests.rs | 13 +- crates/ffi/nemo_relay.h | 2 +- crates/ffi/src/api/llm.rs | 8 +- crates/ffi/tests/unit/types_tests.rs | 4 +- crates/node/pii_redaction.d.ts | 2 +- crates/node/plugin.d.ts | 2 +- crates/node/src/types/mod.rs | 16 +- crates/node/tests/types_tests.mjs | 22 +- crates/pii-redaction/README.md | 2 +- crates/pii-redaction/src/builtin.rs | 6 +- crates/pii-redaction/src/component.rs | 4 +- crates/pii-redaction/src/overlay.rs | 6 +- .../tests/coverage/overlay_tests.rs | 15 +- .../tests/unit/component_tests.rs | 26 +- crates/plugin/src/lib.rs | 6 +- crates/python/src/py_api/mod.rs | 12 +- crates/python/src/py_types/codecs.rs | 20 +- crates/python/src/py_types/mod.rs | 2 +- crates/types/src/codec/response.rs | 20 +- crates/worker/src/lib.rs | 17 +- docs/about-nemo-relay/concepts/codecs.mdx | 2 +- docs/about-nemo-relay/concepts/middleware.mdx | 2 +- .../dynamic-plugins/native-dynamic/about.mdx | 2 +- .../nemo-guardrails/configuration.mdx | 4 +- .../pii-redaction/configuration.mdx | 6 +- .../provider-codecs.mdx | 4 +- .../provider-response-codecs.mdx | 6 +- go/nemo_relay/callbacks_test.go | 4 +- go/nemo_relay/coverage_gap_test.go | 14 +- go/nemo_relay/nemo_relay.go | 12 +- python/nemo_relay/_native.pyi | 4 +- python/nemo_relay/codecs.py | 4 +- python/nemo_relay/codecs.pyi | 4 +- python/nemo_relay/pii_redaction.py | 2 +- python/nemo_relay/pii_redaction.pyi | 2 +- python/tests/plugin/test_worker_sdk.py | 6 +- python/tests/test_builtin_codecs.py | 40 +-- 50 files changed, 448 insertions(+), 410 deletions(-) rename crates/core/src/codec/{gemini.rs => gemini_generate_content.rs} (96%) rename crates/core/tests/unit/codec/{gemini_tests.rs => gemini_generate_content_tests.rs} (94%) diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index bc0dc6a65..a2dd5f029 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -337,7 +337,7 @@ fn lossy_request_shape(surface: ProviderSurface, content: &Json) -> bool { .and_then(Json::as_array) .is_some_and(|blocks| blocks.iter().any(lossy_system_block)) } - ProviderSurface::OpenAIResponses | ProviderSurface::Gemini => false, + ProviderSurface::OpenAIResponses | ProviderSurface::GeminiGenerateContent => false, } } diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs index 7f191246f..5e87b4784 100644 --- a/crates/adaptive/src/response_cache/replay.rs +++ b/crates/adaptive/src/response_cache/replay.rs @@ -83,7 +83,7 @@ fn synthesize_replay_chunks(aggregate: &Json) -> Option> { // aggregate is a valid single native chunk. `replay_is_lossy` still // re-aggregates it and rejects shapes the streaming collector cannot // preserve exactly. - ProviderSurface::Gemini => vec![aggregate.clone()], + ProviderSurface::GeminiGenerateContent => vec![aggregate.clone()], }) } diff --git a/crates/adaptive/tests/unit/response_cache/key_tests.rs b/crates/adaptive/tests/unit/response_cache/key_tests.rs index 781545db5..2c341c178 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -600,8 +600,8 @@ fn gemini_shaped_requests_key_on_the_detected_decode() { "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": {"temperature": 0.0} })); - let (body, effective_codec) = resolved_body("gemini", &request); - assert_eq!(effective_codec, Some("gemini")); + 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" diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index 0c2e6a7fb..5adc048e3 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -161,7 +161,7 @@ pub enum BuiltinLlmCodec { /// Anthropic Messages request and response payloads. AnthropicMessages, /// Gemini generateContent request and response payloads. - Gemini, + GeminiGenerateContent, } impl BuiltinLlmCodec { @@ -172,7 +172,7 @@ impl BuiltinLlmCodec { Self::OpenAiChat => "openai_chat", Self::OpenAiResponses => "openai_responses", Self::AnthropicMessages => "anthropic_messages", - Self::Gemini => "gemini", + Self::GeminiGenerateContent => "gemini_generate_content", } } } diff --git a/crates/core/src/codec/gemini.rs b/crates/core/src/codec/gemini_generate_content.rs similarity index 96% rename from crates/core/src/codec/gemini.rs rename to crates/core/src/codec/gemini_generate_content.rs index 40da37ce5..6d179e3ab 100644 --- a/crates/core/src/codec/gemini.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -33,18 +33,18 @@ const GEMINI_PROVIDER: &str = "gemini"; // --------------------------------------------------------------------------- /// Built-in codec for the Gemini generateContent API. -pub struct GeminiCodec; +pub struct GeminiGenerateContentCodec; pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDescriptor { - surface: ProviderSurface::Gemini, + surface: ProviderSurface::GeminiGenerateContent, detect_request: |obj, _hint| obj.contains_key("contents"), detect_response: detect_gemini_response, - decode_request: |request| GeminiCodec.decode(request), - decode_response: |raw| GeminiCodec.decode_response(raw), - codec_name: "gemini", - request_codec: || std::sync::Arc::new(GeminiCodec), - response_codec: || std::sync::Arc::new(GeminiCodec), - streaming_codec: || Box::new(GeminiStreamingCodec::new()), + 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()), }; // --------------------------------------------------------------------------- @@ -52,7 +52,7 @@ pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDe // --------------------------------------------------------------------------- #[derive(Deserialize)] -struct RawGeminiResponse { +struct RawGeminiGenerateContentResponse { candidates: Option>, #[serde(rename = "usageMetadata")] usage_metadata: Option, @@ -223,10 +223,10 @@ fn parse_optional_id(obj: &serde_json::Map, context: &str) -> Resu None => Ok(None), Some(Json::String(s)) if !s.is_empty() => Ok(Some(s.clone())), Some(Json::String(_)) => Err(FlowError::InvalidArgument(format!( - "Gemini {context}.id must be a non-empty string" + "Gemini generateContent {context}.id must be a non-empty string" ))), Some(_) => Err(FlowError::InvalidArgument(format!( - "Gemini {context}.id must be a string" + "Gemini generateContent {context}.id must be a string" ))), } } @@ -258,7 +258,7 @@ fn validate_single_gemini_part_data_field<'a>( let keys = gemini_part_data_keys(obj); if keys.len() > 1 { return Err(FlowError::InvalidArgument(format!( - "Gemini {context} part must not contain multiple data fields: {}", + "Gemini generateContent {context} part must not contain multiple data fields: {}", keys.join(", ") ))); } @@ -281,7 +281,9 @@ fn gemini_parts_to_message_content( for part in parts { let obj = part.as_object().ok_or_else(|| { - FlowError::InvalidArgument(format!("Gemini {context} parts entry must be an object")) + FlowError::InvalidArgument(format!( + "Gemini generateContent {context} parts entry must be an object" + )) })?; if obj.get("thought").and_then(Json::as_bool) == Some(true) { continue; @@ -302,7 +304,7 @@ fn gemini_parts_to_message_content( "text" => { let text = obj.get("text").and_then(Json::as_str).ok_or_else(|| { FlowError::InvalidArgument(format!( - "Gemini {context} parts[].text must be a string" + "Gemini generateContent {context} parts[].text must be a string" )) })?; let extra: serde_json::Map = obj @@ -460,7 +462,7 @@ fn extract_parts_tool_calls(parts: &[Json]) -> Result, model: Option<&str>) -> (Option, Option) { let model_provider = infer_model_provider("google", model); @@ -1047,7 +1049,7 @@ fn ensure_object_response(content_str: String) -> Json { } } -struct GeminiFunctionResponsePayload { +struct GeminiGenerateContentFunctionResponsePayload { response: Json, parts: Option>, } @@ -1061,7 +1063,7 @@ fn provider_native_gemini_part_value(obj: &serde_json::Map) -> Res if provider != GEMINI_PROVIDER { return Err(FlowError::InvalidArgument(format!( "Gemini encoder: provider_native content part for provider \ - '{provider}' cannot be encoded by the Gemini codec" + '{provider}' cannot be encoded by the Gemini generateContent codec" ))); } match obj.get("kind") { @@ -1103,9 +1105,9 @@ fn provider_native_gemini_part_value(obj: &serde_json::Map) -> Res fn function_response_payload_from_tool_content( content: &Json, -) -> Result { +) -> Result { match content { - Json::Null | Json::String(_) => Ok(GeminiFunctionResponsePayload { + Json::Null | Json::String(_) => Ok(GeminiGenerateContentFunctionResponsePayload { response: ensure_object_response(extract_content_text(content)), parts: None, }), @@ -1154,7 +1156,7 @@ fn function_response_payload_from_tool_content( } } } - Ok(GeminiFunctionResponsePayload { + Ok(GeminiGenerateContentFunctionResponsePayload { response: ensure_object_response(response_texts.join("\n")), parts: Some(native_parts), }) @@ -1193,8 +1195,9 @@ fn insert_serialized( value: &T, context: &str, ) -> Result<()> { - let json = serde_json::to_value(value) - .map_err(|e| FlowError::Internal(format!("Gemini {context} encode: {e}")))?; + let json = serde_json::to_value(value).map_err(|e| { + FlowError::Internal(format!("Gemini generateContent {context} encode: {e}")) + })?; obj.insert(key.into(), json); Ok(()) } @@ -1266,13 +1269,13 @@ fn take_matching_native_group( // LlmResponseCodec // --------------------------------------------------------------------------- -impl LlmResponseCodec for GeminiCodec { +impl LlmResponseCodec for GeminiGenerateContentCodec { fn codec_identity(&self) -> LlmCodecIdentity { - LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini) + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) } fn decode_response(&self, response: &Json) -> Result { - let raw: RawGeminiResponse = serde_json::from_value(response.clone()) + 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()); @@ -1311,13 +1314,15 @@ impl LlmResponseCodec for GeminiCodec { let safety_ratings = extra.remove("safetyRatings"); let grounding_metadata = extra.remove("groundingMetadata"); let citation_metadata = extra.remove("citationMetadata"); - Some(super::response::ApiSpecificResponse::Gemini { - thoughts_tokens, - safety_ratings, - grounding_metadata, - citation_metadata, - extra, - }) + Some( + super::response::ApiSpecificResponse::GeminiGenerateContent { + thoughts_tokens, + safety_ratings, + grounding_metadata, + citation_metadata, + extra, + }, + ) } }; @@ -1339,9 +1344,9 @@ impl LlmResponseCodec for GeminiCodec { // LlmCodec // --------------------------------------------------------------------------- -impl LlmCodec for GeminiCodec { +impl LlmCodec for GeminiGenerateContentCodec { fn codec_identity(&self) -> LlmCodecIdentity { - LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini) + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) } fn decode(&self, request: &LlmRequest) -> Result { @@ -1801,7 +1806,7 @@ fn rebuild_gemini_system_parts(orig_parts: &[Json], text: String) -> Vec { } #[derive(Debug)] -struct GeminiMessageAlignment { +struct GeminiGenerateContentMessageAlignment { prefix_len: usize, ann_gap_end: usize, base_gap_end: usize, @@ -1809,7 +1814,7 @@ struct GeminiMessageAlignment { base_gap_len: usize, } -impl GeminiMessageAlignment { +impl GeminiGenerateContentMessageAlignment { fn new(ann_non_sys: &[&Json], base_non_sys: &[&Json]) -> Self { let prefix_len = ann_non_sys .iter() @@ -1916,7 +1921,7 @@ fn gemini_run_end( start: usize, ann_len: usize, orig_cidx: usize, - alignment: &GeminiMessageAlignment, + alignment: &GeminiGenerateContentMessageAlignment, content_idx_of_base_msg: &[usize], ) -> usize { let mut run_end = start + 1; @@ -1933,7 +1938,7 @@ fn gemini_run_is_unchanged( start: usize, run_end: usize, expected_len: usize, - alignment: &GeminiMessageAlignment, + alignment: &GeminiGenerateContentMessageAlignment, ann_non_sys: &[&Json], base_non_sys: &[&Json], ) -> bool { @@ -1978,7 +1983,7 @@ fn patch_gemini_non_system_contents( .and_then(|v| v.as_array()) .cloned() .unwrap_or_default(); - let alignment = GeminiMessageAlignment::new(ann_non_sys, base_non_sys); + 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 = @@ -2106,7 +2111,7 @@ fn patch_changed_gemini_content( fn gemini_function_response_updates<'a>( ann_msgs: &'a [&Json], ) -> Result<( - HashMap<&'a str, GeminiFunctionResponsePayload>, + HashMap<&'a str, GeminiGenerateContentFunctionResponsePayload>, Vec<&'a str>, )> { let mut updates = HashMap::new(); @@ -2137,7 +2142,7 @@ fn function_response_call_id(fr: &Json) -> Option<&str> { fn build_gemini_function_response_part( call_id: &str, name: &str, - payload: GeminiFunctionResponsePayload, + payload: GeminiGenerateContentFunctionResponsePayload, original_fr: Option<&Json>, ) -> Json { let mut new_fr = original_fr @@ -2641,30 +2646,32 @@ fn patch_extra_fields( // Streaming codec // --------------------------------------------------------------------------- -/// Streaming counterpart to [`GeminiCodec`]. +/// Streaming counterpart to [`GeminiGenerateContentCodec`]. /// /// Accumulates Gemini server-sent event chunks and assembles a complete response -/// that [`GeminiCodec::decode_response`] can consume. -pub struct GeminiStreamingCodec { - state: std::sync::Arc>, +/// that [`GeminiGenerateContentCodec::decode_response`] can consume. +pub struct GeminiGenerateContentStreamingCodec { + state: std::sync::Arc>, } -impl GeminiStreamingCodec { +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(GeminiStreamingState::default())), + state: std::sync::Arc::new(std::sync::Mutex::new( + GeminiGenerateContentStreamingState::default(), + )), } } } -impl Default for GeminiStreamingCodec { +impl Default for GeminiGenerateContentStreamingCodec { fn default() -> Self { Self::new() } } -impl super::streaming::StreamingCodec for GeminiStreamingCodec { +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<()> { @@ -2688,7 +2695,7 @@ impl super::streaming::StreamingCodec for GeminiStreamingCodec { } #[derive(Debug, Default)] -struct GeminiStreamingState { +struct GeminiGenerateContentStreamingState { parts: Vec, finish_reason: Option, usage_metadata: Option, @@ -2699,7 +2706,7 @@ struct GeminiStreamingState { candidate_extra: serde_json::Map, } -impl GeminiStreamingState { +impl GeminiGenerateContentStreamingState { fn push_text_part(&mut self, text: &str, part_obj: &serde_json::Map) { let mut extra: serde_json::Map = part_obj .iter() @@ -2791,7 +2798,7 @@ impl GeminiStreamingState { } // Collect candidate-level metadata fields (safetyRatings, groundingMetadata, // citationMetadata, avgLogprobs, etc.) that non-streaming decode preserves - // in ApiSpecificResponse::Gemini. Later chunks overwrite earlier ones for + // in ApiSpecificResponse::GeminiGenerateContent. Later chunks overwrite earlier ones for // the same key (last-wins), matching the non-streaming behaviour. if let Some(cand_obj) = candidate.as_object() { for (k, v) in cand_obj { @@ -2838,7 +2845,7 @@ impl GeminiStreamingState { } candidate_obj.insert("index".into(), Json::from(0u32)); // Merge accumulated candidate-level metadata so that decode_response can - // populate ApiSpecificResponse::Gemini with the same fields as non-streaming. + // populate ApiSpecificResponse::GeminiGenerateContent with the same fields as non-streaming. for (k, v) in self.candidate_extra { candidate_obj.entry(k).or_insert(v); } @@ -2864,5 +2871,5 @@ impl GeminiStreamingState { // --------------------------------------------------------------------------- #[cfg(test)] -#[path = "../../tests/unit/codec/gemini_tests.rs"] +#[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 eb84a0240..4a587533e 100644 --- a/crates/core/src/codec/mod.rs +++ b/crates/core/src/codec/mod.rs @@ -15,7 +15,7 @@ //! provider codec from a raw payload when no codec annotation is present. pub mod anthropic; -pub mod gemini; +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 6bf03af90..b775ec8fa 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, gemini, 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)] @@ -26,7 +26,7 @@ pub enum ProviderSurface { /// Anthropic Messages. AnthropicMessages, /// Gemini generateContent. - Gemini, + GeminiGenerateContent, } /// Request shape detector; the optional `&str` is a provider hint a codec may use @@ -69,14 +69,14 @@ pub(crate) static BUILTIN_PROVIDER_SURFACES: &[ProviderSurfaceDescriptor] = &[ openai_responses::PROVIDER_SURFACE, anthropic::PROVIDER_SURFACE, openai_chat::PROVIDER_SURFACE, - gemini::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`) > Gemini (`contents`). `None` when no -/// key matches or `body` is not an object. This is a best-effort heuristic: an +/// (`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] @@ -155,7 +155,7 @@ 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::Gemini => &gemini::PROVIDER_SURFACE, + ProviderSurface::GeminiGenerateContent => &gemini_generate_content::PROVIDER_SURFACE, } } diff --git a/crates/core/src/plugins/nemo_guardrails/component.rs b/crates/core/src/plugins/nemo_guardrails/component.rs index fe58df310..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", "gemini"], + values: ["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"], optional: true, }, input => { label: "input", kind: Boolean }, @@ -434,7 +434,7 @@ fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::s "openai_chat", "openai_responses", "anthropic_messages", - "gemini", + "gemini_generate_content", ], None, ) diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs index 9aaffdf30..78ebf92a2 100644 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ b/crates/core/src/plugins/nemo_guardrails/python.rs @@ -891,7 +891,7 @@ enum LocalGuardrailsCodec { OpenAIChat, OpenAIResponses, AnthropicMessages, - Gemini, + GeminiGenerateContent, } impl LocalGuardrailsCodec { @@ -900,7 +900,7 @@ impl LocalGuardrailsCodec { Self::OpenAIChat => ProviderSurface::OpenAIChat, Self::OpenAIResponses => ProviderSurface::OpenAIResponses, Self::AnthropicMessages => ProviderSurface::AnthropicMessages, - Self::Gemini => ProviderSurface::Gemini, + Self::GeminiGenerateContent => ProviderSurface::GeminiGenerateContent, } } @@ -909,7 +909,7 @@ impl LocalGuardrailsCodec { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, - ProviderSurface::Gemini => Self::Gemini, + ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, } } @@ -1350,7 +1350,7 @@ fn extract_stream_text(codec: LocalGuardrailsCodec, chunk: &Json) -> Option { + LocalGuardrailsCodec::GeminiGenerateContent => { let candidates = chunk.get("candidates")?.as_array()?; let parts = candidates .first()? diff --git a/crates/core/tests/unit/codec/gemini_tests.rs b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs similarity index 94% rename from crates/core/tests/unit/codec/gemini_tests.rs rename to crates/core/tests/unit/codec/gemini_generate_content_tests.rs index cef096ed7..f87fdea57 100644 --- a/crates/core/tests/unit/codec/gemini_tests.rs +++ b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Unit tests for GeminiCodec in the NeMo Relay core crate. +//! Unit tests for GeminiGenerateContentCodec in the NeMo Relay core crate. use super::*; use serde_json::json; @@ -30,21 +30,21 @@ fn make_request(content: Json) -> LlmRequest { #[test] fn test_codec_identity_is_gemini_builtin() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; assert_eq!( LlmCodec::codec_identity(&codec), - LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), - "GeminiCodec must not return Opaque; PII sanitization depends on a known identity" + 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; assert_eq!( - ::codec_identity(&codec), - LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), - "GeminiCodec response codec must not return Opaque" + ::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent), + "GeminiGenerateContentCodec response codec must not return Opaque" ); } @@ -54,7 +54,7 @@ fn test_response_codec_identity_is_gemini_builtin() { #[test] fn test_decode_response_text() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -88,7 +88,7 @@ fn test_decode_response_text() { #[test] fn test_decode_response_native_part_as_provider_native_content() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -130,7 +130,7 @@ fn test_decode_response_native_part_as_provider_native_content() { #[test] fn test_decode_response_response_id() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}, "finishReason": "STOP"}], "responseId": "resp-abc-123", @@ -146,7 +146,7 @@ fn test_decode_response_response_id() { #[test] fn test_decode_response_function_call() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -182,7 +182,7 @@ fn test_decode_response_function_call() { #[test] fn test_decode_response_finish_reason_table() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; // (finishReason string, expected FinishReason) // NOTE: these cases have no functionCall parts in the response, so has_tool_calls = false. @@ -242,7 +242,7 @@ fn test_decode_response_finish_reason_table() { #[test] fn test_decode_response_prompt_feedback_block_reason_content_filter() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "promptFeedback": { "blockReason": "SAFETY", @@ -268,7 +268,7 @@ fn test_decode_response_prompt_feedback_block_reason_content_filter() { #[test] fn test_decode_response_prompt_feedback_block_reason_must_be_string() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "promptFeedback": {"blockReason": 123} }); @@ -282,7 +282,7 @@ fn test_decode_response_prompt_feedback_block_reason_must_be_string() { #[test] fn test_decode_response_cached_content_token_count() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": {"role": "model", "parts": [{"text": "cached response"}]}, @@ -307,7 +307,7 @@ fn test_decode_response_cached_content_token_count() { #[test] fn test_decode_response_no_candidates() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 0, "totalTokenCount": 5} @@ -321,7 +321,7 @@ fn test_decode_response_no_candidates() { #[test] fn test_decode_response_extra_fields_preserved() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": {"role": "model", "parts": [{"text": "hi"}]}, @@ -342,7 +342,7 @@ fn test_decode_response_extra_fields_preserved() { #[test] fn test_decode_contents_with_system_instruction() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "Hello"}]}, @@ -373,7 +373,7 @@ fn test_decode_contents_with_system_instruction() { #[test] fn test_decode_function_declarations() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{ @@ -408,7 +408,7 @@ fn test_decode_function_declarations() { #[test] fn test_decode_function_declaration_preserves_provider_fields() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{ @@ -449,7 +449,7 @@ fn test_decode_function_declaration_preserves_provider_fields() { #[test] fn test_decode_generation_config() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": { @@ -473,7 +473,7 @@ fn test_decode_generation_config() { #[test] fn test_decode_extra_fields_captured() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "safetySettings": [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}], @@ -492,7 +492,7 @@ fn test_decode_extra_fields_captured() { #[test] fn test_encode_round_trip_preserves_extra_fields() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original_json = json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "safetySettings": [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}] @@ -509,7 +509,7 @@ fn test_encode_round_trip_preserves_extra_fields() { /// is encoded (not decoded from an existing Gemini request). #[test] fn test_encode_role_assistant_becomes_model() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; // Decode a single-user-turn request to get a valid baseline. let original_json = json!({ "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] @@ -539,7 +539,7 @@ fn test_encode_role_assistant_becomes_model() { /// Decode-then-encode preserves systemInstruction when unchanged. #[test] fn test_encode_preserves_system_instruction_when_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original_json = json!({ "contents": [{"role": "user", "parts": [{"text": "hello"}]}], "systemInstruction": {"parts": [{"text": "You are an assistant."}]} @@ -560,7 +560,7 @@ fn test_encode_preserves_system_instruction_when_unchanged() { /// Decode-then-encode preserves tools when unchanged (round-trip). #[test] fn test_encode_preserves_tools_when_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original_json = json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{ @@ -591,7 +591,7 @@ fn test_encode_preserves_tools_when_unchanged() { /// Decode-then-encode preserves generationConfig when unchanged (round-trip). #[test] fn test_encode_preserves_generation_config_when_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original_json = json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": { @@ -612,7 +612,7 @@ fn test_encode_preserves_generation_config_when_unchanged() { /// Interceptor can edit the system message; the new text appears in the output. #[test] fn test_encode_system_instruction_edit() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original_json = json!({ "contents": [{"role": "user", "parts": [{"text": "hello"}]}], "systemInstruction": {"parts": [{"text": "old prompt"}]} @@ -645,7 +645,7 @@ fn test_encode_system_instruction_edit() { #[test] fn test_streaming_two_chunks_accumulated() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -674,7 +674,7 @@ fn test_streaming_two_chunks_accumulated() { let assembled = finalizer(); - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); assert_eq!( resp.message, @@ -690,7 +690,7 @@ fn test_streaming_two_chunks_accumulated() { #[test] fn test_streaming_preserves_text_part_metadata_from_empty_final_chunk() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -723,7 +723,7 @@ fn test_streaming_preserves_text_part_metadata_from_empty_final_chunk() { "streamed text metadata must survive finalization" ); - let codec = GeminiCodec; + 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"); @@ -739,7 +739,7 @@ fn test_streaming_preserves_text_part_metadata_from_empty_final_chunk() { #[test] fn test_streaming_keeps_non_empty_signed_text_part_separate_from_plain_text() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -779,7 +779,7 @@ fn test_streaming_keeps_non_empty_signed_text_part_separate_from_plain_text() { #[test] fn test_streaming_finalize_valid_response_shape() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -807,7 +807,7 @@ fn test_streaming_finalize_valid_response_shape() { Some("gemini-1.5-pro") ); - let codec = GeminiCodec; + 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())); @@ -815,7 +815,7 @@ fn test_streaming_finalize_valid_response_shape() { #[test] fn test_streaming_last_usage_metadata_wins() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -832,7 +832,7 @@ fn test_streaming_last_usage_metadata_wins() { .unwrap(); let assembled = finalizer(); - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); let usage = resp.usage.unwrap(); @@ -843,7 +843,7 @@ fn test_streaming_last_usage_metadata_wins() { #[test] fn test_streaming_preserves_native_response_parts() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -863,7 +863,7 @@ fn test_streaming_preserves_native_response_parts() { .unwrap(); let assembled = finalizer(); - let codec = GeminiCodec; + 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"); @@ -892,7 +892,7 @@ fn test_streaming_preserves_native_response_parts() { #[test] fn test_streaming_preserves_native_before_text_order() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -925,7 +925,7 @@ fn test_streaming_preserves_native_before_text_order() { #[test] fn test_streaming_does_not_merge_text_across_native_part() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -960,7 +960,7 @@ fn test_streaming_does_not_merge_text_across_native_part() { #[test] fn test_streaming_preserves_adjacent_signed_text_parts() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -995,7 +995,7 @@ fn test_streaming_preserves_adjacent_signed_text_parts() { #[test] fn test_encode_lossless_when_messages_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ { @@ -1030,7 +1030,7 @@ fn test_encode_lossless_when_messages_unchanged() { #[test] fn test_encode_thought_signature_preserved_when_system_message_changes() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "call something"}]}, @@ -1078,7 +1078,7 @@ fn test_encode_thought_signature_preserved_when_system_message_changes() { /// the correct id, name, and response fields. #[test] fn test_encode_thought_signature_preserved_in_multi_turn_continuation() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "call the tool"}]}, @@ -1163,7 +1163,7 @@ fn test_encode_thought_signature_preserved_in_multi_turn_continuation() { /// edits the function call arguments (triggers the rebuild path). #[test] fn test_encode_thought_signature_on_function_call_part_survives_edit() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "run the tool"}]}, @@ -1222,7 +1222,7 @@ fn test_encode_thought_signature_on_function_call_part_survives_edit() { #[test] fn test_encode_inline_data_preserved_when_messages_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let image_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQ=="; let original = make_request(json!({ "contents": [ @@ -1262,7 +1262,7 @@ fn test_encode_inline_data_preserved_when_messages_unchanged() { #[test] fn test_encode_inline_data_preserved_when_text_changes() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let image_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQ=="; let original = make_request(json!({ "contents": [{ @@ -1308,7 +1308,7 @@ fn test_encode_inline_data_preserved_when_text_changes() { #[test] fn test_decode_inline_data_as_provider_native_content_part() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{ "role": "user", @@ -1352,7 +1352,7 @@ fn test_decode_inline_data_as_provider_native_content_part() { #[test] fn test_encode_patches_provider_native_inline_data_content_part() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{ "role": "user", @@ -1392,7 +1392,7 @@ fn test_encode_patches_provider_native_inline_data_content_part() { #[test] fn test_decode_response_filters_thought_parts() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -1416,7 +1416,7 @@ fn test_decode_response_filters_thought_parts() { #[test] fn test_streaming_filters_thought_parts() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -1439,7 +1439,7 @@ fn test_streaming_filters_thought_parts() { .unwrap(); let assembled = finalizer(); - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); assert_eq!( resp.message, @@ -1453,7 +1453,7 @@ fn test_streaming_filters_thought_parts() { /// 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "search for it"}]}, @@ -1492,7 +1492,7 @@ fn test_decode_request_thought_text_does_not_leak_into_assistant_content() { /// content, thought text does not. #[test] fn test_decode_request_visible_text_survives_when_thought_present() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{ "role": "model", @@ -1528,7 +1528,7 @@ fn test_decode_request_visible_text_survives_when_thought_present() { #[test] fn test_encode_preserves_native_tool_groups_when_tools_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [ @@ -1579,7 +1579,7 @@ fn test_encode_preserves_native_tool_groups_when_tools_unchanged() { #[test] fn test_decode_native_only_tool_group_as_provider_native() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"googleSearch": {"apiKey": "sk-tool-secret"}}] @@ -1607,7 +1607,7 @@ fn test_decode_native_only_tool_group_as_provider_native() { #[test] fn test_decode_mixed_tool_group_exposes_native_siblings() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{ @@ -1639,7 +1639,7 @@ fn test_decode_mixed_tool_group_exposes_native_siblings() { #[test] fn test_encode_preserves_native_tool_groups_when_functions_change() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [ @@ -1685,7 +1685,7 @@ fn test_encode_preserves_native_tool_groups_when_functions_change() { /// response, behavior) must survive a decode → edit description → encode round-trip. #[test] fn test_encode_preserves_provider_tool_fields_when_description_changes() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{ @@ -1746,7 +1746,7 @@ fn test_encode_preserves_provider_tool_fields_when_description_changes() { #[test] fn test_encode_patches_gemini_provider_native_tool_group() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"googleSearch": {"apiKey": "sk-tool-secret"}}] @@ -1770,7 +1770,7 @@ fn test_encode_patches_gemini_provider_native_tool_group() { #[test] fn test_encode_patches_mixed_gemini_provider_native_tool_group() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{ @@ -1800,7 +1800,7 @@ fn test_encode_patches_mixed_gemini_provider_native_tool_group() { #[test] fn test_encode_deleting_mixed_native_sibling_does_not_rehome_later_native_group() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [ @@ -1839,7 +1839,7 @@ fn test_encode_deleting_mixed_native_sibling_does_not_rehome_later_native_group( #[test] fn test_decode_response_function_call_uses_id_field() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -1865,7 +1865,7 @@ fn test_decode_response_function_call_uses_id_field() { #[test] fn test_decode_response_function_call_fallback_to_name_when_no_id() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -1888,7 +1888,7 @@ fn test_decode_response_function_call_fallback_to_name_when_no_id() { /// Two simultaneous calls to the same function must have distinct IDs. #[test] fn test_decode_response_multi_call_same_function_distinct_ids() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -1920,7 +1920,7 @@ fn test_decode_response_multi_call_same_function_distinct_ids() { /// 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; // A multi-turn request where Gemini provided a functionCall with an explicit id. let original = make_request(json!({ "contents": [ @@ -1980,7 +1980,7 @@ fn test_encode_function_response_id_not_name() { /// (not the name) is stored in tool_call_id. #[test] fn test_decode_function_response_uses_id_not_name() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "hi"}]}, @@ -2013,7 +2013,7 @@ fn test_decode_function_response_uses_id_not_name() { /// original contents items are left untouched. #[test] fn test_encode_system_message_added_leaves_contents_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "original user"}]}, @@ -2052,7 +2052,7 @@ fn test_encode_system_message_added_leaves_contents_unchanged() { /// normalized_to_gemini_content (not a wrong original as base). #[test] fn test_encode_new_user_message_appended_encodes_correctly() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "hello"}]}, @@ -2103,7 +2103,7 @@ fn test_encode_new_user_message_appended_encodes_correctly() { #[test] fn test_streaming_accumulates_function_call_parts() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -2128,7 +2128,7 @@ fn test_streaming_accumulates_function_call_parts() { .unwrap(); let assembled = finalizer(); - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); assert_eq!(resp.finish_reason, Some(FinishReason::ToolUse)); @@ -2144,7 +2144,7 @@ fn test_streaming_accumulates_function_call_parts() { /// both encoded parts with the second part's signature. #[test] fn test_encode_thought_signature_preserved_for_parallel_same_function_calls() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "run two searches"}]}, @@ -2219,7 +2219,7 @@ fn test_encode_thought_signature_preserved_for_parallel_same_function_calls() { /// 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "run two searches"}]}, @@ -2282,7 +2282,7 @@ fn test_encode_reordered_same_name_function_calls_match_signatures_by_id() { /// index 1; only the message-list index changes from 1 to 2. #[test] fn test_encode_insertion_does_not_bleed_thought_signature() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "user message"}]}, @@ -2387,7 +2387,7 @@ fn test_encode_insertion_does_not_bleed_thought_signature() { /// and model turns must be preserved byte-identically. #[test] fn test_encode_prepend_user_message_preserves_thought_signature() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "original question"}]}, @@ -2473,7 +2473,7 @@ fn test_encode_prepend_user_message_preserves_thought_signature() { /// must have that signature preserved when an interceptor edits the text. #[test] fn test_encode_thought_signature_on_text_part_survives_edit() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "say hello"}]}, @@ -2523,7 +2523,7 @@ fn test_encode_thought_signature_on_text_part_survives_edit() { /// per part — not collapse to just the first. #[test] fn test_decode_multiple_function_responses_in_one_turn() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "run two tools"}]}, @@ -2580,7 +2580,7 @@ fn test_decode_multiple_function_responses_in_one_turn() { #[test] fn test_decode_function_response_with_native_sibling_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{ "role": "user", @@ -2599,7 +2599,7 @@ fn test_decode_function_response_with_native_sibling_errors() { #[test] fn test_function_response_nested_parts_are_exposed_and_patchable() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "show my ordered instrument"}]}, @@ -2670,7 +2670,7 @@ fn test_function_response_nested_parts_are_exposed_and_patchable() { #[test] fn test_decode_rejects_missing_contents() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({"model": "gemini-2.0-flash"})); assert!( codec.decode(&request).is_err(), @@ -2680,7 +2680,7 @@ fn test_decode_rejects_missing_contents() { #[test] fn test_decode_rejects_non_array_contents() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({"contents": "not an array"})); assert!( codec.decode(&request).is_err(), @@ -2690,7 +2690,7 @@ fn test_decode_rejects_non_array_contents() { #[test] fn test_decode_rejects_non_object_generation_config() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": "not an object" @@ -2703,7 +2703,7 @@ fn test_decode_rejects_non_object_generation_config() { #[test] fn test_decode_rejects_invalid_stop_sequences() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": {"stopSequences": "not an array"} @@ -2716,7 +2716,7 @@ fn test_decode_rejects_invalid_stop_sequences() { #[test] fn test_decode_rejects_non_array_tools() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": "not an array" @@ -2729,7 +2729,7 @@ fn test_decode_rejects_non_array_tools() { #[test] fn test_decode_rejects_function_declaration_without_name() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": [{"description": "no name"}]}] @@ -2748,7 +2748,7 @@ fn test_decode_rejects_function_declaration_without_name() { /// must return an error — silently dropping it would be data loss. #[test] fn test_encode_unsupported_role_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hello"}]}] })); @@ -2772,7 +2772,7 @@ fn test_encode_unsupported_role_returns_error() { /// functionResponse content item — not resurrect the deleted result. #[test] fn test_encode_delete_one_parallel_tool_response() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "run two tools"}]}, @@ -2860,7 +2860,7 @@ fn test_encode_delete_one_parallel_tool_response() { /// not duplicate or resurrect the original item. #[test] fn test_encode_insert_between_parallel_tool_responses() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "run two tools"}]}, @@ -2944,13 +2944,13 @@ fn test_encode_insert_between_parallel_tool_responses() { // thoughtsTokenCount accounting // =================================================================== -/// thoughtsTokenCount must be stored in ApiSpecificResponse::Gemini and +/// 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": {"role": "model", "parts": [{"text": "let me think..."}]}, @@ -2976,9 +2976,9 @@ fn test_decode_response_thoughts_token_count() { "fallback total_tokens must include thoughtsTokenCount" ); - // thoughts_tokens must be in ApiSpecificResponse::Gemini. + // thoughts_tokens must be in ApiSpecificResponse::GeminiGenerateContent. match resp.api_specific { - Some(ApiSpecificResponse::Gemini { + Some(ApiSpecificResponse::GeminiGenerateContent { thoughts_tokens, .. }) => { assert_eq!( @@ -2987,14 +2987,14 @@ fn test_decode_response_thoughts_token_count() { "thoughtsTokenCount must be in api_specific" ); } - other => panic!("expected ApiSpecificResponse::Gemini, got: {other:?}"), + 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [ {"parts": [{"text": "hello without a role"}]}, @@ -3024,7 +3024,7 @@ fn test_decode_roleless_content_treated_as_user() { // Wrong numeric types in generationConfig must return an error. #[test] fn test_decode_rejects_non_numeric_temperature() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": {"temperature": "hot"} @@ -3037,7 +3037,7 @@ fn test_decode_rejects_non_numeric_temperature() { #[test] fn test_decode_rejects_non_numeric_max_output_tokens() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": {"maxOutputTokens": "a lot"} @@ -3051,7 +3051,7 @@ fn test_decode_rejects_non_numeric_max_output_tokens() { // Streaming finalizer must propagate responseId when present. #[test] fn test_streaming_propagates_response_id() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -3075,7 +3075,7 @@ fn test_streaming_propagates_response_id() { "streaming finalizer must propagate responseId from SSE events" ); - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); assert_eq!( resp.id.as_deref(), @@ -3088,7 +3088,7 @@ fn test_streaming_propagates_response_id() { // 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -3111,7 +3111,7 @@ fn test_decode_response_max_tokens_with_function_call_is_length_not_tool_use() { // SAFETY + functionCall parts → ContentFilter, not ToolUse. #[test] fn test_decode_response_safety_with_function_call_is_content_filter_not_tool_use() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -3134,7 +3134,7 @@ fn test_decode_response_safety_with_function_call_is_content_filter_not_tool_use // STOP + functionCall parts → ToolUse (unchanged behaviour). #[test] fn test_decode_response_stop_with_function_call_is_tool_use() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": { @@ -3158,7 +3158,7 @@ fn test_decode_response_stop_with_function_call_is_tool_use() { // be silently dropped by patch_changed_gemini_content. #[test] fn test_encode_roleless_original_item_survives_edit() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; // A Gemini request whose user content item has no "role" field. let original = make_request(json!({ "contents": [{"parts": [{"text": "original text"}]}] @@ -3197,7 +3197,7 @@ fn test_encode_roleless_original_item_survives_edit() { // Multiple system messages must be merged into one systemInstruction. #[test] fn test_encode_multiple_system_messages_merged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -3237,7 +3237,7 @@ fn test_encode_multiple_system_messages_merged() { // Encoding unsupported fields must error, not silently ignore. #[test] fn test_encode_unsupported_field_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -3254,7 +3254,7 @@ fn test_encode_unsupported_field_returns_error() { // Unparsable tool-call arguments must return an error rather than becoming {}. #[test] fn test_encode_invalid_tool_call_arguments_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -3281,7 +3281,7 @@ fn test_encode_invalid_tool_call_arguments_returns_error() { // Thinking-only partial usage (no candidatesTokenCount). #[test] fn test_decode_response_thoughts_only_usage_computes_fallback_total() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": {"role": "model", "parts": []}, @@ -3309,7 +3309,7 @@ fn test_decode_response_thoughts_only_usage_computes_fallback_total() { // Decode validation: table-driven negative tests for malformed contents items. #[test] fn test_decode_contents_item_validation() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let cases: &[(&str, Json)] = &[ ( @@ -3367,7 +3367,7 @@ fn test_decode_contents_item_validation() { // Valid roleless contents item must decode successfully as a user message. #[test] fn test_decode_roleless_contents_item_is_valid() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"parts": [{"text": "hello"}]}] })); @@ -3381,7 +3381,7 @@ fn test_decode_roleless_contents_item_is_valid() { // Encoding a normalized message whose content has non-text parts must error. #[test] fn test_encode_non_text_content_part_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -3405,7 +3405,7 @@ fn test_encode_non_text_content_part_returns_error() { // 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "original"}]}] })); @@ -3430,7 +3430,7 @@ fn test_encode_inserted_non_text_content_returns_error() { // groups must return an error rather than silently collapsing them. #[test] fn test_encode_multiple_fn_decl_groups_error_on_edit() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [ @@ -3460,7 +3460,7 @@ fn test_encode_multiple_fn_decl_groups_error_on_edit() { // Streaming: candidate metadata (safetyRatings, groundingMetadata) must survive finalization. #[test] fn test_streaming_candidate_extras_survive_finalize() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let finalizer = streaming_codec.finalizer(); @@ -3480,33 +3480,33 @@ fn test_streaming_candidate_extras_survive_finalize() { let assembled = finalizer(); - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); use super::super::response::ApiSpecificResponse; match &resp.api_specific { - Some(ApiSpecificResponse::Gemini { + Some(ApiSpecificResponse::GeminiGenerateContent { safety_ratings, grounding_metadata, .. }) => { assert!( safety_ratings.is_some(), - "safetyRatings from streaming must survive in ApiSpecificResponse::Gemini" + "safetyRatings from streaming must survive in ApiSpecificResponse::GeminiGenerateContent" ); assert!( grounding_metadata.is_some(), - "groundingMetadata from streaming must survive in ApiSpecificResponse::Gemini" + "groundingMetadata from streaming must survive in ApiSpecificResponse::GeminiGenerateContent" ); } - other => panic!("expected ApiSpecificResponse::Gemini, got: {other:?}"), + other => panic!("expected ApiSpecificResponse::GeminiGenerateContent, got: {other:?}"), } } // Present-but-non-string role must error. #[test] fn test_decode_non_string_role_errors() { - let codec = GeminiCodec; + 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"}]}] @@ -3521,7 +3521,7 @@ fn test_decode_non_string_role_errors() { // functionResponse.response is required; missing must error. #[test] fn test_decode_function_response_missing_response_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "user", "parts": [ {"functionResponse": {"id": "c1", "name": "fn"}} @@ -3536,7 +3536,7 @@ fn test_decode_function_response_missing_response_errors() { // functionResponse.response must be an object. #[test] fn test_decode_function_response_non_object_response_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; for bad in [json!("string"), json!([1, 2]), json!(42)] { let req = make_request(json!({ "contents": [{"role": "user", "parts": [ @@ -3553,7 +3553,7 @@ fn test_decode_function_response_non_object_response_errors() { // Tool encode: non-object content is wrapped in {"output": ...}. #[test] fn test_encode_tool_content_non_object_is_wrapped() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "go"}]}, @@ -3578,7 +3578,7 @@ fn test_encode_tool_content_non_object_is_wrapped() { // System message with non-text content must error. #[test] fn test_encode_system_message_non_text_content_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -3600,7 +3600,7 @@ fn test_encode_system_message_non_text_content_errors() { // Tool message with non-text normalized content must error. #[test] fn test_encode_tool_message_non_text_content_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "go"}]}, @@ -3626,7 +3626,7 @@ fn test_encode_tool_message_non_text_content_errors() { // Refusal content must error (not silently drop). #[test] fn test_encode_refusal_content_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -3647,7 +3647,7 @@ fn test_encode_refusal_content_errors() { // tools[] validation negative cases. #[test] fn test_decode_tools_validation_negative() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let cases: &[(&str, Json)] = &[ ( "tools entry not an object", @@ -3681,7 +3681,7 @@ fn test_decode_tools_validation_negative() { // functionCall.args must be an object on request decode. #[test] fn test_decode_function_call_non_object_args_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; for bad in [json!([1, 2]), json!("string"), json!(42)] { let req = make_request(json!({ "contents": [{ @@ -3699,7 +3699,7 @@ fn test_decode_function_call_non_object_args_errors() { // functionCall arguments must be a JSON object on encode. #[test] fn test_encode_function_call_non_object_args_errors() { - let codec = GeminiCodec; + 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 { @@ -3723,7 +3723,7 @@ fn test_encode_function_call_non_object_args_errors() { // Unchanged payload with native inlineData round-trips byte-identically. #[test] fn test_encode_native_inline_data_round_trips_unchanged() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQ=="; let original = make_request(json!({ "contents": [{ @@ -3746,7 +3746,7 @@ fn test_encode_native_inline_data_round_trips_unchanged() { // generationConfig unmodeled keys are preserved when params are cleared. #[test] fn test_encode_clearing_params_preserves_unmodeled_gen_config_fields() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": { @@ -3781,7 +3781,7 @@ fn test_encode_clearing_params_preserves_unmodeled_gen_config_fields() { // All modeled keys cleared, empty generationConfig is removed. #[test] fn test_encode_clearing_params_removes_empty_gen_config() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": {"temperature": 0.5, "maxOutputTokens": 256} @@ -3798,7 +3798,7 @@ fn test_encode_clearing_params_removes_empty_gen_config() { // Response decode: malformed functionCall parts error rather than silently drop. #[test] fn test_decode_response_malformed_function_call_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let cases: &[(&str, Json)] = &[ ( "functionCall is non-object", @@ -3854,7 +3854,7 @@ fn test_decode_response_malformed_function_call_errors() { // STOP + malformed functionCall must error, not return Complete. #[test] fn test_decode_response_stop_with_malformed_function_call_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{"content": {"role": "model", "parts": [{"functionCall": {"name": "", "args": {}}}]}, @@ -3870,7 +3870,7 @@ fn test_decode_response_stop_with_malformed_function_call_errors() { // Request decode: present functionCall.id must be a non-empty string. #[test] fn test_decode_request_function_call_id_validation() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let cases: &[(&str, Json)] = &[ ( "functionCall.id non-string", @@ -3898,7 +3898,7 @@ fn test_decode_request_function_call_id_validation() { // Request decode: present functionResponse.id must be a non-empty string. #[test] fn test_decode_request_function_response_id_validation() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let cases: &[(&str, Json)] = &[ ( "functionResponse.id non-string", @@ -3926,7 +3926,7 @@ fn test_decode_request_function_response_id_validation() { // systemInstruction validation. #[test] fn test_decode_system_instruction_validation() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let base_contents = json!([{"role": "user", "parts": [{"text": "hi"}]}]); let cases: &[(&str, Json)] = &[ ( @@ -3978,7 +3978,7 @@ fn test_decode_system_instruction_validation() { // model must be a string when present. #[test] fn test_decode_non_string_model_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "model": 42 @@ -3989,7 +3989,7 @@ fn test_decode_non_string_model_errors() { // Encode: empty tool-call name must error. #[test] fn test_encode_empty_tool_call_name_errors() { - let codec = GeminiCodec; + 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 { @@ -4013,7 +4013,7 @@ fn test_encode_empty_tool_call_name_errors() { // Encode: empty tool_call_id must error. #[test] fn test_encode_empty_tool_call_id_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "go"}]}, @@ -4038,7 +4038,7 @@ fn test_encode_empty_tool_call_id_errors() { // Encode: empty FunctionDefinition.name must error. #[test] fn test_encode_empty_function_definition_name_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": [{"name": "good_fn"}]}] @@ -4061,7 +4061,7 @@ fn test_encode_empty_function_definition_name_errors() { // Encode: FunctionDefinition.strict must error when set. #[test] fn test_encode_function_definition_strict_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": [{"name": "fn"}]}] @@ -4084,7 +4084,7 @@ fn test_encode_function_definition_strict_errors() { // Decode: non-string description in functionDeclaration must error. #[test] fn test_decode_function_declaration_non_string_description_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [{"functionDeclarations": [{"name": "fn", "description": 42}]}] @@ -4098,7 +4098,7 @@ fn test_decode_function_declaration_non_string_description_errors() { // Encode: systemInstruction sibling fields are preserved when text changes. #[test] fn test_encode_system_instruction_preserves_sibling_fields() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "systemInstruction": { @@ -4149,7 +4149,7 @@ fn test_extract_content_text_treats_missing_part_type_as_text() { // 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "systemInstruction": {"role": 123, "parts": [{"text": "system"}]} @@ -4163,7 +4163,7 @@ fn test_decode_system_instruction_non_string_role_errors() { // Mixed functionResponse + functionCall in one content item must error. #[test] fn test_decode_mixed_fr_and_fc_parts_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{ "role": "user", @@ -4182,7 +4182,7 @@ fn test_decode_mixed_fr_and_fc_parts_errors() { // Non-numeric topP must error. #[test] fn test_decode_rejects_non_numeric_top_p() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let request = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "generationConfig": {"topP": "high"} @@ -4196,7 +4196,7 @@ fn test_decode_rejects_non_numeric_top_p() { // ProviderNative tool variants owned by a different provider must error on encode. #[test] fn test_encode_mismatched_provider_native_tool_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -4214,7 +4214,7 @@ fn test_encode_mismatched_provider_native_tool_returns_error() { #[test] fn test_encode_provider_native_tool_with_function_declarations_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -4232,7 +4232,7 @@ fn test_encode_provider_native_tool_with_function_declarations_returns_error() { #[test] fn test_encode_system_message_name_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "systemInstruction": {"parts": [{"text": "be helpful"}]}, "contents": [{"role": "user", "parts": [{"text": "hi"}]}] @@ -4251,7 +4251,7 @@ fn test_encode_system_message_name_returns_error() { #[test] fn test_encode_user_message_name_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); @@ -4267,7 +4267,7 @@ fn test_encode_user_message_name_returns_error() { #[test] fn test_encode_assistant_message_name_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "hi"}]}, @@ -4288,7 +4288,7 @@ fn test_encode_assistant_message_name_returns_error() { #[test] fn test_encode_assistant_tool_call_message_name_returns_error() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [ {"role": "user", "parts": [{"text": "use a tool"}]}, @@ -4312,7 +4312,7 @@ fn test_encode_assistant_tool_call_message_name_returns_error() { #[test] fn test_encode_previous_response_id_returns_error() { - let codec = GeminiCodec; + 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()); @@ -4324,7 +4324,7 @@ fn test_encode_previous_response_id_returns_error() { #[test] fn test_decode_non_string_text_part_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": 42}]}] })); @@ -4337,7 +4337,7 @@ fn test_decode_non_string_text_part_errors() { // functionResponse on model-role content must error. #[test] fn test_decode_function_response_on_model_role_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{ "role": "model", @@ -4353,7 +4353,7 @@ fn test_decode_function_response_on_model_role_errors() { // functionCall on user-role content must error. #[test] fn test_decode_function_call_on_user_role_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{ "role": "user", @@ -4369,7 +4369,7 @@ fn test_decode_function_call_on_user_role_errors() { // functionResponse mixed with visible text parts must error. #[test] fn test_decode_function_response_mixed_with_text_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{ "role": "user", @@ -4388,7 +4388,7 @@ fn test_decode_function_response_mixed_with_text_errors() { // Response decode: non-string text value must error. #[test] fn test_decode_response_non_string_text_part_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{ "content": {"role": "model", "parts": [{"text": 42}]}, @@ -4405,7 +4405,7 @@ fn test_decode_response_non_string_text_part_errors() { // systemInstruction part-level metadata is preserved when text changes. #[test] fn test_encode_system_instruction_part_metadata_preserved() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "systemInstruction": { @@ -4431,7 +4431,7 @@ fn test_encode_system_instruction_part_metadata_preserved() { // 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 = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{ "role": "user", @@ -4459,7 +4459,7 @@ fn test_encode_text_part_unknown_metadata_preserved() { // Request decode: part with both text and functionResponse must error. #[test] fn test_decode_part_with_text_and_function_response_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "user", "parts": [ {"text": "visible", "functionResponse": {"name": "fn", "response": {}}} @@ -4474,7 +4474,7 @@ fn test_decode_part_with_text_and_function_response_errors() { // Request decode: part with both text and functionCall must error. #[test] fn test_decode_part_with_text_and_function_call_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let req = make_request(json!({ "contents": [{"role": "model", "parts": [ {"text": "visible", "functionCall": {"name": "fn", "args": {}}} @@ -4489,7 +4489,7 @@ fn test_decode_part_with_text_and_function_call_errors() { // Response decode: non-object part must error. #[test] fn test_decode_response_non_object_part_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{"content": {"role": "model", "parts": ["not an object"]}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 1} @@ -4503,7 +4503,7 @@ fn test_decode_response_non_object_part_errors() { // Response decode: part with both text and functionCall must error. #[test] fn test_decode_response_part_text_and_function_call_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let response = json!({ "candidates": [{"content": {"role": "model", "parts": [ {"text": "hello", "functionCall": {"name": "fn", "args": {}}} @@ -4519,7 +4519,7 @@ fn test_decode_response_part_text_and_function_call_errors() { // Streaming: part with both text and functionCall must error via collector. #[test] fn test_streaming_part_text_and_function_call_errors() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let result = collector(json!({ @@ -4539,7 +4539,7 @@ fn test_streaming_part_text_and_function_call_errors() { // Streaming: non-object part must error. #[test] fn test_streaming_non_object_part_errors() { - let streaming_codec = GeminiStreamingCodec::new(); + 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}] @@ -4553,7 +4553,7 @@ fn test_streaming_non_object_part_errors() { // Streaming: non-string text value must error. #[test] fn test_streaming_non_string_text_errors() { - let streaming_codec = GeminiStreamingCodec::new(); + let streaming_codec = GeminiGenerateContentStreamingCodec::new(); let mut collector = streaming_codec.collector(); let result = collector(json!({ "candidates": [{"content": {"role": "model", "parts": [{"text": 42}]}, "index": 0}] @@ -4567,7 +4567,7 @@ fn test_streaming_non_string_text_errors() { // systemInstruction with multiple text parts must error on edit. #[test] fn test_encode_system_instruction_multiple_text_parts_edit_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "systemInstruction": { @@ -4589,7 +4589,7 @@ fn test_encode_system_instruction_multiple_text_parts_edit_errors() { // systemInstruction with a native non-text part must error on decode. #[test] fn test_decode_system_instruction_non_text_part_errors() { - let codec = GeminiCodec; + let codec = GeminiGenerateContentCodec; let original = make_request(json!({ "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "systemInstruction": { diff --git a/crates/core/tests/unit/codec/parity_tests.rs b/crates/core/tests/unit/codec/parity_tests.rs index 51675d01e..64ab62077 100644 --- a/crates/core/tests/unit/codec/parity_tests.rs +++ b/crates/core/tests/unit/codec/parity_tests.rs @@ -6,9 +6,10 @@ //! 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-specific semantics (contents array, functionResponse -//! role split, thinking tokens, etc.) are covered in gemini_tests.rs. Where -//! the schemas legitimately diverge, the divergence is asserted explicitly: +//! 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. @@ -606,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 454c11902..71a3b4ead 100644 --- a/crates/core/tests/unit/codec/resolve_tests.rs +++ b/crates/core/tests/unit/codec/resolve_tests.rs @@ -26,7 +26,7 @@ fn builtin_provider_surface_registry_keeps_request_priority() { ProviderSurface::OpenAIResponses, ProviderSurface::AnthropicMessages, ProviderSurface::OpenAIChat, - ProviderSurface::Gemini, + ProviderSurface::GeminiGenerateContent, ] ); } @@ -89,11 +89,11 @@ fn detect_request_none_for_unknown_or_non_object() { fn detect_request_gemini_by_contents() { assert_eq!( detect_request_surface(&json!({"contents": []})), - Some(ProviderSurface::Gemini) + Some(ProviderSurface::GeminiGenerateContent) ); assert_eq!( detect_request_surface(&json!({"contents": [{"role": "user", "parts": [{"text": "hi"}]}]})), - Some(ProviderSurface::Gemini) + Some(ProviderSurface::GeminiGenerateContent) ); // Higher-priority surfaces still win when their discriminators are present. assert_eq!( @@ -107,14 +107,14 @@ fn detect_request_gemini_by_contents() { fn detect_response_gemini_by_candidates() { assert_eq!( detect_response_surface(&json!({"candidates": []})), - Some(ProviderSurface::Gemini) + Some(ProviderSurface::GeminiGenerateContent) ); assert_eq!( detect_response_surface(&json!({ "candidates": [{"content": {"parts": [{"text": "hi"}]}, "finishReason": "STOP"}], "usageMetadata": {} })), - Some(ProviderSurface::Gemini) + Some(ProviderSurface::GeminiGenerateContent) ); // A scalar candidates key does not match (must be an array). assert_eq!( @@ -132,7 +132,7 @@ fn detect_response_gemini_by_prompt_feedback_block() { "safetyRatings": [] } })), - Some(ProviderSurface::Gemini) + Some(ProviderSurface::GeminiGenerateContent) ); } @@ -398,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!( @@ -451,7 +451,7 @@ const ALL_SURFACES: [ProviderSurface; 4] = [ ProviderSurface::OpenAIChat, ProviderSurface::OpenAIResponses, ProviderSurface::AnthropicMessages, - ProviderSurface::Gemini, + ProviderSurface::GeminiGenerateContent, ]; #[test] @@ -476,7 +476,10 @@ fn codec_name_uses_canonical_spellings() { ProviderSurface::AnthropicMessages.codec_name(), "anthropic_messages" ); - assert_eq!(ProviderSurface::Gemini.codec_name(), "gemini"); + assert_eq!( + ProviderSurface::GeminiGenerateContent.codec_name(), + "gemini_generate_content" + ); } #[test] @@ -494,7 +497,7 @@ fn supported_codec_names_track_the_builtin_registry() { "openai_responses", "anthropic_messages", "openai_chat", - "gemini" + "gemini_generate_content" ] ); let from_registry: Vec<_> = BUILTIN_PROVIDER_SURFACES @@ -544,7 +547,7 @@ fn request_codec_decodes_each_surface() { "contents": [{"role": "user", "parts": [{"text": "hi"}]}] })); assert!( - !request_codec(ProviderSurface::Gemini) + !request_codec(ProviderSurface::GeminiGenerateContent) .decode(&gemini) .expect("gemini request decodes") .messages @@ -595,7 +598,7 @@ fn response_codec_decodes_each_surface() { "usageMetadata": {"promptTokenCount": 1} }); assert_eq!( - response_codec(ProviderSurface::Gemini) + response_codec(ProviderSurface::GeminiGenerateContent) .decode_response(&gemini) .expect("gemini response decodes") .response_text(), 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 c5d18d523..0013a04c1 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs @@ -403,7 +403,7 @@ fn schema_contains_every_supported_nemo_guardrails_option() { "openai_chat", "openai_responses", "anthropic_messages", - "gemini" + "gemini_generate_content" ] )); assert!(schema_property_has_default( 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 a82d5e32e..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,7 +689,10 @@ fn local_codec_and_rewrite_helpers_cover_all_provider_surfaces() { LocalGuardrailsCodec::AnthropicMessages, ProviderSurface::AnthropicMessages, ), - (LocalGuardrailsCodec::Gemini, ProviderSurface::Gemini), + ( + LocalGuardrailsCodec::GeminiGenerateContent, + ProviderSurface::GeminiGenerateContent, + ), ] { assert_eq!(codec.provider_surface(), surface); assert_eq!( @@ -1078,7 +1081,7 @@ fn stream_text_extraction_handles_supported_codecs() { // Gemini: visible text parts reach the guardrail worker. assert_eq!( extract_stream_text( - LocalGuardrailsCodec::Gemini, + LocalGuardrailsCodec::GeminiGenerateContent, &json!({"candidates": [{"content": {"parts": [{"text": "visible"}]}, "index": 0}]}) ), Some("visible".to_string()) @@ -1097,7 +1100,7 @@ fn stream_text_extraction_handles_supported_codecs() { LocalGuardrailsCodec::AnthropicMessages, json!({"type": "content_block_delta", "delta": {"type": "input_json_delta"}}), ), - (LocalGuardrailsCodec::Gemini, Json::Null), + (LocalGuardrailsCodec::GeminiGenerateContent, Json::Null), ] { assert_eq!(extract_stream_text(codec, &chunk), None); } @@ -1108,7 +1111,7 @@ 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::Gemini, + LocalGuardrailsCodec::GeminiGenerateContent, &json!({"candidates": [{"content": {"parts": [{"thought": true, "text": "internal reasoning"}]}, "index": 0}]}) ), None, @@ -1117,7 +1120,7 @@ fn stream_text_extraction_gemini_skips_thought_parts() { // A chunk with both a thought part and a visible part: only the visible text is forwarded. assert_eq!( extract_stream_text( - LocalGuardrailsCodec::Gemini, + LocalGuardrailsCodec::GeminiGenerateContent, &json!({"candidates": [{"content": {"parts": [ {"thought": true, "text": "reasoning"}, {"text": "answer"} diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 61a7fe26a..0675080f2 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -971,7 +971,7 @@ struct FfiCodecHandle *nemo_relay_anthropic_messages_codec_new(void); * # Safety * Caller must free the returned handle via `nemo_relay_codec_free`. */ -struct FfiCodecHandle *nemo_relay_gemini_codec_new(void); +struct FfiCodecHandle *nemo_relay_gemini_generate_content_codec_new(void); /** * Execute an LLM call end-to-end: run conditional-execution guardrails (on raw diff --git a/crates/ffi/src/api/llm.rs b/crates/ffi/src/api/llm.rs index 43890baf3..31ba4f8e9 100644 --- a/crates/ffi/src/api/llm.rs +++ b/crates/ffi/src/api/llm.rs @@ -380,10 +380,12 @@ pub extern "C" fn nemo_relay_anthropic_messages_codec_new() -> *mut FfiCodecHand /// # Safety /// Caller must free the returned handle via `nemo_relay_codec_free`. #[unsafe(no_mangle)] -pub extern "C" fn nemo_relay_gemini_codec_new() -> *mut FfiCodecHandle { +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::GeminiCodec), - response_codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), + codec: Arc::new(nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec), + response_codec: Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), })) } diff --git a/crates/ffi/tests/unit/types_tests.rs b/crates/ffi/tests/unit/types_tests.rs index a6a09ecc0..15f0aa0cf 100644 --- a/crates/ffi/tests/unit/types_tests.rs +++ b/crates/ffi/tests/unit/types_tests.rs @@ -760,13 +760,13 @@ 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_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(), - "GeminiCodec FFI constructor must return a non-null handle" + "GeminiGenerateContentCodec FFI constructor must return a non-null handle" ); unsafe { diff --git a/crates/node/pii_redaction.d.ts b/crates/node/pii_redaction.d.ts index f546a99dc..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' | 'gemini' | 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 1afb84457..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' | 'gemini' } + | { 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 e9b58eb7f..4c8095991 100644 --- a/crates/node/src/types/mod.rs +++ b/crates/node/src/types/mod.rs @@ -483,20 +483,24 @@ impl OpenAIResponsesCodec { /// Built-in codec for the Gemini generateContent API. /// /// Implements both request codec (decode/encode) and response codec -/// (decodeResponse). Construct with `new GeminiCodec()`. -#[napi(js_name = "GeminiCodec")] -pub struct GeminiCodec { +/// (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 GeminiCodec { +impl GeminiGenerateContentCodec { #[napi(constructor)] pub fn new() -> Self { Self { - inner_codec: std::sync::Arc::new(nemo_relay::codec::gemini::GeminiCodec), - inner_response_codec: std::sync::Arc::new(nemo_relay::codec::gemini::GeminiCodec), + 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, + ), } } diff --git a/crates/node/tests/types_tests.mjs b/crates/node/tests/types_tests.mjs index 1108ed27d..e562b2575 100644 --- a/crates/node/tests/types_tests.mjs +++ b/crates/node/tests/types_tests.mjs @@ -23,7 +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.GeminiCodec, 'function'); + assert.equal(typeof lib.GeminiGenerateContentCodec, 'function'); }); it('scope type enum values', () => { @@ -82,19 +82,19 @@ describe('ScopeStack', () => { }); // =========================================================================== -// GeminiCodec +// GeminiGenerateContentCodec // =========================================================================== -describe('GeminiCodec', () => { - const { GeminiCodec } = lib; +describe('GeminiGenerateContentCodec', () => { + const { GeminiGenerateContentCodec } = lib; it('instantiates', () => { - const codec = new GeminiCodec(); - assert.ok(codec instanceof GeminiCodec); + const codec = new GeminiGenerateContentCodec(); + assert.ok(codec instanceof GeminiGenerateContentCodec); }); it('decode returns an AnnotatedLLMRequest with messages and no params', () => { - const codec = new GeminiCodec(); + const codec = new GeminiGenerateContentCodec(); const req = new LlmRequest({}, { contents: [ { role: 'user', parts: [{ text: 'hello' }] }, @@ -111,7 +111,7 @@ describe('GeminiCodec', () => { }); it('decode captures generationConfig into params', () => { - const codec = new GeminiCodec(); + const codec = new GeminiGenerateContentCodec(); const req = new LlmRequest({}, { contents: [{ role: 'user', parts: [{ text: 'hi' }] }], generationConfig: { temperature: 0.5, maxOutputTokens: 256 }, @@ -124,7 +124,7 @@ describe('GeminiCodec', () => { }); it('encode round-trips extra fields', () => { - const codec = new GeminiCodec(); + const codec = new GeminiGenerateContentCodec(); const req = new LlmRequest({}, { contents: [{ role: 'user', parts: [{ text: 'hi' }] }], safetySettings: [{ category: 'HARM_CATEGORY_HATE_SPEECH', threshold: 'BLOCK_NONE' }], @@ -138,7 +138,7 @@ describe('GeminiCodec', () => { }); it('decodeResponse extracts text and finish reason', () => { - const codec = new GeminiCodec(); + const codec = new GeminiGenerateContentCodec(); const raw = { candidates: [{ content: { role: 'model', parts: [{ text: 'Hello!' }] }, @@ -159,7 +159,7 @@ describe('GeminiCodec', () => { }); it('decodeResponse maps functionCall id correctly', () => { - const codec = new GeminiCodec(); + const codec = new GeminiGenerateContentCodec(); const raw = { candidates: [{ content: { diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 4d8904137..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`, `anthropic_messages`, and `gemini`. + `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 c0d4db900..8d27b69ea 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -304,7 +304,9 @@ impl CompiledBuiltinBackend { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) => { Some(ProviderSurface::AnthropicMessages) } - LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini) => Some(ProviderSurface::Gemini), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) => { + Some(ProviderSurface::GeminiGenerateContent) + } LlmCodecIdentity::Runtime(_) | LlmCodecIdentity::Opaque => None, } } @@ -404,7 +406,7 @@ impl CompiledBuiltinBackend { // 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::Gemini + if surface == ProviderSurface::GeminiGenerateContent && payload .get("candidates") .and_then(Json::as_array) diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 72fa57af7..4de162641 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -274,7 +274,7 @@ nemo_relay::editor_config! { codec => { label: "codec", kind: Enum, - values: ["openai_chat", "openai_responses", "anthropic_messages", "gemini"], + values: ["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"], optional: true, }, profiles => { label: "profiles", kind: List, list: &PII_REDACTION_PROFILE_LIST_ITEM }, @@ -504,7 +504,7 @@ fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::s "openai_chat", "openai_responses", "anthropic_messages", - "gemini", + "gemini_generate_content", ], None, ) diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index 071335cef..e0694cfde 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -12,7 +12,7 @@ pub(crate) enum BuiltinCodecName { OpenAIChat, OpenAIResponses, AnthropicMessages, - Gemini, + GeminiGenerateContent, } impl BuiltinCodecName { @@ -21,7 +21,7 @@ impl BuiltinCodecName { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, - ProviderSurface::Gemini => Self::Gemini, + ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, } } @@ -34,7 +34,7 @@ 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::Gemini => overlay_gemini_response(payload, annotated), + Self::GeminiGenerateContent => overlay_gemini_response(payload, annotated), } } } diff --git a/crates/pii-redaction/tests/coverage/overlay_tests.rs b/crates/pii-redaction/tests/coverage/overlay_tests.rs index 13d675de1..cae8a2ca3 100644 --- a/crates/pii-redaction/tests/coverage/overlay_tests.rs +++ b/crates/pii-redaction/tests/coverage/overlay_tests.rs @@ -165,7 +165,8 @@ fn gemini_overlay_redacts_candidate_text() { }); let annotated = gemini_annotated(Some("[REDACTED]"), None, None, None); - let result = BuiltinCodecName::Gemini.overlay_response_payload(payload, &annotated); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); let text = result["candidates"][0]["content"]["parts"][0]["text"] .as_str() @@ -212,7 +213,8 @@ fn gemini_overlay_redacts_provider_native_candidate_part() { ..Default::default() }; - let result = BuiltinCodecName::Gemini.overlay_response_payload(payload, &annotated); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); assert_eq!( result["candidates"][0]["content"]["parts"][1]["codeExecutionResult"]["output"], json!("[REDACTED]"), @@ -245,7 +247,8 @@ fn gemini_overlay_updates_tool_call_args() { None, None, ); - let result = BuiltinCodecName::Gemini.overlay_response_payload(payload, &annotated); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); let args = &result["candidates"][0]["content"]["parts"][0]["functionCall"]["args"]; assert_eq!( @@ -269,7 +272,8 @@ fn gemini_overlay_updates_response_id_and_model_version() { // 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::Gemini.overlay_response_payload(payload, &annotated); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); assert_eq!( result["responseId"], @@ -304,7 +308,8 @@ fn gemini_overlay_does_not_overwrite_finish_reason() { ..Default::default() }; - let result = BuiltinCodecName::Gemini.overlay_response_payload(payload, &annotated); + let result = + BuiltinCodecName::GeminiGenerateContent.overlay_response_payload(payload, &annotated); assert_eq!( result["candidates"][0]["finishReason"].as_str(), diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index 410aa6c98..cf64ca8cf 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -4355,7 +4355,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_tools_via_codec() { initialize_plugins(plugin_config(json!({ "mode": "builtin", - "codec": "gemini", + "codec": "gemini_generate_content", "input": true, "output": false, "tool_input": false, @@ -4384,7 +4384,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_tools_via_codec() { let _handle = llm_call( LlmCallParams::builder() - .name("gemini") + .name("gemini_generate_content") .request(&request) .build(), ) @@ -4418,7 +4418,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_request_content_via_co initialize_plugins(plugin_config(json!({ "mode": "builtin", - "codec": "gemini", + "codec": "gemini_generate_content", "input": true, "output": false, "tool_input": false, @@ -4454,7 +4454,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_request_content_via_co let _handle = llm_call( LlmCallParams::builder() - .name("gemini") + .name("gemini_generate_content") .request(&request) .build(), ) @@ -4484,7 +4484,7 @@ async fn builtin_backend_sanitizes_gemini_function_response_nested_parts_via_cod initialize_plugins(plugin_config(json!({ "mode": "builtin", - "codec": "gemini", + "codec": "gemini_generate_content", "input": true, "output": false, "tool_input": false, @@ -4529,7 +4529,7 @@ async fn builtin_backend_sanitizes_gemini_function_response_nested_parts_via_cod let _handle = llm_call( LlmCallParams::builder() - .name("gemini") + .name("gemini_generate_content") .request(&request) .build(), ) @@ -4554,14 +4554,14 @@ async fn builtin_backend_sanitizes_gemini_function_response_nested_parts_via_cod #[tokio::test] async fn builtin_backend_sanitizes_gemini_provider_native_response_content_via_codec() { - use crate::codec::gemini::GeminiCodec; + 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", + "codec": "gemini_generate_content", "input": false, "output": true, "tool_input": false, @@ -4593,7 +4593,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_response_content_via_c let _ = llm_call_execute( LlmCallExecuteParams::builder() - .name("gemini") + .name("gemini_generate_content") .request(LlmRequest { headers: serde_json::Map::new(), content: json!({ @@ -4601,7 +4601,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_response_content_via_c }), }) .func(noop_openai_chat_exec_fn(response)) - .response_codec(Arc::new(GeminiCodec)) + .response_codec(Arc::new(GeminiGenerateContentCodec)) .build(), ) .await @@ -5462,14 +5462,14 @@ async fn builtin_backend_sanitizes_openai_responses_output_text_alias_on_stream_ #[tokio::test] async fn builtin_backend_omits_multi_candidate_gemini_normalized_response() { - use crate::codec::gemini::GeminiCodec; + 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", + "codec": "gemini_generate_content", "input": false, "output": true, "tool_input": false, @@ -5512,7 +5512,7 @@ async fn builtin_backend_omits_multi_candidate_gemini_normalized_response() { }) // 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(GeminiCodec)) + .response_codec(Arc::new(GeminiGenerateContentCodec)) .build(), ) .await diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 18d0fff37..6302e3c14 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -60,8 +60,8 @@ pub enum BuiltinLlmCodec { #[serde(rename = "anthropic_messages")] AnthropicMessages, /// Gemini generateContent. - #[serde(rename = "gemini")] - Gemini, + #[serde(rename = "gemini_generate_content")] + GeminiGenerateContent, } /// Per-call LLM codec identity delivered to native plugins. @@ -3448,7 +3448,7 @@ fn llm_codec_identity_from_native( "openai_chat" => BuiltinLlmCodec::OpenAiChat, "openai_responses" => BuiltinLlmCodec::OpenAiResponses, "anthropic_messages" => BuiltinLlmCodec::AnthropicMessages, - "gemini" => BuiltinLlmCodec::Gemini, + "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 2dcd292d6..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, PyGeminiCodec, 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,7 +258,7 @@ fn py_llm_response_codec( if let Ok(builtin) = c.extract::>() { return Some(builtin.inner_response_codec.clone()); } - if let Ok(builtin) = c.extract::>() { + 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 @@ -282,7 +282,7 @@ 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::>() { + if let Ok(builtin) = codec.extract::>() { return Some(builtin.inner_codec.clone()); } Some(Arc::new(py_callable::PyLlmCodecWrapper { diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index 8ee050174..a7a56238d 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -1023,24 +1023,28 @@ impl PyAnthropicMessagesCodec { /// /// Example: /// ```python -/// from nemo_relay.codecs import GeminiCodec -/// codec = GeminiCodec() +/// from nemo_relay.codecs import GeminiGenerateContentCodec +/// codec = GeminiGenerateContentCodec() /// annotated_req = codec.decode(request) /// annotated_resp = codec.decode_response(response) /// ``` -#[pyclass(name = "GeminiCodec")] -pub struct PyGeminiCodec { +#[pyclass(name = "GeminiGenerateContentCodec")] +pub struct PyGeminiGenerateContentCodec { pub(crate) inner_codec: Arc, pub(crate) inner_response_codec: Arc, } #[pymethods] -impl PyGeminiCodec { +impl PyGeminiGenerateContentCodec { #[new] pub(crate) fn new() -> Self { Self { - inner_codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), - inner_response_codec: Arc::new(nemo_relay::codec::gemini::GeminiCodec), + inner_codec: Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), + inner_response_codec: Arc::new( + nemo_relay::codec::gemini_generate_content::GeminiGenerateContentCodec, + ), } } @@ -1077,6 +1081,6 @@ impl PyGeminiCodec { } 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 ac716fca8..f559c39c8 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -186,7 +186,7 @@ fn register_codec_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; 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 d15c5936d..b8932496e 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -191,15 +191,15 @@ fn default_cost_currency() -> String { /// /// Maps from provider-specific stop reasons: /// - **Complete**: OpenAI Chat `"stop"`, Anthropic `"end_turn"`, Responses `"completed"`, -/// Gemini `"STOP"` (without function-call parts) +/// Gemini generateContent `"STOP"` (without function-call parts) /// - **Length**: OpenAI Chat `"length"`, Anthropic `"max_tokens"`, Responses incomplete+max_output_tokens, -/// Gemini `"MAX_TOKENS"` +/// Gemini generateContent `"MAX_TOKENS"` /// - **ToolUse**: OpenAI Chat `"tool_calls"`, Anthropic `"tool_use"`, -/// Gemini `"TOOL_CODE"` (unconditional) or `"STOP"` when function-call parts are present +/// Gemini generateContent `"TOOL_CODE"` (unconditional) or `"STOP"` when function-call parts are present /// - **ContentFilter**: OpenAI Chat `"content_filter"`, Responses incomplete+content_filter, -/// Gemini `"SAFETY"` / `"RECITATION"` / `"BLOCKLIST"` / `"PROHIBITED_CONTENT"` and other policy codes +/// Gemini generateContent `"SAFETY"` / `"RECITATION"` / `"BLOCKLIST"` / `"PROHIBITED_CONTENT"` and other policy codes /// - **Unknown**: Forward-compatible catch-all for unrecognized reasons -/// (e.g. Gemini `"MALFORMED_FUNCTION_CALL"`, `"UNEXPECTED_TOOL_CALL"`) +/// (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 { @@ -234,7 +234,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` and Gemini's `args` are already parsed JSON objects. +/// 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. @@ -332,12 +332,12 @@ pub enum ApiSpecificResponse { }, /// Gemini generateContent API-specific fields. - #[serde(rename = "gemini")] - Gemini { - /// Tokens consumed by the model's internal reasoning (Gemini thinking). + #[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 response. + /// 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.). diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 0029e3ee9..47bc4afc7 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -155,7 +155,7 @@ pub enum BuiltinLlmCodec { /// Anthropic Messages. AnthropicMessages, /// Gemini generateContent request and response payloads. - Gemini, + GeminiGenerateContent, } /// Per-call LLM codec identity supplied to worker sanitizers. @@ -2159,7 +2159,9 @@ fn codec_identity_from_proto( Some("anthropic_messages") => { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) } - Some("gemini") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), + Some("gemini_generate_content") => { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) + } _ => LlmCodecIdentity::Opaque, }, Some(LlmCodecKind::Runtime) => codec_id @@ -2683,8 +2685,9 @@ fn rustc_version_runtime() -> String { mod tests { use super::*; - /// Verify that BuiltinLlmCodec has a Gemini variant and that codec_identity_from_proto - /// decodes the "gemini" id to the correct built-in identity (not Opaque). + /// Verify that BuiltinLlmCodec has a Gemini generateContent variant and that + /// codec_identity_from_proto decodes the "gemini_generate_content" id to the + /// correct built-in identity (not Opaque). #[test] fn test_gemini_codec_identity_decoded_as_builtin_not_opaque() { use nemo_relay_worker_proto::v1::LlmCodecIdentity as ProtoIdentity; @@ -2692,13 +2695,13 @@ mod tests { let proto = ProtoIdentity { kind: LlmCodecKind::Builtin as i32, - id: Some("gemini".to_string()), + id: Some("gemini_generate_content".to_string()), }; let identity = codec_identity_from_proto(Some(&proto)); assert_eq!( identity, - LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::Gemini), - "Gemini codec id must decode to BuiltIn(Gemini), not Opaque" + 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 8bed91b9d..e37d0b540 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -90,7 +90,7 @@ changing a key overlays it. - `api_specific` is a tagged, mutable surface for modeled Anthropic Messages, OpenAI Chat Completions, or OpenAI Responses request controls. Response annotations use the same pattern for provider-only metadata, including Gemini - thinking tokens. + 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` is reserved for unknown future fields. diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index b46637ca8..eeb38b833 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -299,7 +299,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`, `anthropic_messages`, and `gemini` +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 ebd09f2c7..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`, `anthropic_messages`, or `gemini`) 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/nemo-guardrails/configuration.mdx b/docs/configure-plugins/nemo-guardrails/configuration.mdx index b9141495b..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`, `gemini` | +| 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,7 +340,7 @@ The current built-in local mode supports managed LLM execution with: - `openai_chat` - `openai_responses` - `anthropic_messages` -- `gemini` +- `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 faa88702a..b2904c24d 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, Anthropic Messages, and Gemini -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`, `gemini` | 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 708b7f1d6..9df647b4f 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -91,7 +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. -- `GeminiCodec`: Gemini generateContent-compatible requests and responses. +- `GeminiGenerateContentCodec`: Gemini generateContent-compatible requests and responses. ## Provider Codec Roles @@ -107,7 +107,7 @@ 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 | `nemo_relay.codecs.GeminiCodec` | `GeminiCodec` 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. diff --git a/docs/integrate-into-frameworks/provider-response-codecs.mdx b/docs/integrate-into-frameworks/provider-response-codecs.mdx index c4ea2e3cc..0825cd6fd 100644 --- a/docs/integrate-into-frameworks/provider-response-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-response-codecs.mdx @@ -361,7 +361,7 @@ Relay does not model are dropped. Built-in codecs normalize provider field names as follows: -| Normalized field | OpenAI Chat | OpenAI Responses | Anthropic Messages | Gemini | +| 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` | @@ -369,7 +369,7 @@ Built-in codecs normalize provider field names as follows: | `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 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. +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 @@ -445,7 +445,7 @@ The built-in provider codecs also implement response decoding: - `OpenAIChatCodec` — OpenAI Chat Completions API - `OpenAIResponsesCodec` — OpenAI Responses API - `AnthropicMessagesCodec` — Anthropic Messages API -- `GeminiCodec` — Google Gemini generateContent 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/go/nemo_relay/callbacks_test.go b/go/nemo_relay/callbacks_test.go index ddf71a51c..2c19279ec 100644 --- a/go/nemo_relay/callbacks_test.go +++ b/go/nemo_relay/callbacks_test.go @@ -96,7 +96,7 @@ func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) openAIChat := "openai_chat" openAIResponses := "openai_responses" anthropicMessages := "anthropic_messages" - gemini := "gemini" + gemini := "gemini_generate_content" runtimeCodec := "com.example.chat.v1" cases := []codecIdentityTestCase{ @@ -105,7 +105,7 @@ func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) {"openai responses", 1, &openAIResponses, LLMCodecBuiltin}, {"anthropic messages", 1, &anthropicMessages, LLMCodecBuiltin}, // Gemini must be LLMCodecBuiltin, not Opaque — PII sanitization depends on this. - {"gemini", 1, &gemini, 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 ef3c1165b..da56733e2 100644 --- a/go/nemo_relay/coverage_gap_test.go +++ b/go/nemo_relay/coverage_gap_test.go @@ -473,7 +473,7 @@ func testWrapperAndCodecFinalizersRun(t *testing.T) { chatCodec := NewOpenAIChatCodec() responsesCodec := NewOpenAIResponsesCodec() anthropicCodec := NewAnthropicMessagesCodec() - geminiCodec := NewGeminiCodec() + geminiCodec := NewGeminiGenerateContentCodec() if chatCodec == nil || responsesCodec == nil || anthropicCodec == nil || geminiCodec == nil { t.Fatal("expected non-nil codec handles") } @@ -494,12 +494,12 @@ func testWrapperAndCodecFinalizersRun(t *testing.T) { } } -func TestGeminiCodecFunctionCallID(t *testing.T) { - runTestWithScopeStack(t, testGeminiCodecFunctionCallID) +func TestGeminiGenerateContentCodecFunctionCallID(t *testing.T) { + runTestWithScopeStack(t, testGeminiGenerateContentCodecFunctionCallID) } -func testGeminiCodecFunctionCallID(t *testing.T) { - // Verify that NewGeminiCodec returns a usable handle and that, +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(`{ @@ -537,10 +537,10 @@ func testGeminiCodecFunctionCallID(t *testing.T) { }, }, executor, - WithLLMResponseCodec(NewGeminiCodec()), + WithLLMResponseCodec(NewGeminiGenerateContentCodec()), ) if err != nil { - t.Fatalf("LlmCallExecute with GeminiCodec failed: %v", err) + t.Fatalf("LlmCallExecute with GeminiGenerateContentCodec failed: %v", err) } events := capturedEvents() diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index f26524c95..85ddb45fd 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -116,7 +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_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); @@ -924,7 +924,7 @@ 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], -// [NewAnthropicMessagesCodec], or [NewGeminiCodec]. The handle is +// [NewAnthropicMessagesCodec], or [NewGeminiGenerateContentCodec]. The handle is // automatically freed when garbage collected. type CodecHandle struct { ptr *C.FfiCodecHandle @@ -978,13 +978,13 @@ func NewAnthropicMessagesCodec() *CodecHandle { return h } -// NewGeminiCodec creates a codec for the Gemini generateContent API. +// 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 NewGeminiCodec() *CodecHandle { - h := &CodecHandle{ptr: C.nemo_relay_gemini_codec_new()} +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) @@ -996,7 +996,7 @@ func NewGeminiCodec() *CodecHandle { // WithLLMResponseCodec sets the response codec for this LLM call. // Pass a CodecHandle created by [NewOpenAIChatCodec], [NewOpenAIResponsesCodec], -// [NewAnthropicMessagesCodec], or [NewGeminiCodec]. +// [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 cba8889a2..6481de01e 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1157,7 +1157,7 @@ class AnthropicMessagesCodec: """Decode an Anthropic response into a normalized response view.""" ... -class GeminiCodec: +class GeminiGenerateContentCodec: """Built-in codec for Gemini generateContent requests and responses. Summary: @@ -1165,7 +1165,7 @@ class GeminiCodec: """ def __init__(self) -> None: - """Create a Gemini codec.""" + """Create a Gemini generateContent codec.""" ... def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: """Decode a Gemini generateContent request into a normalized request view.""" diff --git a/python/nemo_relay/codecs.py b/python/nemo_relay/codecs.py index 22ac68b7b..98cdb2e32 100644 --- a/python/nemo_relay/codecs.py +++ b/python/nemo_relay/codecs.py @@ -47,7 +47,7 @@ async def impl(request: LLMRequest): from nemo_relay._native import ( AnnotatedLLMRequest, AnthropicMessagesCodec, - GeminiCodec, + GeminiGenerateContentCodec, LLMRequest, OpenAIChatCodec, OpenAIResponsesCodec, @@ -162,7 +162,7 @@ def decode_response(self, response: Json) -> "AnnotatedLLMResponse": __all__ = [ "AnnotatedLLMRequest", "AnthropicMessagesCodec", - "GeminiCodec", + "GeminiGenerateContentCodec", "LlmCodec", "LlmResponseCodec", "OpenAIChatCodec", diff --git a/python/nemo_relay/codecs.pyi b/python/nemo_relay/codecs.pyi index c76d15c0f..d77ed6168 100644 --- a/python/nemo_relay/codecs.pyi +++ b/python/nemo_relay/codecs.pyi @@ -163,7 +163,7 @@ class AnthropicMessagesCodec: """ ... -class GeminiCodec: +class GeminiGenerateContentCodec: """Built-in codec for Gemini generateContent requests and responses.""" def __init__(self) -> None: ... @@ -204,7 +204,7 @@ class GeminiCodec: __all__ = [ "AnnotatedLLMRequest", "AnthropicMessagesCodec", - "GeminiCodec", + "GeminiGenerateContentCodec", "LlmCodec", "LlmResponseCodec", "OpenAIChatCodec", diff --git a/python/nemo_relay/pii_redaction.py b/python/nemo_relay/pii_redaction.py index 4c48d87e2..4c0f91235 100644 --- a/python/nemo_relay/pii_redaction.py +++ b/python/nemo_relay/pii_redaction.py @@ -136,7 +136,7 @@ class PiiRedactionConfig: tool_output: bool = True mark: bool = True priority: int = 100 - codec: Literal["openai_chat", "openai_responses", "anthropic_messages", "gemini"] | 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 1815f054c..8da9ebddc 100644 --- a/python/nemo_relay/pii_redaction.pyi +++ b/python/nemo_relay/pii_redaction.pyi @@ -58,7 +58,7 @@ class PiiRedactionConfig: tool_output: bool = ... mark: bool = ... priority: int = ... - codec: Literal["openai_chat", "openai_responses", "anthropic_messages", "gemini"] | 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 29d49d7e5..0094ed8e3 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -779,7 +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"), + (pb.LLM_CODEC_KIND_BUILTIN, "gemini_generate_content"), (pb.LLM_CODEC_KIND_RUNTIME, "com.example.chat.v1"), (pb.LLM_CODEC_KIND_OPAQUE, None), ]: @@ -831,8 +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"))), - ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "gemini"))), + ("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 b32df1418..07031c402 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -4,7 +4,7 @@ """Tests for built-in codec Python classes and LlmResponseCodec protocol. Covers: -- Built-in codec construction (OpenAIChatCodec, OpenAIResponsesCodec, AnthropicMessagesCodec, GeminiCodec) +- Built-in codec construction (OpenAIChatCodec, OpenAIResponsesCodec, AnthropicMessagesCodec, GeminiGenerateContentCodec) - Built-in codec decode/encode/decode_response methods for all four providers - LlmResponseCodec protocol - response_codec parameter accepts object (not string) @@ -22,7 +22,7 @@ llm, subscribers, ) -from nemo_relay.codecs import AnthropicMessagesCodec, GeminiCodec, OpenAIChatCodec, OpenAIResponsesCodec +from nemo_relay.codecs import AnthropicMessagesCodec, GeminiGenerateContentCodec, OpenAIChatCodec, OpenAIResponsesCodec # --------------------------------------------------------------------------- # 1. Built-in codec construction @@ -67,13 +67,13 @@ def test_anthropic_messages_codec_has_methods(self): assert hasattr(codec, "decode_response") def test_gemini_codec_constructable(self): - """GeminiCodec() is constructable.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec() is constructable.""" + codec = GeminiGenerateContentCodec() assert codec is not None def test_gemini_codec_has_methods(self): - """GeminiCodec has decode, encode, decode_response methods.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec has decode, encode, decode_response methods.""" + codec = GeminiGenerateContentCodec() assert hasattr(codec, "decode") assert hasattr(codec, "encode") assert hasattr(codec, "decode_response") @@ -285,8 +285,8 @@ def test_anthropic_messages_decode_response(self): assert annotated.response_text() == "Hello!" def test_gemini_codec_decode(self): - """GeminiCodec.decode() returns AnnotatedLLMRequest with messages and params.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec.decode() returns AnnotatedLLMRequest with messages and params.""" + codec = GeminiGenerateContentCodec() request = LLMRequest( {}, { @@ -306,8 +306,8 @@ def test_gemini_codec_decode(self): assert annotated.params is not None def test_gemini_codec_encode_round_trip(self): - """GeminiCodec.encode(decode(req), req) is idempotent when nothing changes.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec.encode(decode(req), req) is idempotent when nothing changes.""" + codec = GeminiGenerateContentCodec() original = LLMRequest( {}, { @@ -323,8 +323,8 @@ def test_gemini_codec_encode_round_trip(self): ) def test_gemini_codec_decode_response_text(self): - """GeminiCodec.decode_response() extracts text and usage from a generateContent response.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec.decode_response() extracts text and usage from a generateContent response.""" + codec = GeminiGenerateContentCodec() response = { "candidates": [ { @@ -349,8 +349,8 @@ def test_gemini_codec_decode_response_text(self): assert annotated.usage["prompt_tokens"] == 8 def test_gemini_codec_decode_response_safety_finish_reason(self): - """GeminiCodec maps SAFETY finish reason to 'content_filter', not 'unknown'.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec maps SAFETY finish reason to 'content_filter', not 'unknown'.""" + codec = GeminiGenerateContentCodec() response = { "candidates": [ { @@ -367,8 +367,8 @@ def test_gemini_codec_decode_response_safety_finish_reason(self): ) def test_gemini_codec_decode_response_function_call(self): - """GeminiCodec.decode_response() extracts functionCall parts as tool_calls.""" - codec = GeminiCodec() + """GeminiGenerateContentCodec.decode_response() extracts functionCall parts as tool_calls.""" + codec = GeminiGenerateContentCodec() response = { "candidates": [ { @@ -417,7 +417,7 @@ def test_builtin_codecs_satisfy_protocol(self): assert isinstance(OpenAIChatCodec(), LlmResponseCodec) assert isinstance(OpenAIResponsesCodec(), LlmResponseCodec) assert isinstance(AnthropicMessagesCodec(), LlmResponseCodec) - assert isinstance(GeminiCodec(), LlmResponseCodec) + assert isinstance(GeminiGenerateContentCodec(), LlmResponseCodec) # --------------------------------------------------------------------------- @@ -692,7 +692,7 @@ def test_importable_from_codecs_module(self): """Built-in codecs are importable from nemo_relay.codecs.""" from nemo_relay.codecs import ( AnthropicMessagesCodec, - GeminiCodec, + GeminiGenerateContentCodec, OpenAIChatCodec, OpenAIResponsesCodec, ) @@ -700,11 +700,11 @@ def test_importable_from_codecs_module(self): assert OpenAIChatCodec is not None assert OpenAIResponsesCodec is not None assert AnthropicMessagesCodec is not None - assert GeminiCodec 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, "GeminiCodec") + assert not hasattr(nemo_relay, "GeminiGenerateContentCodec") From 57ab92e5c4f05544fa5b65a48bc08dc1600443db Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 22:15:40 -0700 Subject: [PATCH 03/11] fix(codec): address Gemini review follow-ups Signed-off-by: Sahiti Bulusu --- crates/adaptive/src/response_cache/key.rs | 13 +- .../tests/unit/response_cache/key_tests.rs | 33 ++++ .../tests/unit/response_cache/replay_tests.rs | 8 +- crates/core/src/codec/resolve.rs | 11 ++ crates/core/tests/unit/codec/resolve_tests.rs | 23 +++ crates/node/tests/types_tests.mjs | 39 +++++ crates/pii-redaction/src/builtin.rs | 19 +-- crates/pii-redaction/src/component.rs | 22 +-- crates/pii-redaction/src/overlay.rs | 49 +++--- .../tests/coverage/overlay_tests.rs | 123 ++++++++++++++ .../tests/unit/component_tests.rs | 157 +++++++++++++++++- 11 files changed, 432 insertions(+), 65 deletions(-) diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index a2dd5f029..f7a2b0cca 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -337,7 +337,18 @@ fn lossy_request_shape(surface: ProviderSurface, content: &Json) -> bool { .and_then(Json::as_array) .is_some_and(|blocks| blocks.iter().any(lossy_system_block)) } - ProviderSurface::OpenAIResponses | ProviderSurface::GeminiGenerateContent => false, + ProviderSurface::OpenAIResponses => false, + ProviderSurface::GeminiGenerateContent => object + .get("generationConfig") + .and_then(Json::as_object) + .is_some_and(|generation_config| { + generation_config.keys().any(|key| { + !matches!( + key.as_str(), + "temperature" | "topP" | "maxOutputTokens" | "stopSequences" + ) + }) + }), } } diff --git a/crates/adaptive/tests/unit/response_cache/key_tests.rs b/crates/adaptive/tests/unit/response_cache/key_tests.rs index 2c341c178..90a2ce665 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -608,6 +608,39 @@ fn gemini_shaped_requests_key_on_the_detected_decode() { ); } +#[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 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 a87a26c56..921e172fc 100644 --- a/crates/adaptive/tests/unit/response_cache/replay_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/replay_tests.rs @@ -83,7 +83,13 @@ fn gemini_replay_uses_a_valid_native_stream_event() { let aggregate = json!({ "candidates": [{ "index": 0, - "content": {"role": "model", "parts": [{"text": "hello"}]}, + "content": { + "role": "model", + "parts": [ + {"text": "hello"}, + {"functionCall": {"id": "call_1", "name": "lookup", "args": {"q": "x"}}} + ] + }, "finishReason": "STOP" }], "usageMetadata": { diff --git a/crates/core/src/codec/resolve.rs b/crates/core/src/codec/resolve.rs index b775ec8fa..e6182c168 100644 --- a/crates/core/src/codec/resolve.rs +++ b/crates/core/src/codec/resolve.rs @@ -171,6 +171,9 @@ impl ProviderSurface { /// not a built-in provider codec. #[must_use] pub fn from_codec_name(name: &str) -> Option { + if name == "gemini" { + return Some(Self::GeminiGenerateContent); + } BUILTIN_PROVIDER_SURFACES .iter() .find(|descriptor| descriptor.codec_name == name) @@ -188,6 +191,14 @@ pub fn supported_codec_names() -> Vec<&'static str> { .collect() } +/// Codec names accepted at config boundaries, including compatibility aliases. +#[must_use] +pub fn supported_codec_names_with_aliases() -> Vec<&'static str> { + let mut names = supported_codec_names(); + names.push("gemini"); + names +} + /// Constructs the built-in bidirectional request codec ([`LlmCodec`]) for a surface. #[must_use] pub fn request_codec(surface: ProviderSurface) -> Arc { diff --git a/crates/core/tests/unit/codec/resolve_tests.rs b/crates/core/tests/unit/codec/resolve_tests.rs index 71a3b4ead..44f38644a 100644 --- a/crates/core/tests/unit/codec/resolve_tests.rs +++ b/crates/core/tests/unit/codec/resolve_tests.rs @@ -482,6 +482,19 @@ fn codec_name_uses_canonical_spellings() { ); } +#[test] +fn from_codec_name_accepts_gemini_compatibility_alias() { + assert_eq!( + ProviderSurface::from_codec_name("gemini"), + Some(ProviderSurface::GeminiGenerateContent) + ); + assert_eq!( + ProviderSurface::GeminiGenerateContent.codec_name(), + "gemini_generate_content", + "the legacy alias must not become the canonical codec spelling" + ); +} + #[test] fn from_codec_name_is_none_for_unknown_names() { assert_eq!(ProviderSurface::from_codec_name("generate_content"), None); @@ -505,6 +518,16 @@ fn supported_codec_names_track_the_builtin_registry() { .map(|descriptor| descriptor.codec_name) .collect(); assert_eq!(supported_codec_names(), from_registry); + assert_eq!( + supported_codec_names_with_aliases(), + vec![ + "openai_responses", + "anthropic_messages", + "openai_chat", + "gemini_generate_content", + "gemini" + ] + ); } #[test] diff --git a/crates/node/tests/types_tests.mjs b/crates/node/tests/types_tests.mjs index e562b2575..75c3aa9f0 100644 --- a/crates/node/tests/types_tests.mjs +++ b/crates/node/tests/types_tests.mjs @@ -179,4 +179,43 @@ describe('GeminiGenerateContentCodec', () => { // 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(new LlmRequest({}, { 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 = new LlmRequest({}, { + 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/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index 8d27b69ea..bc52cf8d9 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -399,7 +399,7 @@ 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; } @@ -411,7 +411,7 @@ impl CompiledBuiltinBackend { .get("candidates") .and_then(Json::as_array) .is_some_and(|candidates| candidates.len() > 1) - && self.targets_normalized_gemini_candidate() + && self.targets_normalized_single_projected_response() { return None; } @@ -461,20 +461,7 @@ impl CompiledBuiltinBackend { }) } - fn targets_normalized_openai_chat_choice(&self) -> bool { - self.target_paths.iter().any(|path| { - json_pointer_segments(path) - .and_then(|segments| segments.into_iter().next()) - .is_some_and(|segment| { - matches!( - segment.as_str(), - "message" | "tool_calls" | "finish_reason" | "api_specific" - ) - }) - }) - } - - fn targets_normalized_gemini_candidate(&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 4de162641..3da51e61b 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_with_aliases}; use nemo_relay::plugin::{ ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError, PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior, @@ -498,16 +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", - "gemini_generate_content", - ], - None, - ) + let codec_names = supported_codec_names_with_aliases(); + string_enum_schema(generator, &codec_names, None) } #[cfg(feature = "schema")] @@ -1135,17 +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_with_aliases(); push_policy_diag( diagnostics, policy.unsupported_value, "pii_redaction.unsupported_value", Some(PII_REDACTION_PLUGIN_KIND.to_string()), Some("codec".to_string()), - format!( - "codec must be one of: {}", - supported_codec_names().join(", ") - ), + 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 e0694cfde..dc896053d 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -115,42 +115,29 @@ fn overlay_gemini_response(mut payload: Json, annotated: &AnnotatedLlmResponse) }); parts.extend(sanitized); } else { - // Overlay text: update all non-thought text parts, splitting multi-line normalized - // text across them when there are multiple. + // 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 text_part_count = parts - .iter() - .filter(|p| { - p.get("text").is_some() && p.get("thought").and_then(Json::as_bool) != Some(true) - }) - .count(); - let mut non_thought_text_count = 0usize; - for part in parts.iter_mut() { + 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 { - continue; + return true; } let Some(p) = part.as_object_mut() else { - continue; + return false; }; - let value = if text_part_count <= 1 { - message_text.as_deref() - } else { - // Multiple text parts: split on newline and assign one chunk each. - message_text - .as_deref() - .map(|t| t.split('\n').collect::>()) - .as_ref() - .and_then(|chunks| chunks.get(non_thought_text_count).copied()) - .or_else(|| { - (non_thought_text_count == 0) - .then_some(message_text.as_deref()) - .flatten() - }) + if wrote_text { + return false; + } + let Some(text) = message_text.as_deref() else { + return false; }; - set_optional_string_field(p, "text", value); - non_thought_text_count += 1; - } + set_optional_string_field(p, "text", Some(text)); + wrote_text = true; + true + }); } // Overlay functionCall parts. @@ -167,7 +154,9 @@ fn overlay_gemini_response(mut payload: Json, annotated: &AnnotatedLlmResponse) let Some(sc) = sanitized.next() else { return false; // no sanitized call left — remove this part }; - set_optional_string_field(fc, "id", Some(sc.id.as_str())); + 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 diff --git a/crates/pii-redaction/tests/coverage/overlay_tests.rs b/crates/pii-redaction/tests/coverage/overlay_tests.rs index cae8a2ca3..621eaa5f7 100644 --- a/crates/pii-redaction/tests/coverage/overlay_tests.rs +++ b/crates/pii-redaction/tests/coverage/overlay_tests.rs @@ -177,6 +177,53 @@ fn gemini_overlay_redacts_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!({ @@ -258,6 +305,82 @@ fn gemini_overlay_updates_tool_call_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!({ diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index cf64ca8cf..7668b9a66 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -1794,6 +1794,27 @@ fn validate_allows_llm_surfaces_without_codec() { assert!(report.diagnostics.is_empty(), "{report:?}"); } +#[test] +fn validate_accepts_legacy_gemini_codec_alias() { + 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.is_empty(), "{report:?}"); +} + #[test] fn validate_rejects_regex_replace_without_pattern() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -4347,6 +4368,61 @@ async fn builtin_backend_removes_targeted_message_names_and_ignores_missing_norm clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn builtin_backend_accepts_legacy_gemini_codec_alias() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "gemini", + "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-legacy-alias"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "contents": [{"role": "user", "parts": [{"text": "sk-legacy-secret"}]}], + }), + }; + + let _handle = llm_call( + LlmCallParams::builder() + .name("gemini-legacy-alias") + .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-legacy-alias").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(); @@ -4600,7 +4676,7 @@ async fn builtin_backend_sanitizes_gemini_provider_native_response_content_via_c "contents": [{"role": "user", "parts": [{"text": "run code"}]}] }), }) - .func(noop_openai_chat_exec_fn(response)) + .func(noop_openai_chat_exec_fn(response.clone())) .response_codec(Arc::new(GeminiGenerateContentCodec)) .build(), ) @@ -5533,3 +5609,82 @@ async fn builtin_backend_omits_multi_candidate_gemini_normalized_response() { 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(); +} From f187dda2f9df2724f63e6ee074d2dfac40fb79be Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 22:45:01 -0700 Subject: [PATCH 04/11] fix(codec): address Gemini generateContent review feedback Signed-off-by: Sahiti Bulusu --- .../adaptive/src/acg/request_surfaces/mod.rs | 20 ++- crates/adaptive/src/response_cache/key.rs | 154 ++++++++++++++++-- .../tests/unit/acg/request_surface_tests.rs | 12 ++ .../tests/unit/response_cache/key_tests.rs | 101 ++++++++++++ .../tests/unit/response_cache/replay_tests.rs | 39 ++++- .../core/src/codec/gemini_generate_content.rs | 10 +- crates/core/src/codec/resolve.rs | 8 +- .../codec/gemini_generate_content_tests.rs | 27 +++ .../core/tests/unit/codec/response_tests.rs | 58 +++++++ crates/node/tests/types_tests.mjs | 50 +++--- crates/types/src/codec/response.rs | 36 +++- crates/worker/src/lib.rs | 25 +-- .../worker/tests/unit/codec_identity_tests.rs | 21 +++ docs/about-nemo-relay/concepts/codecs.mdx | 4 +- .../pii-redaction/configuration.mdx | 2 +- .../provider-codecs.mdx | 4 +- .../provider-response-codecs.mdx | 6 +- .../using-codecs.mdx | 2 +- 18 files changed, 500 insertions(+), 79 deletions(-) create mode 100644 crates/worker/tests/unit/codec_identity_tests.rs diff --git a/crates/adaptive/src/acg/request_surfaces/mod.rs b/crates/adaptive/src/acg/request_surfaces/mod.rs index e590dd4a6..b4834aa56 100644 --- a/crates/adaptive/src/acg/request_surfaces/mod.rs +++ b/crates/adaptive/src/acg/request_surfaces/mod.rs @@ -48,7 +48,8 @@ impl RequestSurface { ProviderSurface::OpenAIChat => Some(Self::OpenAIChat), ProviderSurface::OpenAIResponses => Some(Self::OpenAIResponses), ProviderSurface::AnthropicMessages => Some(Self::AnthropicMessages), - _ => None, + // Gemini generateContent ACG request editing is intentionally unsupported. + ProviderSurface::GeminiGenerateContent => None, } } @@ -91,13 +92,16 @@ impl RequestSurface { pub(crate) fn resolve_request_surface_from_request( request: &LlmRequest, ) -> crate::acg::Result { - detect_request_surface(&request.content) - .and_then(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 f7a2b0cca..29f1b6555 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -338,17 +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") - .and_then(Json::as_object) - .is_some_and(|generation_config| { - generation_config.keys().any(|key| { - !matches!( - key.as_str(), - "temperature" | "topP" | "maxOutputTokens" | "stopSequences" - ) + 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)) + } } } @@ -373,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/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 90a2ce665..54912deb6 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -641,6 +641,107 @@ fn gemini_unmodeled_generation_config_fields_do_not_collide() { ); } +#[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 921e172fc..7d5fa65a8 100644 --- a/crates/adaptive/tests/unit/response_cache/replay_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/replay_tests.rs @@ -86,11 +86,22 @@ fn gemini_replay_uses_a_valid_native_stream_event() { "content": { "role": "model", "parts": [ - {"text": "hello"}, - {"functionCall": {"id": "call_1", "name": "lookup", "args": {"q": "x"}}} + {"text": "hello", "thoughtSignature": "sig_TEXT=="}, + { + "functionCall": { + "id": "call_1", + "name": "lookup", + "args": {"q": "x"} + }, + "thoughtSignature": "sig_CALL==" + } ] }, - "finishReason": "STOP" + "finishReason": "STOP", + "safetyRatings": [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"} + ], + "groundingMetadata": {"webSearchQueries": ["example query"]} }], "usageMetadata": { "promptTokenCount": 9, @@ -112,6 +123,28 @@ fn gemini_replay_uses_a_valid_native_stream_event() { ); } +#[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/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs index 6d179e3ab..c8be44b28 100644 --- a/crates/core/src/codec/gemini_generate_content.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -361,7 +361,14 @@ fn validate_gemini_nested_function_response_part(part: &Json) -> Result<&str> { } fn gemini_function_response_to_message_content(fr: &Json) -> Result { - let response = fr.get("response").unwrap().clone(); + 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)); @@ -1314,6 +1321,7 @@ impl LlmResponseCodec for GeminiGenerateContentCodec { let safety_ratings = extra.remove("safetyRatings"); let grounding_metadata = extra.remove("groundingMetadata"); let citation_metadata = extra.remove("citationMetadata"); + extra.remove("api"); Some( super::response::ApiSpecificResponse::GeminiGenerateContent { thoughts_tokens, diff --git a/crates/core/src/codec/resolve.rs b/crates/core/src/codec/resolve.rs index e6182c168..6a766100e 100644 --- a/crates/core/src/codec/resolve.rs +++ b/crates/core/src/codec/resolve.rs @@ -160,15 +160,15 @@ fn descriptor_for(surface: ProviderSurface) -> &'static ProviderSurfaceDescripto } 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 or compatibility alias to its surface, or + /// `None` when `name` is not a built-in provider codec. `"gemini"` is accepted + /// as a compatibility alias for the canonical `"gemini_generate_content"`. #[must_use] pub fn from_codec_name(name: &str) -> Option { if name == "gemini" { diff --git a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs index f87fdea57..6c8f1a067 100644 --- a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs +++ b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs @@ -336,6 +336,28 @@ fn test_decode_response_extra_fields_preserved() { 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))); +} + // =================================================================== // Request decode tests // =================================================================== @@ -3533,6 +3555,11 @@ fn test_decode_function_response_missing_response_errors() { ); } +#[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() { 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/node/tests/types_tests.mjs b/crates/node/tests/types_tests.mjs index 75c3aa9f0..f4812c16d 100644 --- a/crates/node/tests/types_tests.mjs +++ b/crates/node/tests/types_tests.mjs @@ -95,13 +95,16 @@ describe('GeminiGenerateContentCodec', () => { it('decode returns an AnnotatedLLMRequest with messages and no params', () => { const codec = new GeminiGenerateContentCodec(); - const req = new LlmRequest({}, { - contents: [ - { role: 'user', parts: [{ text: 'hello' }] }, - { role: 'model', parts: [{ text: 'hi' }] }, - ], - systemInstruction: { parts: [{ text: 'Be helpful.' }] }, - }); + 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'); @@ -112,10 +115,13 @@ describe('GeminiGenerateContentCodec', () => { it('decode captures generationConfig into params', () => { const codec = new GeminiGenerateContentCodec(); - const req = new LlmRequest({}, { - contents: [{ role: 'user', parts: [{ text: 'hi' }] }], - generationConfig: { temperature: 0.5, maxOutputTokens: 256 }, - }); + 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); @@ -125,10 +131,13 @@ describe('GeminiGenerateContentCodec', () => { it('encode round-trips extra fields', () => { const codec = new GeminiGenerateContentCodec(); - const req = new LlmRequest({}, { - contents: [{ role: 'user', parts: [{ text: 'hi' }] }], - safetySettings: [{ category: 'HARM_CATEGORY_HATE_SPEECH', threshold: 'BLOCK_NONE' }], - }); + 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( @@ -183,16 +192,19 @@ describe('GeminiGenerateContentCodec', () => { it('decode throws for malformed Gemini requests', () => { const codec = new GeminiGenerateContentCodec(); assert.throws( - () => codec.decode(new LlmRequest({}, { contents: 'not an array' })), + () => 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 = new LlmRequest({}, { - contents: [{ role: 'user', parts: [{ text: 'hi' }] }], - }); + const original = { + headers: {}, + content: { + contents: [{ role: 'user', parts: [{ text: 'hi' }] }], + }, + }; assert.throws( () => codec.encode({ messages: 'not an array' }, original), /invalid AnnotatedLlmRequest/, diff --git a/crates/types/src/codec/response.rs b/crates/types/src/codec/response.rs index b8932496e..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; @@ -347,7 +348,13 @@ pub enum ApiSpecificResponse { #[serde(skip_serializing_if = "Option::is_none")] citation_metadata: Option, /// Any remaining candidate-level fields not modeled above. - #[serde(flatten, skip_serializing_if = "serde_json::Map::is_empty")] + #[serde( + flatten, + default, + skip_serializing_if = "gemini_extra_is_empty", + serialize_with = "serialize_gemini_extra", + deserialize_with = "deserialize_gemini_extra" + )] extra: serde_json::Map, }, @@ -361,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 47bc4afc7..92133ae4a 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -2682,26 +2682,5 @@ fn rustc_version_runtime() -> String { } #[cfg(test)] -mod tests { - use super::*; - - /// Verify that BuiltinLlmCodec has a Gemini generateContent variant and that - /// codec_identity_from_proto decodes the "gemini_generate_content" id to the - /// correct built-in identity (not Opaque). - #[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" - ); - } -} +#[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 e37d0b540..dda6c48ff 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -90,7 +90,7 @@ changing a key overlays it. - `api_specific` is a tagged, mutable surface for modeled Anthropic Messages, OpenAI Chat Completions, or OpenAI Responses request controls. Response annotations use the same pattern for provider-only metadata, including Gemini - generateContent thinking tokens. + `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` is reserved for unknown future fields. @@ -121,7 +121,7 @@ normalization logic without changing runtime ownership or public binding APIs. Provider schema extraction is codec-owned. Built-in provider surfaces, such as OpenAI Chat Completions, OpenAI Responses, Anthropic Messages, and Gemini -generateContent, each own the logic that recognizes their request and response +`generateContent`, each own the logic that recognizes their request and response shapes and maps them into `AnnotatedLlmRequest` or `AnnotatedLlmResponse`. When a managed LLM event already has an annotation, subscribers and exporters diff --git a/docs/configure-plugins/pii-redaction/configuration.mdx b/docs/configure-plugins/pii-redaction/configuration.mdx index b2904c24d..e1a77cf12 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -95,7 +95,7 @@ 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, Anthropic Messages, and -Gemini generateContent traffic: +Gemini `generateContent` traffic: ```toml [[components]] diff --git a/docs/integrate-into-frameworks/provider-codecs.mdx b/docs/integrate-into-frameworks/provider-codecs.mdx index 9df647b4f..75c099bec 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -91,7 +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. +- `GeminiGenerateContentCodec`: Gemini `generateContent`-compatible requests and responses. ## Provider Codec Roles @@ -107,7 +107,7 @@ 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` | +| 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. diff --git a/docs/integrate-into-frameworks/provider-response-codecs.mdx b/docs/integrate-into-frameworks/provider-response-codecs.mdx index 0825cd6fd..8b6a7dd2a 100644 --- a/docs/integrate-into-frameworks/provider-response-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-response-codecs.mdx @@ -361,7 +361,7 @@ Relay does not model are dropped. Built-in codecs normalize provider field names as follows: -| Normalized field | OpenAI Chat | OpenAI Responses | Anthropic Messages | Gemini generateContent | +| 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` | @@ -369,7 +369,7 @@ Built-in codecs normalize provider field names as follows: | `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. +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 @@ -445,7 +445,7 @@ The built-in provider codecs also implement response decoding: - `OpenAIChatCodec` — OpenAI Chat Completions API - `OpenAIResponsesCodec` — OpenAI Responses API - `AnthropicMessagesCodec` — Anthropic Messages API -- `GeminiGenerateContentCodec` — Google Gemini generateContent 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 6e4887094..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, Gemini generateContent, 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. From 79c97429f116ed34047a4c27fde533f3960532bd Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 22:53:37 -0700 Subject: [PATCH 05/11] fix(codec): reject lossy Gemini streaming candidates Signed-off-by: Sahiti Bulusu --- .../core/src/codec/gemini_generate_content.rs | 52 +++++++++-- .../codec/gemini_generate_content_tests.rs | 91 +++++++++++++++++++ 2 files changed, 136 insertions(+), 7 deletions(-) diff --git a/crates/core/src/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs index c8be44b28..0a06e4ef8 100644 --- a/crates/core/src/codec/gemini_generate_content.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -2705,6 +2705,7 @@ impl super::streaming::StreamingCodec for GeminiGenerateContentStreamingCodec { #[derive(Debug, Default)] struct GeminiGenerateContentStreamingState { parts: Vec, + candidate_index: Option, finish_reason: Option, usage_metadata: Option, model_version: Option, @@ -2747,7 +2748,43 @@ impl GeminiGenerateContentStreamingState { if let Some(candidates) = event.get("candidates").and_then(Json::as_array) && let Some(candidate) = candidates.first() { - if let Some(parts) = candidate + 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) @@ -2808,11 +2845,9 @@ impl GeminiGenerateContentStreamingState { // 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. - if let Some(cand_obj) = candidate.as_object() { - for (k, v) in cand_obj { - if !matches!(k.as_str(), "content" | "finishReason" | "index") { - self.candidate_extra.insert(k.clone(), v.clone()); - } + for (k, v) in candidate_obj { + if !matches!(k.as_str(), "content" | "finishReason" | "index") { + self.candidate_extra.insert(k.clone(), v.clone()); } } } @@ -2851,7 +2886,10 @@ impl GeminiGenerateContentStreamingState { if let Some(reason) = self.finish_reason { candidate_obj.insert("finishReason".into(), Json::String(reason)); } - candidate_obj.insert("index".into(), Json::from(0u32)); + 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 { diff --git a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs index 6c8f1a067..bd486893a 100644 --- a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs +++ b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs @@ -695,6 +695,7 @@ fn test_streaming_two_chunks_accumulated() { .unwrap(); let assembled = finalizer(); + assert_eq!(assembled["candidates"][0]["index"].as_u64(), Some(0)); let codec = GeminiGenerateContentCodec; let resp = codec.decode_response(&assembled).unwrap(); @@ -710,6 +711,96 @@ fn test_streaming_two_chunks_accumulated() { 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(); From f031d0cd742d3558142e25f14abd62499cc64fe3 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 22:55:11 -0700 Subject: [PATCH 06/11] docs(codec): clarify Gemini codec gateway selection Signed-off-by: Sahiti Bulusu --- docs/about-nemo-relay/concepts/codecs.mdx | 13 +++++++------ docs/integrate-into-frameworks/provider-codecs.mdx | 10 ++++++---- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/docs/about-nemo-relay/concepts/codecs.mdx b/docs/about-nemo-relay/concepts/codecs.mdx index dda6c48ff..d02ff08e6 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -166,12 +166,13 @@ hints let Anthropic Messages requests without a top-level `system` field decode through the Anthropic provider surface instead of being treated as shape-only OpenAI Chat payloads. -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 `/v1/messages`, +`/v1/chat/completions`, and `/v1/responses`, for both buffered and streaming +execution. Count-token, model, probe, Gemini `generateContent`, 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/integrate-into-frameworks/provider-codecs.mdx b/docs/integrate-into-frameworks/provider-codecs.mdx index 75c099bec..e8da08e84 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -111,10 +111,12 @@ The built-in provider codecs expose the same core methods: 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 OpenAI Chat, OpenAI Responses, and +Anthropic Messages request codecs automatically for `/v1/messages`, +`/v1/chat/completions`, and `/v1/responses`, for both buffered and streaming +calls. `GeminiGenerateContentCodec` is available for direct use with Gemini +`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 From 31bf7350d46fc964335bf641928cc08264beb1d0 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 22:59:04 -0700 Subject: [PATCH 07/11] fix(codec): ignore reserved Gemini api metadata Signed-off-by: Sahiti Bulusu --- .../core/src/codec/gemini_generate_content.rs | 7 +++---- .../unit/codec/gemini_generate_content_tests.rs | 17 +++++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/crates/core/src/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs index 0a06e4ef8..00ce5ccd1 100644 --- a/crates/core/src/codec/gemini_generate_content.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -1312,16 +1312,15 @@ impl LlmResponseCodec for GeminiGenerateContentCodec { // citationMetadata, thinking token count, …) that cannot be normalized // across providers. thoughts_tokens lives here, not in shared Usage. let api_specific = { - let cand_extra = candidate.map(|c| c.extra.clone()).unwrap_or_default(); - let has_candidate_data = !cand_extra.is_empty() || thoughts_tokens.is_some(); + 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 mut extra = cand_extra; let safety_ratings = extra.remove("safetyRatings"); let grounding_metadata = extra.remove("groundingMetadata"); let citation_metadata = extra.remove("citationMetadata"); - extra.remove("api"); Some( super::response::ApiSpecificResponse::GeminiGenerateContent { thoughts_tokens, diff --git a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs index bd486893a..9e5fe7ec4 100644 --- a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs +++ b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs @@ -358,6 +358,23 @@ fn test_decode_response_candidate_extra_removes_reserved_api_key() { 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 // =================================================================== From 7d99b5151d6c5744777404b59e7649b86d60b028 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 23:11:38 -0700 Subject: [PATCH 08/11] test(nemo-guardrails): cover Gemini codec validation Signed-off-by: Sahiti Bulusu --- .../nemo_guardrails/component_tests.rs | 34 ++++++++++++++++--- 1 file changed, 29 insertions(+), 5 deletions(-) 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 0013a04c1..f838f1da2 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs @@ -644,12 +644,17 @@ fn assert_invalid_remote_identity_and_codec() { } }))); assert!(bad_codec.has_errors()); - assert!( - bad_codec - .diagnostics + assert!(bad_codec.diagnostics.iter().any(|diag| { + diag.message.contains("codec must be one of:") + && [ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate_content", + ] .iter() - .any(|diag| { diag.message.contains("codec must be one of:") }) - ); + .all(|name| diag.message.contains(name)) + })); let unsupported_remote_codec = validate_plugin_config(&plugin_config(json!({ "mode": "remote", @@ -683,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() { From 810353373a2e38198447aa4bf8bc329d94520299 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Wed, 5 Aug 2026 23:34:49 -0700 Subject: [PATCH 09/11] fix(codec): preserve Gemini streaming signed parts Signed-off-by: Sahiti Bulusu --- crates/adaptive/src/response_cache/replay.rs | 5 +-- .../core/src/codec/gemini_generate_content.rs | 7 ++-- .../codec/gemini_generate_content_tests.rs | 34 +++++++++++++++---- .../adaptive/response-cache.mdx | 11 +++--- 4 files changed, 42 insertions(+), 15 deletions(-) diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs index 5e87b4784..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 diff --git a/crates/core/src/codec/gemini_generate_content.rs b/crates/core/src/codec/gemini_generate_content.rs index 00ce5ccd1..c4137f23d 100644 --- a/crates/core/src/codec/gemini_generate_content.rs +++ b/crates/core/src/codec/gemini_generate_content.rs @@ -2726,7 +2726,8 @@ impl GeminiGenerateContentStreamingState { 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) - && (text.is_empty() || (!current_has_metadata && !previous_has_metadata)); + && !current_has_metadata + && !previous_has_metadata; if can_merge { let Some(Json::String(existing)) = last_obj.get_mut("text") else { return; @@ -2796,8 +2797,10 @@ impl GeminiGenerateContentStreamingState { } let part_obj = part.as_object().unwrap(); let data_key = validate_single_gemini_part_data_field(part_obj, "streaming")?; - // Skip thought parts before inspecting text, matching the non-streaming path. + // 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; } diff --git a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs index 9e5fe7ec4..f1216562f 100644 --- a/crates/core/tests/unit/codec/gemini_generate_content_tests.rs +++ b/crates/core/tests/unit/codec/gemini_generate_content_tests.rs @@ -845,10 +845,18 @@ fn test_streaming_preserves_text_part_metadata_from_empty_final_chunk() { .unwrap(); let assembled = finalizer(); - let part = &assembled["candidates"][0]["content"]["parts"][0]; - assert_eq!(part["text"].as_str(), Some("answer")); + 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!( - part["thoughtSignature"].as_str(), + parts[1]["thoughtSignature"].as_str(), Some("sig_STREAM=="), "streamed text metadata must survive finalization" ); @@ -860,10 +868,13 @@ fn test_streaming_preserves_text_part_metadata_from_empty_final_chunk() { }; assert!(matches!( &parts[0], + ContentPart::Text { text, extra } if text == "answer" && extra.is_empty() + )); + assert!(matches!( + &parts[1], ContentPart::Text { text, extra } - if text == "answer" - && extra.get("thoughtSignature").and_then(Json::as_str) - == Some("sig_STREAM==") + if text.is_empty() + && extra.get("thoughtSignature").and_then(Json::as_str) == Some("sig_STREAM==") )); } @@ -1569,6 +1580,17 @@ fn test_streaming_filters_thought_parts() { .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!( 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. From fd23f501b2742c300421ffcb3c45427340589399 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Thu, 6 Aug 2026 13:39:40 -0700 Subject: [PATCH 10/11] fix(codec): remove ambiguous Gemini alias Signed-off-by: Sahiti Bulusu --- crates/core/src/codec/resolve.rs | 16 +- .../tests/fixtures/native_plugin/.gitignore | 5 - .../tests/fixtures/native_plugin/Cargo.lock | 360 ++++++++++++++++++ crates/core/tests/unit/codec/resolve_tests.rs | 17 +- crates/pii-redaction/src/component.rs | 6 +- .../tests/unit/component_tests.rs | 21 +- docs/about-nemo-relay/concepts/codecs.mdx | 8 +- .../provider-codecs.mdx | 9 +- go/nemo_relay/callbacks_test.go | 1 - 9 files changed, 392 insertions(+), 51 deletions(-) delete mode 100644 crates/core/tests/fixtures/native_plugin/.gitignore create mode 100644 crates/core/tests/fixtures/native_plugin/Cargo.lock diff --git a/crates/core/src/codec/resolve.rs b/crates/core/src/codec/resolve.rs index 6a766100e..f54652291 100644 --- a/crates/core/src/codec/resolve.rs +++ b/crates/core/src/codec/resolve.rs @@ -166,14 +166,10 @@ impl ProviderSurface { descriptor_for(self).codec_name } - /// Resolves a canonical codec name or compatibility alias to its surface, or - /// `None` when `name` is not a built-in provider codec. `"gemini"` is accepted - /// as a compatibility alias for the canonical `"gemini_generate_content"`. + /// 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 { - if name == "gemini" { - return Some(Self::GeminiGenerateContent); - } BUILTIN_PROVIDER_SURFACES .iter() .find(|descriptor| descriptor.codec_name == name) @@ -191,14 +187,6 @@ pub fn supported_codec_names() -> Vec<&'static str> { .collect() } -/// Codec names accepted at config boundaries, including compatibility aliases. -#[must_use] -pub fn supported_codec_names_with_aliases() -> Vec<&'static str> { - let mut names = supported_codec_names(); - names.push("gemini"); - names -} - /// Constructs the built-in bidirectional request codec ([`LlmCodec`]) for a surface. #[must_use] pub fn request_codec(surface: ProviderSurface) -> Arc { diff --git a/crates/core/tests/fixtures/native_plugin/.gitignore b/crates/core/tests/fixtures/native_plugin/.gitignore deleted file mode 100644 index c160acc34..000000000 --- a/crates/core/tests/fixtures/native_plugin/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Library fixture — Cargo.lock is not committed for library crates. -Cargo.lock diff --git a/crates/core/tests/fixtures/native_plugin/Cargo.lock b/crates/core/tests/fixtures/native_plugin/Cargo.lock new file mode 100644 index 000000000..33325b40d --- /dev/null +++ b/crates/core/tests/fixtures/native_plugin/Cargo.lock @@ -0,0 +1,360 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" +dependencies = [ + "serde_core", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "num-traits", + "serde", +] + +[[package]] +name = "futures-core" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "nemo-relay-plugin" +version = "0.8.0" +dependencies = [ + "nemo-relay-types", + "serde", + "serde_json", +] + +[[package]] +name = "nemo-relay-plugin-fixture" +version = "0.0.0" +dependencies = [ + "nemo-relay-plugin", + "serde_json", +] + +[[package]] +name = "nemo-relay-types" +version = "0.8.0" +dependencies = [ + "bitflags", + "chrono", + "serde", + "serde_json", + "typed-builder", + "uuid", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "typed-builder" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31aa81521b70f94402501d848ccc0ecaa8f93c8eb6999eb9747e72287757ffda" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "076a02dc54dd46795c2e9c8282ed40bcfb1e22747e955de9389a1de28190fb26" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "uuid" +version = "1.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2" +dependencies = [ + "getrandom", + "js-sys", + "serde", + "wasm-bindgen", +] + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/crates/core/tests/unit/codec/resolve_tests.rs b/crates/core/tests/unit/codec/resolve_tests.rs index 44f38644a..58fce0f97 100644 --- a/crates/core/tests/unit/codec/resolve_tests.rs +++ b/crates/core/tests/unit/codec/resolve_tests.rs @@ -483,15 +483,16 @@ fn codec_name_uses_canonical_spellings() { } #[test] -fn from_codec_name_accepts_gemini_compatibility_alias() { +fn from_codec_name_rejects_ambiguous_gemini_name() { assert_eq!( ProviderSurface::from_codec_name("gemini"), - Some(ProviderSurface::GeminiGenerateContent) + None, + "Gemini codec names must name the concrete API surface" ); assert_eq!( ProviderSurface::GeminiGenerateContent.codec_name(), "gemini_generate_content", - "the legacy alias must not become the canonical codec spelling" + "the canonical Gemini codec spelling names generateContent explicitly" ); } @@ -518,16 +519,6 @@ fn supported_codec_names_track_the_builtin_registry() { .map(|descriptor| descriptor.codec_name) .collect(); assert_eq!(supported_codec_names(), from_registry); - assert_eq!( - supported_codec_names_with_aliases(), - vec![ - "openai_responses", - "anthropic_messages", - "openai_chat", - "gemini_generate_content", - "gemini" - ] - ); } #[test] diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 3da51e61b..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::{ProviderSurface, supported_codec_names_with_aliases}; +use nemo_relay::codec::resolve::{ProviderSurface, supported_codec_names}; use nemo_relay::plugin::{ ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError, PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior, @@ -498,7 +498,7 @@ fn custom_mark_payload_policy_schema( #[cfg(feature = "schema")] fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema { - let codec_names = supported_codec_names_with_aliases(); + let codec_names = supported_codec_names(); string_enum_schema(generator, &codec_names, None) } @@ -1128,7 +1128,7 @@ fn validate_codec_requirements( }; if ProviderSurface::from_codec_name(codec).is_none() { - let supported = supported_codec_names_with_aliases(); + let supported = supported_codec_names(); push_policy_diag( diagnostics, policy.unsupported_value, diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index 7668b9a66..81276be95 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -1795,7 +1795,7 @@ fn validate_allows_llm_surfaces_without_codec() { } #[test] -fn validate_accepts_legacy_gemini_codec_alias() { +fn validate_rejects_ambiguous_gemini_codec_name() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); @@ -1812,7 +1812,12 @@ fn validate_accepts_legacy_gemini_codec_alias() { } }))); - assert!(report.diagnostics.is_empty(), "{report:?}"); + 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] @@ -4369,14 +4374,14 @@ async fn builtin_backend_removes_targeted_message_names_and_ignores_missing_norm } #[tokio::test] -async fn builtin_backend_accepts_legacy_gemini_codec_alias() { +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", + "codec": "gemini_generate_content", "input": true, "output": false, "tool_input": false, @@ -4391,17 +4396,17 @@ async fn builtin_backend_accepts_legacy_gemini_codec_alias() { .await .unwrap(); - let events = capture_events("pii-redaction-gemini-legacy-alias"); + 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-legacy-secret"}]}], + "contents": [{"role": "user", "parts": [{"text": "sk-gemini-secret"}]}], }), }; let _handle = llm_call( LlmCallParams::builder() - .name("gemini-legacy-alias") + .name("gemini-generate-content") .request(&request) .build(), ) @@ -4419,7 +4424,7 @@ async fn builtin_backend_accepts_legacy_gemini_codec_alias() { })) ); - deregister_subscriber("pii-redaction-gemini-legacy-alias").unwrap(); + deregister_subscriber("pii-redaction-gemini-generate-content").unwrap(); clear_plugin_configuration().unwrap(); } diff --git a/docs/about-nemo-relay/concepts/codecs.mdx b/docs/about-nemo-relay/concepts/codecs.mdx index d02ff08e6..4afd6435f 100644 --- a/docs/about-nemo-relay/concepts/codecs.mdx +++ b/docs/about-nemo-relay/concepts/codecs.mdx @@ -167,9 +167,11 @@ through the Anthropic provider surface instead of being treated as shape-only OpenAI Chat payloads. The `nemo-relay` gateway always enables matching OpenAI Chat, OpenAI Responses, -and Anthropic Messages request codecs for `/v1/messages`, -`/v1/chat/completions`, and `/v1/responses`, for both buffered and streaming -execution. Count-token, model, probe, Gemini `generateContent`, and non-LLM +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. diff --git a/docs/integrate-into-frameworks/provider-codecs.mdx b/docs/integrate-into-frameworks/provider-codecs.mdx index e8da08e84..5df68707e 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -111,10 +111,11 @@ The built-in provider codecs expose the same core methods: 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 the OpenAI Chat, OpenAI Responses, and -Anthropic Messages request codecs automatically for `/v1/messages`, -`/v1/chat/completions`, and `/v1/responses`, for both buffered and streaming -calls. `GeminiGenerateContentCodec` is available for direct use with Gemini +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. diff --git a/go/nemo_relay/callbacks_test.go b/go/nemo_relay/callbacks_test.go index 2c19279ec..8905e7b7e 100644 --- a/go/nemo_relay/callbacks_test.go +++ b/go/nemo_relay/callbacks_test.go @@ -104,7 +104,6 @@ func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) {"openai chat", 1, &openAIChat, LLMCodecBuiltin}, {"openai responses", 1, &openAIResponses, LLMCodecBuiltin}, {"anthropic messages", 1, &anthropicMessages, LLMCodecBuiltin}, - // Gemini must be LLMCodecBuiltin, not Opaque — PII sanitization depends on this. {"gemini_generate_content", 1, &gemini, LLMCodecBuiltin}, {"runtime", 2, &runtimeCodec, LLMCodecRuntime}, {"opaque", 3, nil, LLMCodecOpaque}, From 9c2071689be5ab4d13943bbbcbf1deba557e0fc2 Mon Sep 17 00:00:00 2001 From: Sahiti Bulusu Date: Fri, 7 Aug 2026 10:56:56 -0700 Subject: [PATCH 11/11] style(python): fix Gemini codec formatting Signed-off-by: Sahiti Bulusu --- python/nemo_relay/pii_redaction.py | 4 +++- python/nemo_relay/pii_redaction.pyi | 4 +++- python/tests/test_builtin_codecs.py | 3 ++- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/python/nemo_relay/pii_redaction.py b/python/nemo_relay/pii_redaction.py index 4c0f91235..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", "gemini_generate_content"] | 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 8da9ebddc..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", "gemini_generate_content"] | 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/test_builtin_codecs.py b/python/tests/test_builtin_codecs.py index 07031c402..00458be45 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -4,7 +4,8 @@ """Tests for built-in codec Python classes and LlmResponseCodec protocol. Covers: -- Built-in codec construction (OpenAIChatCodec, OpenAIResponsesCodec, AnthropicMessagesCodec, GeminiGenerateContentCodec) +- 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)