diff --git a/Cargo.toml b/Cargo.toml index c62ce5393..4d711cf0a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,3 +45,6 @@ uuid = "=1.18.1" [workspace.lints.rust] missing_docs = "deny" + +[workspace.lints.clippy] +cognitive_complexity = "deny" diff --git a/clippy.toml b/clippy.toml new file mode 100644 index 000000000..e7c65218e --- /dev/null +++ b/clippy.toml @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +cognitive-complexity-threshold = 18 diff --git a/crates/adaptive/src/plugin_component.rs b/crates/adaptive/src/plugin_component.rs index eae7759d1..1d1b530a9 100644 --- a/crates/adaptive/src/plugin_component.rs +++ b/crates/adaptive/src/plugin_component.rs @@ -215,37 +215,7 @@ fn validate_adaptive_plugin_config_with_policy( ); } - if let Some(state_json) = plugin_config.get("state").and_then(Json::as_object) { - validate_unknown_fields( - &mut diagnostics, - &config.policy, - Some("state".to_string()), - state_json, - &["backend"], - ); - if let Some(backend_json) = state_json.get("backend").and_then(Json::as_object) { - validate_unknown_fields( - &mut diagnostics, - &config.policy, - Some("backend".to_string()), - backend_json, - &["kind", "config"], - ); - let backend_kind = backend_json - .get("kind") - .and_then(Json::as_str) - .unwrap_or_default(); - if let Some(backend_config_json) = backend_json.get("config").and_then(Json::as_object) - { - validate_backend_config_fields( - &mut diagnostics, - &config.policy, - backend_kind, - backend_config_json, - ); - } - } - } + validate_adaptive_state_section(&mut diagnostics, &config.policy, plugin_config); if let Some(telemetry_json) = plugin_config.get("telemetry").and_then(Json::as_object) { validate_unknown_fields( @@ -303,52 +273,94 @@ fn validate_adaptive_plugin_config_with_policy( ); } - if let Some(response_cache_json) = plugin_config + validate_response_cache_section(&mut diagnostics, &config.policy, plugin_config); + + diagnostics.extend(AdaptiveRuntime::validate_config(&config).diagnostics); + diagnostics +} + +fn validate_adaptive_state_section( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + plugin_config: &Map, +) { + let Some(state_json) = plugin_config.get("state").and_then(Json::as_object) else { + return; + }; + validate_unknown_fields( + diagnostics, + policy, + Some("state".to_string()), + state_json, + &["backend"], + ); + let Some(backend_json) = state_json.get("backend").and_then(Json::as_object) else { + return; + }; + validate_unknown_fields( + diagnostics, + policy, + Some("backend".to_string()), + backend_json, + &["kind", "config"], + ); + let backend_kind = backend_json + .get("kind") + .and_then(Json::as_str) + .unwrap_or_default(); + if let Some(backend_config_json) = backend_json.get("config").and_then(Json::as_object) { + validate_backend_config_fields(diagnostics, policy, backend_kind, backend_config_json); + } +} + +fn validate_response_cache_section( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + plugin_config: &Map, +) { + let Some(response_cache_json) = plugin_config .get("response_cache") .and_then(Json::as_object) - { + else { + return; + }; + validate_unknown_fields( + diagnostics, + policy, + Some("response_cache".to_string()), + response_cache_json, + &[ + "ttl_seconds", + "namespace", + "priority", + "bypass_rate", + "cache_nondeterministic", + "key_strategy", + "header_allowlist", + "backend", + ], + ); + if let Some(backend_json) = response_cache_json.get("backend").and_then(Json::as_object) { validate_unknown_fields( - &mut diagnostics, - &config.policy, - Some("response_cache".to_string()), - response_cache_json, - &[ - "ttl_seconds", - "namespace", - "priority", - "bypass_rate", - "cache_nondeterministic", - "key_strategy", - "header_allowlist", - "backend", - ], + diagnostics, + policy, + Some("response_cache.backend".to_string()), + backend_json, + &["kind", "config"], ); - if let Some(backend_json) = response_cache_json.get("backend").and_then(Json::as_object) { - validate_unknown_fields( - &mut diagnostics, - &config.policy, - Some("response_cache.backend".to_string()), - backend_json, - &["kind", "config"], + let backend_kind = backend_json + .get("kind") + .and_then(Json::as_str) + .unwrap_or("in_memory"); + if let Some(backend_config_json) = backend_json.get("config").and_then(Json::as_object) { + validate_response_cache_backend_config_fields( + diagnostics, + policy, + backend_kind, + backend_config_json, ); - let backend_kind = backend_json - .get("kind") - .and_then(Json::as_str) - .unwrap_or("in_memory"); - if let Some(backend_config_json) = backend_json.get("config").and_then(Json::as_object) - { - validate_response_cache_backend_config_fields( - &mut diagnostics, - &config.policy, - backend_kind, - backend_config_json, - ); - } } } - - diagnostics.extend(AdaptiveRuntime::validate_config(&config).diagnostics); - diagnostics } fn validate_response_cache_backend_config_fields( diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index 895d54102..36eeca3d1 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -56,48 +56,8 @@ pub fn build_cache_key( request: &LlmRequest, config: &ResponseCacheConfig, ) -> KeyOutcome { - // Unparseable bodies arrive as null; they would all share one key. - if request.content.is_null() { - return KeyOutcome::Bypass("unparseable_body"); - } - // Cacheability gates run on the RAW request, so they are correct regardless - // of which codec (if any) decodes the body — a chat codec may park `store` - // in `extra` rather than the typed field, so we must not rely on the decode. - if let Some(object) = request.content.as_object() { - // Any present, non-`false` `store` opts into server-side persistence — - // bypass even a malformed non-boolean rather than risk caching a stateful - // call (whose result is otherwise keyed with `store` stripped). - if object - .get("store") - .is_some_and(|value| !matches!(value, Json::Bool(false) | Json::Null)) - { - return KeyOutcome::Bypass("stateful_store"); - } - if object.contains_key("previous_response_id") { - return KeyOutcome::Bypass("stateful_previous_response_id"); - } - // Server-side conversation state the key cannot see. - if object.contains_key("conversation") || object.contains_key("container") { - return KeyOutcome::Bypass("stateful_conversation"); - } - // Responses persists by default; only an explicit opt-out is stateless. - // A `prompt` object is the Responses prompt-template reference; a bare - // string `prompt` is a completions body with no server-side state. - if (object.contains_key("input") - || object.contains_key("instructions") - || object.get("prompt").is_some_and(Json::is_object)) - && !object - .get("store") - .is_some_and(|store| store == &Json::Bool(false)) - { - return KeyOutcome::Bypass("stateful_store"); - } - } - // Toggle off = explicit temperature 0 only; absent defaults to sampling. - if !config.cache_nondeterministic - && request_temperature(&request.content).is_none_or(|temperature| temperature > 0.0) - { - return KeyOutcome::Bypass("nondeterministic_temperature"); + if let Some(reason) = cache_bypass_reason(request, config) { + return KeyOutcome::Bypass(reason); } // Body to fingerprint: the decoded/normalized form when a surface resolves @@ -136,6 +96,44 @@ pub fn build_cache_key( } } +fn cache_bypass_reason(request: &LlmRequest, config: &ResponseCacheConfig) -> Option<&'static str> { + if request.content.is_null() { + return Some("unparseable_body"); + } + if let Some(reason) = request + .content + .as_object() + .and_then(stateful_request_bypass_reason) + { + return Some(reason); + } + (!config.cache_nondeterministic + && request_temperature(&request.content).is_none_or(|temperature| temperature > 0.0)) + .then_some("nondeterministic_temperature") +} + +fn stateful_request_bypass_reason(object: &Map) -> Option<&'static str> { + if object + .get("store") + .is_some_and(|value| !matches!(value, Json::Bool(false) | Json::Null)) + { + return Some("stateful_store"); + } + if object.contains_key("previous_response_id") { + return Some("stateful_previous_response_id"); + } + if object.contains_key("conversation") || object.contains_key("container") { + return Some("stateful_conversation"); + } + let responses_surface = object.contains_key("input") + || object.contains_key("instructions") + || object.get("prompt").is_some_and(Json::is_object); + let explicitly_stateless = object + .get("store") + .is_some_and(|store| store == &Json::Bool(false)); + (responses_surface && !explicitly_stateless).then_some("stateful_store") +} + /// Preserves which OpenAI Chat token-cap field the caller sent. /// /// The Chat codec normalizes both spellings into `GenerationParams.max_tokens`, diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs index ed3d1afe6..02a8bf0d4 100644 --- a/crates/adaptive/src/response_cache/replay.rs +++ b/crates/adaptive/src/response_cache/replay.rs @@ -192,44 +192,7 @@ fn synthesize_chat_chunks(aggregate: &Json) -> Vec { .cloned() .unwrap_or_default(); for (position, choice) in choices.iter().enumerate() { - let index = choice - .get("index") - .and_then(Json::as_u64) - .unwrap_or(position as u64); - let message = choice.get("message").cloned().unwrap_or(json!({})); - if let Some(role) = message.get("role") { - chunks.push(base( - json!([{"index": index, "delta": {"role": role}, "finish_reason": null}]), - )); - } - if let Some(content) = message.get("content").and_then(Json::as_str) - && !content.is_empty() - { - chunks.push(base( - json!([{"index": index, "delta": {"content": content}, "finish_reason": null}]), - )); - } - if let Some(tool_calls) = message.get("tool_calls").and_then(Json::as_array) { - let deltas: Vec = tool_calls - .iter() - .enumerate() - .map(|(call_index, call)| { - let mut delta = call.clone(); - if let Some(map) = delta.as_object_mut() { - map.entry("index".to_string()) - .or_insert(json!(call_index as u64)); - } - delta - }) - .collect(); - chunks.push(base( - json!([{"index": index, "delta": {"tool_calls": deltas}, "finish_reason": null}]), - )); - } - let finish = choice.get("finish_reason").cloned().unwrap_or(Json::Null); - chunks.push(base( - json!([{"index": index, "delta": {}, "finish_reason": finish}]), - )); + chunks.extend(synthesize_chat_choice_chunks(&base, position, choice)); } if let Some(usage) = aggregate.get("usage") { let mut usage_chunk = base(json!([])); @@ -241,6 +204,53 @@ fn synthesize_chat_chunks(aggregate: &Json) -> Vec { chunks } +fn synthesize_chat_choice_chunks( + base: &impl Fn(Json) -> Json, + position: usize, + choice: &Json, +) -> Vec { + let index = choice + .get("index") + .and_then(Json::as_u64) + .unwrap_or(position as u64); + let message = choice.get("message").cloned().unwrap_or(json!({})); + let mut chunks = Vec::new(); + if let Some(role) = message.get("role") { + chunks.push(base( + json!([{"index": index, "delta": {"role": role}, "finish_reason": null}]), + )); + } + if let Some(content) = message.get("content").and_then(Json::as_str) + && !content.is_empty() + { + chunks.push(base( + json!([{"index": index, "delta": {"content": content}, "finish_reason": null}]), + )); + } + if let Some(tool_calls) = message.get("tool_calls").and_then(Json::as_array) { + let deltas = tool_calls + .iter() + .enumerate() + .map(|(call_index, call)| { + let mut delta = call.clone(); + if let Some(map) = delta.as_object_mut() { + map.entry("index".to_string()) + .or_insert(json!(call_index as u64)); + } + delta + }) + .collect::>(); + chunks.push(base( + json!([{"index": index, "delta": {"tool_calls": deltas}, "finish_reason": null}]), + )); + } + let finish = choice.get("finish_reason").cloned().unwrap_or(Json::Null); + chunks.push(base( + json!([{"index": index, "delta": {}, "finish_reason": finish}]), + )); + chunks +} + /// OpenAI Responses: a `response.created` snapshot, one `response.output_item.done` /// per output item, and a `response.completed` carrying the full stored aggregate /// (the collector keeps the last snapshot wholesale, so reassembly is exact). diff --git a/crates/core/src/observability/otel_genai.rs b/crates/core/src/observability/otel_genai.rs index aab9698e2..481dbf4e0 100644 --- a/crates/core/src/observability/otel_genai.rs +++ b/crates/core/src/observability/otel_genai.rs @@ -434,59 +434,31 @@ fn push_error_attributes(attributes: &mut Vec, event: &Event) { } fn scalar_string(event: &Event, keys: &[&str]) -> Option { - if let Some(profile) = event.category_profile() { - for key in keys { - if let Some(value) = profile.extra.get(*key) { - if let Some(value) = value.as_str() { - return Some(value.to_string()); - } - if value.is_number() || value.is_boolean() { - return Some(value.to_string()); - } - } - } - } - for object in event_objects(event) { - for key in keys { - if let Some(value) = object_value(object, key) { - if let Some(value) = value.as_str() { - return Some(value.to_string()); - } - if value.is_number() || value.is_boolean() { - return Some(value.to_string()); - } - } - } - } - None + find_scalar(event, keys, |value| { + value + .as_str() + .map(str::to_string) + .or_else(|| (value.is_number() || value.is_boolean()).then(|| value.to_string())) + }) } fn scalar_i64(event: &Event, keys: &[&str]) -> Option { - if let Some(profile) = event.category_profile() { - for key in keys { - if let Some(value) = profile.extra.get(*key) { - if let Some(value) = value.as_i64() { - return Some(value); - } - if let Some(value) = value.as_u64().and_then(to_i64) { - return Some(value); - } - } - } - } - for object in event_objects(event) { - for key in keys { - if let Some(value) = object_value(object, key) { - if let Some(value) = value.as_i64() { - return Some(value); - } - if let Some(value) = value.as_u64().and_then(to_i64) { - return Some(value); - } - } - } - } - None + find_scalar(event, keys, |value| { + value.as_i64().or_else(|| value.as_u64().and_then(to_i64)) + }) +} + +fn find_scalar(event: &Event, keys: &[&str], convert: impl Fn(&Json) -> Option) -> Option { + let profile_value = event.category_profile().and_then(|profile| { + keys.iter() + .find_map(|key| profile.extra.get(*key).and_then(&convert)) + }); + profile_value.or_else(|| { + event_objects(event).into_iter().find_map(|object| { + keys.iter() + .find_map(|key| object_value(object, key).and_then(&convert)) + }) + }) } fn object_value<'a>(object: &'a Map, key: &str) -> Option<&'a Json> { diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 0dc46285e..cdf00e726 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -1807,22 +1807,7 @@ fn build_otel_config( ))); } }; - for (header, variable) in §ion.header_env { - if variable.trim().is_empty() || variable.trim() != variable { - return Err(PluginError::InvalidConfig(format!( - "OpenTelemetry endpoints[{index}].header_env.{header} must name a nonblank environment variable without surrounding whitespace" - ))); - } - if section - .headers - .keys() - .any(|configured| configured.eq_ignore_ascii_case(header)) - { - return Err(PluginError::InvalidConfig(format!( - "OpenTelemetry endpoints[{index}] header {header:?} cannot appear in both headers and header_env" - ))); - } - } + validate_otel_header_env(index, §ion)?; let mut config = CoreOpenTelemetryConfig::new(section.otel_type, section.endpoint) .with_transport(transport) .with_service_name(section.service_name) @@ -1840,7 +1825,42 @@ fn build_otel_config( for (key, value) in section.headers { config = config.with_header(key, value); } - for (key, variable) in section.header_env { + config = apply_otel_environment_headers(config, index, section.header_env)?; + for (key, value) in section.resource_attributes { + config = config.with_resource_attribute(key, value); + } + Ok(config) +} + +fn validate_otel_header_env( + index: usize, + section: &OpenTelemetryEndpointConfig, +) -> PluginResult<()> { + for (header, variable) in §ion.header_env { + if variable.trim().is_empty() || variable.trim() != variable { + return Err(PluginError::InvalidConfig(format!( + "OpenTelemetry endpoints[{index}].header_env.{header} must name a nonblank environment variable without surrounding whitespace" + ))); + } + if section + .headers + .keys() + .any(|configured| configured.eq_ignore_ascii_case(header)) + { + return Err(PluginError::InvalidConfig(format!( + "OpenTelemetry endpoints[{index}] header {header:?} cannot appear in both headers and header_env" + ))); + } + } + Ok(()) +} + +fn apply_otel_environment_headers( + mut config: CoreOpenTelemetryConfig, + index: usize, + header_env: HashMap, +) -> PluginResult { + for (key, variable) in header_env { let value = std::env::var(&variable).map_err(|error| { PluginError::InvalidConfig(format!( "OpenTelemetry endpoints[{index}].header_env.{key} could not read environment variable {variable:?}: {error}" @@ -1853,9 +1873,6 @@ fn build_otel_config( } config = config.with_header(key, value); } - for (key, value) in section.resource_attributes { - config = config.with_resource_attribute(key, value); - } Ok(config) } diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 97bab10bb..b46d4c8d0 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -1634,33 +1634,54 @@ fn plugin_config_overlay_value(config: &PluginConfig) -> Result { root.remove("version"); } - if let Some(Json::Object(policy)) = root.get_mut("policy") { - let defaults = ConfigPolicy::default(); - if config.policy.unknown_component == defaults.unknown_component { - policy.remove("unknown_component"); - } - if config.policy.unknown_field == defaults.unknown_field { - policy.remove("unknown_field"); - } - if config.policy.unsupported_value == defaults.unsupported_value { - policy.remove("unsupported_value"); - } - if policy.is_empty() { - root.remove("policy"); + remove_default_policy_overlay(root, &config.policy); + remove_default_component_enabled_overlays(root, &config.components); + + Ok(overlay) +} + +fn remove_default_policy_overlay(root: &mut Map, config: &ConfigPolicy) { + let Some(Json::Object(policy)) = root.get_mut("policy") else { + return; + }; + let defaults = ConfigPolicy::default(); + for (field, is_default) in [ + ( + "unknown_component", + config.unknown_component == defaults.unknown_component, + ), + ( + "unknown_field", + config.unknown_field == defaults.unknown_field, + ), + ( + "unsupported_value", + config.unsupported_value == defaults.unsupported_value, + ), + ] { + if is_default { + policy.remove(field); } } + if policy.is_empty() { + root.remove("policy"); + } +} - if let Some(Json::Array(components)) = root.get_mut("components") { - for (component, typed) in components.iter_mut().zip(&config.components) { - if typed.enabled == default_enabled() - && let Json::Object(component) = component - { - component.remove("enabled"); - } +fn remove_default_component_enabled_overlays( + root: &mut Map, + configured: &[PluginComponentSpec], +) { + let Some(Json::Array(components)) = root.get_mut("components") else { + return; + }; + for (component, typed) in components.iter_mut().zip(configured) { + if typed.enabled == default_enabled() + && let Json::Object(component) = component + { + component.remove("enabled"); } } - - Ok(overlay) } /// Resolves the default `plugins.toml` layering into one JSON document, or an diff --git a/crates/core/src/plugins/nemo_guardrails/python.rs b/crates/core/src/plugins/nemo_guardrails/python.rs index 04b99c01c..504a798cb 100644 --- a/crates/core/src/plugins/nemo_guardrails/python.rs +++ b/crates/core/src/plugins/nemo_guardrails/python.rs @@ -1107,7 +1107,10 @@ impl LlmStreamInner for GuardedProviderStream { } } -#[allow(clippy::too_many_arguments)] +#[allow( + clippy::too_many_arguments, + reason = "stream cancellation, monitoring, delivery, and cleanup must remain ordered in one coordinator" +)] async fn forward_guarded_provider_stream( mut provider_stream: LlmJsonStream, codec: LocalGuardrailsCodec, @@ -1130,51 +1133,132 @@ async fn forward_guarded_provider_stream( let Some(item) = item else { break; }; - let chunk = match item { - Ok(chunk) => chunk, - Err(err) => { - let _ = chunk_tx.send(Err(err)).await; - let _ = text_tx.send(None).await; - let _ = monitor.take().expect("monitor available").await; - break; - } + let Some(chunk) = + receive_guarded_provider_chunk(item, &text_tx, &chunk_tx, &mut monitor).await + else { + break; }; - if let Some(message) = blocked_message(&blocked) { - let _ = chunk_tx.send(Err(streaming_output_blocked(message))).await; - let _ = text_tx.send(None).await; - let _ = monitor.take().expect("monitor available").await; + if stop_blocked_provider_stream(&text_tx, &chunk_tx, &blocked, &mut monitor).await { break; } - if let Some(text) = extract_stream_text(codec, &chunk) - && text_tx.send(Some(text)).await.is_err() + if !forward_guarded_stream_text(codec, &chunk, &text_tx, &chunk_tx, &blocked, &mut monitor) + .await { - send_stream_monitor_error( - monitor.take().expect("monitor available"), - &chunk_tx, - &blocked, - ) - .await; break; } - let sent = tokio::select! { - _ = cancel.changed() => break, - sent = chunk_tx.send(Ok(chunk)) => sent, - }; - if sent.is_err() { + if !send_guarded_provider_chunk(chunk, &text_tx, &chunk_tx, &mut monitor, &mut cancel).await + { + break; + } + } + finish_guarded_provider_stream( + &mut provider_stream, + &text_tx, + &chunk_tx, + &blocked, + &mut monitor, + &cancel, + &closed, + ) + .await; +} + +async fn receive_guarded_provider_chunk( + item: FlowResult, + text_tx: &mpsc::Sender>, + chunk_tx: &mpsc::Sender>, + monitor: &mut Option>>, +) -> Option { + match item { + Ok(chunk) => Some(chunk), + Err(err) => { + let _ = chunk_tx.send(Err(err)).await; let _ = text_tx.send(None).await; let _ = monitor.take().expect("monitor available").await; - break; + None } } +} + +async fn stop_blocked_provider_stream( + text_tx: &mpsc::Sender>, + chunk_tx: &mpsc::Sender>, + blocked: &Arc>>, + monitor: &mut Option>>, +) -> bool { + let Some(message) = blocked_message(blocked) else { + return false; + }; + let _ = chunk_tx.send(Err(streaming_output_blocked(message))).await; + let _ = text_tx.send(None).await; + let _ = monitor.take().expect("monitor available").await; + true +} + +async fn forward_guarded_stream_text( + codec: LocalGuardrailsCodec, + chunk: &Json, + text_tx: &mpsc::Sender>, + chunk_tx: &mpsc::Sender>, + blocked: &Arc>>, + monitor: &mut Option>>, +) -> bool { + let Some(text) = extract_stream_text(codec, chunk) else { + return true; + }; + if text_tx.send(Some(text)).await.is_ok() { + return true; + } + send_stream_monitor_error( + monitor.take().expect("monitor available"), + chunk_tx, + blocked, + ) + .await; + false +} + +async fn send_guarded_provider_chunk( + chunk: Json, + text_tx: &mpsc::Sender>, + chunk_tx: &mpsc::Sender>, + monitor: &mut Option>>, + cancel: &mut watch::Receiver, +) -> bool { + let sent = tokio::select! { + _ = cancel.changed() => return false, + sent = chunk_tx.send(Ok(chunk)) => sent, + }; + if sent.is_ok() { + return true; + } + let _ = text_tx.send(None).await; + let _ = monitor.take().expect("monitor available").await; + false +} + +#[allow( + clippy::too_many_arguments, + reason = "stream cleanup needs all channels and lifecycle handles" +)] +async fn finish_guarded_provider_stream( + provider_stream: &mut LlmJsonStream, + text_tx: &mpsc::Sender>, + chunk_tx: &mpsc::Sender>, + blocked: &Arc>>, + monitor: &mut Option>>, + cancel: &watch::Receiver, + closed: &watch::Sender>>, +) { let _ = text_tx.send(None).await; if *cancel.borrow() { if let Some(monitor) = monitor.take() { monitor.abort(); } } else if let Some(monitor) = monitor.take() { - let _ = send_stream_monitor_error(monitor, &chunk_tx, &blocked).await; + let _ = send_stream_monitor_error(monitor, chunk_tx, blocked).await; } closed.send_replace(Some(provider_stream.close().await)); } diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index dd99fde19..4679b2751 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -36,7 +36,7 @@ use crate::api::event::{BaseEvent, MarkEvent}; use crate::api::llm::LlmHandle; use crate::api::llm::emit_reserved_optimization_marks; use crate::api::optimization::finalize_optimization_summary; -use crate::api::runtime::LlmSanitizeResponseContext; +use crate::api::registry::Guardrail; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::subscriber_dispatcher; @@ -44,6 +44,7 @@ use crate::api::runtime::{ EventSubscriberFn, LlmJsonStream, LlmStreamInner, ScopeStackHandle, TASK_SCOPE_STACK, current_scope_stack, }; +use crate::api::runtime::{LlmSanitizeResponseContext, LlmSanitizeResponseFn}; use crate::api::shared::{ metadata_with_otel_error, metadata_with_otel_status, snapshot_event_sanitizers, }; @@ -249,31 +250,8 @@ impl LlmStreamWrapper { aggregated }; - let (entries, sanitizer_snapshot_failed) = match self.scope_stack.read() { - Ok(scope_guard) => { - let scope_locals = scope_guard - .collect_scope_local_registries(|r| &r.llm_sanitize_response_guardrails); - match global_context().read() { - Ok(state) => (state.llm_sanitize_response_entries(&scope_locals), false), - Err(error) => { - log::error!( - target: "nemo_relay.runtime", - event = "stream_end_sanitizer_snapshot_failed"; - "LLM stream END sanitizer snapshot failed; omitting the observability payload: {error}" - ); - (Vec::new(), true) - } - } - } - Err(error) => { - log::error!( - target: "nemo_relay.runtime", - event = "stream_end_sanitizer_snapshot_failed"; - "LLM stream END sanitizer snapshot failed; omitting the observability payload: {error}" - ); - (Vec::new(), true) - } - }; + let (entries, sanitizer_snapshot_failed) = + snapshot_stream_end_sanitizers(&self.scope_stack); let handle = self.handle.clone(); let scope_stack = self.scope_stack.clone(); let finalization_scope_stack = scope_stack.clone(); @@ -412,6 +390,30 @@ impl LlmStreamWrapper { } } +fn snapshot_stream_end_sanitizers( + scope_stack: &ScopeStackHandle, +) -> (Vec>, bool) { + let entries = scope_stack.read().ok().and_then(|scope_guard| { + let scope_locals = scope_guard + .collect_scope_local_registries(|registry| ®istry.llm_sanitize_response_guardrails); + global_context() + .read() + .ok() + .map(|state| state.llm_sanitize_response_entries(&scope_locals)) + }); + match entries { + Some(entries) => (entries, false), + None => { + log::error!( + target: "nemo_relay.runtime", + event = "stream_end_sanitizer_snapshot_failed"; + "LLM stream END sanitizer snapshot failed; omitting the observability payload" + ); + (Vec::new(), true) + } + } +} + impl Stream for LlmStreamWrapper { type Item = Result; diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index db7b727cb..50772a3e6 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -1727,6 +1727,115 @@ pub fn wrap_py_event_subscriber(py_fn: Py) -> EventSubscriberFn { }) } +fn prepare_event_sanitizer_invocation<'py>( + py: Python<'py>, + publication_context: Option<&PythonPublicationContext>, + task_locals: Option, + publication_buffer: Option, +) -> FlowResult<(Option>, Option)> { + match publication_context { + Some(context) => { + let (context, task_locals) = copy_publication_invocation_with_buffer( + py, + context, + task_locals, + publication_buffer, + ) + .map_err(|error| FlowError::Internal(error.to_string()))?; + Ok((Some(context), task_locals)) + } + None => copy_middleware_invocation(py, task_locals) + .map_err(|error| FlowError::Internal(error.to_string())), + } +} + +fn py_event_object(py: Python<'_>, event: &Event) -> PyResult> { + match event { + Event::Scope(inner) => Py::new( + py, + crate::py_types::PyScopeEvent { + inner: inner.clone(), + }, + ) + .map(|value| value.into_any()), + Event::Mark(inner) => Py::new( + py, + crate::py_types::PyMarkEvent { + inner: inner.clone(), + }, + ) + .map(|value| value.into_any()), + } +} + +fn call_event_sanitizer( + py: Python<'_>, + invoke: &Bound<'_, PyAny>, + callback: &Py, + invocation_context: Option<&Bound<'_, PyAny>>, + loop_affine: bool, + py_event: Py, + py_fields: Py, +) -> PyResult> { + let result = match (invocation_context, loop_affine) { + (Some(context), false) => { + context.call_method1("run", (invoke, callback.bind(py), py_event, py_fields)) + } + (None, false) => invoke.call1((callback.bind(py), py_event, py_fields)), + (Some(context), true) => { + context.call_method1("run", (callback.bind(py), py_event, py_fields)) + } + (None, true) => callback.bind(py).call1((py_event, py_fields)), + }?; + Ok(result.unbind()) +} + +fn start_py_event_sanitizer( + py: Python<'_>, + py_fn: &Py, + event: &Event, + fields: &EventSanitizeFields, + publication_context: Option<&PythonPublicationContext>, + task_locals: Option, + publication_buffer: Option, +) -> FlowResult, PyValueFuture>> { + let (invocation_context, task_locals) = prepare_event_sanitizer_invocation( + py, + publication_context, + task_locals, + publication_buffer, + )?; + let py_event = + py_event_object(py, event).map_err(|error| FlowError::Internal(error.to_string()))?; + let fields_json = + serde_json::to_value(fields).map_err(|error| FlowError::Internal(error.to_string()))?; + let py_fields = + json_to_py(py, &fields_json).map_err(|error| FlowError::Internal(error.to_string()))?; + let invoke = py + .import("nemo_relay._event_sanitizer_context") + .and_then(|module| module.getattr("invoke")) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let loop_affine = task_locals.is_some(); + let callback = loop_affine_callback(py, py_fn.bind(py), task_locals.as_ref(), true) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let result = call_event_sanitizer( + py, + &invoke, + &callback, + invocation_context.as_ref(), + loop_affine, + py_event, + py_fields, + ) + .map_err(|error| FlowError::Internal(error.to_string()))?; + split_py_object_or_future_with_locals( + py, + result, + task_locals.as_ref(), + invocation_context.as_ref(), + ) +} + /// Wrap a Python callable ``(Event, EventSanitizeFields) -> EventSanitizeFields``. pub fn wrap_py_event_sanitize_fn(py_fn: Py) -> EventSanitizeFn { let py_fn = Arc::new(py_fn); @@ -1737,83 +1846,17 @@ pub fn wrap_py_event_sanitize_fn(py_fn: Py) -> EventSanitizeFn { let publication_context = publication_context::(); let publication_buffer = capture_nested_publication_buffer(); Box::pin(async move { - let result = Python::attach( - |py| -> FlowResult, PyValueFuture>> { - let (invocation_context, task_locals) = match publication_context.as_ref() { - Some(context) => { - let (context, publication_task_locals) = - copy_publication_invocation_with_buffer( - py, - context, - task_locals, - publication_buffer.clone(), - ) - .map_err(|error| FlowError::Internal(error.to_string()))?; - (Some(context), publication_task_locals) - } - None => { copy_middleware_invocation(py, task_locals) } - .map_err(|error| FlowError::Internal(error.to_string()))?, - }; - let py_event = match event.as_ref() { - Event::Scope(inner) => Py::new( - py, - crate::py_types::PyScopeEvent { - inner: inner.clone(), - }, - ) - .map(|value| value.into_any()), - Event::Mark(inner) => Py::new( - py, - crate::py_types::PyMarkEvent { - inner: inner.clone(), - }, - ) - .map(|value| value.into_any()), - }; - let py_event = match py_event { - Ok(value) => value, - Err(error) => { - return Err(FlowError::Internal(error.to_string())); - } - }; - let fields_json = match serde_json::to_value(&fields) { - Ok(value) => value, - Err(error) => { - return Err(FlowError::Internal(error.to_string())); - } - }; - let py_fields = match json_to_py(py, &fields_json) { - Ok(value) => value, - Err(error) => { - return Err(FlowError::Internal(error.to_string())); - } - }; - let invoke = py - .import("nemo_relay._event_sanitizer_context") - .and_then(|module| module.getattr("invoke")) - .map_err(|error| FlowError::Internal(error.to_string()))?; - let loop_affine = task_locals.is_some(); - let callback = - loop_affine_callback(py, py_fn.bind(py), task_locals.as_ref(), true) - .map_err(|error| FlowError::Internal(error.to_string()))?; - let result = match (invocation_context.as_ref(), !loop_affine) { - (Some(context), true) => context - .call_method1("run", (invoke, callback.bind(py), py_event, py_fields)), - (None, true) => invoke.call1((callback.bind(py), py_event, py_fields)), - (Some(context), false) => { - context.call_method1("run", (callback.bind(py), py_event, py_fields)) - } - (None, false) => callback.bind(py).call1((py_event, py_fields)), - } - .map_err(|error| FlowError::Internal(error.to_string()))?; - split_py_object_or_future_with_locals( - py, - result.unbind(), - task_locals.as_ref(), - invocation_context.as_ref(), - ) - }, - ); + let result = Python::attach(|py| { + start_py_event_sanitizer( + py, + py_fn.as_ref(), + event.as_ref(), + &fields, + publication_context.as_deref(), + task_locals, + publication_buffer, + ) + }); let result = resolve_py_object_or_future(result) .await .and_then(|result| { diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index a3e80fd18..3046120f5 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -2010,16 +2010,15 @@ type OpenTelemetrySubscriber struct { ptr unsafe.Pointer } -// NewOpenTelemetrySubscriber creates a new OpenTelemetry subscriber from config. -func NewOpenTelemetrySubscriber(config OpenTelemetryConfig) (*OpenTelemetrySubscriber, error) { +func normalizeOpenTelemetryConfig(config OpenTelemetryConfig) (OpenTelemetryConfig, error) { if config.Transport == "" { config.Transport = OpenTelemetryTransportHTTPBinary } if config.Type == "" { - return nil, fmt.Errorf("type is required") + return config, fmt.Errorf("type is required") } if config.Endpoint == "" { - return nil, fmt.Errorf("endpoint is required") + return config, fmt.Errorf("endpoint is required") } if config.ServiceName == "" { config.ServiceName = "unknown_service" @@ -2045,17 +2044,30 @@ func NewOpenTelemetrySubscriber(config OpenTelemetryConfig) (*OpenTelemetrySubsc if config.AttributeMappings == nil { config.AttributeMappings = []OtlpAttributeMapping{} } + return config, nil +} + +func optionalCString(value string) *C.char { + if value == "" { + return nil + } + return C.CString(value) +} + +// NewOpenTelemetrySubscriber creates a new OpenTelemetry subscriber from config. +func NewOpenTelemetrySubscriber(config OpenTelemetryConfig) (*OpenTelemetrySubscriber, error) { + config, err := normalizeOpenTelemetryConfig(config) + if err != nil { + return nil, err + } cTransport := C.CString(string(config.Transport)) defer C.free(unsafe.Pointer(cTransport)) cType := C.CString(string(config.Type)) defer C.free(unsafe.Pointer(cType)) - var cEndpoint *C.char - if config.Endpoint != "" { - cEndpoint = C.CString(config.Endpoint) - defer C.free(unsafe.Pointer(cEndpoint)) - } + cEndpoint := C.CString(config.Endpoint) + defer C.free(unsafe.Pointer(cEndpoint)) headersJSON, err := jsonMarshal(config.Headers) if err != nil { @@ -2074,17 +2086,11 @@ func NewOpenTelemetrySubscriber(config OpenTelemetryConfig) (*OpenTelemetrySubsc cServiceName := C.CString(config.ServiceName) defer C.free(unsafe.Pointer(cServiceName)) - var cServiceNamespace *C.char - if config.ServiceNamespace != "" { - cServiceNamespace = C.CString(config.ServiceNamespace) - defer C.free(unsafe.Pointer(cServiceNamespace)) - } + cServiceNamespace := optionalCString(config.ServiceNamespace) + defer C.free(unsafe.Pointer(cServiceNamespace)) - var cServiceVersion *C.char - if config.ServiceVersion != "" { - cServiceVersion = C.CString(config.ServiceVersion) - defer C.free(unsafe.Pointer(cServiceVersion)) - } + cServiceVersion := optionalCString(config.ServiceVersion) + defer C.free(unsafe.Pointer(cServiceVersion)) cInstrumentationScope := C.CString(config.InstrumentationScope) defer C.free(unsafe.Pointer(cInstrumentationScope)) diff --git a/scripts/package_node_musllinux.mjs b/scripts/package_node_musllinux.mjs index 8428370e7..1b2c3146b 100755 --- a/scripts/package_node_musllinux.mjs +++ b/scripts/package_node_musllinux.mjs @@ -20,13 +20,14 @@ function argumentsFrom(args) { let version; let output; let platform; - for (let index = 0; index < args.length; index += 1) { + for (let index = 0; index < args.length; index += 2) { + const value = args[index + 1]; if (args[index] === "--version") { - version = args[++index]; + version = value; } else if (args[index] === "--out") { - output = args[++index]; + output = value; } else if (args[index] === "--platform") { - platform = args[++index]; + platform = value; } else { throw new Error(`Unexpected argument: ${args[index]}`); }