diff --git a/crates/core/src/codec/mod.rs b/crates/core/src/codec/mod.rs index a92f00826..cd688f907 100644 --- a/crates/core/src/codec/mod.rs +++ b/crates/core/src/codec/mod.rs @@ -16,6 +16,7 @@ pub mod anthropic; 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..82d6cfd21 --- /dev/null +++ b/crates/core/src/codec/oci_genai.rs @@ -0,0 +1,353 @@ +// 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 [`LlmResponseCodec`] (response decode) for the OCI Generative AI +//! chat format. +//! +//! # OCI-specific patterns handled +//! +//! - **Two API formats** selected by the response `apiFormat`: `GENERIC` +//! (`choices`-based, OpenAI-style) and `COHERE` (`text`-based). +//! - **Key conventions**: The OCI SDKs emit camelCase JSON while the OCI CLI +//! emits kebab-case; decode tolerates camelCase, kebab-case, and snake_case. +//! - **Responses**: `ChatResult` payloads (`modelId`, `chatResponse`) where the +//! chat response is `choices`-based for `GENERIC` and `text`-based for +//! `COHERE`; `usage` counters are `promptTokens`/`completionTokens`/`totalTokens`. + +use crate::error::{FlowError, Result}; +use crate::json::Json; + +use super::request::{ContentPart, MessageContent, ProviderNativeComponent}; +use super::response::{ + AnnotatedLlmResponse, ApiSpecificResponse, FinishReason, ResponseToolCall, Usage, +}; +use super::traits::LlmResponseCodec; + +// --------------------------------------------------------------------------- +// Public codec struct +// --------------------------------------------------------------------------- + +/// Built-in codec for the OCI Generative AI chat API. +pub struct OCIGenAIChatCodec; + +// --------------------------------------------------------------------------- +// Key-convention helpers +// --------------------------------------------------------------------------- + +/// Convert a camelCase key to its kebab-case form (`maxTokens` -> `max-tokens`). +fn camel_to_kebab(key: &str) -> String { + camel_with_separator(key, '-') +} + +/// Convert a camelCase key to its snake_case form (`maxTokens` -> `max_tokens`). +fn camel_to_snake(key: &str) -> String { + camel_with_separator(key, '_') +} + +fn camel_with_separator(key: &str, separator: char) -> String { + let mut out = String::with_capacity(key.len() + 4); + for c in key.chars() { + if c.is_ascii_uppercase() { + out.push(separator); + out.push(c.to_ascii_lowercase()); + } else { + out.push(c); + } + } + out +} + +/// Return the value for the first present key spelling across naming conventions. +/// +/// The OCI SDKs emit camelCase JSON while the OCI CLI emits kebab-case; callers +/// pass camelCase keys and kebab-case/snake_case fallbacks are derived. +fn get_first<'a>(obj: &'a serde_json::Map, key: &str) -> Option<&'a Json> { + present_key(obj, key).and_then(|present| obj.get(&present)) +} + +/// Return the concrete key spelling present in `obj` for a camelCase `key`. +fn present_key(obj: &serde_json::Map, key: &str) -> Option { + if obj.contains_key(key) { + return Some(key.to_string()); + } + let kebab = camel_to_kebab(key); + if obj.contains_key(&kebab) { + return Some(kebab); + } + let snake = camel_to_snake(key); + if obj.contains_key(&snake) { + return Some(snake); + } + None +} + +// --------------------------------------------------------------------------- +// Shared helpers +// --------------------------------------------------------------------------- + +/// Map an OCI finish reason string to normalized [`FinishReason`]. +/// +/// GENERIC responses use OpenAI-style lowercase reasons; COHERE responses use +/// UPPERCASE Cohere reasons. +fn map_oci_finish_reason(reason: &str) -> FinishReason { + match reason { + "stop" | "COMPLETE" => FinishReason::Complete, + "length" | "MAX_TOKENS" => FinishReason::Length, + "tool_calls" => FinishReason::ToolUse, + other => FinishReason::Unknown(other.to_string()), + } +} + +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 get_first(obj, "type").and_then(Json::as_str) != Some("TEXT") { + return None; + } + match get_first(obj, "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 get_first(obj, "type").and_then(Json::as_str) { + Some("TEXT") => Ok(ContentPart::Text { + text: get_first(obj, "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, + }) + } + } +} + +// --------------------------------------------------------------------------- +// LlmResponseCodec implementation +// --------------------------------------------------------------------------- + +impl LlmResponseCodec for OCIGenAIChatCodec { + 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 get_first(obj, "chatResponse").and_then(Json::as_object) { + Some(chat_response) => (Some(obj), chat_response), + None => (None, obj), + }; + + let model = envelope + .and_then(|envelope| get_first(envelope, "modelId")) + .and_then(Json::as_str) + .map(str::to_string); + let model_version = envelope + .and_then(|envelope| get_first(envelope, "modelVersion")) + .and_then(Json::as_str) + .map(str::to_string); + let api_format = get_first(chat_response, "apiFormat") + .and_then(Json::as_str) + .unwrap_or("GENERIC") + .to_uppercase(); + + let (message, tool_calls, finish_reason) = if api_format == "COHERE" { + decode_cohere_response_body(chat_response) + } else { + decode_generic_response_body(chat_response)? + }; + + let usage = get_first(chat_response, "usage") + .and_then(Json::as_object) + .map(decode_oci_usage); + + Ok(AnnotatedLlmResponse { + id: None, + 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: serde_json::Map::new(), + }) + } +} + +type ResponseBody = ( + Option, + Option>, + Option, +); + +fn decode_generic_response_body( + chat_response: &serde_json::Map, +) -> Result { + let Some(first_choice) = get_first(chat_response, "choices") + .and_then(Json::as_array) + .and_then(|choices| choices.first()) + .and_then(Json::as_object) + else { + return Ok((None, None, None)); + }; + let finish_reason = get_first(first_choice, "finishReason") + .and_then(Json::as_str) + .map(str::to_string); + let Some(raw_message) = get_first(first_choice, "message").and_then(Json::as_object) else { + return Ok((None, None, finish_reason)); + }; + let message = decode_generic_content(get_first(raw_message, "content"))?; + let tool_calls = get_first(raw_message, "toolCalls") + .and_then(Json::as_array) + .map(|calls| calls.iter().filter_map(decode_response_tool_call).collect()) + .filter(|calls: &Vec| !calls.is_empty()); + Ok((message, tool_calls, finish_reason)) +} + +fn decode_cohere_response_body(chat_response: &serde_json::Map) -> ResponseBody { + let message = get_first(chat_response, "text") + .and_then(Json::as_str) + .map(|text| MessageContent::Text(text.to_string())); + let tool_calls = get_first(chat_response, "toolCalls") + .and_then(Json::as_array) + .map(|calls| { + calls + .iter() + .filter_map(decode_response_tool_call) + .collect::>() + }) + .filter(|calls| !calls.is_empty()); + let finish_reason = get_first(chat_response, "finishReason") + .and_then(Json::as_str) + .map(str::to_string); + (message, tool_calls, finish_reason) +} + +/// 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`. +fn decode_response_tool_call(value: &Json) -> Option { + let obj = value.as_object()?; + let name = get_first(obj, "name")?.as_str()?.to_string(); + let arguments = match get_first(obj, "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 => get_first(obj, "parameters").cloned().unwrap_or(Json::Null), + }; + Some(ResponseToolCall { + id: get_first(obj, "id") + .and_then(Json::as_str) + .unwrap_or_default() + .to_string(), + name, + arguments, + }) +} + +/// Map OCI usage counters onto the normalized [`Usage`] field names. +fn decode_oci_usage(usage: &serde_json::Map) -> Usage { + Usage { + prompt_tokens: get_first(usage, "promptTokens").and_then(Json::as_u64), + completion_tokens: get_first(usage, "completionTokens").and_then(Json::as_u64), + total_tokens: get_first(usage, "totalTokens").and_then(Json::as_u64), + cache_read_tokens: None, + cache_write_tokens: None, + cost: None, + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +#[path = "../../tests/unit/codec/oci_genai_tests.rs"] +mod tests; 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..bbaf26818 --- /dev/null +++ b/crates/core/tests/unit/codec/oci_genai_tests.rs @@ -0,0 +1,247 @@ +// 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, MessageContent}; +use super::super::response::{ApiSpecificResponse, FinishReason}; + +// ------------------------------------------------------------------- +// Helpers and fixtures +// ------------------------------------------------------------------- + +const DEDICATED_ENDPOINT: &str = "ocid1.generativeaiendpoint.oc1.us-chicago-1.example"; + +/// 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} + } + }) +} + +// =================================================================== +// 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_kebab_case_cli_shape() { + let cli_shaped = json!({ + "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.as_deref(), Some(DEDICATED_ENDPOINT)); + assert_eq!( + annotated.message, + Some(MessageContent::Text("hello".into())) + ); + assert_eq!(annotated.finish_reason, Some(FinishReason::Complete)); + assert_eq!(annotated.usage.as_ref().unwrap().total_tokens, Some(9)); +} + +#[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_finish_reason_mapping() { + for (raw, expected) in [ + ("stop", FinishReason::Complete), + ("length", FinishReason::Length), + ("tool_calls", FinishReason::ToolUse), + ("COMPLETE", FinishReason::Complete), + ("MAX_TOKENS", FinishReason::Length), + ("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}"); + } +} diff --git a/crates/types/src/codec/request.rs b/crates/types/src/codec/request.rs index 1167b92e6..68bdca806 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` or `COHERE`). + #[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 731617966..17bc8752c 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -326,6 +326,17 @@ pub enum ApiSpecificResponse { content_blocks: Option>, }, + /// OCI Generative AI-specific fields. + #[serde(rename = "oci_genai")] + OCIGenAI { + /// Chat response API format (`GENERIC` or `COHERE`). + #[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, + }, + /// Custom/unknown API -- catch-all for user-implemented codecs. #[serde(rename = "custom")] Custom {