diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index cb585c72d..ae7b5fc09 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -276,6 +276,7 @@ fn emit_llm_start( handle: &LlmHandle, request: &LlmRequest, annotated_request: Option>, + request_codec: Option<&dyn LlmCodec>, ) -> Result<()> { ensure_runtime_owner()?; let (event, subscribers) = { @@ -292,6 +293,15 @@ fn emit_llm_start( .map_err(|error| FlowError::Internal(error.to_string()))?; let sanitized_request = state.llm_sanitize_request_chain(request.clone(), &scope_locals); + let annotated_request = match request_codec { + Some(codec) + if sanitized_request.headers != request.headers + || sanitized_request.content != request.content => + { + codec.decode(&sanitized_request).ok().map(Arc::new) + } + _ => annotated_request, + }; let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null); let event = state.build_llm_start_event(handle, Some(input), annotated_request); (event, subscribers) @@ -341,10 +351,16 @@ pub fn llm_call(params: LlmCallParams<'_>) -> Result { .timestamp_opt(params.timestamp) .build(); let handle = create_llm_handle(handle_params)?; - emit_llm_start(&handle, params.request, params.annotated_request)?; + emit_llm_start(&handle, params.request, params.annotated_request, None)?; Ok(handle) } +#[derive(Clone, Copy)] +struct LlmCallEndBehavior { + response_codec_errors_fatal: bool, + attach_estimated_cost: bool, +} + /// Finish a manual LLM lifecycle span. /// /// This emits an LLM-end event for a handle previously returned by @@ -377,6 +393,19 @@ pub fn llm_call(params: LlmCallParams<'_>) -> Result { /// Sanitize-response guardrails affect only the emitted end-event payload, not /// the caller-owned `response` value. pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> { + llm_call_end_with_behavior( + params, + LlmCallEndBehavior { + response_codec_errors_fatal: true, + attach_estimated_cost: false, + }, + ) +} + +fn llm_call_end_with_behavior( + params: LlmCallEndParams<'_>, + behavior: LlmCallEndBehavior, +) -> Result<()> { let LlmCallEndParams { handle, response, @@ -411,7 +440,12 @@ pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> { Some(annotated_response) => Some(annotated_response), None => match (response_codec.as_ref(), data.as_ref()) { (Some(codec), Some(response)) => match codec.decode_response(response) { - Ok(decoded) => Some(Arc::new(decoded)), + Ok(mut decoded) => { + if behavior.attach_estimated_cost { + attach_estimated_cost_for_provider(&mut decoded, Some(&handle.name)); + } + Some(Arc::new(decoded)) + } Err(error) => { decode_error = Some(error); None @@ -432,7 +466,9 @@ pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> { (event, subscribers) }; NemoRelayContextState::emit_event(&event, &subscribers); - if let Some(error) = decode_error { + if let Some(error) = decode_error + && behavior.response_codec_errors_fatal + { Err(error) } else { Ok(()) @@ -554,6 +590,7 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { } } + let request_codec = codec.clone(); let (intercepted_request, annotated_request) = run_request_intercepts_with_codec(&name, request, codec)?; @@ -567,7 +604,12 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { .model_name_opt(model_name) .build(), )?; - emit_llm_start(&handle, &intercepted_request, annotated_request.clone())?; + emit_llm_start( + &handle, + &intercepted_request, + annotated_request.clone(), + request_codec.as_deref(), + )?; let execution = { let scope_stack = current_scope_stack(); @@ -583,22 +625,18 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { match execution(intercepted_request).await { Ok(response) => { - let annotated_response = response_codec - .as_ref() - .and_then(|codec| { - let mut decoded = codec.decode_response(&response).ok()?; - attach_estimated_cost_for_provider(&mut decoded, Some(&name)); - Some(decoded) - }) - .map(Arc::new); - llm_call_end( + llm_call_end_with_behavior( LlmCallEndParams::builder() .handle(&handle) .response(response.clone()) .data_opt(data) .metadata_opt(metadata) - .annotated_response_opt(annotated_response) + .response_codec_opt(response_codec) .build(), + LlmCallEndBehavior { + response_codec_errors_fatal: false, + attach_estimated_cost: true, + }, )?; Ok(response) } @@ -706,6 +744,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu } } + let request_codec = codec.clone(); let (intercepted_request, annotated_request) = run_request_intercepts_with_codec(&name, request, codec)?; @@ -719,7 +758,12 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu .model_name_opt(model_name) .build(), )?; - emit_llm_start(&handle, &intercepted_request, annotated_request)?; + emit_llm_start( + &handle, + &intercepted_request, + annotated_request, + request_codec.as_deref(), + )?; let execution = { let scope_stack = current_scope_stack(); diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 057d8cdb4..1c4fcc454 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -140,17 +140,6 @@ impl LlmStreamWrapper { None => Json::Null, }; - // Decode aggregated response if response codec is present (non-fatal) - let annotated_response: Option> = self - .response_codec - .as_ref() - .and_then(|c| { - let mut decoded = c.decode_response(&aggregated).ok()?; - attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); - Some(decoded) - }) - .map(Arc::new); - let event_snapshot = { let ss_guard = self.scope_stack.read().expect("scope stack lock poisoned"); let sl = @@ -167,6 +156,18 @@ impl LlmStreamWrapper { } else { Some(sanitized) }; + let annotated_response: Option> = self + .response_codec + .as_ref() + .and_then(|codec| { + let mut decoded = codec.decode_response(data.as_ref()?).ok()?; + attach_estimated_cost_for_provider( + &mut decoded, + Some(&self.handle.name), + ); + Some(decoded) + }) + .map(Arc::new); let event = state.end_llm_handle( &self.handle, data, diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 6b9da3ef4..eefe05cd3 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -19,15 +19,19 @@ use nemo_relay::api::llm::LlmRequest; use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, }; -use nemo_relay::api::registry::{deregister_llm_request_intercept, register_llm_request_intercept}; +use nemo_relay::api::registry::{ + deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, + deregister_llm_sanitize_response_guardrail, register_llm_request_intercept, + register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, +}; use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn}; use nemo_relay::api::runtime::{create_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::ScopeType; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; -use nemo_relay::codec::request::AnnotatedLlmRequest; -use nemo_relay::codec::request::MessageContent; +use nemo_relay::codec::openai_chat::OpenAIChatCodec; +use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; use nemo_relay::codec::response::FinishReason; use nemo_relay::codec::response::{ AnnotatedLlmResponse, PricingCatalog, PricingResolver, Usage, reset_active_pricing_resolver, @@ -226,6 +230,50 @@ fn make_llm_request(content: Json) -> LlmRequest { } } +fn make_openai_chat_request(content: &str) -> LlmRequest { + make_llm_request(json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": content}], + })) +} + +fn make_openai_chat_response(content: &str) -> Json { + json!({ + "id": "chatcmpl-test", + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) +} + +fn first_message_text(request: &AnnotatedLlmRequest) -> Option<&str> { + match request.messages.first()? { + Message::System { + content: MessageContent::Text(text), + .. + } + | Message::User { + content: MessageContent::Text(text), + .. + } + | Message::Tool { + content: MessageContent::Text(text), + .. + } => Some(text.as_str()), + Message::Assistant { + content: Some(MessageContent::Text(text)), + .. + } => Some(text.as_str()), + _ => None, + } +} + fn noop_exec_fn() -> LlmExecutionNextFn { Arc::new(|_req| Box::pin(async move { Ok(json!({"response": "ok"})) })) } @@ -1001,6 +1049,66 @@ async fn test_response_codec_populates_annotated_response() { reset_active_pricing_resolver().unwrap(); } +#[tokio::test] +async fn test_response_codec_annotation_uses_sanitized_managed_response() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::new())); + let ec = events.clone(); + register_subscriber( + "sanitized_resp_codec_sub", + Arc::new(move |e: &Event| { + ec.lock().unwrap().push(e.clone()); + }), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "sanitize_resp_codec_annotation", + 1, + Arc::new(|_response| make_openai_chat_response("Sanitized")), + ) + .unwrap(); + + let func: LlmExecutionNextFn = + Arc::new(|_req| Box::pin(async move { Ok(make_openai_chat_response("SECRET")) })); + let response_codec: Arc = Arc::new(OpenAIChatCodec); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(make_openai_chat_request("hello")) + .func(func) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + assert_eq!(result["choices"][0]["message"]["content"], json!("SECRET")); + + let captured = captured_events_snapshot(&events); + let end_event = captured + .iter() + .find(|e| is_scope_event(e, ScopeType::Llm, ScopeCategory::End)) + .expect("expected LlmEnd event"); + assert_eq!( + end_event.output().unwrap()["choices"][0]["message"]["content"], + json!("Sanitized") + ); + let annotated = end_event + .annotated_response() + .expect("annotated_response should be decoded from sanitized output"); + assert_eq!(annotated.response_text(), Some("Sanitized")); + assert!( + !serde_json::to_string(end_event).unwrap().contains("SECRET"), + "end event should not retain raw response content" + ); + + deregister_subscriber("sanitized_resp_codec_sub").unwrap(); + deregister_llm_sanitize_response_guardrail("sanitize_resp_codec_annotation").unwrap(); +} + #[tokio::test] async fn test_response_codec_none_when_no_codec() { let _lock = TEST_MUTEX.lock().unwrap(); @@ -1141,6 +1249,67 @@ async fn test_request_codec_populates_annotated_request() { deregister_subscriber("req_codec_ann_sub").unwrap(); } +#[tokio::test] +async fn test_request_codec_annotation_uses_sanitized_start_payload() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::new())); + let ec = events.clone(); + register_subscriber( + "sanitized_req_codec_sub", + Arc::new(move |e: &Event| { + ec.lock().unwrap().push(e.clone()); + }), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "sanitize_req_codec_annotation", + 1, + Arc::new(|request| LlmRequest { + headers: request.headers, + content: make_openai_chat_request("Sanitized").content, + }), + ) + .unwrap(); + + let request_codec: Arc = Arc::new(OpenAIChatCodec); + let _result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(make_openai_chat_request("SECRET")) + .func(noop_exec_fn()) + .codec(request_codec) + .build(), + ) + .await + .unwrap(); + + let captured = captured_events_snapshot(&events); + let start_event = captured + .iter() + .find(|e| is_scope_event(e, ScopeType::Llm, ScopeCategory::Start)) + .expect("expected LlmStart event"); + assert_eq!( + start_event.input().unwrap()["content"]["messages"][0]["content"], + json!("Sanitized") + ); + let annotated = start_event + .annotated_request() + .expect("annotated_request should be decoded from sanitized input"); + assert_eq!(first_message_text(annotated), Some("Sanitized")); + assert!( + !serde_json::to_string(start_event) + .unwrap() + .contains("SECRET"), + "start event should not retain raw request content" + ); + + deregister_subscriber("sanitized_req_codec_sub").unwrap(); + deregister_llm_sanitize_request_guardrail("sanitize_req_codec_annotation").unwrap(); +} + #[tokio::test] async fn test_stream_response_codec_populates_annotated_response() { let _lock = TEST_MUTEX.lock().unwrap(); @@ -1203,3 +1372,67 @@ async fn test_stream_response_codec_populates_annotated_response() { deregister_subscriber("stream_resp_codec_sub").unwrap(); reset_active_pricing_resolver().unwrap(); } + +#[tokio::test] +async fn test_stream_response_codec_annotation_uses_sanitized_aggregated_response() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::new())); + let ec = events.clone(); + register_subscriber( + "stream_sanitized_resp_codec_sub", + Arc::new(move |e: &Event| { + ec.lock().unwrap().push(e.clone()); + }), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "stream_sanitize_resp_codec_annotation", + 1, + Arc::new(|_response| make_openai_chat_response("Sanitized")), + ) + .unwrap(); + + let collector: Box Result<()> + Send> = Box::new(|_chunk| Ok(())); + let finalizer: Box Json + Send> = + Box::new(|| make_openai_chat_response("SECRET")); + let response_codec: Arc = Arc::new(OpenAIChatCodec); + + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(make_openai_chat_request("stream me")) + .func(noop_stream_exec_fn()) + .collector(collector) + .finalizer(finalizer) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + while let Some(_chunk) = stream.next().await {} + + let captured = captured_events_snapshot(&events); + let end_event = captured + .iter() + .find(|e| is_scope_event(e, ScopeType::Llm, ScopeCategory::End)) + .expect("expected LlmEnd event after stream drain"); + assert_eq!( + end_event.output().unwrap()["choices"][0]["message"]["content"], + json!("Sanitized") + ); + let annotated = end_event + .annotated_response() + .expect("annotated_response should be decoded from sanitized stream output"); + assert_eq!(annotated.response_text(), Some("Sanitized")); + assert!( + !serde_json::to_string(end_event).unwrap().contains("SECRET"), + "stream end event should not retain raw aggregated response content" + ); + + deregister_subscriber("stream_sanitized_resp_codec_sub").unwrap(); + deregister_llm_sanitize_response_guardrail("stream_sanitize_resp_codec_annotation").unwrap(); +}