Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/adaptive/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }

Expand Down
12 changes: 11 additions & 1 deletion crates/adaptive/src/acg/request_surfaces/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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),
Expand Down
26 changes: 7 additions & 19 deletions crates/adaptive/src/acg_component.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -91,21 +88,12 @@ fn decode_request_for_surface(
request_surface: RequestSurface,
request: &LlmRequest,
) -> Result<AnnotatedLlmRequest> {
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<SemanticRequestView> {
Expand Down
6 changes: 3 additions & 3 deletions crates/adaptive/tests/unit/acg_component_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")
));
}

Expand Down
34 changes: 18 additions & 16 deletions crates/cli/src/gateway.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<dyn LlmResponseCodec>),
},
ProviderRoute::OpenAiResponses => RouteCodecs {
streaming: Some(Box::new(OpenAIResponsesStreamingCodec::new())),
response: Some(Arc::new(OpenAIResponsesCodec) as Arc<dyn LlmResponseCodec>),
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<dyn LlmResponseCodec>),
},
ProviderRoute::AnthropicCountTokens | ProviderRoute::OpenAiModels => RouteCodecs {
None => RouteCodecs {
streaming: None,
response: None,
},
Expand Down Expand Up @@ -955,6 +948,15 @@ impl ProviderRoute {
}
}

const fn provider_surface(self) -> Option<ProviderSurface> {
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.
Expand Down
4 changes: 4 additions & 0 deletions crates/core/src/codec/anthropic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()),
};

// ---------------------------------------------------------------------------
Expand Down
4 changes: 4 additions & 0 deletions crates/core/src/codec/openai_chat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()),
};

// ---------------------------------------------------------------------------
Expand Down
4 changes: 4 additions & 0 deletions crates/core/src/codec/openai_responses.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()),
};

// ---------------------------------------------------------------------------
Expand Down
69 changes: 66 additions & 3 deletions crates/core/src/codec/resolve.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
// 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;
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.
Expand Down Expand Up @@ -48,6 +49,10 @@ pub(crate) struct ProviderSurfaceDescriptor {
pub(crate) detect_response: ResponseSurfaceDetector,
pub(crate) decode_request: fn(&LlmRequest) -> Result<AnnotatedLlmRequest>,
pub(crate) decode_response: fn(&Json) -> Result<AnnotatedLlmResponse>,
pub(crate) codec_name: &'static str,
pub(crate) request_codec: fn() -> Box<dyn LlmCodec>,
pub(crate) response_codec: fn() -> Box<dyn LlmResponseCodec>,
pub(crate) streaming_codec: fn() -> Box<dyn StreamingCodec>,
Comment thread
zhongxuanwang-nv marked this conversation as resolved.
Resolved
}

/// Built-in provider surfaces in request-detection priority order.
Expand Down Expand Up @@ -140,6 +145,64 @@ pub fn normalize_response(raw: &Json) -> Option<AnnotatedLlmResponse> {
(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<Self> {
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<dyn LlmCodec> {
(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<dyn LlmResponseCodec> {
(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<dyn StreamingCodec> {
(descriptor_for(surface).streaming_codec)()
}

#[cfg(test)]
#[path = "../../tests/unit/codec/resolve_tests.rs"]
mod tests;
6 changes: 2 additions & 4 deletions crates/core/src/plugins/nemo_guardrails/component.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
49 changes: 25 additions & 24 deletions crates/core/src/plugins/nemo_guardrails/python.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -815,35 +812,39 @@ enum LocalGuardrailsCodec {
}

impl LocalGuardrailsCodec {
fn decode(&self, request: &LlmRequest) -> FlowResult<AnnotatedLlmRequest> {
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<AnnotatedLlmRequest> {
request_codec(self.provider_surface()).decode(request)
}

fn encode(
&self,
annotated: &AnnotatedLlmRequest,
original: &LlmRequest,
) -> FlowResult<LlmRequest> {
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<crate::codec::response::AnnotatedLlmResponse> {
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)
}
}

Expand All @@ -853,12 +854,12 @@ fn resolve_codec(config: &NeMoGuardrailsConfig) -> PluginResult<Option<LocalGuar
}

match config.codec.as_deref() {
Some("openai_chat") => 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(),
)),
Expand Down
Loading
Loading