diff --git a/crates/adaptive/src/acg/request_surfaces/mod.rs b/crates/adaptive/src/acg/request_surfaces/mod.rs index b4834aa56..b0e3ec4eb 100644 --- a/crates/adaptive/src/acg/request_surfaces/mod.rs +++ b/crates/adaptive/src/acg/request_surfaces/mod.rs @@ -48,6 +48,8 @@ impl RequestSurface { ProviderSurface::OpenAIChat => Some(Self::OpenAIChat), ProviderSurface::OpenAIResponses => Some(Self::OpenAIResponses), ProviderSurface::AnthropicMessages => Some(Self::AnthropicMessages), + // No semantic ACG applier exists for OCI GenAI request shapes yet. + ProviderSurface::OCIGenAI => None, // Gemini generateContent ACG request editing is intentionally unsupported. ProviderSurface::GeminiGenerateContent => None, } diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index dbdb1726b..4b6ee6a5a 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -338,6 +338,12 @@ fn lossy_request_shape(surface: ProviderSurface, content: &Json) -> bool { .is_some_and(|blocks| blocks.iter().any(lossy_system_block)) } ProviderSurface::OpenAIResponses => false, + // OCI GenAI requests carry an envelope (`compartmentId`, `servingMode`) + // whose unmodeled fields the decode does not preserve in `extra`, and + // the generic scalar checks above target OpenAI-shaped keys rather + // than OCI camelCase. Keep OCI raw-keyed — a fallback only ever costs + // a miss. + ProviderSurface::OCIGenAI => true, ProviderSurface::GeminiGenerateContent => { object .get("generationConfig") diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs index 95c405c78..43d7fde1d 100644 --- a/crates/adaptive/src/response_cache/replay.rs +++ b/crates/adaptive/src/response_cache/replay.rs @@ -80,6 +80,11 @@ fn synthesize_replay_chunks(aggregate: &Json) -> Option> { ProviderSurface::AnthropicMessages => synthesize_anthropic_chunks(aggregate), ProviderSurface::OpenAIChat => synthesize_chat_chunks(aggregate), ProviderSurface::OpenAIResponses => synthesize_responses_chunks(aggregate), + // OCI GenAI streaming events are ChatResult-shaped deltas; the OCI + // stream collector accepts a full `chatResponse` aggregate as a single + // native chunk. `replay_is_lossy` still re-aggregates it and rejects + // shapes the streaming collector cannot preserve exactly. + ProviderSurface::OCIGenAI => vec![aggregate.clone()], // 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 diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index 5adc048e3..30040f017 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, + /// OCI Generative AI chat request and response payloads. + OCIGenAI, /// Gemini generateContent request and response payloads. GeminiGenerateContent, } @@ -172,6 +174,7 @@ impl BuiltinLlmCodec { Self::OpenAiChat => "openai_chat", Self::OpenAiResponses => "openai_responses", Self::AnthropicMessages => "anthropic_messages", + Self::OCIGenAI => "oci_genai", Self::GeminiGenerateContent => "gemini_generate_content", } } diff --git a/crates/core/src/codec/mod.rs b/crates/core/src/codec/mod.rs index 4a587533e..dc5ca237e 100644 --- a/crates/core/src/codec/mod.rs +++ b/crates/core/src/codec/mod.rs @@ -17,6 +17,7 @@ pub mod anthropic; pub mod gemini_generate_content; pub mod model_pricing; +pub mod oci_genai; pub mod openai_chat; pub mod openai_responses; pub mod optimization; diff --git a/crates/core/src/codec/oci_genai.rs b/crates/core/src/codec/oci_genai.rs new file mode 100644 index 000000000..6a6698393 --- /dev/null +++ b/crates/core/src/codec/oci_genai.rs @@ -0,0 +1,1876 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Built-in codec for the Oracle Cloud Infrastructure (OCI) Generative AI chat API. +//! +//! Implements [`LlmCodec`] (request decode/encode) and [`LlmResponseCodec`] +//! (response decode) for the OCI Generative AI chat format. +//! +//! # OCI-specific patterns handled +//! +//! - **ChatDetails envelope**: Requests may arrive as a full envelope +//! (`compartmentId`, `servingMode`, `chatRequest`) or as a bare `chatRequest` +//! payload; both are accepted and the envelope is preserved on encode. +//! - **API formats** selected by `apiFormat`: +//! - `GENERIC`: OpenAI-style `messages` with UPPERCASE roles +//! (`USER`/`ASSISTANT`/`SYSTEM`/`TOOL`) whose `content` is a list of typed +//! parts (`{"type": "TEXT", "text": ...}`), flat `toolCalls` +//! (`{id, type: "FUNCTION", name, arguments}`), and `toolCallId` on tool +//! messages. Used by Meta Llama, Google, xAI, OpenAI, and imported +//! open-weights models hosted on dedicated AI clusters. +//! - `COHERE`: a single `message` string plus `chatHistory` turns with +//! `USER`/`CHATBOT`/`SYSTEM` roles and an optional `preambleOverride`. +//! Used by Cohere Command models. +//! - `COHEREV2`: responses are a single assistant `message` with typed +//! content parts and nested `function` tool calls, per the OCI +//! `CohereChatResponseV2` schema. Requests follow the GENERIC `messages` +//! path with COHERE-style `stopSequences`; V2-only request fields the +//! normalized shape does not model (`citationOptions`, `documents`, ...) +//! ride along in `extra` and survive edits untouched. +//! - **Model identity**: Carried in `servingMode.modelId` (on-demand) or +//! `servingMode.endpointId` (dedicated), not in the chat request body. +//! - **Responses**: `ChatResult` payloads (`modelId`, `chatResponse`); `usage` +//! counters are `promptTokens`/`completionTokens`/`totalTokens`. +//! - **Unmodeled fields are preserved**: envelope and chat-response fields the +//! normalized shape does not model (`timeCreated`, future provider fields) +//! are carried in `extra` rather than discarded, consistent with the other +//! response codecs. Unmodeled fields of the decoded choice and assistant +//! message — `logprobs`, `serviceTier`, `groundingMetadata`, +//! `reasoningContent`, `refusal` (GENERIC) and `toolPlan`, `citations` +//! (COHEREV2) per the OCI schema — are namespaced in `extra` under +//! `"choice"` and `"message"`. +//! +//! The codec accepts the REST wire format only: camelCase keys, as documented +//! in the OCI API reference. Alternate renderings produced by Oracle tooling +//! (the CLI's kebab-case `data` envelope, `oci.util.to_dict()` snake_case +//! dicts) are the caller's responsibility to convert. + +use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; +use crate::error::{FlowError, Result}; +use crate::json::Json; + +use super::request::{ + AnnotatedLlmRequest, ApiSpecificRequest, ContentPart, FunctionCall, GenerationParams, Message, + MessageContent, ProviderNativeComponent, ToolCall, ToolChoice, ToolDefinition, +}; +use super::resolve::{ProviderSurface, ProviderSurfaceDescriptor}; +use super::response::{ + AnnotatedLlmResponse, ApiSpecificResponse, FinishReason, ResponseToolCall, Usage, +}; +use super::traits::{LlmCodec, LlmResponseCodec}; + +// --------------------------------------------------------------------------- +// Public codec struct +// --------------------------------------------------------------------------- + +/// Built-in codec for the OCI Generative AI chat API. +pub struct OCIGenAIChatCodec; + +pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDescriptor { + surface: ProviderSurface::OCIGenAI, + detect_request: |obj, hint| { + // The ChatDetails envelope (chatRequest + servingMode/compartmentId) and + // the apiFormat discriminator are unique to OCI Generative AI; a bare + // chatRequest without apiFormat needs the provider hint to classify. + let has_chat_request = obj.get("chatRequest").is_some_and(Json::is_object); + let has_envelope_marker = + obj.get("servingMode").is_some() || obj.get("compartmentId").is_some(); + let hinted_oci = + hint.is_some_and(|hint_value| hint_value == "oci" || hint_value == "oci.genai"); + (has_chat_request && has_envelope_marker) + || obj.get("apiFormat").is_some() + || (hinted_oci && has_chat_request) + }, + detect_response: |obj| match obj.get("chatResponse") { + Some(Json::Object(chat_response)) => chat_response.get("apiFormat").is_some(), + _ => obj.get("apiFormat").is_some(), + }, + decode_request: |request| OCIGenAIChatCodec.decode(request), + decode_response: |raw| OCIGenAIChatCodec.decode_response(raw), + codec_name: "oci_genai", + request_codec: || std::sync::Arc::new(OCIGenAIChatCodec), + response_codec: || std::sync::Arc::new(OCIGenAIChatCodec), + streaming_codec: || Box::new(OCIGenAIStreamingCodec::new()), +}; + +// --------------------------------------------------------------------------- +// Optional-field helpers +// --------------------------------------------------------------------------- + +/// Lookup of an optional list of strings (stop sequences). +fn optional_string_list( + obj: &serde_json::Map, + key: &str, + surface: &str, +) -> Result>> { + let Some(value) = obj.get(key) else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + serde_json::from_value::>(value.clone()) + .map(Some) + .map_err(|error| { + FlowError::InvalidArgument(format!("{surface} {key} must be a string array: {error}")) + }) +} + +// --------------------------------------------------------------------------- +// Modeled-key bookkeeping +// --------------------------------------------------------------------------- + +/// Chat-request keys modeled in [`AnnotatedLlmRequest`] for the GENERIC format. +const MODELED_GENERIC_REQUEST_KEYS: &[&str] = &[ + "apiFormat", + "messages", + "maxTokens", + "temperature", + "topP", + "stop", + "tools", + "toolChoice", +]; + +/// Chat-request keys modeled in [`AnnotatedLlmRequest`] for the COHERE format. +const MODELED_COHERE_REQUEST_KEYS: &[&str] = &[ + "apiFormat", + "message", + "chatHistory", + "preambleOverride", + "maxTokens", + "temperature", + "topP", + "stopSequences", + "tools", + "toolChoice", +]; + +/// Chat-request keys modeled in [`AnnotatedLlmRequest`] for the COHEREV2 +/// format: GENERIC-style `messages` with COHERE-style `stopSequences`. +const MODELED_COHERE_V2_REQUEST_KEYS: &[&str] = &[ + "apiFormat", + "messages", + "maxTokens", + "temperature", + "topP", + "stopSequences", + "tools", + "toolChoice", +]; + +/// Whether `key` is one of the modeled keys. +fn is_modeled_key(key: &str, modeled: &[&str]) -> bool { + modeled.contains(&key) +} + +// --------------------------------------------------------------------------- +// Shared helpers +// --------------------------------------------------------------------------- + +/// Map an OCI finish reason string to normalized [`FinishReason`]. +/// +/// GENERIC responses use OpenAI-style lowercase reasons (Gemini models emit +/// `max_tokens` for the length stop); COHERE and COHEREV2 responses use +/// UPPERCASE Cohere reasons (`TOOL_CALL` and `STOP_SEQUENCE` are V2-only). +fn map_oci_finish_reason(reason: &str) -> FinishReason { + match reason { + "stop" | "COMPLETE" | "STOP_SEQUENCE" => FinishReason::Complete, + "length" | "max_tokens" | "MAX_TOKENS" => FinishReason::Length, + "tool_calls" | "TOOL_CALL" => FinishReason::ToolUse, + "content_filter" => FinishReason::ContentFilter, + other => FinishReason::Unknown(other.to_string()), + } +} + +/// Collect the fields of `obj` that are not in `modeled` for `extra` carriage. +fn unmodeled_fields( + obj: &serde_json::Map, + modeled: &[&str], +) -> serde_json::Map { + obj.iter() + .filter(|(key, _)| !modeled.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +/// Helper to construct a [`Json`] number from an `f64`. +fn json_f64(v: f64) -> Json { + serde_json::Number::from_f64(v) + .map(Json::Number) + .unwrap_or(Json::Null) +} + +fn insert_json(obj: &mut serde_json::Map, key: &str, value: Json) { + obj.insert(key.to_string(), value); +} + +fn set_or_remove_json(obj: &mut serde_json::Map, key: &str, value: Option) { + if let Some(value) = value { + obj.insert(key.into(), value); + } else { + obj.remove(key); + } +} + +fn patch_extra_fields( + obj: &mut serde_json::Map, + baseline: &serde_json::Map, + edited: &serde_json::Map, +) { + for key in baseline.keys().filter(|key| !edited.contains_key(*key)) { + obj.remove(key); + } + for (key, value) in edited { + if baseline.get(key) != Some(value) { + obj.insert(key.clone(), value.clone()); + } + } +} + +fn native_component(value: &Json) -> ProviderNativeComponent { + ProviderNativeComponent { + provider: "oci_genai".to_string(), + kind: value + .get("type") + .and_then(Json::as_str) + .unwrap_or("unknown") + .to_string(), + value: value.clone(), + } +} + +// --------------------------------------------------------------------------- +// GENERIC content conversion +// --------------------------------------------------------------------------- + +/// Flatten a GENERIC content value into normalized [`MessageContent`]. +/// +/// A content-part list whose parts are all `{"type": "TEXT", "text": ...}` is +/// flattened to plain text; lists carrying any non-text part are preserved as +/// typed parts so image or future block types survive losslessly. +fn decode_generic_content(value: Option<&Json>) -> Result> { + let value = match value { + None | Some(Json::Null) => return Ok(None), + Some(value) => value, + }; + if let Some(text) = value.as_str() { + return Ok(Some(MessageContent::Text(text.to_string()))); + } + let parts = value.as_array().ok_or_else(|| { + FlowError::InvalidArgument( + "OCI GenAI GENERIC message content must be a string, an array, or null".into(), + ) + })?; + if parts.is_empty() { + // Tool-call-only messages carry `"content": []`; there is no content. + return Ok(None); + } + if let Some(text) = flatten_all_text_parts(parts) { + return Ok(Some(MessageContent::Text(text))); + } + let parts = parts + .iter() + .map(decode_generic_content_part) + .collect::>>()?; + Ok(Some(MessageContent::Parts(parts))) +} + +/// Join a part list into plain text when every part is a `TEXT` part. +fn flatten_all_text_parts(parts: &[Json]) -> Option { + let mut text = String::new(); + for part in parts { + let obj = part.as_object()?; + if obj.get("type").and_then(Json::as_str) != Some("TEXT") { + return None; + } + match obj.get("text") { + None | Some(Json::Null) => {} + Some(Json::String(part_text)) => text.push_str(part_text), + Some(_) => return None, + } + } + Some(text) +} + +fn decode_generic_content_part(value: &Json) -> Result { + let Some(obj) = value.as_object() else { + return Err(FlowError::InvalidArgument( + "OCI GenAI GENERIC content part must be an object".into(), + )); + }; + match obj.get("type").and_then(Json::as_str) { + // A TEXT part whose `text` is not a string falls through to the + // provider-native branch so the raw value survives the round trip + // instead of collapsing to an empty string. + Some("TEXT") if obj.get("text").is_none_or(Json::is_string) => Ok(ContentPart::Text { + text: obj + .get("text") + .and_then(Json::as_str) + .unwrap_or_default() + .to_string(), + extra: obj + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "text")) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + }), + _ => { + let native = native_component(value); + Ok(ContentPart::ProviderNative { + provider: native.provider, + kind: native.kind, + value: native.value, + }) + } + } +} + +/// Wrap normalized content back into the GENERIC typed content-part list. +fn encode_generic_content(content: &MessageContent) -> Result { + match content { + MessageContent::Text(text) => Ok(serde_json::json!([{"type": "TEXT", "text": text}])), + MessageContent::Parts(parts) => Ok(Json::Array( + parts + .iter() + .map(encode_generic_content_part) + .collect::>>()?, + )), + } +} + +fn encode_generic_content_part(part: &ContentPart) -> Result { + match part { + ContentPart::Text { text, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("TEXT".into())); + obj.insert("text".into(), Json::String(text.clone())); + Ok(Json::Object(obj)) + } + ContentPart::ProviderNative { + provider, value, .. + } if provider == "oci_genai" => Ok(value.clone()), + other => Err(FlowError::InvalidArgument(format!( + "content part {other:?} cannot be encoded for OCI GenAI" + ))), + } +} + +// --------------------------------------------------------------------------- +// Tool call conversion +// --------------------------------------------------------------------------- + +/// Convert a flat OCI `toolCalls` entry into the normalized nested [`ToolCall`]. +fn decode_oci_tool_call(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OCI GenAI toolCalls entry must be an object".into()) + })?; + // A nested `function` object means the entry is already normalized. + let function = obj.get("function").and_then(Json::as_object); + let (name, arguments) = match function { + Some(function) => (function.get("name"), function.get("arguments")), + None => (obj.get("name"), obj.get("arguments")), + }; + Ok(ToolCall { + id: obj + .get("id") + .and_then(Json::as_str) + .unwrap_or_default() + .to_string(), + call_type: "function".to_string(), + function: FunctionCall { + name: name.and_then(Json::as_str).unwrap_or_default().to_string(), + arguments: match arguments { + Some(Json::String(text)) => text.clone(), + Some(other) => other.to_string(), + None => String::new(), + }, + }, + }) +} + +/// Convert a normalized [`ToolCall`] into the COHEREV2 nested-function shape. +fn encode_oci_tool_call_nested(tool_call: &ToolCall) -> Json { + let mut function = serde_json::Map::new(); + function.insert("name".into(), Json::String(tool_call.function.name.clone())); + function.insert( + "arguments".into(), + Json::String(tool_call.function.arguments.clone()), + ); + let mut obj = serde_json::Map::new(); + if !tool_call.id.is_empty() { + obj.insert("id".into(), Json::String(tool_call.id.clone())); + } + obj.insert("type".into(), Json::String("FUNCTION".into())); + obj.insert("function".into(), Json::Object(function)); + Json::Object(obj) +} + +/// Convert a normalized nested [`ToolCall`] back into the flat OCI shape. +/// +/// A missing wire `id` decodes to an empty string, so an empty id is omitted +/// on re-encode rather than materializing an `"id": ""` field. +fn encode_oci_tool_call(tool_call: &ToolCall) -> Json { + let mut obj = serde_json::Map::new(); + if !tool_call.id.is_empty() { + obj.insert("id".into(), Json::String(tool_call.id.clone())); + } + obj.insert("type".into(), Json::String("FUNCTION".into())); + obj.insert("name".into(), Json::String(tool_call.function.name.clone())); + obj.insert( + "arguments".into(), + Json::String(tool_call.function.arguments.clone()), + ); + Json::Object(obj) +} + +// --------------------------------------------------------------------------- +// GENERIC message decode/encode +// --------------------------------------------------------------------------- + +fn decode_generic_message(value: &Json) -> Result { + let obj = value.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OCI GenAI GENERIC message must be an object".into()) + })?; + let role = obj + .get("role") + .and_then(Json::as_str) + .unwrap_or("USER") + .to_lowercase(); + let content = decode_generic_content(obj.get("content"))?; + let tool_calls = match obj.get("toolCalls") { + None | Some(Json::Null) => None, + Some(Json::Array(calls)) => Some( + calls + .iter() + .map(decode_oci_tool_call) + .collect::>>()?, + ), + Some(_) => { + return Err(FlowError::InvalidArgument( + "OCI GenAI GENERIC toolCalls must be an array".into(), + )); + } + }; + let tool_call_id = obj + .get("toolCallId") + .and_then(Json::as_str) + .map(str::to_string); + match role.as_str() { + "system" => match content { + Some(content) => Ok(Message::System { + content, + name: None, + }), + None => Ok(provider_native_message(&role, value)), + }, + "user" => match content { + Some(content) => Ok(Message::User { + content, + name: None, + }), + None => Ok(provider_native_message(&role, value)), + }, + "assistant" => Ok(Message::Assistant { + content, + tool_calls, + name: None, + }), + "tool" => match (content, tool_call_id) { + (Some(content), Some(tool_call_id)) => Ok(Message::Tool { + content, + tool_call_id, + }), + _ => Ok(provider_native_message(&role, value)), + }, + _ => Ok(provider_native_message(&role, value)), + } +} + +fn provider_native_message(kind: &str, value: &Json) -> Message { + Message::ProviderNative { + provider: "oci_genai".into(), + kind: kind.to_string(), + value: value.clone(), + } +} + +fn encode_generic_message(message: &Message, api_format: &str) -> Result { + let mut obj = serde_json::Map::new(); + match message { + Message::System { content, .. } => { + obj.insert("role".into(), Json::String("SYSTEM".into())); + obj.insert("content".into(), encode_generic_content(content)?); + } + Message::User { content, .. } => { + obj.insert("role".into(), Json::String("USER".into())); + obj.insert("content".into(), encode_generic_content(content)?); + } + Message::Assistant { + content, + tool_calls, + .. + } => { + obj.insert("role".into(), Json::String("ASSISTANT".into())); + // `None` round-trips a tool-call-only message: an empty part list + // decodes to `None`, so re-encode as `[]` rather than `null` to + // keep the OCI typed-part-list shape. + obj.insert( + "content".into(), + match content { + Some(content) => encode_generic_content(content)?, + None => Json::Array(Vec::new()), + }, + ); + if let Some(tool_calls) = tool_calls { + // COHEREV2 nests name/arguments under `function`; GENERIC is flat. + let encode = if api_format == "COHEREV2" { + encode_oci_tool_call_nested + } else { + encode_oci_tool_call + }; + obj.insert( + "toolCalls".into(), + Json::Array(tool_calls.iter().map(encode).collect()), + ); + } + } + Message::Tool { + content, + tool_call_id, + } => { + obj.insert("role".into(), Json::String("TOOL".into())); + obj.insert("content".into(), encode_generic_content(content)?); + obj.insert("toolCallId".into(), Json::String(tool_call_id.clone())); + } + Message::ProviderNative { + provider, value, .. + } if provider == "oci_genai" => return Ok(value.clone()), + other => { + return Err(FlowError::InvalidArgument(format!( + "message {other:?} cannot be encoded for OCI GenAI" + ))); + } + } + Ok(Json::Object(obj)) +} + +/// Rewrite only the GENERIC messages that intercepts actually changed. +/// +/// Unchanged messages are carried over from the raw payload verbatim so +/// per-message provider fields without a normalized equivalent survive. +fn patch_generic_messages( + chat_request: &mut serde_json::Map, + edited: &[Message], + baseline: &[Message], + api_format: &str, +) -> Result<()> { + let raw_messages: Vec = chat_request + .get("messages") + .and_then(Json::as_array) + .cloned() + .unwrap_or_default(); + let patched = edited + .iter() + .enumerate() + .map(|(index, message)| { + let unchanged = baseline.get(index) == Some(message); + match raw_messages.get(index) { + Some(raw) if unchanged => Ok(raw.clone()), + _ => encode_generic_message(message, api_format), + } + }) + .collect::>>()?; + insert_json(chat_request, "messages", Json::Array(patched)); + Ok(()) +} + +// --------------------------------------------------------------------------- +// COHERE message decode/encode +// --------------------------------------------------------------------------- + +fn decode_cohere_messages(chat_request: &serde_json::Map) -> Result> { + let mut messages = Vec::new(); + + if let Some(preamble) = chat_request.get("preambleOverride").and_then(Json::as_str) + && !preamble.is_empty() + { + messages.push(Message::System { + content: MessageContent::Text(preamble.to_string()), + name: None, + }); + } + + if let Some(history) = chat_request.get("chatHistory") { + let turns = history.as_array().ok_or_else(|| { + FlowError::InvalidArgument("OCI GenAI COHERE chatHistory must be an array".into()) + })?; + for turn in turns { + messages.push(decode_cohere_turn(turn)?); + } + } + + if let Some(current) = chat_request.get("message").and_then(Json::as_str) { + messages.push(Message::User { + content: MessageContent::Text(current.to_string()), + name: None, + }); + } + + Ok(messages) +} + +fn decode_cohere_turn(turn: &Json) -> Result { + let obj = turn.as_object().ok_or_else(|| { + FlowError::InvalidArgument("OCI GenAI COHERE chatHistory turn must be an object".into()) + })?; + let role = obj + .get("role") + .and_then(Json::as_str) + .unwrap_or("USER") + .to_uppercase(); + let Some(text) = obj.get("message").and_then(Json::as_str) else { + return Ok(provider_native_message(&role, turn)); + }; + let content = MessageContent::Text(text.to_string()); + match role.as_str() { + "USER" => Ok(Message::User { + content, + name: None, + }), + "CHATBOT" => Ok(Message::Assistant { + content: Some(content), + tool_calls: None, + name: None, + }), + "SYSTEM" => Ok(Message::System { + content, + name: None, + }), + _ => Ok(provider_native_message(&role, turn)), + } +} + +/// Extract the plain-text body of a normalized message for COHERE encoding. +fn cohere_text(content: &MessageContent) -> Result { + match content { + MessageContent::Text(text) => Ok(text.clone()), + MessageContent::Parts(_) => Err(FlowError::InvalidArgument( + "multimodal content cannot be encoded for the OCI GenAI COHERE format".into(), + )), + } +} + +/// Rebuild the COHERE `preambleOverride`/`chatHistory`/`message` fields from +/// edited messages. COHERE turns are plain strings, so edits rebuild the +/// modeled fields rather than patching individual turns. +fn encode_cohere_messages( + chat_request: &mut serde_json::Map, + messages: &[Message], +) -> Result<()> { + let mut remaining = messages; + + if let Some(Message::System { content, .. }) = remaining.first() { + insert_json( + chat_request, + "preambleOverride", + Json::String(cohere_text(content)?), + ); + remaining = &remaining[1..]; + } else { + // The encoder merges into the raw request, so without this removal a + // preamble deleted (or re-roled) by an intercept would survive on the + // wire while the normalized annotation no longer contains it. + chat_request.remove("preambleOverride"); + } + + // The COHERE chat request requires a non-empty `message` prompt; an edited + // list without a trailing user turn would otherwise silently send `""`. + let Some(Message::User { content, .. }) = remaining.last() else { + return Err(FlowError::InvalidArgument( + "OCI GenAI COHERE requests require the last message to be a user message".into(), + )); + }; + let current = cohere_text(content)?; + remaining = &remaining[..remaining.len() - 1]; + + let history = remaining + .iter() + .map(encode_cohere_turn) + .collect::>>()?; + + chat_request.insert("message".into(), Json::String(current)); + if !history.is_empty() || chat_request.contains_key("chatHistory") { + chat_request.insert("chatHistory".into(), Json::Array(history)); + } + Ok(()) +} + +fn encode_cohere_turn(message: &Message) -> Result { + let (role, content) = match message { + Message::User { content, .. } => ("USER", content), + Message::Assistant { + content: Some(content), + .. + } => ("CHATBOT", content), + Message::System { content, .. } => ("SYSTEM", content), + // OCI's CohereToolMessage carries structured `toolResults`, not a plain + // `message` string, and has no field for the normalized tool_call_id. + // Reject rather than silently dropping the identifier; wire-shaped TOOL + // turns survive untouched as ProviderNative messages below. + Message::Tool { .. } => { + return Err(FlowError::InvalidArgument( + "OCI GenAI COHERE tool results cannot be rebuilt from a normalized tool message" + .into(), + )); + } + Message::ProviderNative { + provider, value, .. + } if provider == "oci_genai" => return Ok(value.clone()), + other => { + return Err(FlowError::InvalidArgument(format!( + "message {other:?} cannot be encoded as an OCI GenAI COHERE chatHistory turn" + ))); + } + }; + Ok(serde_json::json!({"role": role, "message": cohere_text(content)?})) +} + +// --------------------------------------------------------------------------- +// Params, tools, and envelope helpers +// --------------------------------------------------------------------------- + +/// Decode the normalized generation params for one API format. +fn decode_params( + chat_request: &serde_json::Map, + api_format: &str, +) -> Result> { + const SURFACE: &str = "OCI GenAI"; + let temperature = super::optional_f64(chat_request, "temperature", SURFACE)?; + let max_tokens = super::optional_u64(chat_request, "maxTokens", SURFACE)?; + let top_p = super::optional_f64(chat_request, "topP", SURFACE)?; + // Both Cohere formats spell the stop list `stopSequences`. + let stop_key = if api_format.starts_with("COHERE") { + "stopSequences" + } else { + "stop" + }; + let stop = optional_string_list(chat_request, stop_key, SURFACE)?; + if temperature.is_some() || max_tokens.is_some() || top_p.is_some() || stop.is_some() { + Ok(Some(GenerationParams { + temperature, + max_tokens, + top_p, + stop, + })) + } else { + Ok(None) + } +} + +/// Patch only the generation params an intercept actually changed. +/// +/// A param cleared to `None` removes the raw key, matching the set-or-remove +/// semantics of the other provider codecs. +fn patch_params( + chat_request: &mut serde_json::Map, + edited: Option<&GenerationParams>, + baseline: Option<&GenerationParams>, + api_format: &str, +) { + if edited == baseline { + return; + } + let temperature = edited.and_then(|params| params.temperature); + if temperature != baseline.and_then(|params| params.temperature) { + set_or_remove_json(chat_request, "temperature", temperature.map(json_f64)); + } + let top_p = edited.and_then(|params| params.top_p); + if top_p != baseline.and_then(|params| params.top_p) { + set_or_remove_json(chat_request, "topP", top_p.map(json_f64)); + } + let max_tokens = edited.and_then(|params| params.max_tokens); + if max_tokens != baseline.and_then(|params| params.max_tokens) { + set_or_remove_json(chat_request, "maxTokens", max_tokens.map(Json::from)); + } + let stop = edited.and_then(|params| params.stop.as_ref()); + if stop != baseline.and_then(|params| params.stop.as_ref()) { + let stop_key = if api_format.starts_with("COHERE") { + "stopSequences" + } else { + "stop" + }; + set_or_remove_json( + chat_request, + stop_key, + stop.map(|values| serde_json::json!(values)), + ); + } +} + +fn decode_tools( + chat_request: &serde_json::Map, +) -> Result>> { + match chat_request.get("tools") { + None | Some(Json::Null) => Ok(None), + Some(Json::Array(tools)) => Ok(Some( + tools + .iter() + .map(|tool| { + let native = native_component(tool); + ToolDefinition::ProviderNative { + provider: native.provider, + kind: native.kind, + value: native.value, + } + }) + .collect(), + )), + Some(_) => Err(FlowError::InvalidArgument( + "OCI GenAI tools must be an array".into(), + )), + } +} + +fn encode_oci_tool(tool: &ToolDefinition) -> Result { + match tool { + ToolDefinition::ProviderNative { + provider, value, .. + } if provider == "oci_genai" => Ok(value.clone()), + ToolDefinition::Function { function, extra } => { + let mut obj = extra.clone(); + obj.insert("type".into(), Json::String("FUNCTION".into())); + obj.insert("name".into(), Json::String(function.name.clone())); + if let Some(description) = &function.description { + obj.insert("description".into(), Json::String(description.clone())); + } + if let Some(parameters) = &function.parameters { + obj.insert("parameters".into(), parameters.clone()); + } + obj.extend(function.extra.clone()); + Ok(Json::Object(obj)) + } + other => Err(FlowError::InvalidArgument(format!( + "tool {other:?} cannot be encoded for OCI GenAI" + ))), + } +} + +fn encode_oci_tool_choice(tool_choice: &ToolChoice) -> Result { + match tool_choice { + ToolChoice::ProviderNative(native) if native.provider == "oci_genai" => { + Ok(native.value.clone()) + } + other => Err(FlowError::InvalidArgument(format!( + "tool choice {other:?} cannot be encoded for OCI GenAI" + ))), + } +} + +/// Extract the model identity from the `servingMode` envelope object. +fn model_from_envelope(envelope: &serde_json::Map) -> Option { + let serving_mode = envelope.get("servingMode")?.as_object()?; + serving_mode + .get("modelId") + .or_else(|| serving_mode.get("endpointId")) + .and_then(Json::as_str) + .map(str::to_string) +} + +/// Split the request content into the optional ChatDetails envelope and the +/// chat request object. +fn split_envelope( + obj: &serde_json::Map, +) -> ( + Option<&serde_json::Map>, + &serde_json::Map, +) { + match obj.get("chatRequest").and_then(Json::as_object) { + Some(chat_request) => (Some(obj), chat_request), + None => (None, obj), + } +} + +/// Resolve the request API format (uppercased), defaulting to `GENERIC`. +fn request_api_format(chat_request: &serde_json::Map) -> String { + chat_request + .get("apiFormat") + .and_then(Json::as_str) + .unwrap_or("GENERIC") + .to_uppercase() +} + +fn validate_oci_supported_fields( + annotated: &AnnotatedLlmRequest, + baseline: &AnnotatedLlmRequest, +) -> Result<()> { + let unsupported = [ + annotated.model != baseline.model, + annotated.instructions != baseline.instructions, + annotated.store != baseline.store, + annotated.previous_response_id != baseline.previous_response_id, + annotated.truncation != baseline.truncation, + annotated.reasoning != baseline.reasoning, + annotated.include != baseline.include, + annotated.user != baseline.user, + annotated.metadata != baseline.metadata, + annotated.service_tier != baseline.service_tier, + annotated.parallel_tool_calls != baseline.parallel_tool_calls, + annotated.max_output_tokens != baseline.max_output_tokens, + annotated.max_tool_calls != baseline.max_tool_calls, + annotated.top_logprobs != baseline.top_logprobs, + annotated.stream != baseline.stream, + ] + .into_iter() + .any(|changed| changed); + if unsupported { + return Err(FlowError::InvalidArgument( + "request contains fields that cannot be encoded for OCI GenAI".into(), + )); + } + Ok(()) +} + +/// Patch envelope-level fields (`compartmentId`, `servingMode`) when the +/// api-specific annotation changed them. +/// +/// `api_format` is read-only: the encoder patches the raw payload in place, so +/// switching formats cannot rebuild the body without leaving the other +/// format's modeled fields behind, and an edit is rejected instead. +fn patch_oci_api_specific( + envelope: Option<&mut serde_json::Map>, + edited: &Option, + baseline: &Option, +) -> Result<()> { + let (compartment_id, serving_mode, old_compartment_id, old_serving_mode) = + match (edited, baseline) { + ( + Some(ApiSpecificRequest::OCIGenAI { + compartment_id, + serving_mode, + api_format, + }), + Some(ApiSpecificRequest::OCIGenAI { + compartment_id: old_compartment_id, + serving_mode: old_serving_mode, + api_format: old_api_format, + }), + ) => { + if api_format != old_api_format { + return Err(FlowError::InvalidArgument( + "the OCI GenAI api_format cannot be edited".into(), + )); + } + ( + compartment_id, + serving_mode, + old_compartment_id, + old_serving_mode, + ) + } + // A dropped api_specific annotation leaves the envelope untouched; + // the raw payload keeps serving as the source of truth. + (None, _) => return Ok(()), + (Some(_), _) => { + return Err(FlowError::InvalidArgument( + "api_specific provider does not match OCI GenAI".into(), + )); + } + }; + if compartment_id == old_compartment_id && serving_mode == old_serving_mode { + return Ok(()); + } + let Some(envelope) = envelope else { + return Err(FlowError::InvalidArgument( + "compartmentId and servingMode edits require a ChatDetails envelope".into(), + )); + }; + if compartment_id != old_compartment_id { + set_or_remove_json( + envelope, + "compartmentId", + compartment_id.clone().map(Json::String), + ); + } + if serving_mode != old_serving_mode { + set_or_remove_json(envelope, "servingMode", serving_mode.clone()); + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// LlmCodec implementation +// --------------------------------------------------------------------------- + +impl LlmCodec for OCIGenAIChatCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI) + } + + 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 (envelope, chat_request) = split_envelope(obj); + let api_format = request_api_format(chat_request); + + let messages = if api_format == "COHERE" { + decode_cohere_messages(chat_request)? + } else { + match chat_request.get("messages") { + None | Some(Json::Null) => Vec::new(), + Some(Json::Array(messages)) => messages + .iter() + .map(decode_generic_message) + .collect::>>()?, + Some(_) => { + return Err(FlowError::InvalidArgument( + "OCI GenAI GENERIC messages must be an array".into(), + )); + } + } + }; + let params = decode_params(chat_request, &api_format)?; + let tools = decode_tools(chat_request)?; + let tool_choice = chat_request + .get("toolChoice") + .filter(|value| !value.is_null()) + .map(|value| ToolChoice::ProviderNative(native_component(value))); + + let modeled = match api_format.as_str() { + "COHERE" => MODELED_COHERE_REQUEST_KEYS, + "COHEREV2" => MODELED_COHERE_V2_REQUEST_KEYS, + _ => MODELED_GENERIC_REQUEST_KEYS, + }; + let extra: serde_json::Map = chat_request + .iter() + .filter(|(key, _)| !is_modeled_key(key, modeled)) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + + Ok(AnnotatedLlmRequest { + messages, + instructions: None, + model: envelope.and_then(model_from_envelope), + params, + tools, + tool_choice, + 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: Some(ApiSpecificRequest::OCIGenAI { + compartment_id: envelope + .and_then(|envelope| envelope.get("compartmentId")) + .and_then(Json::as_str) + .map(str::to_string), + serving_mode: envelope + .and_then(|envelope| envelope.get("servingMode")) + .cloned(), + api_format: Some(api_format), + }), + extra, + }) + } + + fn encode(&self, annotated: &AnnotatedLlmRequest, original: &LlmRequest) -> Result { + let baseline = self.decode(original)?; + let mut content = original.content.clone(); + let obj = content + .as_object_mut() + .ok_or_else(|| FlowError::Internal("original content is not an object".into()))?; + + // Split the mutable envelope from a working copy of the chat request. + let chat_request_key = obj + .get("chatRequest") + .is_some_and(Json::is_object) + .then(|| "chatRequest".to_string()); + let mut chat_request = match &chat_request_key { + Some(key) => obj + .get(key) + .and_then(Json::as_object) + .cloned() + .unwrap_or_default(), + None => obj.clone(), + }; + let api_format = request_api_format(&chat_request); + + validate_oci_supported_fields(annotated, &baseline)?; + + if annotated.messages != baseline.messages { + if api_format == "COHERE" { + encode_cohere_messages(&mut chat_request, &annotated.messages)?; + } else { + patch_generic_messages( + &mut chat_request, + &annotated.messages, + &baseline.messages, + &api_format, + )?; + } + } + + patch_params( + &mut chat_request, + annotated.params.as_ref(), + baseline.params.as_ref(), + &api_format, + ); + + if annotated.tools != baseline.tools { + let tools = annotated + .tools + .as_deref() + .map(|tools| { + tools + .iter() + .map(encode_oci_tool) + .collect::>>() + }) + .transpose()? + .map(Json::Array); + set_or_remove_json(&mut chat_request, "tools", tools); + } + if annotated.tool_choice != baseline.tool_choice { + let tool_choice = annotated + .tool_choice + .as_ref() + .map(encode_oci_tool_choice) + .transpose()?; + set_or_remove_json(&mut chat_request, "toolChoice", tool_choice); + } + + patch_extra_fields(&mut chat_request, &baseline.extra, &annotated.extra); + + match chat_request_key { + Some(key) => { + obj.insert(key, Json::Object(chat_request)); + patch_oci_api_specific(Some(obj), &annotated.api_specific, &baseline.api_specific)?; + Ok(LlmRequest { + headers: original.headers.clone(), + content, + }) + } + None => { + patch_oci_api_specific(None, &annotated.api_specific, &baseline.api_specific)?; + Ok(LlmRequest { + headers: original.headers.clone(), + content: Json::Object(chat_request), + }) + } + } + } +} + +// --------------------------------------------------------------------------- +// LlmResponseCodec implementation +// --------------------------------------------------------------------------- + +impl LlmResponseCodec for OCIGenAIChatCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI) + } + + fn decode_response(&self, response: &Json) -> Result { + let Some(obj) = response.as_object() else { + // Non-object responses are preserved raw so observability still + // captures whatever the provider path produced. + let mut extra = serde_json::Map::new(); + extra.insert("raw".to_string(), response.clone()); + return Ok(AnnotatedLlmResponse { + extra, + ..AnnotatedLlmResponse::default() + }); + }; + + let (envelope, chat_response) = match obj.get("chatResponse").and_then(Json::as_object) { + Some(chat_response) => (Some(obj), chat_response), + None => (None, obj), + }; + + let model = envelope + .and_then(|envelope| envelope.get("modelId")) + .and_then(Json::as_str) + .map(str::to_string); + let model_version = envelope + .and_then(|envelope| envelope.get("modelVersion")) + .and_then(Json::as_str) + .map(str::to_string); + let api_format = chat_response + .get("apiFormat") + .and_then(Json::as_str) + .unwrap_or("GENERIC") + .to_uppercase(); + + let (message, tool_calls, finish_reason, nested_extra) = match api_format.as_str() { + "COHERE" => decode_cohere_response_body(chat_response), + "COHEREV2" => decode_cohere_v2_response_body(chat_response)?, + _ => decode_generic_response_body(chat_response)?, + }; + + let id = if api_format == "COHEREV2" { + chat_response + .get("id") + .and_then(Json::as_str) + .map(str::to_string) + } else { + None + }; + + let usage = chat_response + .get("usage") + .and_then(Json::as_object) + .map(decode_oci_usage); + + // Preserve fields the normalized shape does not model so observability + // keeps timeCreated, service tiers, grounding metadata, and future + // provider fields. + let modeled_response_keys: &[&str] = match api_format.as_str() { + "COHERE" => &["apiFormat", "text", "finishReason", "toolCalls", "usage"], + "COHEREV2" => &["apiFormat", "id", "message", "finishReason", "usage"], + _ => &["apiFormat", "choices", "usage"], + }; + let mut extra = match envelope { + Some(envelope) => { + unmodeled_fields(envelope, &["chatResponse", "modelId", "modelVersion"]) + } + None => serde_json::Map::new(), + }; + extra.extend(unmodeled_fields(chat_response, modeled_response_keys)); + extra.extend(nested_extra); + + Ok(AnnotatedLlmResponse { + id, + model, + message, + tool_calls, + finish_reason: finish_reason.as_deref().map(map_oci_finish_reason), + usage, + optimization_summary: None, + api_specific: Some(ApiSpecificResponse::OCIGenAI { + api_format: Some(api_format), + model_version, + }), + extra, + }) + } +} + +type ResponseBody = ( + Option, + Option>, + Option, + serde_json::Map, +); + +/// Keys of the decoded GENERIC choice consumed by the normalized shape. +/// +/// `index` is excluded from `extra` carriage as well: it is positional trivia +/// (always `0` for the single decoded choice) rather than provider data. +const MODELED_CHOICE_KEYS: &[&str] = &["message", "finishReason", "index"]; + +/// Keys of a decoded assistant message consumed by the normalized shape. +const MODELED_MESSAGE_KEYS: &[&str] = &["role", "content", "toolCalls"]; + +/// Namespace unmodeled fields of a decoded nested container under `key`. +/// +/// The choice-level fields of GENERIC responses (`logprobs`, `usage`, +/// `groundingMetadata`, `serviceTier`) and the message-level fields of +/// GENERIC (`refusal`, `annotations`, `reasoningContent`) and COHEREV2 +/// (`toolPlan`, `citations`) responses are documented in the OCI schema but +/// not normalized; they are carried in `extra` under the container's wire key +/// so their origin stays unambiguous. +fn nest_unmodeled_fields( + extra: &mut serde_json::Map, + key: &str, + obj: &serde_json::Map, + modeled: &[&str], +) { + let unmodeled = unmodeled_fields(obj, modeled); + if !unmodeled.is_empty() { + extra.insert(key.to_string(), Json::Object(unmodeled)); + } +} + +fn decode_generic_response_body( + chat_response: &serde_json::Map, +) -> Result { + let mut nested_extra = serde_json::Map::new(); + let Some(first_choice) = chat_response + .get("choices") + .and_then(Json::as_array) + .and_then(|choices| choices.first()) + .and_then(Json::as_object) + else { + return Ok((None, None, None, nested_extra)); + }; + nest_unmodeled_fields( + &mut nested_extra, + "choice", + first_choice, + MODELED_CHOICE_KEYS, + ); + let finish_reason = first_choice + .get("finishReason") + .and_then(Json::as_str) + .map(str::to_string); + let Some(raw_message) = first_choice.get("message").and_then(Json::as_object) else { + return Ok((None, None, finish_reason, nested_extra)); + }; + nest_unmodeled_fields( + &mut nested_extra, + "message", + raw_message, + MODELED_MESSAGE_KEYS, + ); + let message = decode_generic_content(raw_message.get("content"))?; + let tool_calls = raw_message + .get("toolCalls") + .and_then(Json::as_array) + .map(|calls| decode_response_tool_calls(calls)) + .filter(|calls: &Vec| !calls.is_empty()); + Ok((message, tool_calls, finish_reason, nested_extra)) +} + +fn decode_cohere_response_body(chat_response: &serde_json::Map) -> ResponseBody { + let message = chat_response + .get("text") + .and_then(Json::as_str) + .map(|text| MessageContent::Text(text.to_string())); + let tool_calls = chat_response + .get("toolCalls") + .and_then(Json::as_array) + .map(|calls| decode_response_tool_calls(calls)) + .filter(|calls| !calls.is_empty()); + let finish_reason = chat_response + .get("finishReason") + .and_then(Json::as_str) + .map(str::to_string); + // COHERE (v1) is flat: unmodeled fields live directly on the chat + // response and are already carried by the chat-response-level pass. + (message, tool_calls, finish_reason, serde_json::Map::new()) +} + +/// Decode a COHEREV2 (`CohereChatResponseV2`) body: a single assistant +/// `message` whose `content` is a typed part list (`TEXT`, `THINKING`, +/// `IMAGE_URL`, `DOCUMENT`) and whose tool calls nest an OpenAI-style +/// `function` object. +fn decode_cohere_v2_response_body( + chat_response: &serde_json::Map, +) -> Result { + let mut nested_extra = serde_json::Map::new(); + let finish_reason = chat_response + .get("finishReason") + .and_then(Json::as_str) + .map(str::to_string); + let Some(raw_message) = chat_response.get("message").and_then(Json::as_object) else { + return Ok((None, None, finish_reason, nested_extra)); + }; + nest_unmodeled_fields( + &mut nested_extra, + "message", + raw_message, + MODELED_MESSAGE_KEYS, + ); + let message = decode_generic_content(raw_message.get("content"))?; + let tool_calls = raw_message + .get("toolCalls") + .and_then(Json::as_array) + .map(|calls| decode_response_tool_calls(calls)) + .filter(|calls: &Vec| !calls.is_empty()); + Ok((message, tool_calls, finish_reason, nested_extra)) +} + +/// Convert an OCI response tool-call list into [`ResponseToolCall`]s. +fn decode_response_tool_calls(calls: &[Json]) -> Vec { + calls + .iter() + .enumerate() + .filter_map(|(index, call)| decode_response_tool_call(index, call)) + .collect() +} + +/// Convert an OCI response tool call into [`ResponseToolCall`]. +/// +/// GENERIC calls are flat (`{id, type, name, arguments}`) with `arguments` as a +/// JSON-encoded string; COHERE calls carry `name` plus parsed `parameters` and +/// no `id`, so a positional `call_{index}` id is synthesized to keep parallel +/// calls distinguishable; COHEREV2 calls nest `name`/`arguments` under an +/// OpenAI-style `function` object next to the `id`. +fn decode_response_tool_call(index: usize, value: &Json) -> Option { + let obj = value.as_object()?; + let body = obj.get("function").and_then(Json::as_object).unwrap_or(obj); + let name = body.get("name")?.as_str()?.to_string(); + let arguments = match body.get("arguments") { + Some(Json::String(text)) => { + // CRITICAL: GENERIC arguments arrive JSON-encoded; parse for the + // normalized shape, preserving the raw string when unparseable. + serde_json::from_str::(text).unwrap_or_else(|_| Json::String(text.clone())) + } + Some(other) => other.clone(), + None => body.get("parameters").cloned().unwrap_or(Json::Null), + }; + let id = match obj.get("id").and_then(Json::as_str) { + Some(id) => id.to_string(), + None => format!("call_{index}"), + }; + Some(ResponseToolCall { + id, + name, + arguments, + }) +} + +/// Map OCI usage counters onto the normalized [`Usage`] field names. +/// +/// OpenAI and xAI models report cache hits under +/// `promptTokensDetails.cachedTokens`. +fn decode_oci_usage(usage: &serde_json::Map) -> Usage { + let cache_read_tokens = usage + .get("promptTokensDetails") + .and_then(Json::as_object) + .and_then(|details| details.get("cachedTokens")) + .and_then(Json::as_u64); + Usage { + prompt_tokens: usage.get("promptTokens").and_then(Json::as_u64), + completion_tokens: usage.get("completionTokens").and_then(Json::as_u64), + total_tokens: usage.get("totalTokens").and_then(Json::as_u64), + cache_read_tokens, + cache_write_tokens: None, + cost: None, + } +} + +// --------------------------------------------------------------------------- +// Streaming codec +// --------------------------------------------------------------------------- + +/// Streaming counterpart to [`OCIGenAIChatCodec`]. +/// +/// Replays the OCI Generative AI SSE event sequence into the same JSON shape a +/// non-streaming `ChatResult` carries (`{modelId, chatResponse: {apiFormat, +/// ...}}`). Once finalized, the assembled JSON can be fed back through +/// [`OCIGenAIChatCodec::decode_response`] to produce an +/// [`AnnotatedLlmResponse`] — meaning streaming and non-streaming OCI requests +/// converge on the same observability output. +/// +/// # Strategy +/// +/// OCI streams untagged chat-response deltas. `GENERIC` events carry +/// `{index, message: {role, content: [{type: "TEXT", text}], toolCalls}, finishReason}` +/// fragments whose text and tool-call `arguments` accumulate per choice index; +/// `COHERE` events carry incremental `{apiFormat: "COHERE", text}` fragments +/// with `finishReason` on the terminal event. Events wrapped in a +/// `chatResponse` envelope are unwrapped first, and `modelId`/`usage` are +/// captured whenever a chunk supplies them. +/// +/// Internal state lives behind `Arc>` so the `&self`-produced +/// collector and finalizer closures share access. Each instance is single-use +/// because [`LlmFinalizerFn`] consumes the finalize step. +/// +/// [`LlmFinalizerFn`]: crate::api::runtime::LlmFinalizerFn +pub struct OCIGenAIStreamingCodec { + state: std::sync::Arc>, +} + +impl OCIGenAIStreamingCodec { + /// Creates a fresh streaming codec with empty accumulator state. + pub fn new() -> Self { + Self { + state: std::sync::Arc::new(std::sync::Mutex::new(OCIGenAIStreamingState::default())), + } + } +} + +impl Default for OCIGenAIStreamingCodec { + fn default() -> Self { + Self::new() + } +} + +impl super::streaming::StreamingCodec for OCIGenAIStreamingCodec { + 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()); + // Move state out so finalize can consume it; the codec is single-use, so leaving a + // default behind is intentional and never observed by another caller. + std::mem::take(&mut *guard).finalize() + }) + } +} + +#[derive(Debug, Default)] +struct OCIGenAIStreamingState { + model_id: Option, + /// Resolved from the first event's `apiFormat`, or inferred from the event + /// shape (`message`/`index` => GENERIC, bare `text` => COHERE). + api_format: Option, + /// Latest non-null usage snapshot; the terminal event's counters win. + usage: Option, + /// Per-choice accumulators keyed by `index`. BTreeMap so finalize emits + /// choices in stable order. + choices: std::collections::BTreeMap, + cohere_text: String, + cohere_finish_reason: Option, +} + +#[derive(Debug, Default)] +struct OCIChoiceState { + role: Option, + /// Typed content parts in arrival order; consecutive TEXT fragments merge + /// into one part, non-TEXT typed parts (THINKING, IMAGE_URL, DOCUMENT, + /// future kinds) are preserved verbatim. + parts: Vec, + /// Tool-call accumulators in first-seen order. Fragments that carry an + /// `id` are matched to the accumulator with that id (OCI provides no + /// per-call `index`, and parallel calls can each arrive at event-local + /// position 0); id-less fragments fall back to their array position. + tool_calls: Vec, + finish_reason: Option, +} + +#[derive(Debug, Default)] +struct OCIToolCallState { + id: Option, + type_: Option, + name: Option, + arguments: String, +} + +impl OCIGenAIStreamingState { + fn observe(&mut self, event: &Json) { + let Some(obj) = event.as_object() else { + return; + }; + // Some transports wrap each delta in the ChatResult envelope; unwrap it. + let inner = obj.get("chatResponse").and_then(Json::as_object); + if let Some(model_id) = obj.get("modelId").and_then(Json::as_str) { + self.model_id = Some(model_id.to_string()); + } + let obj = inner.unwrap_or(obj); + + if self.api_format.is_none() { + self.api_format = obj + .get("apiFormat") + .and_then(Json::as_str) + .map(str::to_uppercase) + .or_else(|| self.infer_api_format(obj)); + } + if let Some(usage) = obj.get("usage") + && !usage.is_null() + { + self.usage = Some(usage.clone()); + } + + if self.api_format.as_deref() == Some("COHERE") { + let finish_reason = obj.get("finishReason").and_then(Json::as_str); + if let Some(text) = obj.get("text").and_then(Json::as_str) { + if finish_reason.is_some() && !text.is_empty() { + // The live service's terminal COHERE event repeats the + // complete response text (alongside chatHistory and + // finishReason); take it as authoritative rather than + // appending, which would double the assembled text. + self.cohere_text = text.to_string(); + } else { + self.cohere_text.push_str(text); + } + } + if let Some(reason) = finish_reason { + self.cohere_finish_reason = Some(reason.to_string()); + } + return; + } + + // GENERIC: the event is either a bare choice delta or carries a + // `choices` array of deltas. + match obj.get("choices").and_then(Json::as_array) { + Some(choices) => { + for choice in choices { + if let Some(choice) = choice.as_object() { + self.observe_generic_choice(choice); + } + } + } + None => self.observe_generic_choice(obj), + } + } + + fn infer_api_format(&self, obj: &serde_json::Map) -> Option { + if obj.get("message").is_some() + || obj.get("choices").is_some() + || obj.get("index").is_some() + { + Some("GENERIC".to_string()) + } else if obj.get("text").is_some() { + Some("COHERE".to_string()) + } else { + None + } + } + + fn observe_generic_choice(&mut self, choice: &serde_json::Map) { + let index = choice.get("index").and_then(Json::as_u64).unwrap_or(0); + let entry = self.choices.entry(index).or_default(); + if let Some(reason) = choice.get("finishReason").and_then(Json::as_str) { + entry.finish_reason = Some(reason.to_string()); + } + let Some(message) = choice.get("message").and_then(Json::as_object) else { + return; + }; + if let Some(role) = message.get("role").and_then(Json::as_str) { + entry.role = Some(role.to_string()); + } + if let Some(parts) = message.get("content").and_then(Json::as_array) { + for part in parts { + entry.observe_content_part(part); + } + } + if let Some(tool_calls) = message.get("toolCalls").and_then(Json::as_array) { + for (position, tool_call) in tool_calls.iter().enumerate() { + if let Some(tool_call) = tool_call.as_object() { + entry.observe_tool_call(position, tool_call); + } + } + } + } + + fn finalize(self) -> Json { + let api_format = self.api_format.unwrap_or_else(|| "GENERIC".to_string()); + let mut chat_response = serde_json::Map::new(); + chat_response.insert("apiFormat".to_string(), Json::String(api_format.clone())); + if api_format == "COHERE" { + chat_response.insert("text".to_string(), Json::String(self.cohere_text)); + if let Some(reason) = self.cohere_finish_reason { + chat_response.insert("finishReason".to_string(), Json::String(reason)); + } + } else if api_format == "COHEREV2" { + // The COHEREV2 response decoder reads a single root-level + // `message` (nested-function tool calls, root finishReason) + // rather than a `choices` array. + let choice = self.choices.into_values().next().unwrap_or_default(); + let finish_reason = choice.finish_reason.clone(); + chat_response.insert("message".to_string(), choice.finalize_v2_message()); + if let Some(reason) = finish_reason { + chat_response.insert("finishReason".to_string(), Json::String(reason)); + } + } else { + let choices: Vec = self + .choices + .into_iter() + .map(|(index, choice)| choice.finalize(index)) + .collect(); + chat_response.insert("choices".to_string(), Json::Array(choices)); + } + if let Some(usage) = self.usage { + chat_response.insert("usage".to_string(), usage); + } + let mut output = serde_json::Map::new(); + if let Some(model_id) = self.model_id { + output.insert("modelId".to_string(), Json::String(model_id)); + } + output.insert("chatResponse".to_string(), Json::Object(chat_response)); + Json::Object(output) + } +} + +impl OCIChoiceState { + fn observe_content_part(&mut self, part: &Json) { + if part.get("type").and_then(Json::as_str) == Some("TEXT") { + let Some(text) = part.get("text").and_then(Json::as_str) else { + return; + }; + if let Some(Json::Object(last)) = self.parts.last_mut() + && last.get("type").and_then(Json::as_str) == Some("TEXT") + { + let merged = format!( + "{}{}", + last.get("text").and_then(Json::as_str).unwrap_or_default(), + text + ); + last.insert("text".to_string(), Json::String(merged)); + return; + } + self.parts + .push(serde_json::json!({"type": "TEXT", "text": text})); + return; + } + if part.is_object() { + self.parts.push(part.clone()); + } + } + + /// The accumulated parts with empty TEXT placeholders removed; an + /// all-empty stream yields `[]` so the response decode reports no + /// assistant message, matching the non-streaming path. + fn content_parts(parts: Vec) -> Vec { + parts + .into_iter() + .filter(|part| { + part.get("type").and_then(Json::as_str) != Some("TEXT") + || part + .get("text") + .and_then(Json::as_str) + .is_some_and(|text| !text.is_empty()) + }) + .collect() + } + + fn observe_tool_call(&mut self, position: usize, tool_call: &serde_json::Map) { + let slot = match tool_call.get("id").and_then(Json::as_str) { + Some(id) => self + .tool_calls + .iter() + .position(|state| state.id.as_deref() == Some(id)) + .unwrap_or_else(|| { + self.tool_calls.push(OCIToolCallState::default()); + self.tool_calls.len() - 1 + }), + None if position < self.tool_calls.len() => position, + None => { + self.tool_calls.push(OCIToolCallState::default()); + self.tool_calls.len() - 1 + } + }; + let state = &mut self.tool_calls[slot]; + if let Some(id) = tool_call.get("id").and_then(Json::as_str) { + state.id = Some(id.to_string()); + } + if let Some(type_) = tool_call.get("type").and_then(Json::as_str) { + state.type_ = Some(type_.to_string()); + } + // COHEREV2 fragments nest name/arguments under a `function` object; + // GENERIC fragments are flat. + let body = tool_call + .get("function") + .and_then(Json::as_object) + .unwrap_or(tool_call); + if let Some(name) = body.get("name").and_then(Json::as_str) { + state.name = Some(name.to_string()); + } + if let Some(arguments) = body.get("arguments").and_then(Json::as_str) { + state.arguments.push_str(arguments); + } + } + + /// Assemble the accumulated choice as a COHEREV2 root `message` + /// (typed content parts, nested-function tool calls). + fn finalize_v2_message(self) -> Json { + let mut message = serde_json::Map::new(); + message.insert( + "role".to_string(), + Json::String(self.role.unwrap_or_else(|| "ASSISTANT".to_string())), + ); + message.insert( + "content".to_string(), + Json::Array(Self::content_parts(self.parts)), + ); + if !self.tool_calls.is_empty() { + let tool_calls: Vec = self + .tool_calls + .into_iter() + .map(OCIToolCallState::finalize_v2) + .collect(); + message.insert("toolCalls".to_string(), Json::Array(tool_calls)); + } + Json::Object(message) + } + + fn finalize(self, index: u64) -> Json { + let mut message = serde_json::Map::new(); + message.insert( + "role".to_string(), + Json::String(self.role.unwrap_or_else(|| "ASSISTANT".to_string())), + ); + message.insert( + "content".to_string(), + Json::Array(Self::content_parts(self.parts)), + ); + if !self.tool_calls.is_empty() { + let tool_calls: Vec = self + .tool_calls + .into_iter() + .map(OCIToolCallState::finalize) + .collect(); + message.insert("toolCalls".to_string(), Json::Array(tool_calls)); + } + let mut choice = serde_json::Map::new(); + choice.insert("index".to_string(), Json::Number(index.into())); + choice.insert("message".to_string(), Json::Object(message)); + if let Some(reason) = self.finish_reason { + choice.insert("finishReason".to_string(), Json::String(reason)); + } + Json::Object(choice) + } +} + +impl OCIToolCallState { + /// Assemble the call in the COHEREV2 nested-function wire shape. + fn finalize_v2(self) -> Json { + let mut function = serde_json::Map::new(); + function.insert( + "name".to_string(), + Json::String(self.name.unwrap_or_default()), + ); + function.insert("arguments".to_string(), Json::String(self.arguments)); + let mut call = serde_json::Map::new(); + if let Some(id) = self.id { + call.insert("id".to_string(), Json::String(id)); + } + call.insert( + "type".to_string(), + Json::String(self.type_.unwrap_or_else(|| "FUNCTION".to_string())), + ); + call.insert("function".to_string(), Json::Object(function)); + Json::Object(call) + } + + fn finalize(self) -> Json { + let mut call = serde_json::Map::new(); + if let Some(id) = self.id { + call.insert("id".to_string(), Json::String(id)); + } + call.insert( + "type".to_string(), + Json::String(self.type_.unwrap_or_else(|| "FUNCTION".to_string())), + ); + call.insert( + "name".to_string(), + Json::String(self.name.unwrap_or_default()), + ); + call.insert("arguments".to_string(), Json::String(self.arguments)); + Json::Object(call) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +#[path = "../../tests/unit/codec/oci_genai_tests.rs"] +mod tests; diff --git a/crates/core/src/codec/resolve.rs b/crates/core/src/codec/resolve.rs index f54652291..30a585520 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_generate_content, openai_chat, openai_responses}; +use super::{anthropic, gemini_generate_content, oci_genai, 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, + /// OCI Generative AI chat. + OCIGenAI, /// Gemini generateContent. GeminiGenerateContent, } @@ -68,6 +70,10 @@ pub(crate) struct ProviderSurfaceDescriptor { pub(crate) static BUILTIN_PROVIDER_SURFACES: &[ProviderSurfaceDescriptor] = &[ openai_responses::PROVIDER_SURFACE, anthropic::PROVIDER_SURFACE, + // OCI GenAI must precede OpenAI Chat: a bare OCI GENERIC chatRequest body + // carries `messages`, which the looser OpenAI Chat detector would claim + // under first-match-wins. It has no overlap with the surfaces above. + oci_genai::PROVIDER_SURFACE, openai_chat::PROVIDER_SURFACE, gemini_generate_content::PROVIDER_SURFACE, ]; @@ -155,6 +161,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::OCIGenAI => &oci_genai::PROVIDER_SURFACE, ProviderSurface::GeminiGenerateContent => &gemini_generate_content::PROVIDER_SURFACE, } } diff --git a/crates/core/src/observability/otel_genai.rs b/crates/core/src/observability/otel_genai.rs index 527170f3f..17aaf4034 100644 --- a/crates/core/src/observability/otel_genai.rs +++ b/crates/core/src/observability/otel_genai.rs @@ -620,6 +620,9 @@ fn provider_from_normalized_request(event: &Event) -> Option<&'static str> { ApiSpecificRequest::OpenAIChat { .. } | ApiSpecificRequest::OpenAIResponses { .. } => { Some("openai") } + // Not an OTel well-known value yet; follows the dotted cloud-provider + // convention (`aws.bedrock`, `gcp.gemini`). + ApiSpecificRequest::OCIGenAI { .. } => Some("oci.genai"), ApiSpecificRequest::Custom { .. } => None, } } diff --git a/crates/core/src/plugins/nemo_guardrails/component.rs b/crates/core/src/plugins/nemo_guardrails/component.rs index 9689c5bab..d23ab59fe 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_generate_content"], + values: ["openai_chat", "openai_responses", "anthropic_messages", "oci_genai", "gemini_generate_content"], optional: true, }, input => { label: "input", kind: Boolean }, @@ -434,6 +434,7 @@ fn codec_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::s "openai_chat", "openai_responses", "anthropic_messages", + "oci_genai", "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 baa2ef915..6557d4c14 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, + OCIGenAI, GeminiGenerateContent, } @@ -900,6 +901,7 @@ impl LocalGuardrailsCodec { Self::OpenAIChat => ProviderSurface::OpenAIChat, Self::OpenAIResponses => ProviderSurface::OpenAIResponses, Self::AnthropicMessages => ProviderSurface::AnthropicMessages, + Self::OCIGenAI => ProviderSurface::OCIGenAI, Self::GeminiGenerateContent => ProviderSurface::GeminiGenerateContent, } } @@ -909,6 +911,7 @@ impl LocalGuardrailsCodec { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::OCIGenAI => Self::OCIGenAI, ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, } } @@ -1311,10 +1314,62 @@ fn extract_stream_text(codec: LocalGuardrailsCodec, chunk: &Json) -> Option extract_openai_chat_stream_text(chunk), LocalGuardrailsCodec::OpenAIResponses => extract_openai_response_stream_text(chunk), LocalGuardrailsCodec::AnthropicMessages => extract_anthropic_stream_text(chunk), + LocalGuardrailsCodec::OCIGenAI => extract_oci_genai_stream_text(chunk), LocalGuardrailsCodec::GeminiGenerateContent => extract_gemini_stream_text(chunk), } } +/// Collect the concatenated TEXT-part text from OCI GENERIC stream deltas or +/// the bare `text` fragment of COHERE deltas. Events may arrive wrapped in a +/// `chatResponse` envelope, and GENERIC deltas are either a bare choice +/// (`message` at the top level) or carry a `choices` array of deltas, +/// mirroring the stream shapes the OCI streaming codec accepts. +/// +/// The live service's terminal COHERE event (the one carrying `finishReason`) +/// repeats the complete response text already delivered by earlier deltas. +/// Forwarding it would double the text the output rails evaluate, so it is +/// suppressed here, mirroring the deduplication in the OCI streaming codec. +fn extract_oci_genai_stream_text(chunk: &serde_json::Map) -> Option { + let chunk = chunk + .get("chatResponse") + .and_then(Json::as_object) + .unwrap_or(chunk); + if let Some(text) = chunk.get("text").and_then(Json::as_str) { + if chunk.get("finishReason").and_then(Json::as_str).is_some() { + return None; + } + return (!text.is_empty()).then(|| text.to_string()); + } + fn collect_generic_text(message: &Json, collected: &mut String) { + let Some(parts) = message.get("content").and_then(Json::as_array) else { + return; + }; + for part in parts { + if part.get("type").and_then(Json::as_str) == Some("TEXT") + && let Some(text) = part.get("text").and_then(Json::as_str) + { + collected.push_str(text); + } + } + } + let mut collected = String::new(); + match chunk.get("choices").and_then(Json::as_array) { + Some(choices) => { + for choice in choices { + if let Some(message) = choice.get("message") { + collect_generic_text(message, &mut collected); + } + } + } + None => { + if let Some(message) = chunk.get("message") { + collect_generic_text(message, &mut collected); + } + } + } + (!collected.is_empty()).then_some(collected) +} + fn extract_openai_chat_stream_text(chunk: &serde_json::Map) -> Option { let choices = chunk.get("choices")?.as_array()?; let parts = choices diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index c0c836684..705107dba 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -34,6 +34,7 @@ use nemo_relay::api::runtime::{create_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::{EmitMarkEventParams, ScopeType, event}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::codec::anthropic::AnthropicMessagesCodec; +use nemo_relay::codec::oci_genai::OCIGenAIChatCodec; use nemo_relay::codec::openai_chat::OpenAIChatCodec; use nemo_relay::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationKind, LlmOptimizationModel, @@ -1456,6 +1457,90 @@ async fn test_response_codec_populates_annotated_response() { deregister_subscriber("resp_codec_sub").unwrap(); } +#[tokio::test] +async fn test_oci_genai_response_codec_populates_annotated_response() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::new())); + let ec = events.clone(); + register_subscriber( + "oci_resp_codec_sub", + Arc::new(move |e: &Event| { + ec.lock().unwrap().push(e.clone()); + }), + ) + .unwrap(); + + // Shape observed from a live OCI Generative AI GENERIC tool-call response. + let func: LlmExecutionNextFn = Arc::new(|_req| { + Box::pin(async move { + Ok(json!({ + "modelId": "meta.llama-4-maverick-17b-128e-instruct-fp8", + "modelVersion": "1.0.0", + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "toolCalls": [{ + "type": "FUNCTION", + "id": "chatcmpl-tool-bda9d62eab5cea3c", + "name": "get_weather", + "arguments": "{\"city\": \"Paris\"}" + }] + }, + "finishReason": "tool_calls" + }], + "usage": {"promptTokens": 627, "completionTokens": 13, "totalTokens": 640} + } + })) + }) + }); + let response_codec: Arc = Arc::new(OCIGenAIChatCodec); + + let _result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("oci_genai") + .request(make_llm_request( + json!({"messages": [{"role": "USER", "content": [{"type": "TEXT", "text": "hi"}]}]}), + )) + .func(func) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + let captured = captured_events_snapshot(&events); + let end_event = captured + .iter() + .find(|e| is_scope_event(e, ScopeType::Llm, ScopeCategory::End)) + .expect("expected LlmEnd event"); + + let annotated = end_event + .annotated_response() + .expect("annotated_response should be Some when the OCI codec is active"); + assert_eq!( + annotated.model.as_deref(), + Some("meta.llama-4-maverick-17b-128e-instruct-fp8") + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::ToolUse)); + assert_eq!(annotated.message, None); + let tool_calls = annotated.tool_calls.as_ref().expect("tool calls decoded"); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].arguments, json!({"city": "Paris"})); + let usage = annotated.usage.as_ref().expect("usage decoded"); + assert_eq!(usage.prompt_tokens, Some(627)); + assert_eq!(usage.completion_tokens, Some(13)); + assert_eq!(usage.total_tokens, Some(640)); + + deregister_subscriber("oci_resp_codec_sub").unwrap(); +} + #[tokio::test] async fn test_response_codec_annotation_uses_sanitized_managed_response() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/unit/codec/oci_genai_tests.rs b/crates/core/tests/unit/codec/oci_genai_tests.rs new file mode 100644 index 000000000..aa6384fbe --- /dev/null +++ b/crates/core/tests/unit/codec/oci_genai_tests.rs @@ -0,0 +1,2066 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Unit tests for the OCI Generative AI codec in the NeMo Relay core crate. + +use super::*; +use serde_json::json; + +use super::super::request::{ContentPart, Message, MessageContent, ToolChoice}; +use super::super::resolve::{ + detect_request_surface, detect_request_surface_with_hint, detect_response_surface, +}; +use super::super::response::{ApiSpecificResponse, FinishReason}; +use super::super::streaming::StreamingCodec; + +// ------------------------------------------------------------------- +// Helpers and fixtures +// ------------------------------------------------------------------- + +const DEDICATED_ENDPOINT: &str = "ocid1.generativeaiendpoint.oc1.us-chicago-1.example"; + +fn make_request(content: Json) -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content, + } +} + +fn generic_chat_details() -> Json { + json!({ + "compartmentId": "ocid1.compartment.oc1..example", + "servingMode": {"servingType": "DEDICATED", "endpointId": DEDICATED_ENDPOINT}, + "chatRequest": { + "apiFormat": "GENERIC", + "messages": [ + {"role": "SYSTEM", "content": [{"type": "TEXT", "text": "You are terse."}]}, + {"role": "USER", "content": [{"type": "TEXT", "text": "My SSN is 111-22-3333."}]} + ], + "maxTokens": 600, + "temperature": 0.0 + } + }) +} + +fn cohere_chat_details() -> Json { + json!({ + "compartmentId": "ocid1.compartment.oc1..example", + "servingMode": {"servingType": "ON_DEMAND", "modelId": "cohere.command-a-03-2025"}, + "chatRequest": { + "apiFormat": "COHERE", + "preambleOverride": "You are terse.", + "chatHistory": [ + {"role": "USER", "message": "hello"}, + {"role": "CHATBOT", "message": "hi"} + ], + "message": "What is the weather?", + "maxTokens": 100 + } + }) +} + +/// Shape observed from a live dedicated-endpoint chat (imported NVIDIA Nemotron 3). +fn generic_chat_result() -> Json { + json!({ + "modelId": DEDICATED_ENDPOINT, + "modelVersion": "1.0", + "chatResponse": { + "apiFormat": "GENERIC", + "timeCreated": "2026-07-23T22:59:00.000Z", + "choices": [ + { + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "NEMOTRON3_OK"}] + }, + "finishReason": "stop" + } + ], + "usage": {"promptTokens": 18, "completionTokens": 5, "totalTokens": 23} + } + }) +} + +fn cohere_chat_result() -> Json { + json!({ + "modelId": "cohere.command-a-03-2025", + "chatResponse": { + "apiFormat": "COHERE", + "text": "Sunny and 72.", + "finishReason": "COMPLETE", + "usage": {"promptTokens": 12, "completionTokens": 4, "totalTokens": 16} + } + }) +} + +/// Envelope, request, and per-message levels all carry unmodeled fields. +fn unmodeled_generic() -> Json { + json!({ + "compartmentId": "ocid1.compartment.oc1..example", + "opcRetryToken": "retry-abc", + "servingMode": {"servingType": "DEDICATED", "endpointId": DEDICATED_ENDPOINT, "futureFlag": true}, + "chatRequest": { + "apiFormat": "GENERIC", + "messages": [ + {"role": "SYSTEM", "content": [{"type": "TEXT", "text": "Be terse."}], "name": "sys-1"}, + {"role": "USER", "content": [{"type": "TEXT", "text": "hello"}], "unknownPerMessage": 7} + ], + "maxTokens": 64, + "topK": 40, + "seed": 7, + "unknownFutureField": {"nested": true} + } + }) +} + +fn message_role(message: &Message) -> &'static str { + match message { + Message::System { .. } => "system", + Message::User { .. } => "user", + Message::Developer { .. } => "developer", + Message::Assistant { .. } => "assistant", + Message::Tool { .. } => "tool", + Message::Function { .. } => "function", + Message::ToolCallItem { .. } => "tool_call", + Message::ToolResultItem { .. } => "tool_result", + Message::ProviderNative { .. } => "provider_native", + } +} + +fn message_text(message: &Message) -> Option<&str> { + let content = match message { + Message::System { content, .. } + | Message::User { content, .. } + | Message::Tool { content, .. } => content, + Message::Assistant { + content: Some(content), + .. + } => content, + _ => return None, + }; + match content { + MessageContent::Text(text) => Some(text.as_str()), + MessageContent::Parts(_) => None, + } +} + +// =================================================================== +// codec_identity +// =================================================================== + +#[test] +fn test_codec_identity_is_oci_genai_builtin() { + let codec = OCIGenAIChatCodec; + assert_eq!( + LlmCodec::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI), + "OCIGenAIChatCodec must not return Opaque; PII sanitization depends on a known identity" + ); +} + +#[test] +fn test_response_codec_identity_is_oci_genai_builtin() { + let codec = OCIGenAIChatCodec; + assert_eq!( + ::codec_identity(&codec), + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI), + "OCIGenAIChatCodec response codec must not return Opaque" + ); +} + +// =================================================================== +// GENERIC request decode tests +// =================================================================== + +#[test] +fn test_generic_decode_envelope() { + let annotated = OCIGenAIChatCodec + .decode(&make_request(generic_chat_details())) + .unwrap(); + + let roles: Vec<_> = annotated.messages.iter().map(message_role).collect(); + assert_eq!(roles, vec!["system", "user"]); + assert_eq!( + message_text(&annotated.messages[1]), + Some("My SSN is 111-22-3333.") + ); + assert_eq!(annotated.model.as_deref(), Some(DEDICATED_ENDPOINT)); + + let params = annotated.params.as_ref().unwrap(); + assert_eq!(params.max_tokens, Some(600)); + assert_eq!(params.temperature, Some(0.0)); + + assert_eq!( + annotated.api_specific, + Some(ApiSpecificRequest::OCIGenAI { + compartment_id: Some("ocid1.compartment.oc1..example".into()), + serving_mode: Some( + json!({"servingType": "DEDICATED", "endpointId": DEDICATED_ENDPOINT}) + ), + api_format: Some("GENERIC".into()), + }) + ); +} + +#[test] +fn test_generic_decode_bare_chat_request() { + let bare = generic_chat_details().get("chatRequest").cloned().unwrap(); + let annotated = OCIGenAIChatCodec.decode(&make_request(bare)).unwrap(); + + let roles: Vec<_> = annotated.messages.iter().map(message_role).collect(); + assert_eq!(roles, vec!["system", "user"]); + assert_eq!(annotated.model, None); + assert_eq!( + annotated.api_specific, + Some(ApiSpecificRequest::OCIGenAI { + compartment_id: None, + serving_mode: None, + api_format: Some("GENERIC".into()), + }) + ); +} + +#[test] +fn test_generic_decode_defaults_missing_api_format_to_generic() { + let annotated = OCIGenAIChatCodec + .decode(&make_request(json!({ + "messages": [{"role": "USER", "content": [{"type": "TEXT", "text": "hi"}]}], + "chatRequest": "not-an-object" + }))) + .unwrap(); + assert_eq!( + annotated.api_specific, + Some(ApiSpecificRequest::OCIGenAI { + compartment_id: None, + serving_mode: None, + api_format: Some("GENERIC".into()), + }) + ); + assert_eq!(message_text(&annotated.messages[0]), Some("hi")); +} + +#[test] +fn test_generic_decode_rejects_non_array_messages() { + let error = OCIGenAIChatCodec + .decode(&make_request(json!({ + "apiFormat": "GENERIC", + "messages": "oops" + }))) + .unwrap_err(); + assert!(matches!(error, FlowError::InvalidArgument(_)), "{error}"); +} + +#[test] +fn test_non_wire_request_renderings_are_not_decoded() { + // The codec accepts the REST wire format only (camelCase). Alternate + // renderings from Oracle tooling (CLI kebab-case, SDK-dict snake_case) + // are the caller's responsibility to convert first. + let annotated = OCIGenAIChatCodec + .decode(&make_request(json!({ + "compartment-id": "ocid1.compartment.oc1..kebab", + "serving-mode": {"serving-type": "ON_DEMAND", "model-id": "meta.llama-3.3-70b-instruct"}, + "chat-request": { + "api-format": "GENERIC", + "messages": [{"role": "USER", "content": [{"type": "TEXT", "text": "hi"}]}], + "max-tokens": 32 + } + }))) + .unwrap(); + assert_eq!(annotated.model, None); + assert!(annotated.messages.is_empty()); + assert_eq!(annotated.params, None); +} + +// =================================================================== +// GENERIC request encode tests +// =================================================================== + +#[test] +fn test_redaction_round_trip_preserves_envelope() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.messages[1] = Message::User { + content: MessageContent::Text("My SSN is [REDACTED].".into()), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = encoded.content.get("chatRequest").unwrap(); + + assert_eq!( + chat_request["messages"][1], + json!({ + "role": "USER", + "content": [{"type": "TEXT", "text": "My SSN is [REDACTED]."}] + }) + ); + // Envelope fields survive untouched. + assert_eq!( + encoded.content["compartmentId"], + json!("ocid1.compartment.oc1..example") + ); + assert_eq!( + encoded.content["servingMode"], + json!({"servingType": "DEDICATED", "endpointId": DEDICATED_ENDPOINT}) + ); + assert_eq!(chat_request["maxTokens"], json!(600)); +} + +#[test] +fn test_tool_calls_round_trip() { + let payload = json!({ + "apiFormat": "GENERIC", + "messages": [ + { + "role": "ASSISTANT", + "content": [], + "toolCalls": [ + {"id": "call-1", "type": "FUNCTION", "name": "get_weather", "arguments": "{}"} + ] + }, + {"role": "TOOL", "content": [{"type": "TEXT", "text": "72F"}], "toolCallId": "call-1"} + ] + }); + let codec = OCIGenAIChatCodec; + let original = make_request(payload.clone()); + let annotated = codec.decode(&original).unwrap(); + + match &annotated.messages[0] { + Message::Assistant { + tool_calls: Some(tool_calls), + .. + } => { + assert_eq!(tool_calls[0].id, "call-1"); + assert_eq!(tool_calls[0].function.name, "get_weather"); + assert_eq!(tool_calls[0].function.arguments, "{}"); + } + other => panic!("expected assistant with tool calls, got {other:?}"), + } + match &annotated.messages[1] { + Message::Tool { tool_call_id, .. } => assert_eq!(tool_call_id, "call-1"), + other => panic!("expected tool message, got {other:?}"), + } + + // Unedited round trip is the identity. + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded.content, payload); + + // Editing the assistant message forces a rebuild through the flat OCI + // tool-call shape. + let mut edited = annotated.clone(); + edited.messages[0] = Message::Assistant { + content: Some(MessageContent::Text("checking".into())), + tool_calls: match &annotated.messages[0] { + Message::Assistant { tool_calls, .. } => tool_calls.clone(), + _ => unreachable!(), + }, + name: None, + }; + let encoded = codec.encode(&edited, &original).unwrap(); + assert_eq!( + encoded.content["messages"][0]["toolCalls"][0], + json!({"id": "call-1", "type": "FUNCTION", "name": "get_weather", "arguments": "{}"}) + ); + assert_eq!( + encoded.content["messages"][1]["toolCallId"], + json!("call-1") + ); +} + +// =================================================================== +// COHERE request tests +// =================================================================== + +#[test] +fn test_cohere_decode() { + let annotated = OCIGenAIChatCodec + .decode(&make_request(cohere_chat_details())) + .unwrap(); + + let roles: Vec<_> = annotated.messages.iter().map(message_role).collect(); + assert_eq!(roles, vec!["system", "user", "assistant", "user"]); + assert_eq!(message_text(&annotated.messages[0]), Some("You are terse.")); + assert_eq!( + message_text(annotated.messages.last().unwrap()), + Some("What is the weather?") + ); + assert_eq!(annotated.model.as_deref(), Some("cohere.command-a-03-2025")); + assert_eq!(annotated.params.as_ref().unwrap().max_tokens, Some(100)); + assert!(matches!( + &annotated.api_specific, + Some(ApiSpecificRequest::OCIGenAI { + api_format: Some(api_format), + .. + }) if api_format == "COHERE" + )); +} + +#[test] +fn test_cohere_round_trip() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_chat_details()); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + + // Unedited COHERE requests round-trip to the identical payload. + assert_eq!(encoded.content, cohere_chat_details()); +} + +#[test] +fn test_cohere_edit_rebuilds_modeled_fields() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + let last = annotated.messages.len() - 1; + annotated.messages[last] = Message::User { + content: MessageContent::Text("What is the weather in [REDACTED]?".into()), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = encoded.content.get("chatRequest").unwrap(); + assert_eq!( + chat_request["message"], + json!("What is the weather in [REDACTED]?") + ); + assert_eq!(chat_request["preambleOverride"], json!("You are terse.")); + assert_eq!( + chat_request["chatHistory"], + json!([ + {"role": "USER", "message": "hello"}, + {"role": "CHATBOT", "message": "hi"} + ]) + ); + assert_eq!( + encoded.content["servingMode"], + json!({"servingType": "ON_DEMAND", "modelId": "cohere.command-a-03-2025"}) + ); + assert_eq!(chat_request["maxTokens"], json!(100)); +} + +#[test] +fn test_cohere_stop_sequences_map_to_stop() { + let mut payload = cohere_chat_details(); + payload["chatRequest"]["stopSequences"] = json!(["END"]); + let annotated = OCIGenAIChatCodec.decode(&make_request(payload)).unwrap(); + assert_eq!( + annotated.params.as_ref().unwrap().stop, + Some(vec!["END".to_string()]) + ); +} + +// =================================================================== +// Identity invariant: encode(decode(original), original) == original +// =================================================================== + +#[test] +fn test_generic_identity() { + let codec = OCIGenAIChatCodec; + let original = make_request(unmodeled_generic()); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + + assert_eq!(encoded.content, unmodeled_generic()); +} + +#[test] +fn test_cohere_identity() { + let mut payload = cohere_chat_details(); + payload["chatRequest"]["isForceSingleStep"] = json!(true); + let codec = OCIGenAIChatCodec; + let original = make_request(payload.clone()); + let annotated = codec.decode(&original).unwrap(); + let encoded = codec.encode(&annotated, &original).unwrap(); + + assert_eq!(encoded.content, payload); +} + +#[test] +fn test_edit_preserves_unmodeled_fields_on_untouched_messages() { + let codec = OCIGenAIChatCodec; + let original = make_request(unmodeled_generic()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.messages[1] = Message::User { + content: MessageContent::Text("redacted".into()), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = encoded.content.get("chatRequest").unwrap(); + + // Untouched system message keeps its unmodeled per-message field. + assert_eq!( + chat_request["messages"][0], + unmodeled_generic()["chatRequest"]["messages"][0] + ); + // Edited message carries the redaction. + assert_eq!( + chat_request["messages"][1]["content"], + json!([{"type": "TEXT", "text": "redacted"}]) + ); + // Unmodeled request-level fields survive. + assert_eq!(chat_request["topK"], json!(40)); + assert_eq!(chat_request["seed"], json!(7)); + assert_eq!(chat_request["unknownFutureField"], json!({"nested": true})); + assert_eq!(encoded.content["opcRetryToken"], json!("retry-abc")); +} + +#[test] +fn test_param_edit_only_touches_changed_param() { + let codec = OCIGenAIChatCodec; + let original = make_request(unmodeled_generic()); + let mut annotated = codec.decode(&original).unwrap(); + + let mut params = annotated.params.clone().unwrap_or_default(); + params.max_tokens = Some(128); + annotated.params = Some(params); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = encoded.content.get("chatRequest").unwrap(); + + assert_eq!(chat_request["maxTokens"], json!(128)); + assert_eq!( + chat_request["messages"], + unmodeled_generic()["chatRequest"]["messages"] + ); +} + +#[test] +fn test_tool_choice_survives_as_provider_native() { + let payload = json!({ + "apiFormat": "GENERIC", + "messages": [{"role": "USER", "content": [{"type": "TEXT", "text": "hi"}]}], + "tools": [{"type": "FUNCTION", "name": "get_weather", "parameters": {"type": "object"}}], + "toolChoice": {"type": "auto"} + }); + let codec = OCIGenAIChatCodec; + let original = make_request(payload.clone()); + let annotated = codec.decode(&original).unwrap(); + + assert!(matches!( + &annotated.tool_choice, + Some(ToolChoice::ProviderNative(native)) if native.provider == "oci_genai" + )); + assert_eq!(annotated.tools.as_ref().map(Vec::len), Some(1)); + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded.content, payload); +} + +#[test] +fn test_model_edit_is_rejected() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + annotated.model = Some("other-model".into()); + + let error = codec.encode(&annotated, &original).unwrap_err(); + assert!(matches!(error, FlowError::InvalidArgument(_)), "{error}"); +} + +// =================================================================== +// Response decode tests +// =================================================================== + +#[test] +fn test_generic_chat_result() { + let annotated = OCIGenAIChatCodec + .decode_response(&generic_chat_result()) + .unwrap(); + + assert_eq!(annotated.model.as_deref(), Some(DEDICATED_ENDPOINT)); + assert_eq!( + annotated.message, + Some(MessageContent::Text("NEMOTRON3_OK".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); + + let usage = annotated.usage.as_ref().unwrap(); + assert_eq!(usage.prompt_tokens, Some(18)); + assert_eq!(usage.completion_tokens, Some(5)); + assert_eq!(usage.total_tokens, Some(23)); + + assert_eq!( + annotated.api_specific, + Some(ApiSpecificResponse::OCIGenAI { + api_format: Some("GENERIC".into()), + model_version: Some("1.0".into()), + }) + ); +} + +#[test] +fn test_cohere_chat_result() { + let annotated = OCIGenAIChatCodec + .decode_response(&cohere_chat_result()) + .unwrap(); + + assert_eq!( + annotated.message, + Some(MessageContent::Text("Sunny and 72.".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); + assert_eq!(annotated.model.as_deref(), Some("cohere.command-a-03-2025")); + assert_eq!( + annotated.api_specific, + Some(ApiSpecificResponse::OCIGenAI { + api_format: Some("COHERE".into()), + model_version: None, + }) + ); +} + +#[test] +fn test_non_wire_renderings_are_not_decoded() { + // The codec accepts the REST wire format only (camelCase). Alternate + // renderings from Oracle tooling (CLI kebab-case, SDK-dict snake_case) + // are the caller's responsibility to convert first. + let snake_cased = json!({ + "model_id": DEDICATED_ENDPOINT, + "chat_response": { + "api_format": "GENERIC", + "choices": [{ + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "hello"}]}, + "finish_reason": "stop" + }] + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&snake_cased).unwrap(); + + assert_eq!(annotated.model, None); + assert_eq!(annotated.message, None); + assert_eq!(annotated.finish_reason, None); + + // CLI output: kebab-case keys wrapped in a `data` envelope. + let cli_shaped = json!({ + "data": { + "model-id": DEDICATED_ENDPOINT, + "chat-response": { + "api-format": "GENERIC", + "choices": [{ + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "hello"}]}, + "finish-reason": "stop" + }], + "usage": {"total-tokens": 9} + } + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&cli_shaped).unwrap(); + + assert_eq!(annotated.model, None); + assert_eq!(annotated.message, None); + assert_eq!(annotated.finish_reason, None); + assert_eq!(annotated.usage, None); +} + +#[test] +fn test_non_dict_response() { + let annotated = OCIGenAIChatCodec + .decode_response(&json!("plain text")) + .unwrap(); + assert_eq!(annotated.extra.get("raw"), Some(&json!("plain text"))); + assert_eq!(annotated.message, None); +} + +#[test] +fn test_response_tool_calls_parse_string_arguments() { + let raw = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "message": { + "role": "ASSISTANT", + "content": [], + "toolCalls": [{ + "id": "call-9", + "type": "FUNCTION", + "name": "get_weather", + "arguments": "{\"city\": \"NYC\"}" + }] + }, + "finishReason": "tool_calls" + }] + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&raw).unwrap(); + + assert_eq!(annotated.finish_reason, Some(FinishReason::ToolUse)); + let tool_calls = annotated.tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls[0].id, "call-9"); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].arguments, json!({"city": "NYC"})); + // A tool-call-only message with `"content": []` has no assistant content. + assert_eq!(annotated.message, None); +} + +#[test] +fn test_non_text_parts_preserved_as_provider_native() { + let response = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "message": { + "role": "ASSISTANT", + "content": [ + {"type": "TEXT", "text": "see image"}, + {"type": "IMAGE", "imageUrl": {"url": "https://example.com/x.png"}} + ] + }, + "finishReason": "stop" + }] + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&response).unwrap(); + + let Some(MessageContent::Parts(parts)) = annotated.message else { + panic!("expected typed parts, got {:?}", annotated.message); + }; + assert_eq!(parts.len(), 2); + assert!(matches!(&parts[0], ContentPart::Text { text, .. } if text == "see image")); + let ContentPart::ProviderNative { + provider, + kind, + value, + } = &parts[1] + else { + panic!("expected ProviderNative part, got {:?}", parts[1]); + }; + assert_eq!(provider, "oci_genai"); + assert_eq!(kind, "IMAGE"); + assert_eq!(value["imageUrl"]["url"], json!("https://example.com/x.png")); +} + +#[test] +fn test_invalid_generic_content_shape_errors() { + for bad_content in [json!(42), json!({"type": "TEXT"}), json!([17])] { + let response = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "message": {"role": "ASSISTANT", "content": bad_content}, + "finishReason": "stop" + }] + } + }); + let error = OCIGenAIChatCodec.decode_response(&response).unwrap_err(); + assert!( + matches!(error, crate::error::FlowError::InvalidArgument(_)), + "expected InvalidArgument, got {error:?}" + ); + } +} + +#[test] +fn test_cohere_parallel_tool_calls_get_positional_ids() { + // Shape observed live: COHERE tool calls carry no `id`, so parallel calls + // must receive distinct synthesized ids. + let response = json!({ + "modelId": "cohere.command-r-08-2024", + "chatResponse": { + "apiFormat": "COHERE", + "text": "I will use the tool for each city.", + "finishReason": "COMPLETE", + "toolCalls": [ + {"name": "get_weather", "parameters": {"city": "Paris"}}, + {"name": "get_weather", "parameters": {"city": "Rome"}} + ] + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&response).unwrap(); + + let tool_calls = annotated.tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls.len(), 2); + assert_eq!(tool_calls[0].id, "call_0"); + assert_eq!(tool_calls[1].id, "call_1"); + assert_eq!(tool_calls[0].arguments, json!({"city": "Paris"})); + assert_eq!(tool_calls[1].arguments, json!({"city": "Rome"})); +} + +#[test] +fn test_usage_cached_tokens_mapped_to_cache_read() { + // Shape observed live from OpenAI and xAI models on OCI: cache hits are + // reported under `promptTokensDetails.cachedTokens`. + let response = json!({ + "modelId": "xai.grok-3-mini", + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "hi"}]}, + "finishReason": "stop" + }], + "usage": { + "promptTokens": 13, + "completionTokens": 8, + "totalTokens": 607, + "promptTokensDetails": {"cachedTokens": 3}, + "completionTokensDetails": {"reasoningTokens": 586} + } + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&response).unwrap(); + + let usage = annotated.usage.as_ref().unwrap(); + assert_eq!(usage.prompt_tokens, Some(13)); + assert_eq!(usage.cache_read_tokens, Some(3)); + assert_eq!(usage.cache_write_tokens, None); +} + +#[test] +fn test_cohere_v2_chat_result() { + // Shape per the OCI `CohereChatResponseV2` schema (apiFormat COHEREV2): + // a single assistant message with typed content parts and nested-function + // tool calls. Confirmed against the live service (us-chicago-1, + // 2026-07-29): the wire matches this schema, including provider-supplied + // nested-function tool-call ids, JSON-encoded string arguments, and + // message-level toolPlan/citations. + let response = json!({ + "modelId": "cohere.command-a-03-2025", + "modelVersion": "2.0", + "chatResponse": { + "apiFormat": "COHEREV2", + "id": "resp-v2-123", + "message": { + "role": "ASSISTANT", + "content": [ + {"type": "THINKING", "thinking": "I should call the tool."}, + {"type": "TEXT", "text": "Checking the weather."} + ], + "toolCalls": [{ + "id": "call-v2-1", + "type": "FUNCTION", + "function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"} + }], + // Message-level per the OCI CohereAssistantMessageV2 schema. + "toolPlan": "I will check the weather.", + "citations": [{"start": 0, "end": 8, "text": "Checking"}] + }, + "finishReason": "TOOL_CALL", + "usage": {"promptTokens": 20, "completionTokens": 15, "totalTokens": 35} + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&response).unwrap(); + + // Grounding metadata and the tool plan are not normalized but must + // survive, namespaced under the message they came from. + assert_eq!( + annotated.extra.get("message"), + Some(&json!({ + "toolPlan": "I will check the weather.", + "citations": [{"start": 0, "end": 8, "text": "Checking"}] + })) + ); + + assert_eq!(annotated.id.as_deref(), Some("resp-v2-123")); + assert_eq!(annotated.model.as_deref(), Some("cohere.command-a-03-2025")); + assert_eq!(annotated.finish_reason, Some(FinishReason::ToolUse)); + + let Some(MessageContent::Parts(parts)) = &annotated.message else { + panic!("expected typed parts, got {:?}", annotated.message); + }; + assert_eq!(parts.len(), 2); + assert!( + matches!(&parts[0], ContentPart::ProviderNative { kind, .. } if kind == "THINKING"), + "THINKING content should be preserved as a provider-native part" + ); + assert!(matches!(&parts[1], ContentPart::Text { text, .. } if text == "Checking the weather.")); + + let tool_calls = annotated.tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].id, "call-v2-1"); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].arguments, json!({"city": "Paris"})); + + let usage = annotated.usage.as_ref().unwrap(); + assert_eq!(usage.total_tokens, Some(35)); + assert_eq!( + annotated.api_specific, + Some(ApiSpecificResponse::OCIGenAI { + api_format: Some("COHEREV2".into()), + model_version: Some("2.0".into()), + }) + ); +} + +#[test] +fn test_cohere_v2_text_only_flattens() { + let response = json!({ + "chatResponse": { + "apiFormat": "COHEREV2", + "id": "resp-v2-456", + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "Sunny and 72."}] + }, + "finishReason": "COMPLETE" + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&response).unwrap(); + + assert_eq!(annotated.id.as_deref(), Some("resp-v2-456")); + assert_eq!( + annotated.message, + Some(MessageContent::Text("Sunny and 72.".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); +} + +#[test] +fn test_unmodeled_response_fields_preserved_in_extra() { + // GENERIC: timeCreated and serviceTier are not normalized but must + // survive; envelope-level unknown fields likewise. + let generic = json!({ + "modelId": DEDICATED_ENDPOINT, + "modelVersion": "1.0", + "futureEnvelopeField": {"nested": true}, + "chatResponse": { + "apiFormat": "GENERIC", + "timeCreated": "2026-07-27T17:27:25.871Z", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "hi"}], + "refusal": null, + "reasoningContent": "chain of thought" + }, + "finishReason": "stop", + // Choice-level per the OCI ChatChoice schema. + "serviceTier": "default", + "groundingMetadata": {"sources": ["doc-1"]}, + "logprobs": {"tokenLogprobs": [-0.1]} + }], + "usage": {"totalTokens": 9} + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&generic).unwrap(); + + assert_eq!( + annotated.extra.get("timeCreated"), + Some(&json!("2026-07-27T17:27:25.871Z")) + ); + assert_eq!( + annotated.extra.get("futureEnvelopeField"), + Some(&json!({"nested": true})) + ); + // Choice- and message-level unmodeled fields are namespaced by origin. + assert_eq!( + annotated.extra.get("choice"), + Some(&json!({ + "serviceTier": "default", + "groundingMetadata": {"sources": ["doc-1"]}, + "logprobs": {"tokenLogprobs": [-0.1]} + })) + ); + assert_eq!( + annotated.extra.get("message"), + Some(&json!({"refusal": null, "reasoningContent": "chain of thought"})) + ); + // Modeled fields stay normalized-only. + for modeled in [ + "apiFormat", + "choices", + "usage", + "chatResponse", + "modelId", + "modelVersion", + ] { + assert!( + !annotated.extra.contains_key(modeled), + "{modeled} should not be duplicated into extra" + ); + } + + // COHERE: chatHistory is not normalized and must survive. + let cohere = json!({ + "modelId": "cohere.command-r-08-2024", + "chatResponse": { + "apiFormat": "COHERE", + "text": "hi", + "chatHistory": [{"role": "USER", "message": "hello"}], + "finishReason": "COMPLETE" + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&cohere).unwrap(); + assert_eq!( + annotated.extra.get("chatHistory"), + Some(&json!([{"role": "USER", "message": "hello"}])) + ); +} + +#[test] +fn test_finish_reason_mapping() { + for (raw, expected) in [ + ("stop", FinishReason::Complete), + ("length", FinishReason::Length), + ("tool_calls", FinishReason::ToolUse), + ("content_filter", FinishReason::ContentFilter), + ("COMPLETE", FinishReason::Complete), + ("MAX_TOKENS", FinishReason::Length), + // Live Gemini-on-OCI responses use the lowercase spelling. + ("max_tokens", FinishReason::Length), + // COHEREV2 reasons per the OCI CohereChatResponseV2 schema. + ("TOOL_CALL", FinishReason::ToolUse), + ("STOP_SEQUENCE", FinishReason::Complete), + ("weird", FinishReason::Unknown("weird".into())), + ] { + let response = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{"message": {"role": "ASSISTANT", "content": []}, "finishReason": raw}] + } + }); + let annotated = OCIGenAIChatCodec.decode_response(&response).unwrap(); + assert_eq!(annotated.finish_reason, Some(expected), "for {raw}"); + } +} + +// =================================================================== +// Surface detection tests +// =================================================================== + +#[test] +fn test_detect_request_envelope_and_api_format() { + assert_eq!( + detect_request_surface(&generic_chat_details()), + Some(ProviderSurface::OCIGenAI) + ); + assert_eq!( + detect_request_surface(&cohere_chat_details()), + Some(ProviderSurface::OCIGenAI) + ); + // A bare chatRequest carries the apiFormat discriminator. + assert_eq!( + detect_request_surface(&generic_chat_details()["chatRequest"]), + Some(ProviderSurface::OCIGenAI) + ); +} + +#[test] +fn test_detect_request_hint_resolves_bare_chat_request() { + let bare = json!({"chatRequest": {"messages": []}}); + assert_eq!(detect_request_surface(&bare), None); + assert_eq!( + detect_request_surface_with_hint(&bare, Some("oci")), + Some(ProviderSurface::OCIGenAI) + ); + assert_eq!( + detect_request_surface_with_hint(&bare, Some("oci.genai")), + Some(ProviderSurface::OCIGenAI) + ); + assert_eq!(detect_request_surface_with_hint(&bare, Some("other")), None); +} + +#[test] +fn test_detect_request_does_not_shadow_other_surfaces() { + assert_eq!( + detect_request_surface(&json!({"messages": []})), + Some(ProviderSurface::OpenAIChat) + ); + assert_eq!( + detect_request_surface(&json!({"system": "x", "messages": []})), + Some(ProviderSurface::AnthropicMessages) + ); + assert_eq!( + detect_request_surface(&json!({"input": []})), + Some(ProviderSurface::OpenAIResponses) + ); +} + +#[test] +fn test_detect_response_chat_result() { + assert_eq!( + detect_response_surface(&generic_chat_result()), + Some(ProviderSurface::OCIGenAI) + ); + assert_eq!( + detect_response_surface(&cohere_chat_result()), + Some(ProviderSurface::OCIGenAI) + ); + // A bare COHERE chat response has no `choices`, so it stays unambiguous. + assert_eq!( + detect_response_surface(&cohere_chat_result()["chatResponse"]), + Some(ProviderSurface::OCIGenAI) + ); +} + +#[test] +fn test_detect_response_bare_generic_is_ambiguous_with_openai_chat() { + // A bare GENERIC chat response carries both `apiFormat` and `choices`; + // strict response detection refuses ambiguous shapes. + assert_eq!( + detect_response_surface(&generic_chat_result()["chatResponse"]), + None + ); +} + +#[test] +fn test_detect_response_does_not_shadow_other_surfaces() { + assert_eq!( + detect_response_surface(&json!({"choices": []})), + Some(ProviderSurface::OpenAIChat) + ); + assert_eq!( + detect_response_surface(&json!({"type": "message", "content": []})), + Some(ProviderSurface::AnthropicMessages) + ); + assert_eq!( + detect_response_surface(&json!({"output": []})), + Some(ProviderSurface::OpenAIResponses) + ); +} + +// =================================================================== +// Streaming codec tests +// =================================================================== + +#[test] +fn oci_streaming_codec_assembles_generic_text_response() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({ + "modelId": DEDICATED_ENDPOINT, + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "Hello, "}]} + }] + } + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [{"type": "TEXT", "text": "world."}]} + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": []}, + "finishReason": "stop", + "usage": {"promptTokens": 12, "completionTokens": 3, "totalTokens": 15} + })) + .unwrap(); + + let assembled = finalizer(); + // Wire-compatible with a ChatResult — feed it back through the decoder. + let annotated = OCIGenAIChatCodec.decode_response(&assembled).unwrap(); + assert_eq!(annotated.model.as_deref(), Some(DEDICATED_ENDPOINT)); + assert_eq!( + annotated.message, + Some(MessageContent::Text("Hello, world.".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); + let usage = annotated.usage.as_ref().unwrap(); + assert_eq!(usage.prompt_tokens, Some(12)); + assert_eq!(usage.completion_tokens, Some(3)); + assert_eq!(usage.total_tokens, Some(15)); +} + +#[test] +fn oci_streaming_codec_accumulates_generic_tool_call_arguments() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({ + "apiFormat": "GENERIC", + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [], + "toolCalls": [{"id": "call-1", "type": "FUNCTION", "name": "get_weather", "arguments": "{\"city\":"}] + } + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [], "toolCalls": [{"arguments": " \"NYC\"}"}]}, + "finishReason": "tool_calls" + })) + .unwrap(); + + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + assert_eq!(annotated.finish_reason, Some(FinishReason::ToolUse)); + let tool_calls = annotated.tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls[0].id, "call-1"); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].arguments, json!({"city": "NYC"})); +} + +#[test] +fn oci_streaming_codec_assembles_cohere_text_response() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({"apiFormat": "COHERE", "text": "Sunny"})).unwrap(); + collector(json!({"apiFormat": "COHERE", "text": " and 72."})).unwrap(); + collector(json!({ + "apiFormat": "COHERE", + "text": "", + "finishReason": "COMPLETE", + "usage": {"promptTokens": 8, "completionTokens": 4, "totalTokens": 12} + })) + .unwrap(); + + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + assert_eq!( + annotated.message, + Some(MessageContent::Text("Sunny and 72.".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); + assert_eq!(annotated.usage.as_ref().unwrap().total_tokens, Some(12)); + assert_eq!( + annotated.api_specific, + Some(ApiSpecificResponse::OCIGenAI { + api_format: Some("COHERE".into()), + model_version: None, + }) + ); +} + +// =================================================================== +// Encode edge cases surfaced in review +// =================================================================== + +#[test] +fn test_cohere_encode_rejects_multimodal_content() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + let last = annotated.messages.len() - 1; + annotated.messages[last] = Message::User { + content: MessageContent::Parts(vec![ContentPart::Text { + text: "described image".into(), + extra: Default::default(), + }]), + name: None, + }; + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!(matches!(err, FlowError::InvalidArgument(_)), "{err:?}"); +} + +#[test] +fn test_cohere_encode_requires_trailing_user_message() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.messages.push(Message::Assistant { + content: Some(MessageContent::Text("appended".into())), + tool_calls: None, + name: None, + }); + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!(matches!(err, FlowError::InvalidArgument(_)), "{err:?}"); +} + +#[test] +fn test_cohere_encode_rejects_normalized_tool_message() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + let last = annotated.messages.len() - 1; + annotated.messages.insert( + last, + Message::Tool { + content: MessageContent::Text("72F".into()), + tool_call_id: "call-1".into(), + }, + ); + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!(matches!(err, FlowError::InvalidArgument(_)), "{err:?}"); +} + +#[test] +fn test_tool_call_only_assistant_reencodes_without_null_content_or_empty_id() { + use super::super::request::{FunctionCall, ToolCall}; + + let codec = OCIGenAIChatCodec; + let mut payload = generic_chat_details(); + payload["chatRequest"]["messages"] + .as_array_mut() + .unwrap() + .push(json!({ + "role": "ASSISTANT", + "content": [], + "toolCalls": [{"type": "FUNCTION", "name": "get_weather", "arguments": "{\"city\": \"NYC\"}"}] + })); + let original = make_request(payload); + let mut annotated = codec.decode(&original).unwrap(); + + let last = annotated.messages.len() - 1; + let Message::Assistant { tool_calls, .. } = &annotated.messages[last] else { + panic!("expected an assistant message"); + }; + assert_eq!(tool_calls.as_ref().unwrap()[0].id, ""); + annotated.messages[last] = Message::Assistant { + content: None, + tool_calls: Some(vec![ToolCall { + id: String::new(), + call_type: "function".into(), + function: FunctionCall { + name: "get_weather".into(), + arguments: "{\"city\": \"[REDACTED]\"}".into(), + }, + }]), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let message = &encoded.content["chatRequest"]["messages"][2]; + assert_eq!(message["content"], json!([])); + let tool_call = &message["toolCalls"][0]; + assert!( + tool_call.get("id").is_none(), + "empty tool-call id must be omitted, got {tool_call}" + ); + assert_eq!(tool_call["arguments"], json!("{\"city\": \"[REDACTED]\"}")); +} + +#[test] +fn test_api_format_edit_is_rejected() { + // Merge-not-replace encoding patches the raw payload in place, so a format + // switch could never remove the previous format's modeled fields; the + // api_format annotation is therefore read-only. + let codec = OCIGenAIChatCodec; + for (payload, new_format) in [ + (generic_chat_details(), "COHERE"), + (generic_chat_details(), "COHEREV2"), + (cohere_chat_details(), "GENERIC"), + ] { + let original = make_request(payload); + let mut annotated = codec.decode(&original).unwrap(); + let Some(ApiSpecificRequest::OCIGenAI { api_format, .. }) = &mut annotated.api_specific + else { + panic!("expected an OCI api_specific annotation"); + }; + *api_format = Some(new_format.into()); + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!( + matches!(&err, FlowError::InvalidArgument(message) + if message.contains("api_format cannot be edited")), + "{new_format}: {err:?}" + ); + } +} + +#[test] +fn oci_streaming_codec_tracks_parallel_tool_calls_by_id() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + // Two parallel calls whose fragments arrive in separate events, each at + // event-local position 0. + collector(json!({ + "apiFormat": "GENERIC", + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [], + "toolCalls": [{"id": "call-a", "type": "FUNCTION", "name": "get_weather", "arguments": "{\"city\":"}] + } + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [], "toolCalls": [{"id": "call-b", "type": "FUNCTION", "name": "get_weather", "arguments": "{\"city\":"}]} + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [], "toolCalls": [{"id": "call-a", "arguments": " \"Paris\"}"}]} + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [], "toolCalls": [{"id": "call-b", "arguments": " \"Rome\"}"}]}, + "finishReason": "tool_calls" + })) + .unwrap(); + + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + let tool_calls = annotated.tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls.len(), 2); + assert_eq!(tool_calls[0].id, "call-a"); + assert_eq!(tool_calls[0].arguments, json!({"city": "Paris"})); + assert_eq!(tool_calls[1].id, "call-b"); + assert_eq!(tool_calls[1].arguments, json!({"city": "Rome"})); +} + +#[test] +fn oci_streaming_codec_tool_call_only_stream_decodes_without_message() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({ + "apiFormat": "GENERIC", + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [], + "toolCalls": [{"id": "call-1", "type": "FUNCTION", "name": "get_weather", "arguments": "{}"}] + }, + "finishReason": "tool_calls" + })) + .unwrap(); + + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + assert_eq!(annotated.message, None); + assert_eq!(annotated.tool_calls.as_ref().unwrap().len(), 1); +} + +#[test] +fn test_cohere_edit_removes_stale_preamble_override() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + // Drop the leading system message (the decoded preambleOverride). + assert!(matches!( + annotated.messages.first(), + Some(Message::System { .. }) + )); + annotated.messages.remove(0); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = encoded.content["chatRequest"].as_object().unwrap(); + assert!( + !chat_request.contains_key("preambleOverride"), + "a removed system message must not leave the original preamble on the wire" + ); +} + +fn cohere_v2_chat_details() -> Json { + json!({ + "compartmentId": "ocid1.compartment.oc1..example", + "servingMode": {"servingType": "ON_DEMAND", "modelId": "cohere.command-a-03-2025"}, + "chatRequest": { + "apiFormat": "COHEREV2", + "messages": [ + {"role": "USER", "content": [{"type": "TEXT", "text": "What is the weather?"}]} + ], + "maxTokens": 100, + "stopSequences": ["END"], + "citationOptions": {"mode": "OFF"} + } + }) +} + +#[test] +fn test_cohere_v2_request_decode_and_identity() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_v2_chat_details()); + let annotated = codec.decode(&original).unwrap(); + + assert_eq!(annotated.messages.len(), 1); + assert!(matches!(&annotated.messages[0], Message::User { .. })); + let params = annotated.params.as_ref().unwrap(); + assert_eq!(params.stop, Some(vec!["END".to_string()])); + assert_eq!(params.max_tokens, Some(100)); + // V2-only request fields ride along in extra rather than being dropped. + assert_eq!(annotated.extra["citationOptions"], json!({"mode": "OFF"})); + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded.content, cohere_v2_chat_details()); +} + +#[test] +fn test_cohere_v2_request_edit_patches_stop_sequences() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_v2_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.params.as_mut().unwrap().stop = Some(vec!["HALT".to_string()]); + annotated.messages[0] = Message::User { + content: MessageContent::Text("What is the weather in [REDACTED]?".into()), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = &encoded.content["chatRequest"]; + assert_eq!(chat_request["stopSequences"], json!(["HALT"])); + assert!( + chat_request.get("stop").is_none(), + "COHEREV2 edits must patch stopSequences, not the GENERIC stop key" + ); + assert_eq!( + chat_request["messages"][0]["content"], + json!([{"type": "TEXT", "text": "What is the weather in [REDACTED]?"}]) + ); + assert_eq!(chat_request["citationOptions"], json!({"mode": "OFF"})); +} + +#[test] +fn oci_streaming_codec_does_not_double_cohere_text_on_full_terminal_event() { + // Live-captured shape: the service's terminal COHERE event repeats the + // complete response text alongside chatHistory and finishReason. + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + for fragment in ["Rome", " is", " the", " capital", " of", " Italy", "."] { + collector(json!({"apiFormat": "COHERE", "text": fragment})).unwrap(); + } + collector(json!({ + "apiFormat": "COHERE", + "text": "Rome is the capital of Italy.", + "chatHistory": [ + {"role": "USER", "message": "What is the capital of Italy?"}, + {"role": "CHATBOT", "message": "Rome is the capital of Italy."} + ], + "finishReason": "COMPLETE" + })) + .unwrap(); + + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + assert_eq!( + annotated.message, + Some(MessageContent::Text("Rome is the capital of Italy.".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); +} + +#[test] +fn oci_streaming_codec_finalizes_cohere_v2_root_message() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({ + "apiFormat": "COHEREV2", + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "Checking "}]} + })) + .unwrap(); + collector(json!({ + "message": {"content": [{"type": "TEXT", "text": "the weather."}]} + })) + .unwrap(); + collector(json!({ + "message": { + "content": [], + "toolCalls": [{ + "id": "call-1", + "type": "FUNCTION", + "function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"} + }] + }, + "finishReason": "TOOL_CALL", + "usage": {"promptTokens": 20, "completionTokens": 9, "totalTokens": 29} + })) + .unwrap(); + + let assembled = finalizer(); + let chat_response = &assembled["chatResponse"]; + assert!( + chat_response.get("message").is_some() && chat_response.get("choices").is_none(), + "COHEREV2 streams must finalize to the root-message shape, got {chat_response}" + ); + + let annotated = OCIGenAIChatCodec.decode_response(&assembled).unwrap(); + assert_eq!( + annotated.message, + Some(MessageContent::Text("Checking the weather.".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::ToolUse)); + let tool_calls = annotated.tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls[0].id, "call-1"); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].arguments, json!({"city": "Paris"})); + assert_eq!(annotated.usage.as_ref().unwrap().total_tokens, Some(29)); + assert_eq!( + annotated.api_specific, + Some(ApiSpecificResponse::OCIGenAI { + api_format: Some("COHEREV2".into()), + model_version: None, + }) + ); +} + +#[test] +fn oci_streaming_codec_preserves_non_text_cohere_v2_parts_in_order() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({ + "apiFormat": "COHEREV2", + "message": {"role": "ASSISTANT", "content": [ + {"type": "THINKING", "thinking": "internal reasoning"} + ]} + })) + .unwrap(); + collector(json!({ + "message": {"content": [{"type": "TEXT", "text": "The answer "}]} + })) + .unwrap(); + collector(json!({ + "message": {"content": [{"type": "TEXT", "text": "is 42."}]}, + "finishReason": "COMPLETE" + })) + .unwrap(); + + let assembled = finalizer(); + // The finalized wire shape keeps the typed parts in arrival order, with + // consecutive TEXT fragments merged into one part. + assert_eq!( + assembled["chatResponse"]["message"]["content"], + json!([ + {"type": "THINKING", "thinking": "internal reasoning"}, + {"type": "TEXT", "text": "The answer is 42."} + ]) + ); + + let annotated = OCIGenAIChatCodec.decode_response(&assembled).unwrap(); + let Some(MessageContent::Parts(parts)) = &annotated.message else { + panic!("expected typed parts, got {:?}", annotated.message); + }; + assert!(matches!( + &parts[0], + ContentPart::ProviderNative { kind, .. } if kind == "THINKING" + )); + assert!(matches!( + &parts[1], + ContentPart::Text { text, .. } if text == "The answer is 42." + )); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); +} + +#[test] +fn oci_streaming_codec_preserves_non_text_generic_parts_in_order() { + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({ + "apiFormat": "GENERIC", + "index": 0, + "message": {"role": "ASSISTANT", "content": [ + {"type": "THINKING", "thinking": "internal reasoning"} + ]} + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [{"type": "TEXT", "text": "The answer "}]} + })) + .unwrap(); + collector(json!({ + "index": 0, + "message": {"content": [{"type": "TEXT", "text": "is 42."}]}, + "finishReason": "stop" + })) + .unwrap(); + + let assembled = finalizer(); + assert_eq!( + assembled["chatResponse"]["choices"][0]["message"]["content"], + json!([ + {"type": "THINKING", "thinking": "internal reasoning"}, + {"type": "TEXT", "text": "The answer is 42."} + ]) + ); + + let annotated = OCIGenAIChatCodec.decode_response(&assembled).unwrap(); + let Some(MessageContent::Parts(parts)) = &annotated.message else { + panic!("expected typed parts, got {:?}", annotated.message); + }; + assert!(matches!( + &parts[0], + ContentPart::ProviderNative { kind, .. } if kind == "THINKING" + )); + assert!(matches!( + &parts[1], + ContentPart::Text { text, .. } if text == "The answer is 42." + )); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); +} + +#[test] +fn test_non_oci_api_specific_is_rejected() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.api_specific = Some(ApiSpecificRequest::Custom { + api_name: "not-oci".into(), + data: Json::Null, + }); + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!( + matches!(&err, FlowError::InvalidArgument(message) + if message.contains("does not match OCI GenAI")), + "{err:?}" + ); +} + +#[test] +fn test_envelope_edit_on_bare_chat_request_is_rejected() { + let codec = OCIGenAIChatCodec; + let bare = generic_chat_details().get("chatRequest").cloned().unwrap(); + let original = make_request(bare); + let mut annotated = codec.decode(&original).unwrap(); + + let Some(ApiSpecificRequest::OCIGenAI { compartment_id, .. }) = &mut annotated.api_specific + else { + panic!("expected an OCI api_specific annotation"); + }; + *compartment_id = Some("ocid1.compartment.oc1..edited".into()); + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!( + matches!(&err, FlowError::InvalidArgument(message) + if message.contains("require a ChatDetails envelope")), + "{err:?}" + ); +} + +#[test] +fn test_cohere_v2_assistant_edit_reencodes_nested_tool_calls() { + use super::super::request::{FunctionCall, ToolCall}; + + let codec = OCIGenAIChatCodec; + let mut payload = cohere_v2_chat_details(); + payload["chatRequest"]["messages"] + .as_array_mut() + .unwrap() + .push(json!({ + "role": "ASSISTANT", + "content": [], + "toolCalls": [{ + "id": "call-1", + "type": "FUNCTION", + "function": {"name": "get_weather", "arguments": "{\"city\": \"NYC\"}"} + }] + })); + let original = make_request(payload); + let mut annotated = codec.decode(&original).unwrap(); + + let last = annotated.messages.len() - 1; + annotated.messages[last] = Message::Assistant { + content: None, + tool_calls: Some(vec![ToolCall { + id: "call-1".into(), + call_type: "function".into(), + function: FunctionCall { + name: "get_weather".into(), + arguments: "{\"city\": \"[REDACTED]\"}".into(), + }, + }]), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let tool_call = &encoded.content["chatRequest"]["messages"][1]["toolCalls"][0]; + assert_eq!( + tool_call["function"], + json!({"name": "get_weather", "arguments": "{\"city\": \"[REDACTED]\"}"}), + "COHEREV2 edits must re-encode nested-function tool calls, got {tool_call}" + ); + assert!( + tool_call.get("name").is_none() && tool_call.get("arguments").is_none(), + "flat GENERIC keys must not appear on a COHEREV2 tool call, got {tool_call}" + ); +} + +#[test] +fn test_non_string_text_part_survives_as_provider_native() { + let codec = OCIGenAIChatCodec; + let mut payload = generic_chat_details(); + payload["chatRequest"]["messages"][1]["content"] = json!([ + {"type": "TEXT", "text": {"unexpected": "object"}} + ]); + let original = make_request(payload.clone()); + let annotated = codec.decode(&original).unwrap(); + + let Some(Message::User { + content: MessageContent::Parts(parts), + .. + }) = annotated.messages.get(1) + else { + panic!("expected a user message with typed parts"); + }; + assert!( + matches!(&parts[0], ContentPart::ProviderNative { .. }), + "non-string TEXT value must survive as provider-native, got {:?}", + parts[0] + ); + + // Identity: the raw value re-encodes untouched. + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded.content, payload); +} + +// =================================================================== +// Coverage: edit paths and error paths not exercised elsewhere +// =================================================================== + +#[test] +fn test_tools_and_tool_choice_edits_reencode() { + use super::super::request::{FunctionDefinition, ProviderNativeComponent, ToolDefinition}; + + let codec = OCIGenAIChatCodec; + let mut payload = generic_chat_details(); + payload["chatRequest"]["tools"] = json!([ + {"type": "FUNCTION", "name": "old_tool", "parameters": {"type": "object"}} + ]); + payload["chatRequest"]["toolChoice"] = json!({"type": "AUTO"}); + let original = make_request(payload); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.tools = Some(vec![ToolDefinition::Function { + function: FunctionDefinition { + name: "get_weather".into(), + description: Some("Get weather".into()), + parameters: Some(json!({"type": "object", "properties": {}})), + strict: None, + extra: Default::default(), + }, + extra: Default::default(), + }]); + annotated.tool_choice = Some(ToolChoice::ProviderNative(ProviderNativeComponent { + provider: "oci_genai".into(), + kind: "tool_choice".into(), + value: json!({"type": "REQUIRED"}), + })); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = &encoded.content["chatRequest"]; + assert_eq!( + chat_request["tools"], + json!([{ + "type": "FUNCTION", + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {}} + }]) + ); + assert_eq!(chat_request["toolChoice"], json!({"type": "REQUIRED"})); +} + +#[test] +fn test_dropping_tools_removes_the_wire_key() { + let codec = OCIGenAIChatCodec; + let mut payload = generic_chat_details(); + payload["chatRequest"]["tools"] = json!([ + {"type": "FUNCTION", "name": "old_tool", "parameters": {"type": "object"}} + ]); + let original = make_request(payload); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.tools = None; + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert!( + encoded.content["chatRequest"].get("tools").is_none(), + "dropping the tools annotation must remove the wire key" + ); +} + +#[test] +fn test_non_native_tool_choice_edit_is_rejected() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.tool_choice = Some(ToolChoice::Auto); + + let err = codec.encode(&annotated, &original).unwrap_err(); + assert!(matches!(err, FlowError::InvalidArgument(_)), "{err:?}"); +} + +#[test] +fn test_parts_content_edit_reencodes_typed_parts() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.messages[1] = Message::User { + content: MessageContent::Parts(vec![ + ContentPart::Text { + text: "look at this".into(), + extra: Default::default(), + }, + ContentPart::ProviderNative { + provider: "oci_genai".into(), + kind: "IMAGE".into(), + value: json!({"type": "IMAGE", "imageUrl": {"url": "data:image/png;base64,AA"}}), + }, + ]), + name: None, + }; + + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!( + encoded.content["chatRequest"]["messages"][1]["content"], + json!([ + {"type": "TEXT", "text": "look at this"}, + {"type": "IMAGE", "imageUrl": {"url": "data:image/png;base64,AA"}} + ]) + ); +} + +#[test] +fn test_top_p_edit_patches_only_top_p() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + let params = annotated.params.as_mut().unwrap(); + params.top_p = Some(0.5); + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = &encoded.content["chatRequest"]; + assert_eq!(chat_request["topP"], json!(0.5)); + assert_eq!(chat_request["temperature"], json!(0.0)); + assert_eq!(chat_request["maxTokens"], json!(600)); +} + +#[test] +fn test_clearing_params_removes_provider_fields() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + let params = annotated.params.as_mut().unwrap(); + params.temperature = None; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = &encoded.content["chatRequest"]; + assert!( + chat_request.get("temperature").is_none(), + "clearing temperature must remove the provider field" + ); + // Untouched params keep their raw values. + assert_eq!(chat_request["maxTokens"], json!(600)); +} + +#[test] +fn test_clearing_all_params_removes_all_provider_fields() { + let codec = OCIGenAIChatCodec; + let original = make_request(generic_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.params = None; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = &encoded.content["chatRequest"]; + assert!(chat_request.get("temperature").is_none()); + assert!(chat_request.get("maxTokens").is_none()); + // Non-param request fields survive a full params clear. + assert_eq!(chat_request["apiFormat"], json!("GENERIC")); + assert!(chat_request.get("messages").is_some()); +} + +#[test] +fn test_cohere_v2_clearing_stop_removes_stop_sequences() { + let codec = OCIGenAIChatCodec; + let original = make_request(cohere_v2_chat_details()); + let mut annotated = codec.decode(&original).unwrap(); + + annotated.params.as_mut().unwrap().stop = None; + + let encoded = codec.encode(&annotated, &original).unwrap(); + let chat_request = &encoded.content["chatRequest"]; + assert!( + chat_request.get("stopSequences").is_none(), + "clearing stop must remove the COHERE stopSequences field" + ); + assert_eq!(chat_request["maxTokens"], json!(100)); + assert_eq!(chat_request["citationOptions"], json!({"mode": "OFF"})); +} + +#[test] +fn test_decode_error_paths() { + let codec = OCIGenAIChatCodec; + + // Request content that is not an object. + assert!(codec.decode(&make_request(json!("nope"))).is_err()); + + // A stop list that is not a string array. + let mut payload = generic_chat_details(); + payload["chatRequest"]["stop"] = json!("HALT"); + assert!(codec.decode(&make_request(payload)).is_err()); + + // A toolCalls entry that is not an object. + let mut payload = generic_chat_details(); + payload["chatRequest"]["messages"] = json!([ + {"role": "ASSISTANT", "content": [], "toolCalls": ["nope"]} + ]); + assert!(codec.decode(&make_request(payload)).is_err()); + + // A GENERIC message that is not an object. + let mut payload = generic_chat_details(); + payload["chatRequest"]["messages"] = json!(["nope"]); + assert!(codec.decode(&make_request(payload)).is_err()); + + // A COHERE chatHistory turn that is not an object. + let mut payload = cohere_chat_details(); + payload["chatRequest"]["chatHistory"] = json!(["nope"]); + assert!(codec.decode(&make_request(payload)).is_err()); +} + +#[test] +fn test_unknown_generic_role_survives_as_provider_native() { + let codec = OCIGenAIChatCodec; + let mut payload = generic_chat_details(); + payload["chatRequest"]["messages"] = json!([ + {"role": "MODERATOR", "content": [{"type": "TEXT", "text": "hi"}]} + ]); + let original = make_request(payload.clone()); + let annotated = codec.decode(&original).unwrap(); + + assert!(matches!( + &annotated.messages[0], + Message::ProviderNative { provider, .. } if provider == "oci_genai" + )); + let encoded = codec.encode(&annotated, &original).unwrap(); + assert_eq!(encoded.content, payload); +} + +#[test] +fn oci_streaming_codec_infers_api_format_from_event_shape() { + // GENERIC inferred from a bare choice delta with no apiFormat anywhere. + let codec = OCIGenAIStreamingCodec::default(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + collector(json!({ + "index": 0, + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "hi"}]}, + "finishReason": "stop" + })) + .unwrap(); + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + assert_eq!(annotated.message, Some(MessageContent::Text("hi".into()))); + + // COHERE inferred from a bare text fragment with no apiFormat anywhere. + let codec = OCIGenAIStreamingCodec::new(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + collector(json!({"text": "hello"})).unwrap(); + collector(json!({"text": "!"})).unwrap(); + // The live terminal event repeats the complete text. + collector(json!({"text": "hello!", "finishReason": "COMPLETE"})).unwrap(); + let annotated = OCIGenAIChatCodec.decode_response(&finalizer()).unwrap(); + assert_eq!( + annotated.message, + Some(MessageContent::Text("hello!".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); +} diff --git a/crates/core/tests/unit/codec/parity_tests.rs b/crates/core/tests/unit/codec/parity_tests.rs index 64ab62077..03ff25735 100644 --- a/crates/core/tests/unit/codec/parity_tests.rs +++ b/crates/core/tests/unit/codec/parity_tests.rs @@ -836,3 +836,234 @@ fn baseline_patching_rejects_multiple_reordered_and_edited_items_without_provena .contains("multiple edited array items without stable identities") ); } + +// =================================================================== +// OCI GenAI parity: the fourth built-in surface agrees with the others +// =================================================================== + +/// The same logical response (one assistant text message) in the OCI GenAI +/// GENERIC `ChatResult` schema. +fn oci_text_response(model: &str) -> Json { + json!({ + "modelId": model, + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "hello"}] + }, + "finishReason": "stop" + }] + } + }) +} + +fn oci_response_with_usage(model: &str, extra_usage: Json) -> Json { + let mut raw = oci_text_response(model); + // OCI reports no cache token counters; only the three basic counters. + let mut usage = json!({ + "promptTokens": 1000, + "completionTokens": 500, + "totalTokens": 1500 + }); + merge_object(&mut usage, extra_usage); + raw["chatResponse"] + .as_object_mut() + .unwrap() + .insert("usage".into(), usage); + raw +} + +#[test] +fn test_oci_response_model_name_and_text_parity() { + let chat = decode(&chat_text_response("parity-shared-model")); + let oci = decode(&oci_text_response("parity-shared-model")); + + assert_eq!(oci.model, chat.model); + assert_eq!(oci.response_text(), chat.response_text()); + assert_eq!(oci.finish_reason, chat.finish_reason); + // OCI ChatResult payloads carry no response id; the other schemas do. + assert_eq!(oci.id, None); +} + +#[test] +fn test_oci_finish_reason_parity() { + // Complete: GENERIC "stop" and COHERE "COMPLETE" agree with the others. + let generic_stop = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{"message": {"role": "ASSISTANT", "content": []}, "finishReason": "stop"}] + } + }); + let cohere_complete = json!({ + "chatResponse": {"apiFormat": "COHERE", "text": "x", "finishReason": "COMPLETE"} + }); + for raw in [&generic_stop, &cohere_complete] { + assert_eq!( + decode(raw).finish_reason, + Some(FinishReason::Complete), + "expected Complete for {raw}", + ); + } + + // Length: GENERIC "length" and COHERE "MAX_TOKENS" agree with the others. + let generic_length = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{"message": {"role": "ASSISTANT", "content": []}, "finishReason": "length"}] + } + }); + let cohere_max_tokens = json!({ + "chatResponse": {"apiFormat": "COHERE", "text": "x", "finishReason": "MAX_TOKENS"} + }); + for raw in [&generic_length, &cohere_max_tokens] { + assert_eq!( + decode(raw).finish_reason, + Some(FinishReason::Length), + "expected Length for {raw}", + ); + } +} + +#[test] +fn test_oci_response_tool_call_parity() { + // The same logical tool invocation as test_response_tool_call_parity, in + // the flat OCI GENERIC shape (string-encoded arguments like OpenAI Chat). + let chat = decode(&json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_parity_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"NYC\",\"units\":\"c\"}" + } + }] + }, + "finish_reason": "tool_calls" + }] + })); + let oci = decode(&json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "message": { + "role": "ASSISTANT", + "content": [], + "toolCalls": [{ + "id": "call_parity_1", + "type": "FUNCTION", + "name": "get_weather", + "arguments": "{\"city\":\"NYC\",\"units\":\"c\"}" + }] + }, + "finishReason": "tool_calls" + }] + } + })); + + assert_eq!(oci.finish_reason, Some(FinishReason::ToolUse)); + assert_eq!(oci.tool_calls, chat.tool_calls); + let calls = oci.tool_calls.expect("oci tool calls"); + assert_eq!(calls[0].arguments, json!({"city": "NYC", "units": "c"})); + assert!(calls[0].arguments.is_object()); +} + +#[test] +fn test_oci_response_usage_parity() { + let chat = decode(&chat_response_with_usage("parity-usage-model", json!({}))); + let oci = decode(&oci_response_with_usage("parity-usage-model", json!({}))); + + let chat_usage = chat.usage.unwrap(); + let oci_usage = oci.usage.unwrap(); + assert_eq!(oci_usage.prompt_tokens, chat_usage.prompt_tokens); + assert_eq!(oci_usage.completion_tokens, chat_usage.completion_tokens); + assert_eq!(oci_usage.total_tokens, chat_usage.total_tokens); + // Divergence: OCI has no prompt-cache counters, so cache_read_tokens is + // None where the OpenAI Chat fixture reports 200 cached tokens. + assert_eq!(oci_usage.cache_read_tokens, None); + assert_eq!(chat_usage.cache_read_tokens, Some(200)); + assert!(matches!( + oci.api_specific, + Some(ApiSpecificResponse::OCIGenAI { .. }) + )); +} + +#[test] +fn test_oci_request_normalization_parity() { + let chat = normalize_request(&req(json!({ + "model": "parity-request-model", + "messages": [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Summarize the docs."} + ], + "temperature": 0.5, + "max_tokens": 256, + "stop": ["END"] + }))) + .expect("chat request decodes"); + + let oci = normalize_request(&req(json!({ + "compartmentId": "ocid1.compartment.oc1..parity", + "servingMode": {"servingType": "ON_DEMAND", "modelId": "parity-request-model"}, + "chatRequest": { + "apiFormat": "GENERIC", + "messages": [ + {"role": "SYSTEM", "content": [{"type": "TEXT", "text": "You are terse."}]}, + {"role": "USER", "content": [{"type": "TEXT", "text": "Summarize the docs."}]} + ], + "temperature": 0.5, + "maxTokens": 256, + "stop": ["END"] + } + }))) + .expect("oci request decodes"); + + // UPPERCASE roles and TEXT content-part lists normalize to the same + // messages OpenAI Chat produces; the model comes from servingMode. + assert_eq!(oci.messages, chat.messages); + assert_eq!(oci.params, chat.params); + assert_eq!(oci.model, chat.model); + assert_eq!(oci.system_prompt(), chat.system_prompt()); + assert_eq!(oci.last_user_message(), chat.last_user_message()); +} + +#[test] +fn test_oci_request_hint_never_overrides_strong_signals() { + // The "oci" hint must not reroute unambiguous non-OCI bodies. + let chat_request = req(json!({ + "model": "gpt-parity", + "messages": [{"role": "user", "content": "hi"}] + })); + let hinted = normalize_request_with_hint(&chat_request, Some("oci")) + .expect("chat request decodes despite oci hint"); + assert_eq!( + hinted, + normalize_request(&chat_request).expect("chat request decodes"), + ); + // And the reverse: a strong OCI envelope decodes as OCI even with a + // wrong hint. + let oci_request = req(json!({ + "compartmentId": "ocid1.compartment.oc1..parity", + "servingMode": {"servingType": "ON_DEMAND", "modelId": "m"}, + "chatRequest": { + "apiFormat": "GENERIC", + "messages": [{"role": "USER", "content": [{"type": "TEXT", "text": "hi"}]}] + } + })); + let hinted = normalize_request_with_hint(&oci_request, Some("anthropic")) + .expect("oci request decodes despite wrong hint"); + assert_eq!( + hinted, + normalize_request(&oci_request).expect("oci request decodes"), + ); + assert!(matches!( + hinted.api_specific, + Some(super::request::ApiSpecificRequest::OCIGenAI { .. }) + )); +} diff --git a/crates/core/tests/unit/codec/resolve_tests.rs b/crates/core/tests/unit/codec/resolve_tests.rs index 58fce0f97..d8a60a3f3 100644 --- a/crates/core/tests/unit/codec/resolve_tests.rs +++ b/crates/core/tests/unit/codec/resolve_tests.rs @@ -25,6 +25,7 @@ fn builtin_provider_surface_registry_keeps_request_priority() { vec![ ProviderSurface::OpenAIResponses, ProviderSurface::AnthropicMessages, + ProviderSurface::OCIGenAI, ProviderSurface::OpenAIChat, ProviderSurface::GeminiGenerateContent, ] @@ -447,10 +448,11 @@ fn hint_does_not_classify_non_object_or_keyless() { // Provider-codec factory (name<->surface mapping + codec construction) // --------------------------------------------------------------------------- -const ALL_SURFACES: [ProviderSurface; 4] = [ +const ALL_SURFACES: [ProviderSurface; 5] = [ ProviderSurface::OpenAIChat, ProviderSurface::OpenAIResponses, ProviderSurface::AnthropicMessages, + ProviderSurface::OCIGenAI, ProviderSurface::GeminiGenerateContent, ]; @@ -476,6 +478,7 @@ fn codec_name_uses_canonical_spellings() { ProviderSurface::AnthropicMessages.codec_name(), "anthropic_messages" ); + assert_eq!(ProviderSurface::OCIGenAI.codec_name(), "oci_genai"); assert_eq!( ProviderSurface::GeminiGenerateContent.codec_name(), "gemini_generate_content" @@ -510,6 +513,7 @@ fn supported_codec_names_track_the_builtin_registry() { vec![ "openai_responses", "anthropic_messages", + "oci_genai", "openai_chat", "gemini_generate_content" ] diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index 6d8ff05ad..ca327909e 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -32,6 +32,7 @@ use crate::api::scope::{COMPACTION_EVENT_NAME, EmitMarkEventParams, event}; use crate::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use crate::codec::anthropic::AnthropicMessagesCodec; +use crate::codec::oci_genai::OCIGenAIChatCodec; use crate::codec::openai_chat::OpenAIChatCodec; use crate::codec::openai_responses::OpenAIResponsesCodec; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; @@ -178,6 +179,10 @@ fn request_sanitizer_context_preserves_all_codec_identity_states() { sanitize_context_for_request_codec(Some(&AnthropicMessagesCodec)).codec(), &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) ); + assert_eq!( + sanitize_context_for_request_codec(Some(&OCIGenAIChatCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI) + ); assert_eq!( sanitize_context_for_request_codec(Some(&RuntimeIdentityCodec)).codec(), &LlmCodecIdentity::Runtime("com.example.chat.v1".into()) @@ -218,6 +223,11 @@ fn response_sanitizer_context_preserves_all_codec_identity_states() { .codec(), &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) ); + assert_eq!( + sanitize_context_for_response_codec(Some(&OCIGenAIChatCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI) + ); assert_eq!( sanitize_context_for_response_codec(Some(&RuntimeIdentityCodec)).codec(), &LlmCodecIdentity::Runtime("com.example.chat.v1".into()) 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 f838f1da2..142b4dee2 100644 --- a/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs +++ b/crates/core/tests/unit/plugins/nemo_guardrails/component_tests.rs @@ -403,6 +403,7 @@ fn schema_contains_every_supported_nemo_guardrails_option() { "openai_chat", "openai_responses", "anthropic_messages", + "oci_genai", "gemini_generate_content" ] )); @@ -650,6 +651,7 @@ fn assert_invalid_remote_identity_and_codec() { "openai_chat", "openai_responses", "anthropic_messages", + "oci_genai", "gemini_generate_content", ] .iter() @@ -689,6 +691,20 @@ fn assert_invalid_remote_identity_and_codec() { }) ); + let unsupported_remote_oci_codec = validate_plugin_config(&plugin_config(json!({ + "mode": "remote", + "codec": "oci_genai", + "remote": { + "endpoint": "http://localhost:8000", + "config_id": "default" + } + }))); + assert!(unsupported_remote_oci_codec.has_errors()); + assert!(unsupported_remote_oci_codec.diagnostics.iter().any(|diag| { + diag.message + .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", 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 ce173c97a..fc69225a8 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 @@ -1078,6 +1078,30 @@ fn stream_text_extraction_handles_supported_codecs() { ), Some("hello".to_string()) ); + // OCI GENERIC: bare choice delta with a top-level `message`. + assert_eq!( + extract_stream_text( + LocalGuardrailsCodec::OCIGenAI, + &json!({"index": 0, "message": {"content": [{"type": "TEXT", "text": "hello"}]}}) + ), + Some("hello".to_string()) + ); + // OCI GENERIC: `choices`-wrapped deltas, optionally inside `chatResponse`. + assert_eq!( + extract_stream_text( + LocalGuardrailsCodec::OCIGenAI, + &json!({"chatResponse": {"choices": [ + {"index": 0, "message": {"content": [{"type": "TEXT", "text": "hel"}]}}, + {"index": 1, "message": {"content": [{"type": "TEXT", "text": "lo"}]}} + ]}}) + ), + Some("hello".to_string()) + ); + // OCI COHERE: bare text fragment. + assert_eq!( + extract_stream_text(LocalGuardrailsCodec::OCIGenAI, &json!({"text": "hello"})), + Some("hello".to_string()) + ); // Gemini: visible text parts reach the guardrail worker. assert_eq!( extract_stream_text( @@ -1101,11 +1125,43 @@ fn stream_text_extraction_handles_supported_codecs() { json!({"type": "content_block_delta", "delta": {"type": "input_json_delta"}}), ), (LocalGuardrailsCodec::GeminiGenerateContent, Json::Null), + // A tool-call-only OCI delta carries no user-visible text and must not + // reach the guardrail worker. + ( + LocalGuardrailsCodec::OCIGenAI, + json!({"index": 0, "message": {"content": [], "toolCalls": [{"arguments": "{"}]}}), + ), + // The terminal COHERE event repeats the full response text alongside + // finishReason; forwarding it would double the rail input. + ( + LocalGuardrailsCodec::OCIGenAI, + json!({"apiFormat": "COHERE", "text": "hello!", "finishReason": "COMPLETE"}), + ), + ( + LocalGuardrailsCodec::OCIGenAI, + json!({"chatResponse": {"apiFormat": "COHERE", "text": "hello!", "finishReason": "COMPLETE"}}), + ), ] { assert_eq!(extract_stream_text(codec, &chunk), None); } } +#[test] +fn stream_text_extraction_oci_cohere_stream_is_not_doubled() { + // Live-shaped COHERE stream: incremental deltas, then a terminal event + // repeating the complete text. The rails must see the text exactly once. + let stream = [ + json!({"apiFormat": "COHERE", "text": "hello"}), + json!({"apiFormat": "COHERE", "text": "!"}), + json!({"apiFormat": "COHERE", "text": "hello!", "finishReason": "COMPLETE"}), + ]; + let forwarded: String = stream + .iter() + .filter_map(|chunk| extract_stream_text(LocalGuardrailsCodec::OCIGenAI, chunk)) + .collect(); + assert_eq!(forwarded, "hello!"); +} + #[test] fn stream_text_extraction_gemini_skips_thought_parts() { // A thought chunk (thought: true) must NOT reach the guardrail worker. diff --git a/crates/node/pii_redaction.d.ts b/crates/node/pii_redaction.d.ts index 969eae122..96065e6d5 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_generate_content' | string; + codec?: 'openai_chat' | 'openai_responses' | 'anthropic_messages' | 'oci_genai' | '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 1ee433eee..b0d7e42a9 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_generate_content' } + | { kind: 'builtin'; id: 'openai_chat' | 'openai_responses' | 'anthropic_messages' | 'oci_genai' | '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 4c8095991..3a3fd38bb 100644 --- a/crates/node/src/types/mod.rs +++ b/crates/node/src/types/mod.rs @@ -599,3 +599,62 @@ impl AnthropicMessagesCodec { serde_json::to_value(&annotated).map_err(|e| napi::Error::from_reason(e.to_string())) } } + +/// Built-in codec for the OCI Generative AI chat API. +/// +/// Implements both request codec (decode/encode) and response codec +/// (decodeResponse). Construct with `new OCIGenAIChatCodec()`. +#[napi(js_name = "OCIGenAIChatCodec")] +pub struct OCIGenAIChatCodec { + pub(crate) inner_codec: std::sync::Arc, + pub(crate) inner_response_codec: std::sync::Arc, +} + +#[napi] +impl OCIGenAIChatCodec { + #[napi(constructor)] + pub fn new() -> Self { + Self { + inner_codec: std::sync::Arc::new(nemo_relay::codec::oci_genai::OCIGenAIChatCodec), + inner_response_codec: std::sync::Arc::new( + nemo_relay::codec::oci_genai::OCIGenAIChatCodec, + ), + } + } + + /// 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())) + } +} diff --git a/crates/node/tests/types_tests.mjs b/crates/node/tests/types_tests.mjs index f4812c16d..4e60824c7 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.OCIGenAIChatCodec, 'function'); assert.equal(typeof lib.GeminiGenerateContentCodec, 'function'); }); @@ -81,6 +82,87 @@ describe('ScopeStack', () => { }); }); +// =========================================================================== +// OCIGenAIChatCodec +// =========================================================================== + +describe('OCIGenAIChatCodec', () => { + const { OCIGenAIChatCodec } = lib; + + const chatDetails = () => ({ + headers: {}, + content: { + compartmentId: 'ocid1.compartment.oc1..example', + servingMode: { servingType: 'ON_DEMAND', modelId: 'meta.llama-3.3-70b-instruct' }, + chatRequest: { + apiFormat: 'GENERIC', + messages: [ + { role: 'USER', content: [{ type: 'TEXT', text: 'My SSN is 111-22-3333.' }] }, + ], + maxTokens: 600, + seed: 7, + }, + }, + }); + + it('instantiates', () => { + const codec = new OCIGenAIChatCodec(); + assert.ok(codec instanceof OCIGenAIChatCodec); + }); + + it('decode returns an AnnotatedLLMRequest with model and params', () => { + const codec = new OCIGenAIChatCodec(); + const annotated = codec.decode(chatDetails()); + assert.equal(annotated.model, 'meta.llama-3.3-70b-instruct'); + assert.equal(annotated.messages.length, 1); + assert.equal(annotated.messages[0].role, 'user'); + // Rust serializes GenerationParams fields in snake_case (max_tokens, not maxTokens) + assert.equal(annotated.params.max_tokens, 600); + }); + + it('encode is an identity for an unedited annotation', () => { + const codec = new OCIGenAIChatCodec(); + const req = chatDetails(); + const annotated = codec.decode(req); + const reEncoded = codec.encode(annotated, req); + assert.deepEqual(reEncoded.content, req.content); + }); + + it('encode applies edited messages and keeps unmodeled fields', () => { + const codec = new OCIGenAIChatCodec(); + const req = chatDetails(); + const annotated = codec.decode(req); + annotated.messages = [{ role: 'user', content: 'My SSN is [REDACTED].' }]; + const reEncoded = codec.encode(annotated, req); + assert.deepEqual(reEncoded.content.chatRequest.messages[0].content, [ + { type: 'TEXT', text: 'My SSN is [REDACTED].' }, + ]); + assert.equal(reEncoded.content.chatRequest.seed, 7, 'unmodeled fields must survive edits'); + }); + + it('decodeResponse extracts text, finish reason, and usage', () => { + const codec = new OCIGenAIChatCodec(); + const raw = { + modelId: 'meta.llama-3.3-70b-instruct', + chatResponse: { + apiFormat: 'GENERIC', + choices: [{ + index: 0, + message: { role: 'ASSISTANT', content: [{ type: 'TEXT', text: 'Hello!' }] }, + finishReason: 'stop', + }], + usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 }, + }, + }; + const resp = codec.decodeResponse(raw); + // message is a plain string (MessageContent::Text serializes to a string, not {text: ...}) + assert.equal(resp.message, 'Hello!'); + assert.equal(resp.finish_reason, 'complete'); + assert.equal(resp.model, 'meta.llama-3.3-70b-instruct'); + assert.equal(resp.usage?.prompt_tokens, 10); + }); +}); + // =========================================================================== // GeminiGenerateContentCodec // =========================================================================== diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 4460e107b..b456d275f 100644 --- a/crates/pii-redaction/README.md +++ b/crates/pii-redaction/README.md @@ -33,7 +33,8 @@ 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_generate_content`. + `openai_responses`, `anthropic_messages`, `oci_genai`, 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 bc52cf8d9..fe8f8f754 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::OCIGenAI) => Some(ProviderSurface::OCIGenAI), LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) => { Some(ProviderSurface::GeminiGenerateContent) } diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 38e581d15..c382eb026 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_generate_content"], + values: ["openai_chat", "openai_responses", "anthropic_messages", "oci_genai", "gemini_generate_content"], optional: true, }, profiles => { label: "profiles", kind: List, list: &PII_REDACTION_PROFILE_LIST_ITEM }, diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index 5f935a834..fcfb26e8d 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, + OCIGenAI, GeminiGenerateContent, } @@ -21,6 +22,7 @@ impl BuiltinCodecName { ProviderSurface::OpenAIChat => Self::OpenAIChat, ProviderSurface::OpenAIResponses => Self::OpenAIResponses, ProviderSurface::AnthropicMessages => Self::AnthropicMessages, + ProviderSurface::OCIGenAI => Self::OCIGenAI, ProviderSurface::GeminiGenerateContent => Self::GeminiGenerateContent, } } @@ -34,6 +36,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::OCIGenAI => overlay_oci_genai_response(payload, annotated), Self::GeminiGenerateContent => overlay_gemini_response(payload, annotated), } } @@ -260,6 +263,233 @@ fn overlay_anthropic_response(mut payload: Json, annotated: &AnnotatedLlmRespons payload } +fn overlay_oci_genai_response(mut payload: Json, annotated: &AnnotatedLlmResponse) -> Json { + let Some(root) = payload.as_object_mut() else { + return payload; + }; + if root.contains_key("modelId") { + set_optional_string_field(root, "modelId", annotated.model.as_deref()); + } + if root.get("chatResponse").is_some_and(Json::is_object) { + if let Some(chat_response) = root.get_mut("chatResponse").and_then(Json::as_object_mut) { + overlay_oci_chat_response(chat_response, annotated); + } + } else { + overlay_oci_chat_response(root, annotated); + } + payload +} + +fn overlay_oci_chat_response( + chat_response: &mut Map, + annotated: &AnnotatedLlmResponse, +) { + let api_format = chat_response + .get("apiFormat") + .and_then(Json::as_str) + .unwrap_or("GENERIC") + .to_uppercase(); + + if api_format == "COHERE" { + set_optional_string_field( + chat_response, + "text", + annotated_message_text(annotated.message.as_ref()).as_deref(), + ); + set_optional_string_field( + chat_response, + "finishReason", + annotated + .finish_reason + .as_ref() + .map(oci_cohere_finish_reason), + ); + overlay_oci_cohere_tool_calls(chat_response, annotated.tool_calls.as_deref()); + return; + } + + if api_format == "COHEREV2" { + // COHEREV2 carries a single root-level assistant `message` (typed + // content parts, nested-function tool calls) instead of `choices`. + set_optional_string_field( + chat_response, + "finishReason", + annotated + .finish_reason + .as_ref() + .map(oci_cohere_v2_finish_reason), + ); + let Some(message) = chat_response + .get_mut("message") + .and_then(Json::as_object_mut) + else { + return; + }; + overlay_oci_message(message, annotated); + return; + } + + let Some(choices) = chat_response + .get_mut("choices") + .and_then(Json::as_array_mut) + else { + return; + }; + // The normalized annotation models a single choice; any additional raw + // choices have no sanitized counterpart and would leak unredacted data. + choices.truncate(1); + let Some(choice) = choices.first_mut().and_then(Json::as_object_mut) else { + return; + }; + set_optional_string_field( + choice, + "finishReason", + annotated + .finish_reason + .as_ref() + .map(oci_generic_finish_reason), + ); + let Some(message) = choice.get_mut("message").and_then(Json::as_object_mut) else { + return; + }; + overlay_oci_message(message, annotated); +} + +/// Sanitize an OCI assistant message: typed TEXT parts or a bare string +/// `content` (both shapes the decoder accepts), plus tool calls. +fn overlay_oci_message(message: &mut Map, annotated: &AnnotatedLlmResponse) { + match message.get_mut("content") { + Some(Json::Array(blocks)) => { + overlay_oci_text_parts(blocks, annotated_message_text(annotated.message.as_ref())); + } + Some(Json::String(_)) => { + set_optional_string_field( + message, + "content", + annotated_message_text(annotated.message.as_ref()).as_deref(), + ); + } + _ => {} + } + overlay_oci_tool_calls(message, annotated.tool_calls.as_deref()); +} + +/// Sanitize flat COHERE (v1) tool calls: `{name, parameters}` entries directly +/// on the chat response, with `parameters` as a parsed JSON object and no `id` +/// on the wire. +fn overlay_oci_cohere_tool_calls( + chat_response: &mut Map, + tool_calls: Option<&[ResponseToolCall]>, +) { + let Some(raw_calls) = chat_response + .get_mut("toolCalls") + .and_then(Json::as_array_mut) + else { + return; + }; + let Some(tool_calls) = tool_calls else { + chat_response.remove("toolCalls"); + return; + }; + // The COHERE wire documents `parameters` as an object; a sanitizer that + // produced any other shape cannot be overlaid faithfully, so drop the + // calls rather than emit an invalid wire shape. + if tool_calls.iter().any(|call| !call.arguments.is_object()) { + chat_response.remove("toolCalls"); + return; + } + + raw_calls.truncate(tool_calls.len()); + + for (raw_call, sanitized_call) in raw_calls.iter_mut().zip(tool_calls.iter()) { + let Some(raw_call) = raw_call.as_object_mut() else { + chat_response.remove("toolCalls"); + return; + }; + set_optional_string_field(raw_call, "name", Some(sanitized_call.name.as_str())); + raw_call.insert("parameters".into(), sanitized_call.arguments.clone()); + } +} + +fn overlay_oci_text_parts(blocks: &mut [Json], message_text: Option) { + let text_part_count = blocks + .iter() + .filter(|block| block.get("type").and_then(Json::as_str) == Some("TEXT")) + .count(); + // `splitn` keeps surplus newline-separated text inside the final fragment + // so a sanitized line that itself contains a newline is never dropped. + let parts = message_text.as_deref().map(|text| { + text.splitn(text_part_count.max(1), '\n') + .collect::>() + }); + let mut text_part_index = 0usize; + + for block in blocks { + if block.get("type").and_then(Json::as_str) != Some("TEXT") { + continue; + } + let Some(block) = block.as_object_mut() else { + continue; + }; + if text_part_count <= 1 { + set_optional_string_field(block, "text", message_text.as_deref()); + text_part_index += 1; + continue; + } + let part = parts + .as_ref() + .and_then(|parts| parts.get(text_part_index).copied()) + .or_else(|| { + (text_part_index == 0) + .then_some(message_text.as_deref()) + .flatten() + }); + set_optional_string_field(block, "text", part); + text_part_index += 1; + } +} + +fn overlay_oci_tool_calls( + message: &mut Map, + tool_calls: Option<&[ResponseToolCall]>, +) { + let Some(raw_calls) = message.get_mut("toolCalls").and_then(Json::as_array_mut) else { + return; + }; + let Some(tool_calls) = tool_calls else { + message.remove("toolCalls"); + return; + }; + + raw_calls.truncate(tool_calls.len()); + + for (raw_call, sanitized_call) in raw_calls.iter_mut().zip(tool_calls.iter()) { + let Some(raw_call) = raw_call.as_object_mut() else { + message.remove("toolCalls"); + return; + }; + set_optional_string_field(raw_call, "id", Some(sanitized_call.id.as_str())); + // The OCI decode reads `function.name`/`function.arguments` when a + // nested `function` object exists, so sanitize that object too; the + // flat fields are the plain OCI wire shape. + if let Some(function) = raw_call.get_mut("function").and_then(Json::as_object_mut) { + set_optional_string_field(function, "name", Some(sanitized_call.name.as_str())); + set_optional_string_field( + function, + "arguments", + Some(json_string(&sanitized_call.arguments).as_str()), + ); + } else { + set_optional_string_field(raw_call, "name", Some(sanitized_call.name.as_str())); + set_optional_string_field( + raw_call, + "arguments", + Some(json_string(&sanitized_call.arguments).as_str()), + ); + } + } +} + fn overlay_openai_chat_tool_calls( message: &mut Map, tool_calls: Option<&[ResponseToolCall]>, @@ -508,6 +738,35 @@ fn anthropic_stop_reason(reason: &FinishReason) -> &str { } } +fn oci_generic_finish_reason(reason: &FinishReason) -> &str { + match reason { + FinishReason::Complete => "stop", + FinishReason::Length => "length", + FinishReason::ToolUse => "tool_calls", + FinishReason::ContentFilter => "content_filter", + FinishReason::Unknown(other) => other.as_str(), + } +} + +fn oci_cohere_finish_reason(reason: &FinishReason) -> &str { + match reason { + FinishReason::Complete => "COMPLETE", + FinishReason::Length => "MAX_TOKENS", + FinishReason::ToolUse | FinishReason::ContentFilter => "COMPLETE", + FinishReason::Unknown(other) => other.as_str(), + } +} + +fn oci_cohere_v2_finish_reason(reason: &FinishReason) -> &str { + match reason { + FinishReason::Complete => "COMPLETE", + FinishReason::Length => "MAX_TOKENS", + FinishReason::ToolUse => "TOOL_CALL", + FinishReason::ContentFilter => "COMPLETE", + FinishReason::Unknown(other) => other.as_str(), + } +} + #[cfg(test)] #[path = "../tests/coverage/overlay_tests.rs"] mod tests; diff --git a/crates/pii-redaction/tests/coverage/overlay_tests.rs b/crates/pii-redaction/tests/coverage/overlay_tests.rs index 621eaa5f7..068ec7869 100644 --- a/crates/pii-redaction/tests/coverage/overlay_tests.rs +++ b/crates/pii-redaction/tests/coverage/overlay_tests.rs @@ -134,6 +134,51 @@ fn anthropic_overlay_preserves_full_multiline_text_in_single_text_block() { assert_eq!(blocks[0]["text"], json!("line one\nline two")); } +#[test] +fn oci_genai_overlay_rewrites_generic_text_and_tool_calls() { + let payload = json!({ + "modelId": "meta.llama-3.3-70b-instruct", + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "raw secret"}], + "toolCalls": [ + {"id": "call_1", "type": "FUNCTION", "name": "one", "arguments": "{\"secret\":\"raw-1\"}"}, + {"id": "call_2", "type": "FUNCTION", "name": "two", "arguments": "{\"secret\":\"raw-2\"}"} + ] + }, + "finishReason": "tool_calls" + }] + } + }); + let annotated = AnnotatedLlmResponse { + model: Some("meta.llama-3.3-70b-instruct".into()), + message: Some(MessageContent::Text("[REDACTED]".into())), + tool_calls: Some(vec![tool_call( + "call_1", + "one", + json!({"secret": "[REDACTED]"}), + )]), + finish_reason: Some(FinishReason::ToolUse), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + let message = &overlaid["chatResponse"]["choices"][0]["message"]; + assert_eq!(message["content"][0]["text"], json!("[REDACTED]")); + let calls = message["toolCalls"].as_array().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0]["arguments"], json!("{\"secret\":\"[REDACTED]\"}")); + assert_eq!( + overlaid["chatResponse"]["choices"][0]["finishReason"], + json!("tool_calls") + ); +} + fn gemini_annotated( message: Option<&str>, tool_calls: Option>, @@ -410,6 +455,106 @@ fn gemini_overlay_updates_response_id_and_model_version() { ); } +#[test] +fn oci_genai_overlay_rewrites_cohere_text() { + let payload = json!({ + "chatResponse": { + "apiFormat": "COHERE", + "text": "raw secret", + "finishReason": "COMPLETE" + } + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + finish_reason: Some(FinishReason::Complete), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + assert_eq!(overlaid["chatResponse"]["text"], json!("[REDACTED]")); + assert_eq!(overlaid["chatResponse"]["finishReason"], json!("COMPLETE")); +} + +#[test] +fn oci_genai_overlay_rewrites_each_text_part_and_keeps_non_text_blocks() { + let payload = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [ + {"type": "TEXT", "text": "raw one"}, + {"type": "IMAGE", "imageUrl": {"url": "data:image/png;base64,AAAA"}}, + {"type": "TEXT", "text": "raw two"} + ] + } + }] + } + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text( + "[REDACTED ONE]\n[REDACTED TWO]\nwith remainder".into(), + )), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + let content = &overlaid["chatResponse"]["choices"][0]["message"]["content"]; + assert_eq!(content[0]["text"], json!("[REDACTED ONE]")); + assert_eq!( + content[1], + json!({"type": "IMAGE", "imageUrl": {"url": "data:image/png;base64,AAAA"}}) + ); + // The final TEXT part keeps any surplus newline-separated text. + assert_eq!(content[2]["text"], json!("[REDACTED TWO]\nwith remainder")); +} + +#[test] +fn oci_genai_overlay_sanitizes_nested_function_tool_calls() { + let payload = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [], + "toolCalls": [{ + "id": "call_1", + "type": "FUNCTION", + "function": {"name": "one", "arguments": "{\"secret\":\"raw-1\"}"} + }] + } + }] + } + }); + let annotated = AnnotatedLlmResponse { + tool_calls: Some(vec![tool_call( + "call_1", + "one", + json!({"secret": "[REDACTED]"}), + )]), + finish_reason: Some(FinishReason::ToolUse), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + let call = &overlaid["chatResponse"]["choices"][0]["message"]["toolCalls"][0]; + assert_eq!( + call["function"]["arguments"], + json!("{\"secret\":\"[REDACTED]\"}") + ); + assert!( + call.get("arguments").is_none(), + "sanitized arguments must land on the nested function object, got {call}" + ); +} + #[test] fn gemini_overlay_does_not_overwrite_finish_reason() { // A STOP response with a functionCall part: normalized finish_reason is ToolUse, @@ -440,3 +585,341 @@ fn gemini_overlay_does_not_overwrite_finish_reason() { "Gemini overlay must not overwrite native finishReason with the derived ToolUse value" ); } + +#[test] +fn oci_genai_overlay_sanitizes_flat_cohere_tool_calls() { + let payload = json!({ + "chatResponse": { + "apiFormat": "COHERE", + "text": "raw secret", + "finishReason": "COMPLETE", + "toolCalls": [ + {"name": "one", "parameters": {"secret": "raw-1"}}, + {"name": "two", "parameters": {"secret": "raw-2"}} + ] + } + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + tool_calls: Some(vec![tool_call( + "call_0", + "one", + json!({"secret": "[REDACTED]"}), + )]), + finish_reason: Some(FinishReason::Complete), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + let chat_response = &overlaid["chatResponse"]; + assert_eq!(chat_response["text"], json!("[REDACTED]")); + let calls = chat_response["toolCalls"].as_array().unwrap(); + assert_eq!(calls.len(), 1, "dropped sanitized calls must be truncated"); + assert_eq!(calls[0]["parameters"], json!({"secret": "[REDACTED]"})); + assert!( + calls[0].get("id").is_none(), + "COHERE wire tool calls carry no id and must not gain one" + ); +} + +#[test] +fn oci_genai_overlay_sanitizes_cohere_v2_root_message() { + let payload = json!({ + "chatResponse": { + "apiFormat": "COHEREV2", + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "raw secret"}], + "toolCalls": [{ + "id": "call_1", + "type": "FUNCTION", + "function": {"name": "one", "arguments": "{\"secret\":\"raw-1\"}"} + }] + }, + "finishReason": "TOOL_CALL" + } + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + tool_calls: Some(vec![tool_call( + "call_1", + "one", + json!({"secret": "[REDACTED]"}), + )]), + finish_reason: Some(FinishReason::ToolUse), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + let message = &overlaid["chatResponse"]["message"]; + assert_eq!(message["content"][0]["text"], json!("[REDACTED]")); + assert_eq!( + message["toolCalls"][0]["function"]["arguments"], + json!("{\"secret\":\"[REDACTED]\"}") + ); + assert_eq!(overlaid["chatResponse"]["finishReason"], json!("TOOL_CALL")); +} + +#[test] +fn oci_genai_overlay_sanitizes_generic_string_content() { + let payload = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": {"role": "ASSISTANT", "content": "raw secret"} + }] + } + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + assert_eq!( + overlaid["chatResponse"]["choices"][0]["message"]["content"], + json!("[REDACTED]") + ); +} + +#[test] +fn oci_genai_overlay_drops_cohere_tool_calls_with_non_object_arguments() { + for arguments in [json!("scalar"), json!([1, 2]), json!(null)] { + let payload = json!({ + "chatResponse": { + "apiFormat": "COHERE", + "text": "ok", + "toolCalls": [{"name": "one", "parameters": {"secret": "raw-1"}}] + } + }); + let annotated = AnnotatedLlmResponse { + tool_calls: Some(vec![tool_call("call_0", "one", arguments.clone())]), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + assert!( + overlaid["chatResponse"].get("toolCalls").is_none(), + "non-object sanitized arguments ({arguments}) must drop toolCalls, got {}", + overlaid["chatResponse"] + ); + } +} + +#[test] +fn oci_genai_overlay_removes_unsanitized_additional_choices() { + let payload = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [ + { + "index": 0, + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "raw secret"}]} + }, + { + "index": 1, + "message": {"role": "ASSISTANT", "content": [{"type": "TEXT", "text": "second raw secret"}]} + } + ] + } + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + + let choices = overlaid["chatResponse"]["choices"].as_array().unwrap(); + assert_eq!( + choices.len(), + 1, + "additional raw choices have no sanitized counterpart and must be removed" + ); + assert_eq!( + choices[0]["message"]["content"][0]["text"], + json!("[REDACTED]") + ); +} + +#[test] +fn oci_genai_overlay_guards_pass_unrecognized_shapes_through() { + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + ..AnnotatedLlmResponse::default() + }; + + // Non-object payloads and shapes without the expected structure pass + // through unchanged instead of panicking or half-sanitizing. + for payload in [ + json!("not an object"), + json!({"chatResponse": {"apiFormat": "GENERIC"}}), + json!({"chatResponse": {"apiFormat": "GENERIC", "choices": ["not-an-object"]}}), + json!({"chatResponse": {"apiFormat": "GENERIC", "choices": [{"index": 0}]}}), + json!({"chatResponse": {"apiFormat": "COHEREV2"}}), + ] { + let overlaid = + BuiltinCodecName::OCIGenAI.overlay_response_payload(payload.clone(), &annotated); + assert_eq!(overlaid, payload); + } +} + +#[test] +fn oci_genai_overlay_reaches_bare_chat_response_via_provider_surface() { + // Envelope-less payload routed through the provider-surface mapping. + let payload = json!({ + "apiFormat": "COHERE", + "text": "raw secret", + "finishReason": "COMPLETE" + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + finish_reason: Some(FinishReason::Complete), + ..AnnotatedLlmResponse::default() + }; + + let overlaid = BuiltinCodecName::from_provider_surface(ProviderSurface::OCIGenAI) + .overlay_response_payload(payload, &annotated); + + assert_eq!(overlaid["text"], json!("[REDACTED]")); +} + +#[test] +fn oci_genai_overlay_removes_tool_calls_without_sanitized_counterparts() { + // GENERIC: sanitized None removes the key; a non-object raw call also + // removes the key rather than leaving unredacted entries behind. + for (payload_calls, sanitized) in [ + ( + json!([{"id": "call_1", "name": "one", "arguments": "{\"secret\":\"raw\"}"}]), + None, + ), + ( + json!(["not-an-object"]), + Some(vec![tool_call("call_1", "one", json!({}))]), + ), + ] { + let payload = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": {"role": "ASSISTANT", "content": [], "toolCalls": payload_calls} + }] + } + }); + let annotated = AnnotatedLlmResponse { + tool_calls: sanitized, + ..AnnotatedLlmResponse::default() + }; + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + assert!( + overlaid["chatResponse"]["choices"][0]["message"] + .get("toolCalls") + .is_none(), + "unsanitizable toolCalls must be removed" + ); + } + + // COHERE: same removal semantics on the flat root-level calls. + for (payload_calls, sanitized) in [ + ( + json!([{"name": "one", "parameters": {"secret": "raw"}}]), + None, + ), + ( + json!(["not-an-object"]), + Some(vec![tool_call("call_0", "one", json!({}))]), + ), + ] { + let payload = json!({ + "chatResponse": {"apiFormat": "COHERE", "text": "ok", "toolCalls": payload_calls} + }); + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("ok".into())), + tool_calls: sanitized, + ..AnnotatedLlmResponse::default() + }; + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + assert!(overlaid["chatResponse"].get("toolCalls").is_none()); + } +} + +#[test] +fn oci_genai_overlay_multi_part_text_handles_short_and_non_object_blocks() { + let payload = json!({ + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [ + {"type": "TEXT", "text": "raw one"}, + {"type": "TEXT", "text": "raw two"} + ] + } + }] + } + }); + // A single sanitized line for two TEXT parts: the first block takes the + // full text (index-0 fallback), the second has no fragment and is + // left untouched. + let annotated = AnnotatedLlmResponse { + message: Some(MessageContent::Text("[REDACTED]".into())), + ..AnnotatedLlmResponse::default() + }; + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(payload, &annotated); + let content = &overlaid["chatResponse"]["choices"][0]["message"]["content"]; + assert_eq!(content[0]["text"], json!("[REDACTED]")); +} + +#[test] +fn oci_genai_overlay_maps_every_finish_reason_variant() { + for (reason, generic, cohere, v2) in [ + (FinishReason::Complete, "stop", "COMPLETE", "COMPLETE"), + (FinishReason::Length, "length", "MAX_TOKENS", "MAX_TOKENS"), + (FinishReason::ToolUse, "tool_calls", "COMPLETE", "TOOL_CALL"), + ( + FinishReason::ContentFilter, + "content_filter", + "COMPLETE", + "COMPLETE", + ), + ( + FinishReason::Unknown("mystery".into()), + "mystery", + "mystery", + "mystery", + ), + ] { + let annotated = AnnotatedLlmResponse { + finish_reason: Some(reason), + ..AnnotatedLlmResponse::default() + }; + + let generic_payload = json!({"chatResponse": {"apiFormat": "GENERIC", + "choices": [{"index": 0, "finishReason": "x", "message": {"role": "ASSISTANT", "content": []}}]}}); + let overlaid = + BuiltinCodecName::OCIGenAI.overlay_response_payload(generic_payload, &annotated); + assert_eq!( + overlaid["chatResponse"]["choices"][0]["finishReason"], + json!(generic) + ); + + let cohere_payload = + json!({"chatResponse": {"apiFormat": "COHERE", "text": "ok", "finishReason": "x"}}); + let overlaid = + BuiltinCodecName::OCIGenAI.overlay_response_payload(cohere_payload, &annotated); + assert_eq!(overlaid["chatResponse"]["finishReason"], json!(cohere)); + + let v2_payload = json!({"chatResponse": {"apiFormat": "COHEREV2", "finishReason": "x", + "message": {"role": "ASSISTANT", "content": []}}}); + let overlaid = BuiltinCodecName::OCIGenAI.overlay_response_payload(v2_payload, &annotated); + assert_eq!(overlaid["chatResponse"]["finishReason"], json!(v2)); + } +} diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index 81276be95..3d464e3af 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -421,6 +421,49 @@ async fn normalized_llm_paths_use_the_active_codec_and_fail_closed_for_unknown_c ); } +#[tokio::test] +async fn normalized_llm_response_paths_use_the_active_oci_genai_codec_identity() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec!["/message".to_string()], + ..BuiltinBackendConfig::default() + }, + None, + ) + .unwrap(); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + + let sanitized = sanitize_response( + json!({ + "modelId": "meta.llama-3.3-70b-instruct", + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [{ + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "sk-oci-secret"}] + }, + "finishReason": "stop" + }] + } + }), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::BuiltIn( + BuiltinLlmCodec::OCIGenAI, + )), + ) + .await + .expect("sanitizer callback must succeed") + .expect("the active OCI GenAI codec identity must retain the payload"); + assert_eq!( + sanitized["chatResponse"]["choices"][0]["message"]["content"][0]["text"], + json!("[REDACTED]") + ); +} + #[tokio::test] async fn normalized_llm_paths_omit_payloads_when_legacy_codec_decode_fails() { let backend = crate::builtin::CompiledBuiltinBackend::new( diff --git a/crates/plugin/src/async_sdk.rs b/crates/plugin/src/async_sdk.rs index 9d33de9d9..57a4752f5 100644 --- a/crates/plugin/src/async_sdk.rs +++ b/crates/plugin/src/async_sdk.rs @@ -1043,6 +1043,7 @@ impl CodecIdentityInvocation { "anthropic_messages" => Ok(LlmCodecIdentity::BuiltIn( BuiltinLlmCodec::AnthropicMessages, )), + "oci_genai" => Ok(LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI)), "gemini_generate_content" => Ok(LlmCodecIdentity::BuiltIn( BuiltinLlmCodec::GeminiGenerateContent, )), diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 2276bed54..97ff278ae 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -65,6 +65,9 @@ pub enum BuiltinLlmCodec { /// Anthropic Messages. #[serde(rename = "anthropic_messages")] AnthropicMessages, + /// OCI Generative AI chat. + #[serde(rename = "oci_genai")] + OCIGenAI, /// Gemini generateContent. #[serde(rename = "gemini_generate_content")] GeminiGenerateContent, diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index ffcbe5658..09aa73c91 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -3320,6 +3320,54 @@ fn typed_async_middleware_registers_and_round_trips_every_surface() { assert_eq!(live_host_strings(), 0); } +#[test] +fn typed_async_llm_sanitize_context_decodes_oci_genai_builtin_identity() { + let _guard = begin_test(); + let host = test_host_v4(); + let mut ctx = test_context(&host.v3.v1); + + ctx.register_llm_sanitize_request_guardrail( + "llm-request-oci-genai", + 0, + |request, context| async move { + assert_eq!( + context.codec, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI) + ); + Ok(Some(request)) + }, + ) + .unwrap(); + + let registration = + take_async_registration(NemoRelayNativeAsyncMiddlewareKind::LlmSanitizeRequest); + assert_eq!( + invoke_async_registration( + &host, + ®istration, + json!({ + "request": test_llm_request(), + "context": { "codec_kind": "builtin", "codec_id": "oci_genai" } + }), + None, + ) + .unwrap()["content"], + json!({ "prompt": "hello" }) + ); + unsafe { registration.free() }; + // The context's retained codec capability is released on the SDK executor + // after the result completion is delivered, so poll instead of asserting + // immediately. + let deadline = Instant::now() + Duration::from_secs(5); + while live_host_strings() != 0 { + assert!( + Instant::now() < deadline, + "host strings were not released after the sanitize invocation" + ); + std::thread::yield_now(); + } +} + #[test] fn typed_async_registration_failure_rolls_back_callback_state() { let _guard = begin_test(); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 81c69d1b4..fce7e33f5 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -51,9 +51,10 @@ 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, PyGeminiGenerateContentCodec, - PyLLMAttributes, PyLLMHandle, PyLLMRequest, PyLlmStream, PyOpenAIChatCodec, - PyOpenAIResponsesCodec, PyPropagationContext, PyScopeAttributes, PyScopeHandle, PyScopeStack, - PyScopeType, PyThreadScopeStackBinding, PyToolAttributes, PyToolHandle, + PyLLMAttributes, PyLLMHandle, PyLLMRequest, PyLlmStream, PyOCIGenAIChatCodec, + PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, PyScopeAttributes, + PyScopeHandle, PyScopeStack, PyScopeType, PyThreadScopeStackBinding, PyToolAttributes, + PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -259,6 +260,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()); + } if let Ok(builtin) = c.extract::>() { return Some(builtin.inner_response_codec.clone()); } @@ -283,6 +287,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()); + } if let Ok(builtin) = codec.extract::>() { return Some(builtin.inner_codec.clone()); } diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index a7a56238d..2b18c1087 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -1016,6 +1016,71 @@ impl PyAnthropicMessagesCodec { } } +/// Built-in codec for the OCI Generative AI chat API. +/// +/// Implements both ``LlmCodec`` (decode/encode for requests) and +/// ``LlmResponseCodec`` (decode_response for responses). +/// +/// Example: +/// ```python +/// from nemo_relay.codecs import OCIGenAIChatCodec +/// codec = OCIGenAIChatCodec() +/// annotated_req = codec.decode(request) +/// annotated_resp = codec.decode_response(response) +/// ``` +#[pyclass(name = "OCIGenAIChatCodec")] +pub struct PyOCIGenAIChatCodec { + pub(crate) inner_codec: Arc, + pub(crate) inner_response_codec: Arc, +} + +#[pymethods] +impl PyOCIGenAIChatCodec { + #[new] + pub(crate) fn new() -> Self { + Self { + inner_codec: Arc::new(nemo_relay::codec::oci_genai::OCIGenAIChatCodec), + inner_response_codec: Arc::new(nemo_relay::codec::oci_genai::OCIGenAIChatCodec), + } + } + + /// 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 { + "" + } +} + /// Built-in codec for the Gemini generateContent API. /// /// Implements both ``LlmCodec`` (decode/encode for requests) and diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index f559c39c8..f7bd7fb11 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::()?; m.add_class::()?; Ok(()) } diff --git a/crates/types/src/codec/request.rs b/crates/types/src/codec/request.rs index 1167b92e6..02c3fc5c1 100644 --- a/crates/types/src/codec/request.rs +++ b/crates/types/src/codec/request.rs @@ -534,6 +534,19 @@ pub enum ApiSpecificRequest { #[serde(skip_serializing_if = "Option::is_none")] text: Option, }, + /// OCI Generative AI-specific request fields. + #[serde(rename = "oci_genai")] + OCIGenAI { + /// Compartment OCID from the `ChatDetails` envelope. + #[serde(skip_serializing_if = "Option::is_none")] + compartment_id: Option, + /// Serving mode object (`servingType` plus `modelId` or `endpointId`). + #[serde(skip_serializing_if = "Option::is_none")] + serving_mode: Option, + /// Chat request API format (`GENERIC`, `COHERE`, or `COHEREV2`). + #[serde(skip_serializing_if = "Option::is_none")] + api_format: Option, + }, /// Custom provider request fields. #[serde(rename = "custom")] Custom { diff --git a/crates/types/src/codec/response.rs b/crates/types/src/codec/response.rs index 8b3a6abf8..7f50d2aab 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -332,6 +332,17 @@ pub enum ApiSpecificResponse { content_blocks: Option>, }, + /// OCI Generative AI-specific fields. + #[serde(rename = "oci_genai")] + OCIGenAI { + /// Chat response API format (`GENERIC`, `COHERE`, or `COHEREV2`). + #[serde(skip_serializing_if = "Option::is_none")] + api_format: Option, + /// Model version reported on the `ChatResult` envelope. + #[serde(skip_serializing_if = "Option::is_none")] + model_version: Option, + }, + /// Gemini generateContent API-specific fields. #[serde(rename = "gemini_generate_content")] GeminiGenerateContent { diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 92133ae4a..d80937fce 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -154,6 +154,8 @@ pub enum BuiltinLlmCodec { OpenAiResponses, /// Anthropic Messages. AnthropicMessages, + /// OCI Generative AI chat request and response payloads. + OCIGenAI, /// Gemini generateContent request and response payloads. GeminiGenerateContent, } @@ -2159,6 +2161,7 @@ fn codec_identity_from_proto( Some("anthropic_messages") => { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) } + Some("oci_genai") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI), Some("gemini_generate_content") => { LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::GeminiGenerateContent) } diff --git a/crates/worker/tests/unit/codec_identity_tests.rs b/crates/worker/tests/unit/codec_identity_tests.rs index 022e6dceb..8dc1d7420 100644 --- a/crates/worker/tests/unit/codec_identity_tests.rs +++ b/crates/worker/tests/unit/codec_identity_tests.rs @@ -19,3 +19,20 @@ fn test_gemini_codec_identity_decoded_as_builtin_not_opaque() { "Gemini generateContent codec id must decode to BuiltIn(GeminiGenerateContent), not Opaque" ); } + +#[test] +fn test_oci_genai_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("oci_genai".to_string()), + }; + let identity = codec_identity_from_proto(Some(&proto)); + assert_eq!( + identity, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OCIGenAI), + "OCI GenAI codec id must decode to BuiltIn(OCIGenAI), not Opaque" + ); +} diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 20e0b52a3..a5d4ad901 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -362,7 +362,8 @@ 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_generate_content` +Relay's built-in `openai_chat`, `openai_responses`, `anthropic_messages`, `oci_genai`, 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 4bccf8b01..d5cec5403 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -271,7 +271,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_generate_content`) and +(one of `openai_chat`, `openai_responses`, `anthropic_messages`, `oci_genai`, 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 f4915c252..c4850a563 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_generate_content` | +| Codec support | `openai_chat` | `openai_chat`, `openai_responses`, `anthropic_messages`, `oci_genai`, `gemini_generate_content` | | Runtime availability | Any runtime that includes the remote backend | Runtimes that can start `python3 >= 3.11` with `nemoguardrails==0.22.0` installed | ## Remote Mode @@ -340,6 +340,7 @@ The current built-in local mode supports managed LLM execution with: - `openai_chat` - `openai_responses` - `anthropic_messages` +- `oci_genai` - `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 e1a77cf12..e0f26b6f3 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -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_generate_content` | Runtime-specific future implementation | +| Codec support | `openai_chat`, `openai_responses`, `anthropic_messages`, `oci_genai`, `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 5df68707e..2ad656399 100644 --- a/docs/integrate-into-frameworks/provider-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-codecs.mdx @@ -73,7 +73,7 @@ Use the annotated request surfaces according to their ownership: - Portable `messages`, content parts, function calls and results, `tools`, and `tool_choice` when the component has shared semantics. - Tagged `api_specific` fields for modeled controls that belong to Anthropic - Messages, OpenAI Chat Completions, or OpenAI Responses. + Messages, OpenAI Chat Completions, OpenAI Responses, or OCI Generative AI. - `{ provider, kind, value }` native components for provider-only input items, blocks, tools, tool choices, and future union members. `value` is the exact provider JSON. @@ -91,6 +91,8 @@ 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. +- `OCIGenAIChatCodec`: OCI Generative AI chat-compatible requests and responses + in the `GENERIC`, `COHERE`, and `COHEREV2` API formats. - `GeminiGenerateContentCodec`: Gemini `generateContent`-compatible requests and responses. ## Provider Codec Roles @@ -107,6 +109,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` | +| OCI Generative AI chat | `nemo_relay.codecs.OCIGenAIChatCodec` | `OCIGenAIChatCodec` 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. @@ -114,9 +117,9 @@ Choose the provider codec that matches the payload shape the framework already s 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 +`generateContent` and OCI Generative AI chat do not yet have gateway routes, +so use `GeminiGenerateContentCodec` or `OCIGenAIChatCodec` directly when a +framework sends those payload shapes. Count-token, model, probe, and non-LLM passthrough routes do not use request codecs. ## Example: Add a System Message with a Provider Codec diff --git a/docs/integrate-into-frameworks/provider-response-codecs.mdx b/docs/integrate-into-frameworks/provider-response-codecs.mdx index 8b6a7dd2a..19c1c65a5 100644 --- a/docs/integrate-into-frameworks/provider-response-codecs.mdx +++ b/docs/integrate-into-frameworks/provider-response-codecs.mdx @@ -361,13 +361,13 @@ 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` | -|---|---|---|---|---| -| `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` | — | +| Normalized Field | OpenAI Chat | OpenAI Responses | Anthropic Messages | OCI Generative AI | Gemini `generateContent` | +|---|---|---|---|---|---| +| `prompt_tokens` | `prompt_tokens` | `input_tokens` | `input_tokens` | `promptTokens` | `promptTokenCount` | +| `completion_tokens` | `completion_tokens` | `output_tokens` | `output_tokens` | `completionTokens` | `candidatesTokenCount` | +| `total_tokens` | `total_tokens` | `total_tokens` | computed | `totalTokens` | `totalTokenCount` (or computed as prompt + candidates + thinking) | +| `cache_read_tokens` | `prompt_tokens_details.cached_tokens` | `input_tokens_details.cached_tokens` | `cache_read_input_tokens` | `promptTokensDetails.cachedTokens` | `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. @@ -445,6 +445,7 @@ The built-in provider codecs also implement response decoding: - `OpenAIChatCodec` — OpenAI Chat Completions API - `OpenAIResponsesCodec` — OpenAI Responses API - `AnthropicMessagesCodec` — Anthropic Messages API +- `OCIGenAIChatCodec` — OCI Generative AI chat API (`GENERIC`, `COHERE`, and `COHEREV2` formats) - `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 5a529a32c..5db471f7d 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, OCI Generative AI chat, 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/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index f9273c34f..d5beaf95f 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 OCIGenAIChatCodec: + """Built-in codec for OCI Generative AI chat requests and responses. + + Summary: + Native codec bridge for OCI Generative AI chat payloads. + """ + + def __init__(self) -> None: + """Create an OCI Generative AI chat codec.""" + ... + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: + """Decode an OCI GenAI chat request into a normalized request view.""" + ... + def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMRequest: + """Encode a normalized request back into OCI GenAI chat shape.""" + ... + def decode_response(self, response: _Json) -> AnnotatedLLMResponse: + """Decode an OCI GenAI chat response into a normalized response view.""" + ... + class GeminiGenerateContentCodec: """Built-in codec for Gemini generateContent requests and responses. diff --git a/python/nemo_relay/codecs.py b/python/nemo_relay/codecs.py index 98cdb2e32..7972c8a75 100644 --- a/python/nemo_relay/codecs.py +++ b/python/nemo_relay/codecs.py @@ -49,6 +49,7 @@ async def impl(request: LLMRequest): AnthropicMessagesCodec, GeminiGenerateContentCodec, LLMRequest, + OCIGenAIChatCodec, OpenAIChatCodec, OpenAIResponsesCodec, ) @@ -165,6 +166,7 @@ def decode_response(self, response: Json) -> "AnnotatedLLMResponse": "GeminiGenerateContentCodec", "LlmCodec", "LlmResponseCodec", + "OCIGenAIChatCodec", "OpenAIChatCodec", "OpenAIResponsesCodec", ] diff --git a/python/nemo_relay/codecs.pyi b/python/nemo_relay/codecs.pyi index d77ed6168..80b3496a1 100644 --- a/python/nemo_relay/codecs.pyi +++ b/python/nemo_relay/codecs.pyi @@ -163,6 +163,44 @@ class AnthropicMessagesCodec: """ ... +class OCIGenAIChatCodec: + """Built-in codec for OCI Generative AI chat requests and responses.""" + + def __init__(self) -> None: ... + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: + """Decode an OCI Generative AI chat request. + + Args: + request: Raw OCI ``ChatDetails`` request payload. + + Returns: + AnnotatedLLMRequest: Normalized request representation. + """ + ... + + def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMRequest: + """Encode a normalized request back into OCI chat format. + + Args: + annotated: Normalized request after intercept edits. + original: Original OCI ``ChatDetails`` request. + + Returns: + LLMRequest: Updated OCI chat request payload. + """ + ... + + def decode_response(self, response: Json) -> AnnotatedLLMResponse: + """Decode an OCI Generative AI chat response. + + Args: + response: Raw OCI ``ChatResult`` response payload. + + Returns: + AnnotatedLLMResponse: Normalized response representation. + """ + ... + class GeminiGenerateContentCodec: """Built-in codec for Gemini generateContent requests and responses.""" @@ -207,6 +245,7 @@ __all__ = [ "GeminiGenerateContentCodec", "LlmCodec", "LlmResponseCodec", + "OCIGenAIChatCodec", "OpenAIChatCodec", "OpenAIResponsesCodec", ] diff --git a/python/nemo_relay/pii_redaction.py b/python/nemo_relay/pii_redaction.py index 99dea9047..f27061fa5 100644 --- a/python/nemo_relay/pii_redaction.py +++ b/python/nemo_relay/pii_redaction.py @@ -136,9 +136,11 @@ 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", "oci_genai", "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 d7719d5e0..c4c4bbf2c 100644 --- a/python/nemo_relay/pii_redaction.pyi +++ b/python/nemo_relay/pii_redaction.pyi @@ -59,7 +59,9 @@ class PiiRedactionConfig: mark: bool = ... priority: int = ... codec: ( - Literal["openai_chat", "openai_responses", "anthropic_messages", "gemini_generate_content"] | str | None + Literal["openai_chat", "openai_responses", "anthropic_messages", "oci_genai", "gemini_generate_content"] + | str + | None ) = ... builtin: BuiltinConfig | None = ... local: LocalModelConfig | None = ... diff --git a/python/tests/test_builtin_codecs.py b/python/tests/test_builtin_codecs.py index 00458be45..ac6e65ddd 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -5,8 +5,8 @@ Covers: - Built-in codec construction for OpenAIChatCodec, OpenAIResponsesCodec, - AnthropicMessagesCodec, and GeminiGenerateContentCodec -- Built-in codec decode/encode/decode_response methods for all four providers + AnthropicMessagesCodec, OCIGenAIChatCodec, and GeminiGenerateContentCodec +- Built-in codec decode/encode/decode_response methods for all five providers - LlmResponseCodec protocol - response_codec parameter accepts object (not string) """ @@ -23,7 +23,13 @@ llm, subscribers, ) -from nemo_relay.codecs import AnthropicMessagesCodec, GeminiGenerateContentCodec, OpenAIChatCodec, OpenAIResponsesCodec +from nemo_relay.codecs import ( + AnthropicMessagesCodec, + GeminiGenerateContentCodec, + OCIGenAIChatCodec, + OpenAIChatCodec, + OpenAIResponsesCodec, +) # --------------------------------------------------------------------------- # 1. Built-in codec construction @@ -67,6 +73,18 @@ def test_anthropic_messages_codec_has_methods(self): assert hasattr(codec, "encode") assert hasattr(codec, "decode_response") + def test_oci_genai_chat_codec_constructable(self): + """OCIGenAIChatCodec() is constructable.""" + codec = OCIGenAIChatCodec() + assert codec is not None + + def test_oci_genai_chat_codec_has_methods(self): + """OCIGenAIChatCodec has decode, encode, decode_response methods.""" + codec = OCIGenAIChatCodec() + assert hasattr(codec, "decode") + assert hasattr(codec, "encode") + assert hasattr(codec, "decode_response") + def test_gemini_codec_constructable(self): """GeminiGenerateContentCodec() is constructable.""" codec = GeminiGenerateContentCodec() @@ -285,6 +303,65 @@ def test_anthropic_messages_decode_response(self): assert annotated.model == "claude-3-sonnet-20240229" assert annotated.response_text() == "Hello!" + def test_oci_genai_request_decode_encode_round_trip(self): + """OCIGenAIChatCodec decodes and re-encodes an OCI ChatDetails request.""" + codec = OCIGenAIChatCodec() + original = LLMRequest( + {}, + { + "compartmentId": "ocid1.compartment.oc1..example", + "servingMode": {"servingType": "ON_DEMAND", "modelId": "meta.llama-3.3-70b-instruct"}, + "chatRequest": { + "apiFormat": "GENERIC", + "messages": [{"role": "USER", "content": [{"type": "TEXT", "text": "My SSN is 111-22-3333."}]}], + "maxTokens": 600, + }, + }, + ) + annotated = codec.decode(original) + assert isinstance(annotated, AnnotatedLLMRequest) + assert annotated.model == "meta.llama-3.3-70b-instruct" + + # Identity: an unedited annotation re-encodes byte-identically. + identical = codec.encode(annotated, original) + assert identical.content == original.content + + annotated.messages = [ + {"role": "user", "content": "My SSN is [REDACTED]."}, + ] + encoded = codec.encode(annotated, original) + encoded_content = cast(JsonObject, encoded.content) + chat_request = cast(JsonObject, encoded_content["chatRequest"]) + messages = cast(list[JsonObject], chat_request["messages"]) + assert messages[0]["content"] == [{"type": "TEXT", "text": "My SSN is [REDACTED]."}] + assert cast(int, chat_request["maxTokens"]) == 600 + + def test_oci_genai_decode_response(self): + """OCIGenAIChatCodec.decode_response() returns AnnotatedLLMResponse.""" + codec = OCIGenAIChatCodec() + response = { + "modelId": "meta.llama-3.3-70b-instruct", + "chatResponse": { + "apiFormat": "GENERIC", + "choices": [ + { + "index": 0, + "message": { + "role": "ASSISTANT", + "content": [{"type": "TEXT", "text": "Hello!"}], + }, + "finishReason": "stop", + } + ], + "usage": {"promptTokens": 10, "completionTokens": 5, "totalTokens": 15}, + }, + } + annotated = codec.decode_response(response) + assert isinstance(annotated, AnnotatedLLMResponse) + assert annotated.model == "meta.llama-3.3-70b-instruct" + assert annotated.response_text() == "Hello!" + assert annotated.finish_reason == "complete" + def test_gemini_codec_decode(self): """GeminiGenerateContentCodec.decode() returns AnnotatedLLMRequest with messages and params.""" codec = GeminiGenerateContentCodec() @@ -418,6 +495,8 @@ def test_builtin_codecs_satisfy_protocol(self): assert isinstance(OpenAIChatCodec(), LlmResponseCodec) assert isinstance(OpenAIResponsesCodec(), LlmResponseCodec) assert isinstance(AnthropicMessagesCodec(), LlmResponseCodec) + assert isinstance(OCIGenAIChatCodec(), LlmResponseCodec) + assert isinstance(GeminiGenerateContentCodec(), LlmResponseCodec)