diff --git a/Cargo.lock b/Cargo.lock index 5927aa621..bc29d5356 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1519,6 +1519,7 @@ dependencies = [ "serde_json", "serde_json_canonicalizer", "sha2", + "strum", "tdigest", "thiserror 2.0.18", "tokio", diff --git a/crates/adaptive/Cargo.toml b/crates/adaptive/Cargo.toml index c59bf27e4..a44349dc6 100644 --- a/crates/adaptive/Cargo.toml +++ b/crates/adaptive/Cargo.toml @@ -25,6 +25,7 @@ tokio = { version = "1", features = ["rt", "macros", "sync", "time"] } regex = "1" serde_json_canonicalizer = "0.3" sha2 = "0.11" +strum = { version = "0.27", features = ["derive"] } redis = { version = "1.1", features = ["tokio-comp", "connection-manager"], optional = true } diff --git a/crates/adaptive/src/acg/request_surfaces/mod.rs b/crates/adaptive/src/acg/request_surfaces/mod.rs index 48172e4aa..922148fa0 100644 --- a/crates/adaptive/src/acg/request_surfaces/mod.rs +++ b/crates/adaptive/src/acg/request_surfaces/mod.rs @@ -18,12 +18,14 @@ use std::collections::HashSet; use nemo_relay::api::llm::LlmRequest; use nemo_relay::codec::resolve::{ProviderSurface, detect_request_surface}; use serde_json::Value; +use strum::AsRefStr; use crate::acg::prompt_ir::PromptIR; use crate::acg::translation::{HintPlan, HintTarget}; #[cfg_attr(not(test), allow(dead_code))] -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, AsRefStr)] +#[strum(serialize_all = "snake_case")] pub(crate) enum RequestSurface { AnthropicMessages, OpenAIChat, @@ -49,6 +51,14 @@ impl RequestSurface { } } + pub(crate) fn provider_surface(self) -> ProviderSurface { + match self { + Self::AnthropicMessages => ProviderSurface::AnthropicMessages, + Self::OpenAIChat => ProviderSurface::OpenAIChat, + Self::OpenAIResponses => ProviderSurface::OpenAIResponses, + } + } + pub(crate) fn supports_provider(self, provider: &str) -> bool { match provider { "anthropic" => matches!(self, Self::AnthropicMessages), diff --git a/crates/adaptive/src/acg_component.rs b/crates/adaptive/src/acg_component.rs index 3d344866a..443e376ff 100644 --- a/crates/adaptive/src/acg_component.rs +++ b/crates/adaptive/src/acg_component.rs @@ -28,11 +28,8 @@ use nemo_relay::api::runtime::{ LlmExecutionFn, LlmExecutionNextFn, LlmRequestInterceptFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, }; -use nemo_relay::codec::anthropic::AnthropicMessagesCodec; -use nemo_relay::codec::openai_chat::OpenAIChatCodec; -use nemo_relay::codec::openai_responses::OpenAIResponsesCodec; use nemo_relay::codec::request::AnnotatedLlmRequest; -use nemo_relay::codec::traits::LlmCodec; +use nemo_relay::codec::resolve::request_codec; use nemo_relay::json::Json; use uuid::Uuid; @@ -91,21 +88,12 @@ fn decode_request_for_surface( request_surface: RequestSurface, request: &LlmRequest, ) -> Result { - match request_surface { - RequestSurface::AnthropicMessages => { - AnthropicMessagesCodec.decode(request).map_err(|error| { - AdaptiveError::Internal(format!("failed to decode anthropic request: {error}")) - }) - } - RequestSurface::OpenAIChat => OpenAIChatCodec.decode(request).map_err(|error| { - AdaptiveError::Internal(format!("failed to decode openai chat request: {error}")) - }), - RequestSurface::OpenAIResponses => OpenAIResponsesCodec.decode(request).map_err(|error| { - AdaptiveError::Internal(format!( - "failed to decode openai responses request: {error}" - )) - }), - } + request_codec(request_surface.provider_surface()) + .decode(request) + .map_err(|error| { + let surface_label: &str = request_surface.as_ref(); + AdaptiveError::Internal(format!("failed to decode {surface_label} request: {error}")) + }) } fn build_semantic_request_view(request: &LlmRequest) -> Result { diff --git a/crates/adaptive/tests/unit/acg_component_tests.rs b/crates/adaptive/tests/unit/acg_component_tests.rs index 3a0194fbb..04e601381 100644 --- a/crates/adaptive/tests/unit/acg_component_tests.rs +++ b/crates/adaptive/tests/unit/acg_component_tests.rs @@ -1124,15 +1124,15 @@ fn acg_component_decode_request_for_surface_reports_codec_specific_errors() { assert!(matches!( decode_request_for_surface(RequestSurface::AnthropicMessages, &invalid_request), - Err(AdaptiveError::Internal(message)) if message.contains("failed to decode anthropic request") + Err(AdaptiveError::Internal(message)) if message.contains("failed to decode anthropic_messages request") )); assert!(matches!( decode_request_for_surface(RequestSurface::OpenAIChat, &invalid_request), - Err(AdaptiveError::Internal(message)) if message.contains("failed to decode openai chat request") + Err(AdaptiveError::Internal(message)) if message.contains("failed to decode open_ai_chat request") )); assert!(matches!( decode_request_for_surface(RequestSurface::OpenAIResponses, &invalid_request), - Err(AdaptiveError::Internal(message)) if message.contains("failed to decode openai responses request") + Err(AdaptiveError::Internal(message)) if message.contains("failed to decode open_ai_responses request") )); } diff --git a/crates/cli/src/gateway.rs b/crates/cli/src/gateway.rs index ddac6837b..c839eca6d 100644 --- a/crates/cli/src/gateway.rs +++ b/crates/cli/src/gateway.rs @@ -17,9 +17,10 @@ use nemo_relay::api::llm::{ use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, }; -use nemo_relay::codec::anthropic::{AnthropicMessagesCodec, AnthropicMessagesStreamingCodec}; -use nemo_relay::codec::openai_chat::{OpenAIChatCodec, OpenAIChatStreamingCodec}; -use nemo_relay::codec::openai_responses::{OpenAIResponsesCodec, OpenAIResponsesStreamingCodec}; +use nemo_relay::codec::resolve::{ + ProviderSurface, response_codec as build_response_codec, + streaming_codec as build_streaming_codec, +}; use nemo_relay::codec::streaming::StreamingCodec; use nemo_relay::codec::traits::LlmResponseCodec; use nemo_relay::error::FlowError; @@ -248,20 +249,12 @@ struct RouteCodecs { } fn codecs_for_route(route: ProviderRoute) -> RouteCodecs { - match route { - ProviderRoute::AnthropicMessages => RouteCodecs { - streaming: Some(Box::new(AnthropicMessagesStreamingCodec::new())), - response: Some(Arc::new(AnthropicMessagesCodec) as Arc), - }, - ProviderRoute::OpenAiResponses => RouteCodecs { - streaming: Some(Box::new(OpenAIResponsesStreamingCodec::new())), - response: Some(Arc::new(OpenAIResponsesCodec) as Arc), + match route.provider_surface() { + Some(surface) => RouteCodecs { + streaming: Some(build_streaming_codec(surface)), + response: Some(Arc::from(build_response_codec(surface))), }, - ProviderRoute::OpenAiChatCompletions => RouteCodecs { - streaming: Some(Box::new(OpenAIChatStreamingCodec::new())), - response: Some(Arc::new(OpenAIChatCodec) as Arc), - }, - ProviderRoute::AnthropicCountTokens | ProviderRoute::OpenAiModels => RouteCodecs { + None => RouteCodecs { streaming: None, response: None, }, @@ -955,6 +948,15 @@ impl ProviderRoute { } } + const fn provider_surface(self) -> Option { + match self { + Self::OpenAiResponses => Some(ProviderSurface::OpenAIResponses), + Self::OpenAiChatCompletions => Some(ProviderSurface::OpenAIChat), + Self::AnthropicMessages => Some(ProviderSurface::AnthropicMessages), + Self::AnthropicCountTokens | Self::OpenAiModels => None, + } + } + // Returns the provider route name recorded on managed LLM events. These names split OpenAI API // variants because their request/response schemas differ even when they share a base URL, and // they double as codec hints for ambiguous provider request shapes. diff --git a/crates/core/src/codec/anthropic.rs b/crates/core/src/codec/anthropic.rs index c82d22925..64456b49e 100644 --- a/crates/core/src/codec/anthropic.rs +++ b/crates/core/src/codec/anthropic.rs @@ -56,6 +56,10 @@ pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDe }, decode_request: |request| AnthropicMessagesCodec.decode(request), decode_response: |raw| AnthropicMessagesCodec.decode_response(raw), + codec_name: "anthropic_messages", + request_codec: || Box::new(AnthropicMessagesCodec), + response_codec: || Box::new(AnthropicMessagesCodec), + streaming_codec: || Box::new(AnthropicMessagesStreamingCodec::new()), }; // --------------------------------------------------------------------------- diff --git a/crates/core/src/codec/openai_chat.rs b/crates/core/src/codec/openai_chat.rs index 2c7c66c7e..2b7dfdce8 100644 --- a/crates/core/src/codec/openai_chat.rs +++ b/crates/core/src/codec/openai_chat.rs @@ -33,6 +33,10 @@ pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDe detect_response: |obj| obj.get("choices").is_some_and(Json::is_array), decode_request: |request| OpenAIChatCodec.decode(request), decode_response: |raw| OpenAIChatCodec.decode_response(raw), + codec_name: "openai_chat", + request_codec: || Box::new(OpenAIChatCodec), + response_codec: || Box::new(OpenAIChatCodec), + streaming_codec: || Box::new(OpenAIChatStreamingCodec::new()), }; // --------------------------------------------------------------------------- diff --git a/crates/core/src/codec/openai_responses.rs b/crates/core/src/codec/openai_responses.rs index 57b05661c..d3adf42cf 100644 --- a/crates/core/src/codec/openai_responses.rs +++ b/crates/core/src/codec/openai_responses.rs @@ -48,6 +48,10 @@ pub(crate) const PROVIDER_SURFACE: ProviderSurfaceDescriptor = ProviderSurfaceDe }, decode_request: |request| OpenAIResponsesCodec.decode(request), decode_response: |raw| OpenAIResponsesCodec.decode_response(raw), + codec_name: "openai_responses", + request_codec: || Box::new(OpenAIResponsesCodec), + response_codec: || Box::new(OpenAIResponsesCodec), + streaming_codec: || Box::new(OpenAIResponsesStreamingCodec::new()), }; // --------------------------------------------------------------------------- diff --git a/crates/core/src/codec/resolve.rs b/crates/core/src/codec/resolve.rs index 1769ea100..7a815ee28 100644 --- a/crates/core/src/codec/resolve.rs +++ b/crates/core/src/codec/resolve.rs @@ -1,9 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Provider-surface detection and best-effort normalization: the preferred path -//! for turning raw provider JSON into normalized types when no codec annotation -//! is present. +//! Provider-surface detection, best-effort normalization, and construction of +//! the matching built-in codecs from a raw payload, surface, or codec name. use crate::api::llm::LlmRequest; use crate::error::Result; @@ -11,6 +10,8 @@ use crate::json::Json; use super::request::AnnotatedLlmRequest; use super::response::AnnotatedLlmResponse; +use super::streaming::StreamingCodec; +use super::traits::{LlmCodec, LlmResponseCodec}; use super::{anthropic, openai_chat, openai_responses}; /// A built-in provider request/response surface. @@ -48,6 +49,10 @@ pub(crate) struct ProviderSurfaceDescriptor { pub(crate) detect_response: ResponseSurfaceDetector, pub(crate) decode_request: fn(&LlmRequest) -> Result, pub(crate) decode_response: fn(&Json) -> Result, + pub(crate) codec_name: &'static str, + pub(crate) request_codec: fn() -> Box, + pub(crate) response_codec: fn() -> Box, + pub(crate) streaming_codec: fn() -> Box, } /// Built-in provider surfaces in request-detection priority order. @@ -140,6 +145,64 @@ pub fn normalize_response(raw: &Json) -> Option { (descriptor.decode_response)(raw).ok() } +fn descriptor_for(surface: ProviderSurface) -> &'static ProviderSurfaceDescriptor { + match surface { + ProviderSurface::OpenAIChat => &openai_chat::PROVIDER_SURFACE, + ProviderSurface::OpenAIResponses => &openai_responses::PROVIDER_SURFACE, + ProviderSurface::AnthropicMessages => &anthropic::PROVIDER_SURFACE, + } +} + +impl ProviderSurface { + /// The canonical codec name for this surface (e.g. `"openai_chat"`), the + /// inverse of [`Self::from_codec_name`]. + #[must_use] + pub fn codec_name(self) -> &'static str { + descriptor_for(self).codec_name + } + + /// Resolves a canonical codec name to its surface, or `None` when `name` is + /// not a built-in provider codec. + #[must_use] + pub fn from_codec_name(name: &str) -> Option { + BUILTIN_PROVIDER_SURFACES + .iter() + .find(|descriptor| descriptor.codec_name == name) + .map(|descriptor| descriptor.surface) + } +} + +/// The canonical codec names of every built-in provider surface, for config +/// validation and "supported codec" messages. +#[must_use] +pub fn supported_codec_names() -> Vec<&'static str> { + BUILTIN_PROVIDER_SURFACES + .iter() + .map(|descriptor| descriptor.codec_name) + .collect() +} + +/// Constructs the built-in bidirectional request codec ([`LlmCodec`]) for a surface. +#[must_use] +pub fn request_codec(surface: ProviderSurface) -> Box { + (descriptor_for(surface).request_codec)() +} + +/// Constructs the built-in decode-only response codec ([`LlmResponseCodec`]) for a surface. +#[must_use] +pub fn response_codec(surface: ProviderSurface) -> Box { + (descriptor_for(surface).response_codec)() +} + +/// Constructs a fresh, single-use streaming codec ([`StreamingCodec`]) for a surface. +/// +/// A [`StreamingCodec`] finalizer consumes its accumulator, so callers must +/// construct one instance per managed streaming call. +#[must_use] +pub fn streaming_codec(surface: ProviderSurface) -> Box { + (descriptor_for(surface).streaming_codec)() +} + #[cfg(test)] #[path = "../../tests/unit/codec/resolve_tests.rs"] mod tests; diff --git a/crates/core/src/plugins/nemo_guardrails/component.rs b/crates/core/src/plugins/nemo_guardrails/component.rs index 7ece7fa3d..3b6f30e90 100644 --- a/crates/core/src/plugins/nemo_guardrails/component.rs +++ b/crates/core/src/plugins/nemo_guardrails/component.rs @@ -11,6 +11,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; +use crate::codec::resolve::supported_codec_names; use crate::plugin::{ ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError, PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior, deregister_plugin, @@ -901,10 +902,7 @@ fn validate_codec_requirements( return; }; - if !matches!( - codec, - "openai_chat" | "openai_responses" | "anthropic_messages" - ) { + if !supported_codec_names().contains(&codec) { push_policy_diag( diagnostics, policy.unsupported_value, diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs index d52200030..63784d102 100644 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ b/crates/core/src/plugins/nemo_guardrails/python.rs @@ -20,11 +20,8 @@ use tokio_stream::wrappers::ReceiverStream; use crate::api::llm::LlmRequest; use crate::api::runtime::{LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn, ToolExecutionFn}; -use crate::codec::anthropic::AnthropicMessagesCodec; -use crate::codec::openai_chat::OpenAIChatCodec; -use crate::codec::openai_responses::OpenAIResponsesCodec; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; -use crate::codec::traits::{LlmCodec, LlmResponseCodec}; +use crate::codec::resolve::{ProviderSurface, request_codec, response_codec}; use crate::error::{FlowError, Result as FlowResult}; use crate::json::Json; use crate::plugin::{PluginError, PluginRegistrationContext, Result as PluginResult}; @@ -815,35 +812,39 @@ enum LocalGuardrailsCodec { } impl LocalGuardrailsCodec { - fn decode(&self, request: &LlmRequest) -> FlowResult { + fn provider_surface(self) -> ProviderSurface { match self { - Self::OpenAIChat => OpenAIChatCodec.decode(request), - Self::OpenAIResponses => OpenAIResponsesCodec.decode(request), - Self::AnthropicMessages => AnthropicMessagesCodec.decode(request), + Self::OpenAIChat => ProviderSurface::OpenAIChat, + Self::OpenAIResponses => ProviderSurface::OpenAIResponses, + Self::AnthropicMessages => ProviderSurface::AnthropicMessages, + } + } + + fn from_provider_surface(surface: ProviderSurface) -> Self { + match surface { + ProviderSurface::OpenAIChat => Self::OpenAIChat, + ProviderSurface::OpenAIResponses => Self::OpenAIResponses, + ProviderSurface::AnthropicMessages => Self::AnthropicMessages, } } + fn decode(&self, request: &LlmRequest) -> FlowResult { + request_codec(self.provider_surface()).decode(request) + } + fn encode( &self, annotated: &AnnotatedLlmRequest, original: &LlmRequest, ) -> FlowResult { - match self { - Self::OpenAIChat => OpenAIChatCodec.encode(annotated, original), - Self::OpenAIResponses => OpenAIResponsesCodec.encode(annotated, original), - Self::AnthropicMessages => AnthropicMessagesCodec.encode(annotated, original), - } + request_codec(self.provider_surface()).encode(annotated, original) } fn decode_response( &self, response: &Json, ) -> FlowResult { - match self { - Self::OpenAIChat => OpenAIChatCodec.decode_response(response), - Self::OpenAIResponses => OpenAIResponsesCodec.decode_response(response), - Self::AnthropicMessages => AnthropicMessagesCodec.decode_response(response), - } + response_codec(self.provider_surface()).decode_response(response) } } @@ -853,12 +854,12 @@ fn resolve_codec(config: &NeMoGuardrailsConfig) -> PluginResult Ok(Some(LocalGuardrailsCodec::OpenAIChat)), - Some("openai_responses") => Ok(Some(LocalGuardrailsCodec::OpenAIResponses)), - Some("anthropic_messages") => Ok(Some(LocalGuardrailsCodec::AnthropicMessages)), - Some(other) => Err(PluginError::InvalidConfig(format!( - "unsupported local NeMo Guardrails codec '{other}'" - ))), + Some(name) => match ProviderSurface::from_codec_name(name) { + Some(surface) => Ok(Some(LocalGuardrailsCodec::from_provider_surface(surface))), + None => Err(PluginError::InvalidConfig(format!( + "unsupported local NeMo Guardrails codec '{name}'" + ))), + }, None => Err(PluginError::InvalidConfig( "local NeMo Guardrails backend requires a supported codec".to_string(), )), diff --git a/crates/core/tests/unit/codec/resolve_tests.rs b/crates/core/tests/unit/codec/resolve_tests.rs index c1460cdbe..70907a933 100644 --- a/crates/core/tests/unit/codec/resolve_tests.rs +++ b/crates/core/tests/unit/codec/resolve_tests.rs @@ -390,3 +390,173 @@ fn hint_does_not_classify_non_object_or_keyless() { None ); } + +// --------------------------------------------------------------------------- +// Provider-codec factory (name<->surface mapping + codec construction) +// --------------------------------------------------------------------------- + +const ALL_SURFACES: [ProviderSurface; 3] = [ + ProviderSurface::OpenAIChat, + ProviderSurface::OpenAIResponses, + ProviderSurface::AnthropicMessages, +]; + +#[test] +fn codec_name_round_trips_for_every_surface() { + for surface in ALL_SURFACES { + assert_eq!( + ProviderSurface::from_codec_name(surface.codec_name()), + Some(surface), + "codec_name/from_codec_name must round-trip for {surface:?}", + ); + } +} + +#[test] +fn codec_name_uses_canonical_spellings() { + assert_eq!(ProviderSurface::OpenAIChat.codec_name(), "openai_chat"); + assert_eq!( + ProviderSurface::OpenAIResponses.codec_name(), + "openai_responses" + ); + assert_eq!( + ProviderSurface::AnthropicMessages.codec_name(), + "anthropic_messages" + ); +} + +#[test] +fn from_codec_name_is_none_for_unknown_names() { + assert_eq!(ProviderSurface::from_codec_name("gemini"), None); + assert_eq!(ProviderSurface::from_codec_name(""), None); + assert_eq!(ProviderSurface::from_codec_name("OpenAIChat"), None); +} + +#[test] +fn supported_codec_names_track_the_builtin_registry() { + assert_eq!( + supported_codec_names(), + vec!["openai_responses", "anthropic_messages", "openai_chat"] + ); + let from_registry: Vec<_> = BUILTIN_PROVIDER_SURFACES + .iter() + .map(|descriptor| descriptor.codec_name) + .collect(); + assert_eq!(supported_codec_names(), from_registry); +} + +#[test] +fn request_codec_decodes_each_surface() { + let chat = req(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}] + })); + assert!( + !request_codec(ProviderSurface::OpenAIChat) + .decode(&chat) + .expect("chat request decodes") + .messages + .is_empty() + ); + + let anthropic = req(json!({ + "model": "claude-3-5-sonnet", + "system": "be terse", + "messages": [{"role": "user", "content": "hi"}] + })); + assert!( + !request_codec(ProviderSurface::AnthropicMessages) + .decode(&anthropic) + .expect("anthropic request decodes") + .messages + .is_empty() + ); + + let responses = req(json!({"model": "gpt-4o", "input": "hi"})); + assert!( + !request_codec(ProviderSurface::OpenAIResponses) + .decode(&responses) + .expect("responses request decodes") + .messages + .is_empty() + ); +} + +#[test] +fn response_codec_decodes_each_surface() { + let chat = json!({ + "choices": [{ + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }] + }); + assert_eq!( + response_codec(ProviderSurface::OpenAIChat) + .decode_response(&chat) + .expect("chat response decodes") + .response_text(), + Some("hello") + ); + + let anthropic = json!({ + "type": "message", + "content": [{"type": "text", "text": "hey"}], + "stop_reason": "end_turn" + }); + assert_eq!( + response_codec(ProviderSurface::AnthropicMessages) + .decode_response(&anthropic) + .expect("anthropic response decodes") + .response_text(), + Some("hey") + ); + + let responses = json!({"output": [], "output_text": "yo"}); + assert_eq!( + response_codec(ProviderSurface::OpenAIResponses) + .decode_response(&responses) + .expect("responses response decodes") + .response_text(), + Some("yo") + ); +} + +#[test] +fn streaming_codec_round_trips_through_its_response_codec() { + let codec = streaming_codec(ProviderSurface::OpenAIChat); + let mut collect = codec.collector(); + collect(json!({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "gpt-4o", + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": null}] + })) + .unwrap(); + for part in ["Hello, ", "world", "."] { + collect(json!({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": part}, "finish_reason": null}] + })) + .unwrap(); + } + collect(json!({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}] + })) + .unwrap(); + + let assembled = codec.finalizer()(); + let decoded = response_codec(ProviderSurface::OpenAIChat) + .decode_response(&assembled) + .expect("assembled stream decodes"); + assert_eq!(decoded.response_text(), Some("Hello, world.")); +} + +#[test] +fn streaming_codec_constructs_a_usable_codec_for_every_surface() { + for surface in ALL_SURFACES { + let assembled = streaming_codec(surface).finalizer()(); + assert!( + assembled.is_object(), + "{surface:?} streaming codec finalizes to a JSON object", + ); + } +} diff --git a/crates/pii-redaction/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index 62646a51c..eb27b4033 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -14,9 +14,9 @@ use nemo_relay::api::llm::LlmRequest; use nemo_relay::api::runtime::{ EventSanitizeFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, ToolSanitizeFn, }; -use nemo_relay::codec::anthropic::AnthropicMessagesCodec; -use nemo_relay::codec::openai_chat::OpenAIChatCodec; -use nemo_relay::codec::openai_responses::OpenAIResponsesCodec; +use nemo_relay::codec::resolve::{ + ProviderSurface, request_codec as build_request_codec, response_codec as build_response_codec, +}; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::plugin::{PluginError, Result as PluginResult}; @@ -28,7 +28,8 @@ use super::overlay::BuiltinCodecName; pub(super) struct CompiledBuiltinBackend { action: BuiltinAction, target_paths: Arc>, - codec: Option>, + request_codec: Option>, + response_codec: Option>, codec_name: Option, } @@ -65,10 +66,6 @@ enum BuiltinMaskStrategy { }, } -trait BuiltinRequestResponseCodec: LlmCodec + LlmResponseCodec + Send + Sync {} - -impl BuiltinRequestResponseCodec for T where T: LlmCodec + LlmResponseCodec + Send + Sync {} - impl CompiledBuiltinBackend { pub(super) fn new( config: BuiltinBackendConfig, @@ -117,14 +114,19 @@ impl CompiledBuiltinBackend { } }; + let surface = match codec_name.as_deref() { + Some(name) => Some(ProviderSurface::from_codec_name(name).ok_or_else(|| { + PluginError::InvalidConfig(format!("unsupported codec '{name}'")) + })?), + None => None, + }; + Ok(Self { action, target_paths: Arc::new(config.target_paths), - codec_name: codec_name.as_deref().and_then(BuiltinCodecName::parse), - codec: codec_name - .as_deref() - .map(instantiate_builtin_codec) - .transpose()?, + request_codec: surface.map(build_request_codec).map(Arc::from), + response_codec: surface.map(build_response_codec).map(Arc::from), + codec_name: surface.map(BuiltinCodecName::from_provider_surface), }) } @@ -240,14 +242,14 @@ impl CompiledBuiltinBackend { } fn sanitize_request_with_codec(&self, request: &LlmRequest) -> Option { - let codec = self.codec.as_ref()?; + let codec = self.request_codec.as_ref()?; let annotated = codec.decode(request).ok()?; let sanitized_annotated = sanitize_serializable_with_backend(self, annotated).ok()?; codec.encode(&sanitized_annotated, request).ok() } fn sanitize_response_with_codec(&self, payload: Json) -> Option { - let codec = self.codec.as_ref()?; + let codec = self.response_codec.as_ref()?; let codec_name = self.codec_name?; let annotated = codec.decode_response(&payload).ok()?; let sanitized_annotated = sanitize_serializable_with_backend(self, annotated).ok()?; @@ -421,22 +423,6 @@ fn compile_builtin_matcher( Ok(Some(Arc::new(pattern))) } -fn instantiate_builtin_codec( - codec_name: &str, -) -> PluginResult> { - let codec: Arc = match codec_name { - "openai_chat" => Arc::new(OpenAIChatCodec), - "openai_responses" => Arc::new(OpenAIResponsesCodec), - "anthropic_messages" => Arc::new(AnthropicMessagesCodec), - other => { - return Err(PluginError::InvalidConfig(format!( - "unsupported codec '{other}'" - ))); - } - }; - Ok(codec) -} - fn sanitize_serializable_with_backend( backend: &CompiledBuiltinBackend, value: T, diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index bb07fd7bb..a68dd8c70 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -7,6 +7,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use nemo_relay::codec::resolve::supported_codec_names; use nemo_relay::plugin::{ ConfigDiagnostic, ConfigPolicy, DiagnosticLevel, Plugin, PluginComponentSpec, PluginError, PluginRegistrationContext, Result as PluginResult, UnsupportedBehavior, deregister_plugin, @@ -707,10 +708,7 @@ fn validate_codec_requirements( return; }; - if !matches!( - codec, - "openai_chat" | "openai_responses" | "anthropic_messages" - ) { + if !supported_codec_names().contains(&codec) { push_policy_diag( diagnostics, policy.unsupported_value, diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index efcce94f3..0f97caef4 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -4,6 +4,7 @@ use serde_json::{Map, Value as Json}; use nemo_relay::codec::request::{ContentPart, MessageContent}; +use nemo_relay::codec::resolve::ProviderSurface; use nemo_relay::codec::response::{AnnotatedLlmResponse, FinishReason, ResponseToolCall}; #[derive(Clone, Copy)] @@ -14,12 +15,11 @@ pub(crate) enum BuiltinCodecName { } impl BuiltinCodecName { - pub(crate) fn parse(value: &str) -> Option { - match value { - "openai_chat" => Some(Self::OpenAIChat), - "openai_responses" => Some(Self::OpenAIResponses), - "anthropic_messages" => Some(Self::AnthropicMessages), - _ => None, + pub(crate) fn from_provider_surface(surface: ProviderSurface) -> Self { + match surface { + ProviderSurface::OpenAIChat => Self::OpenAIChat, + ProviderSurface::OpenAIResponses => Self::OpenAIResponses, + ProviderSurface::AnthropicMessages => Self::AnthropicMessages, } }