From d6b3c549eee54cd7a772add5cda25c0c0f24718d Mon Sep 17 00:00:00 2001 From: maciej-berri Date: Thu, 12 Nov 2026 17:13:55 -0800 Subject: [PATCH] perf(rust): stop re-serializing the Messages request body in Python logging callbacks (#51440) * perf(rust): store the sent wire bytes as the Messages request instead of re-serializing The Rust /v1/messages path handed record_response the body as a Python dict, and Python ran json.dumps on the whole body again to build httpx_response.request. Rust now keeps the bytes it serialized for the final attempt and passes them through, so the stored request is exactly what the provider received and Python does no extra full-body serialization. The hosted transport path sends the same Rust bytes instead of a second Python json.dumps. This also fixes the stored request after a retry with a wide int, NaN or Infinity, which previously carried the internal marked string instead of the real value. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bne5zRJx1T5sEYBcsGtvqx * perf(rust): serialize the Messages logged input in Rust instead of json.dumps messages_input ran json.dumps on the whole pre_call body in Python for every Rust /v1/messages request. Rust now serializes it with the same raw_json writer it uses for the wire, at the same point (after caller substitution, before Logging.pre_call), so callback edits still reach the wire but not the logged input. Python only returns the top-level key order, with the handler signature read once, and top-level scalars keep the caller's own spelling. The logged input string becomes compact JSON with raw UTF-8. Top-level key order and values are unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bne5zRJx1T5sEYBcsGtvqx * fix(rust): keep body keys missing from the order list in the Messages logged input The logged input only wrote keys named by Python's order tuple, so a key left out of that tuple would vanish from the log without an error. Keys the order does not name now follow in their sent order. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bne5zRJx1T5sEYBcsGtvqx --------- Co-authored-by: Claude Opus 5.5 --- litellm/rust/Cargo.lock | 2 + .../crates/callbacks-legacy-python/Cargo.toml | 1 + .../python_contract.json | 3 +- .../callbacks-legacy-python/src/adapter.rs | 72 +++++++- .../callbacks-legacy-python/src/callbacks.rs | 22 ++- .../callbacks-legacy-python/src/error.rs | 6 +- .../crates/callbacks-legacy-python/src/lib.rs | 1 + .../src/messages_input.rs | 135 ++++++++++++++ .../callbacks-legacy-python/src/python.rs | 4 +- .../src/test_support.rs | 4 +- litellm/rust/crates/host/Cargo.toml | 1 + litellm/rust/crates/host/src/interceptors.rs | 5 +- litellm/rust/crates/http/src/outbound.rs | 2 +- .../crates/inference-messages/src/handler.rs | 66 ++++--- .../rust/crates/inference-messages/src/lib.rs | 4 +- .../inference-messages/src/transport.rs | 8 +- .../inference-messages/tests/messages/host.rs | 50 ++++-- .../tests/messages/retry.rs | 24 +++ .../tests/messages/transport.rs | 51 +++++- .../inference-messages/tests/support/mod.rs | 13 ++ litellm/rust/crates/inference/src/outbound.rs | 15 +- .../python-bridge/benches/serialization.rs | 78 +++++++- .../src/routes/messages/transport.rs | 10 +- .../host/callbacks_legacy_python.py | 30 ++-- litellm/rust_bridge/host/routes/messages.py | 5 +- .../messages/test_callbacks.py | 166 +++++++++++++++++- .../messages/test_streaming.py | 16 +- .../rust_bridge/host/routes/test_messages.py | 16 +- .../host/test_callbacks_legacy_python.py | 43 +++-- 29 files changed, 696 insertions(+), 157 deletions(-) create mode 100644 litellm/rust/crates/callbacks-legacy-python/src/messages_input.rs diff --git a/litellm/rust/Cargo.lock b/litellm/rust/Cargo.lock index 1308df05f37..d3cbb960521 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -3691,6 +3691,7 @@ dependencies = [ "proptest", "pyo3", "rstest", + "serde", "serde_json", "strum", "thiserror 2.0.19", @@ -4009,6 +4010,7 @@ dependencies = [ name = "litellm-host" version = "0.1.0" dependencies = [ + "bytes", "futures-util", "litellm-coroutine", "litellm-secret-value", diff --git a/litellm/rust/crates/callbacks-legacy-python/Cargo.toml b/litellm/rust/crates/callbacks-legacy-python/Cargo.toml index 0133af16b80..7819979029c 100644 --- a/litellm/rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm/rust/crates/callbacks-legacy-python/Cargo.toml @@ -11,6 +11,7 @@ litellm-host.workspace = true litellm-host-python.workspace = true pyo3.workspace = true +serde.workspace = true strum.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm/rust/crates/callbacks-legacy-python/python_contract.json b/litellm/rust/crates/callbacks-legacy-python/python_contract.json index a96e9f04bdc..505651b5f16 100644 --- a/litellm/rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm/rust/crates/callbacks-legacy-python/python_contract.json @@ -28,7 +28,7 @@ "api_key", "additional_args" ], - "messages_input": [ + "messages_input_order": [ "body", "kwargs" ], @@ -43,6 +43,7 @@ "status", "headers", "body", + "streamed", "sent" ], "retried_request": [ diff --git a/litellm/rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm/rust/crates/callbacks-legacy-python/src/adapter.rs index 71ca89b0790..33dc9996bb9 100644 --- a/litellm/rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm/rust/crates/callbacks-legacy-python/src/adapter.rs @@ -21,7 +21,7 @@ use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, SentRequest, deferred::{PendingLogging, PendingSuccess}, - finalize, is_internal_call, + finalize, is_internal_call, messages_input, python::{Logging, Streaming}, setup, }; @@ -216,12 +216,18 @@ impl LegacyLogging { &self, py: Python<'py>, body: &Bound<'py, PyDict>, + sent: &Value, + as_called: &serde_json::Map, ) -> PyResult> { let description = match self.operation { LoggingOperation::Completion => "Chat completions", LoggingOperation::Responses => "Responses", LoggingOperation::Messages => { - return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py))); + let order: Vec = Logging::MessagesInputOrder + .call(py, (body, self.call.kwargs().bind(py)))? + .extract()?; + let content = messages_input::serialized(&order, sent, as_called)?; + return messages_input::logged(py, &content); } LoggingOperation::Ocr => "OCR document processing", }; @@ -598,8 +604,12 @@ impl LegacyLogging { to_py(py, &wire.body)?.into_bound(py) } .cast_into::()?; + let mut as_called = serde_json::Map::new(); for (name, sent) in wire.body.as_object().into_iter().flatten() { if let Some(value) = self.caller_value(py, name, sent)? { + if lossless && messages_input::is_caller_scalar(&value) { + as_called.insert(name.clone(), json_from_py(&value)?); + } body.set_item(name, value)?; } } @@ -620,7 +630,7 @@ impl LegacyLogging { url: wire.url.clone(), context: context.clone(), }); - let input = self.logged_input(py, &body)?; + let input = self.logged_input(py, &body, &wire.body, &as_called)?; self.logger()?.pre_call( py, &input, @@ -752,10 +762,16 @@ impl LegacyLogging { let sent = self.request.as_ref().map(|request| SentRequest { url: &request.url, headers: &request.headers, - body: &request.body, + body: &head.sent_body, }); - self.logger()? - .record_response(py, head.status, &head.headers, &raw.body, sent)?; + self.logger()?.record_response( + py, + head.status, + &head.headers, + &raw.body, + head.streamed, + sent, + )?; } return Ok(HookStep::Ready(())); } @@ -1734,7 +1750,7 @@ assert reasons == [{'kind': 'thinking_signature'}, {'kind': 'strict_tools', 'dro #[case::responses(crate::LoggingOperation::Responses, "Responses")] #[case::messages( crate::LoggingOperation::Messages, - r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"# + r#"[{'role': 'user', 'content': '{"pages":[1]}'}]"# )] #[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")] fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases( @@ -1783,6 +1799,48 @@ assert original == [0] }); } + #[rstest] + fn the_messages_input_is_the_body_before_pre_call_edits_with_the_callers_scalars() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + run( + py, + &locals, + c" +kwargs['temperature'] = 1 +def on_pre_call(additional_args): + additional_args['complete_input_dict']['system'] = 'edited' +", + ); + let mut logging = LegacyLogging { + operation: crate::LoggingOperation::Messages, + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + let wire = WireRequest { + body: json!({"model": "m", "system": "caf\u{e9} \u{1f600}", "temperature": 1.0}), + ..route_wire() + }; + + let (_, step) = send_and_receive(py, &mut logging, wire, &route_context()); + + let HookStep::Ready(sent) = step else { + panic!("pre_call answers at once"); + }; + assert_eq!(sent.body["system"], "edited"); + run( + py, + &locals, + c" +assert logger.pre_input == [ + {'role': 'user', 'content': '{\"model\":\"m\",\"system\":\"caf\u{e9} \u{1f600}\",\"temperature\":1}'} +], logger.pre_input +", + ); + }); + } + #[rstest::rstest] fn a_cycle_through_the_retained_headers_is_collected() { Python::initialize(); diff --git a/litellm/rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm/rust/crates/callbacks-legacy-python/src/callbacks.rs index 9eed008910c..ff31a686a8a 100644 --- a/litellm/rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm/rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -4,7 +4,11 @@ use litellm_host::interceptors::{RequestContext, WireRequest}; use litellm_host_python::json_to_py; -use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use pyo3::{ + exceptions::PyBaseException, + prelude::*, + types::{PyBytes, PyDict}, +}; use serde_json::Value; use crate::logger::PythonLogger; @@ -13,7 +17,7 @@ use crate::python::{Logging, Wrapper}; pub struct SentRequest<'a> { pub url: &'a str, pub headers: &'a Py, - pub body: &'a Py, + pub body: &'a [u8], } pub trait LegacyCallbacks { @@ -47,6 +51,7 @@ pub trait LegacyCallbacks { status: u16, headers: &[(String, String)], body: &str, + streamed: bool, request: Option>, ) -> PyResult<()>; @@ -182,12 +187,21 @@ impl LegacyCallbacks for PythonLogger { status: u16, headers: &[(String, String)], body: &str, + streamed: bool, request: Option>, ) -> PyResult<()> { - let request = request.map(|sent| (sent.url, sent.headers.bind(py), sent.body.bind(py))); + let request = + request.map(|sent| (sent.url, sent.headers.bind(py), PyBytes::new(py, sent.body))); Logging::RecordResponse.call( py, - (self.object(py), status, headers.to_vec(), body, request), + ( + self.object(py), + status, + headers.to_vec(), + body, + streamed, + request, + ), )?; Ok(()) } diff --git a/litellm/rust/crates/callbacks-legacy-python/src/error.rs b/litellm/rust/crates/callbacks-legacy-python/src/error.rs index 97ce62ae99c..4be7e31b81f 100644 --- a/litellm/rust/crates/callbacks-legacy-python/src/error.rs +++ b/litellm/rust/crates/callbacks-legacy-python/src/error.rs @@ -5,13 +5,17 @@ use pyo3::{PyErr, exceptions::PyRuntimeError}; pub(crate) enum Error { #[error("call logging is not initialized")] LoggingNotInitialized, + #[error("the Messages input could not be serialized for logging")] + LoggedInput(#[source] serde_json::Error), } impl From for PyErr { fn from(error: Error) -> Self { let message = Report::new(&error).to_string(); match error { - Error::LoggingNotInitialized => PyRuntimeError::new_err(message), + Error::LoggingNotInitialized | Error::LoggedInput(_) => { + PyRuntimeError::new_err(message) + } } } } diff --git a/litellm/rust/crates/callbacks-legacy-python/src/lib.rs b/litellm/rust/crates/callbacks-legacy-python/src/lib.rs index 3a8c15ae54c..8ee096bd189 100644 --- a/litellm/rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm/rust/crates/callbacks-legacy-python/src/lib.rs @@ -14,6 +14,7 @@ mod deferred; mod error; mod logger; mod mapping; +mod messages_input; mod python; pub use adapter::LegacyLogging; pub use call::PublicCall; diff --git a/litellm/rust/crates/callbacks-legacy-python/src/messages_input.rs b/litellm/rust/crates/callbacks-legacy-python/src/messages_input.rs new file mode 100644 index 00000000000..3cf42205cb5 --- /dev/null +++ b/litellm/rust/crates/callbacks-legacy-python/src/messages_input.rs @@ -0,0 +1,135 @@ +use pyo3::{ + prelude::*, + types::{PyBool, PyDict, PyFloat, PyInt, PyList}, +}; +use serde::ser::{Error as _, Serialize, SerializeMap, Serializer}; +use serde_json::{Map, Value}; + +use crate::error::Error; + +struct Ordered<'a> { + order: &'a [String], + sent: &'a Map, + as_called: &'a Map, +} + +impl Serialize for Ordered<'_> { + fn serialize(&self, serializer: S) -> Result { + let mut map = serializer.serialize_map(None)?; + for name in self.order { + if let Some(value) = self.as_called.get(name).or_else(|| self.sent.get(name)) { + map.serialize_entry(name, value)?; + } + } + for (name, value) in self.sent { + if !self.order.contains(name) { + map.serialize_entry(name, value)?; + } + } + map.end() + } +} + +pub(crate) fn is_caller_scalar(value: &Bound<'_, PyAny>) -> bool { + value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() +} + +pub(crate) fn serialized( + order: &[String], + sent: &Value, + as_called: &Map, +) -> Result { + let bytes = match sent { + Value::Object(sent) => litellm_host::raw_json::to_vec(&Ordered { + order, + sent, + as_called, + }), + other => litellm_host::raw_json::to_vec(other), + } + .map_err(Error::LoggedInput)?; + String::from_utf8(bytes).map_err(|error| Error::LoggedInput(serde_json::Error::custom(error))) +} + +pub(crate) fn logged<'py>(py: Python<'py>, content: &str) -> PyResult> { + let entry = PyDict::new(py); + entry.set_item("role", "user")?; + entry.set_item("content", content)?; + Ok(PyList::new(py, [entry])?.into_any()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn names(order: &[&str]) -> Vec { + order.iter().map(|name| (*name).to_owned()).collect() + } + + #[rstest] + #[case::caller_order( + &["messages", "max_tokens", "model"], + r#"{"messages":[],"max_tokens":8,"model":"m"}"# + )] + #[case::unlisted_keys_follow_in_sent_order( + &["model"], + r#"{"model":"m","max_tokens":8,"messages":[]}"# + )] + #[case::unknown_name_skipped( + &["absent", "model", "max_tokens", "messages"], + r#"{"model":"m","max_tokens":8,"messages":[]}"# + )] + fn top_level_keys_follow_the_given_order(#[case] order: &[&str], #[case] expected: &str) { + let sent = json!({"model": "m", "max_tokens": 8, "messages": []}); + + assert_eq!( + serialized(&names(order), &sent, &Map::new()).unwrap(), + expected + ); + } + + #[rstest] + #[case::int_for_a_float(json!(1.0), json!(1), r#"{"temperature":1}"#)] + #[case::bool_for_an_int(json!(1), json!(true), r#"{"temperature":true}"#)] + fn a_caller_scalar_keeps_its_spelling( + #[case] sent: Value, + #[case] called: Value, + #[case] expected: &str, + ) { + let as_called = Map::from_iter([("temperature".to_owned(), called)]); + + assert_eq!( + serialized( + &names(&["temperature"]), + &json!({"temperature": sent}), + &as_called + ) + .unwrap(), + expected + ); + } + + #[rstest] + #[case::non_ascii(json!("caf\u{e9} \u{1f600}"), "\"caf\u{e9} \u{1f600}\"")] + #[case::small_float(json!(1e-7), "1e-7")] + #[case::wide_integer( + json!(litellm_host::raw_json::encode("1180591620717411303424")), + "1180591620717411303424" + )] + fn values_are_written_as_the_wire_writes_them(#[case] leaf: Value, #[case] spelled: &str) { + let sent = json!({"system": leaf}); + + let logged = serialized(&names(&["system"]), &sent, &Map::new()).unwrap(); + + assert_eq!(logged, format!(r#"{{"system":{spelled}}}"#)); + assert_eq!( + logged.as_bytes(), + litellm_host::raw_json::to_vec(&sent).unwrap() + ); + } +} diff --git a/litellm/rust/crates/callbacks-legacy-python/src/python.rs b/litellm/rust/crates/callbacks-legacy-python/src/python.rs index 58ddf7c708a..85ee3b831ad 100644 --- a/litellm/rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm/rust/crates/callbacks-legacy-python/src/python.rs @@ -42,8 +42,8 @@ pub(crate) enum Logging { Update, #[strum(serialize = "pre_call")] PreCall, - #[strum(serialize = "messages_input")] - MessagesInput, + #[strum(serialize = "messages_input_order")] + MessagesInputOrder, #[strum(serialize = "post_call")] PostCall, #[strum(serialize = "record_response")] diff --git a/litellm/rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm/rust/crates/callbacks-legacy-python/src/test_support.rs index 5b8b2bc4f96..ab1a4ab9cdb 100644 --- a/litellm/rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm/rust/crates/callbacks-legacy-python/src/test_support.rs @@ -62,11 +62,11 @@ FAKES = { custom_llm_provider=provider, ), 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), - 'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}], + 'messages_input_order': lambda body, kwargs: tuple(body), 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), - 'record_response': lambda logger, status, headers, body, sent: logger.model_call_details.update( + 'record_response': lambda logger, status, headers, body, streamed, sent: logger.model_call_details.update( response_headers=dict(headers) ), 'retried_request': lambda logger, logged_body, body, reasons: ( diff --git a/litellm/rust/crates/host/Cargo.toml b/litellm/rust/crates/host/Cargo.toml index 2c113bc644c..73392827f31 100644 --- a/litellm/rust/crates/host/Cargo.toml +++ b/litellm/rust/crates/host/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +bytes.workspace = true litellm-secret-value.workspace = true futures-util.workspace = true litellm-coroutine.workspace = true diff --git a/litellm/rust/crates/host/src/interceptors.rs b/litellm/rust/crates/host/src/interceptors.rs index 8146448cd98..40591ec9ccd 100644 --- a/litellm/rust/crates/host/src/interceptors.rs +++ b/litellm/rust/crates/host/src/interceptors.rs @@ -1,5 +1,6 @@ use std::future::Future; +use bytes::Bytes; use serde_json::Value; /// The provider request as it is about to leave, offered to the host for rewriting. @@ -40,7 +41,7 @@ pub struct RequestContext { #[derive(Clone, Debug, PartialEq, Eq)] pub struct RawResponse { pub body: String, - /// The provider's HTTP status and headers, when the route forwards them to logging. + /// The provider's HTTP exchange, when the route forwards it to logging. pub head: Option, } @@ -48,6 +49,8 @@ pub struct RawResponse { pub struct ResponseHead { pub status: u16, pub headers: Vec<(String, String)>, + pub sent_body: Bytes, + pub streamed: bool, } impl RawResponse { diff --git a/litellm/rust/crates/http/src/outbound.rs b/litellm/rust/crates/http/src/outbound.rs index b11041151af..49ff7fa6df0 100644 --- a/litellm/rust/crates/http/src/outbound.rs +++ b/litellm/rust/crates/http/src/outbound.rs @@ -38,7 +38,7 @@ pub fn json( pub fn json_bytes( url: &str, headers: &[(String, String)], - body: Vec, + body: impl Into, timeout: Option, ) -> Result { let url = Url::parse(url).map_err(Repr::ParseUrl)?; diff --git a/litellm/rust/crates/inference-messages/src/handler.rs b/litellm/rust/crates/inference-messages/src/handler.rs index 4796a640e25..d54d4f5f8d1 100644 --- a/litellm/rust/crates/inference-messages/src/handler.rs +++ b/litellm/rust/crates/inference-messages/src/handler.rs @@ -30,8 +30,7 @@ use serde_json::{Map, Value}; use super::{ Delivered, Error, HostAnswer, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, - Transport, TransportBody, TransportRequest, TransportResponse, UnencodableHeader, - WorkloadIdentity, + Transport, TransportRequest, TransportResponse, UnencodableHeader, WorkloadIdentity, common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, constants::{ MESSAGES_HTTP_ERROR_ATTEMPTS, MESSAGES_STREAM_ERROR_BODY_READ_SECS, MESSAGES_TIMEOUT_SECS, @@ -42,7 +41,7 @@ use super::{ use litellm_inference::{ caching::{CacheRequest, CallCache}, context::CallContext, - outbound::raw_leaf_outbound_request, + outbound::{raw_leaf_body, serialized_outbound_request}, }; struct ProviderCall { @@ -505,9 +504,13 @@ async fn call_provider( let timeout = request .timeout .unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); - let response = match send_recovering(http, transport, &request, timeout, context).await? { + let (upstream, sent_body) = + send_recovering(http, transport, &request, timeout, context).await?; + let response = match upstream { Upstream::Native(response) => response, - Upstream::Hosted(response) => return hosted_response(response, request, context).await, + Upstream::Hosted(response) => { + return hosted_response(response, sent_body, request, context).await; + } }; let ProviderCall { identity, @@ -522,6 +525,8 @@ async fn call_provider( let head = ResponseHead { status: response.status().as_u16(), headers: super::encoding::header_pairs(response.headers()), + sent_body, + streamed: stream, }; if stream { context @@ -556,6 +561,7 @@ async fn call_provider( async fn hosted_response( response: TransportResponse, + sent_body: Bytes, request: ProviderCall, context: &CallContext<'_, impl Interceptors>, ) -> Result { @@ -579,6 +585,8 @@ async fn hosted_response( .raw_response_received(RawResponse::new(String::new()).with_head(ResponseHead { status, headers: headers.clone(), + sent_body, + streamed: true, })) .await?; let streamed = @@ -593,9 +601,12 @@ async fn hosted_response( }; log_response_body(&text); context - .raw_response_received( - RawResponse::new(text.clone()).with_head(ResponseHead { status, headers }), - ) + .raw_response_received(RawResponse::new(text.clone()).with_head(ResponseHead { + status, + headers, + sent_body, + streamed: false, + })) .await?; if serde_json::from_str::(&text).is_err() { return Err(litellm_http::Error::upstream(status, text).into()); @@ -606,14 +617,14 @@ async fn hosted_response( fn hosted_request( authenticated: Authenticated, url: &str, - body: &Value, + body: Bytes, stream: bool, ) -> Result { if authenticated.signer.is_none() { return Ok(TransportRequest { url: url.to_owned(), headers: authenticated.headers, - body: TransportBody::Json(body.clone()), + body, stream, }); } @@ -622,7 +633,7 @@ fn hosted_request( .iter() .map(|(name, _)| name.clone()) .collect(); - let signed = raw_leaf_outbound_request(authenticated, url, body, None)?; + let signed = serialized_outbound_request(authenticated, url, body.clone(), None)?; let headers = signed .headers() .iter() @@ -635,15 +646,10 @@ fn hosted_request( Some((spelled, value.to_str().ok()?.to_owned())) }) .collect(); - let bytes = signed - .body() - .and_then(reqwest::Body::as_bytes) - .map(Bytes::copy_from_slice) - .unwrap_or_default(); Ok(TransportRequest { url: signed.url().to_string(), headers, - body: TransportBody::Signed(bytes), + body, stream, }) } @@ -654,7 +660,7 @@ async fn send_recovering( request: &ProviderCall, timeout: Duration, context: &CallContext<'_, impl Interceptors>, -) -> Result { +) -> Result<(Upstream, Bytes), Error> { let config = request.provider.config(); let mut body = std::borrow::Cow::Borrowed(&request.wire.body); let mut attempt = 1; @@ -680,16 +686,24 @@ async fn send_recovering( headers: request.wire.headers.clone(), signer: request.signer.clone(), }; + let sent_body = raw_leaf_body(&body)?; let response = match transport { None => Upstream::Native( - send(http, authenticated, &request.wire.url, &body, timeout).await?, + send( + http, + authenticated, + &request.wire.url, + sent_body.clone(), + timeout, + ) + .await?, ), Some(transport) => Upstream::Hosted( transport .send(hosted_request( authenticated, &request.wire.url, - &body, + sent_body.clone(), request.stream, )?) .await?, @@ -703,7 +717,7 @@ async fn send_recovering( } let failure = match response { Upstream::Native(response) if response.status().is_success() => { - return Ok(Upstream::Native(response)); + return Ok((Upstream::Native(response), sent_body)); } Upstream::Native(response) => { provider_failure(response, timeout, request.stream).await? @@ -721,7 +735,7 @@ async fn send_recovering( raised: Some(error), } } - hosted @ Upstream::Hosted(_) => return Ok(hosted), + hosted @ Upstream::Hosted(_) => return Ok((hosted, sent_body)), }; let retried = if attempt == MESSAGES_HTTP_ERROR_ATTEMPTS { None @@ -756,10 +770,10 @@ async fn send( http: &litellm_http::Client, authenticated: Authenticated, url: &str, - body: &Value, + body: Bytes, timeout: Duration, ) -> Result { - let request = raw_leaf_outbound_request(authenticated, url, body, None)?; + let request = serialized_outbound_request(authenticated, url, body, None)?; tokio::time::timeout(timeout, litellm_inference::outbound::send(request, http)) .await .map_err(litellm_http::Error::read_timeout)? @@ -1746,13 +1760,13 @@ mod tests { .await .unwrap(); assert_eq!(call.wire.url, format!("{base}{path}")); - let signed = raw_leaf_outbound_request( + let signed = serialized_outbound_request( Authenticated { headers: call.wire.headers, signer: call.signer, }, &call.wire.url, - &call.wire.body, + raw_leaf_body(&call.wire.body).unwrap(), None, ) .unwrap(); diff --git a/litellm/rust/crates/inference-messages/src/lib.rs b/litellm/rust/crates/inference-messages/src/lib.rs index cd312ee61b1..e65804ce108 100644 --- a/litellm/rust/crates/inference-messages/src/lib.rs +++ b/litellm/rust/crates/inference-messages/src/lib.rs @@ -28,9 +28,7 @@ pub use host_answer::HostAnswer; pub use litellm_inference::Error; pub use pre_request::PreRequestHooks; pub use safeguards::Safeguards; -pub use transport::{ - Delivered, Transport, TransportBody, TransportRequest, TransportResponse, UnencodableHeader, -}; +pub use transport::{Delivered, Transport, TransportRequest, TransportResponse, UnencodableHeader}; pub use types::{ MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, messages_body, diff --git a/litellm/rust/crates/inference-messages/src/transport.rs b/litellm/rust/crates/inference-messages/src/transport.rs index 1353592447b..7d26cc81348 100644 --- a/litellm/rust/crates/inference-messages/src/transport.rs +++ b/litellm/rust/crates/inference-messages/src/transport.rs @@ -1,18 +1,12 @@ use bytes::Bytes; use futures_util::{future::BoxFuture, stream::BoxStream}; -use serde_json::Value; use crate::Error; -pub enum TransportBody { - Json(Value), - Signed(Bytes), -} - pub struct TransportRequest { pub url: String, pub headers: Vec<(String, String)>, - pub body: TransportBody, + pub body: Bytes, pub stream: bool, } diff --git a/litellm/rust/crates/inference-messages/tests/messages/host.rs b/litellm/rust/crates/inference-messages/tests/messages/host.rs index 7254830a2b0..8667616f279 100644 --- a/litellm/rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm/rust/crates/inference-messages/tests/messages/host.rs @@ -1,10 +1,8 @@ use litellm_core_utils::decline::Decline; -use litellm_host::lifecycle::ExecutionEvent; use std::sync::Mutex; -use litellm_host::{ - interceptors::{ExecutionFacts, RawResponse, RequestContext, ResultSource, WireRequest}, - lifecycle::CallEvent, +use litellm_host::interceptors::{ + ExecutionFacts, RawResponse, RequestContext, ResultSource, WireRequest, }; use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; @@ -47,17 +45,7 @@ impl RecordingHost { } fn raw_heads(&self) -> Vec { - self.events - .lock() - .unwrap() - .iter() - .filter_map(|event| match event { - CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - Some(raw.clone()) - } - _ => None, - }) - .collect() + self.events.raw_heads() } fn raw_responses(&self) -> Vec { @@ -444,6 +432,38 @@ async fn a_stream_emits_its_head_without_a_body_before_any_chunk(call: MessagesC ); } +#[rstest] +#[tokio::test] +async fn the_raw_head_carries_the_exact_bytes_the_provider_received( + call: MessagesCall, + #[values(false, true)] stream: bool, +) { + let response = if stream { + ResponseTemplate::new(200).set_body_raw( + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "text/event-stream", + ) + } else { + message_response() + }; + let upstream = upstream([response]).await; + let host = RecordingHost::passthrough(authenticated( + with_fields( + call, + json!({"stream": stream, "temperature": 1e-7, "system": "caf\u{e9} \u{1f600}"}), + ), + upstream.uri(), + )); + + let _ = run_through(&host).await.expect("messages call succeeds"); + + let [raw] = <[RawResponse; 1]>::try_from(host.raw_heads()) + .unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len())); + let head = raw.head.expect("the head is forwarded"); + assert_eq!(head.sent_body, only_request(&upstream).await.body); + assert_eq!(head.streamed, stream); +} + /// Python logs `optional_params` after its requested-param gate, so a dropped param must /// not resurface in callbacks. #[rstest] diff --git a/litellm/rust/crates/inference-messages/tests/messages/retry.rs b/litellm/rust/crates/inference-messages/tests/messages/retry.rs index e48482dc8b0..7eac27b0b10 100644 --- a/litellm/rust/crates/inference-messages/tests/messages/retry.rs +++ b/litellm/rust/crates/inference-messages/tests/messages/retry.rs @@ -149,6 +149,30 @@ async fn a_rejected_thinking_replay_is_sent_once_more_without_thinking( assert_eq!(chunks.is_empty(), !stream); } +#[rstest] +#[tokio::test] +async fn the_raw_head_of_a_retried_call_carries_the_retried_bytes( + call: MessagesCall, + #[values(false, true)] stream: bool, +) { + let success = if stream { + sse_response() + } else { + message_response() + }; + let upstream = upstream([rejection(400, INVALID_SIGNATURE), success]).await; + let host = RecordingCall::::new(replay(call, upstream.uri(), stream, enabled())); + + execute(&host).await.expect("the retry succeeds"); + + let [_, retried] = <[_; 2]>::try_from(received(&upstream).await).unwrap(); + let [raw] = <[_; 1]>::try_from(host.events.0.raw_heads()).unwrap(); + assert_eq!( + raw.head.expect("the head is forwarded").sent_body, + retried.body + ); +} + #[rstest] #[tokio::test] async fn a_second_rejection_is_the_error_the_caller_sees( diff --git a/litellm/rust/crates/inference-messages/tests/messages/transport.rs b/litellm/rust/crates/inference-messages/tests/messages/transport.rs index f561b09e55f..bb566e2ee1f 100644 --- a/litellm/rust/crates/inference-messages/tests/messages/transport.rs +++ b/litellm/rust/crates/inference-messages/tests/messages/transport.rs @@ -3,7 +3,7 @@ use std::collections::VecDeque; use bytes::Bytes; use futures_util::{StreamExt, TryStreamExt, future::BoxFuture}; use litellm_inference_messages::{ - Delivered, MessagesCallResponse, Transport, TransportBody, TransportRequest, TransportResponse, + Delivered, MessagesCallResponse, Transport, TransportRequest, TransportResponse, }; use litellm_inference_testing::RecordingSecrets; use rstest::rstest; @@ -28,6 +28,7 @@ struct Sent { url: String, headers: Vec<(String, String)>, body: Value, + raw_body: Bytes, stream: bool, } @@ -55,13 +56,11 @@ impl Transport for ScriptedTransport { &self, request: TransportRequest, ) -> BoxFuture<'static, Result> { - let TransportBody::Json(body) = request.body else { - panic!("an unsigned request hands the host its JSON body"); - }; self.sent.lock().unwrap().push(Sent { url: request.url, headers: request.headers, - body, + body: serde_json::from_slice(&request.body).expect("the host is handed JSON bytes"), + raw_body: request.body, stream: request.stream, }); let next = self @@ -113,10 +112,16 @@ async fn execute( transport: Arc, call: MessagesCall, ) -> Result<(MessagesCallResponse, Vec), Error> { - let host = RecordingCall::::new(call); + execute_recorded(transport, &RecordingCall::::new(call)).await +} + +async fn execute_recorded( + transport: Arc, + host: &RecordingCall, +) -> Result<(MessagesCallResponse, Vec), Error> { let route = messages_route(Arc::new(RecordingSecrets::empty())).with_transport(transport); match route - .execute(host.request()?, &host, Some(host.events.0.sender.clone())) + .execute(host.request()?, host, Some(host.events.0.sender.clone())) .await? { MessagesCallResponse::Stream { head, chunks } => Ok(( @@ -230,6 +235,38 @@ async fn a_recoverable_rejection_is_retried_through_the_host_transport(call: Mes assert!(second.body.get("thinking").is_none()); } +#[rstest] +#[tokio::test] +async fn the_raw_head_carries_the_bytes_the_host_transport_last_sent( + call: MessagesCall, + #[values(false, true)] stream: bool, +) { + let frames: &'static [&'static str] = + &["event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"]; + let success = if stream { + Scripted::Stream(frames) + } else { + Scripted::Complete(message_body()) + }; + let transport = ScriptedTransport::new([Scripted::Rejected(400, INVALID_SIGNATURE), success]); + let fields = json!({"stream": stream, "temperature": 1e-7, "system": "caf\u{e9}"}); + let host = + RecordingCall::::new(hosted_call(with_fields(call, replayed_thinking()), fields)); + + execute_recorded(transport.clone(), &host) + .await + .expect("the retry succeeds"); + + let sent = transport.sent(); + let [_, retried] = sent.as_slice() else { + panic!("two requests"); + }; + let [raw] = <[_; 1]>::try_from(host.events.0.raw_heads()).unwrap(); + let head = raw.head.expect("the head is forwarded"); + assert_eq!(head.sent_body, retried.raw_body); + assert_eq!(head.streamed, stream); +} + #[rstest] #[case::final_rejection(vec![Scripted::Rejected(429, "slow down")])] #[case::spent_retry(vec![ diff --git a/litellm/rust/crates/inference-messages/tests/support/mod.rs b/litellm/rust/crates/inference-messages/tests/support/mod.rs index 9115b6e89bb..fc1545796ae 100644 --- a/litellm/rust/crates/inference-messages/tests/support/mod.rs +++ b/litellm/rust/crates/inference-messages/tests/support/mod.rs @@ -153,6 +153,19 @@ impl Observations { } Ok(events) } + + pub fn raw_heads(&self) -> Vec { + self.lock() + .unwrap() + .iter() + .filter_map(|event| match event { + litellm_host::lifecycle::CallEvent::Execution( + litellm_host::lifecycle::ExecutionEvent::ProviderResponseReceived { raw }, + ) => Some(raw.clone()), + _ => None, + }) + .collect() + } } impl litellm_host::lifecycle::CallObserver for CallEvents { diff --git a/litellm/rust/crates/inference/src/outbound.rs b/litellm/rust/crates/inference/src/outbound.rs index 87fae0bbe9e..51636cf550c 100644 --- a/litellm/rust/crates/inference/src/outbound.rs +++ b/litellm/rust/crates/inference/src/outbound.rs @@ -1,5 +1,7 @@ use std::time::Duration; +use bytes::Bytes; + use litellm_http::outbound::{json, json_bytes, signed}; use litellm_llms::base_llm::auth::Authenticated; use serde_json::Value; @@ -35,15 +37,20 @@ pub fn outbound_request( } } -pub fn raw_leaf_outbound_request( +pub fn raw_leaf_body(body: &Value) -> Result { + litellm_host::raw_json::to_vec(body) + .map(Bytes::from) + .map_err(litellm_http::Error::request_body) +} + +pub fn serialized_outbound_request( authenticated: Authenticated, url: &str, - body: &Value, + body: Bytes, timeout: Option, ) -> Result { let Authenticated { headers, signer } = authenticated; - let bytes = litellm_host::raw_json::to_vec(body).map_err(litellm_http::Error::request_body)?; - let request = json_bytes(url, &headers, bytes, timeout)?; + let request = json_bytes(url, &headers, body, timeout)?; match signer { None => Ok(request), Some(signer) => signed(request, &signer), diff --git a/litellm/rust/crates/python-bridge/benches/serialization.rs b/litellm/rust/crates/python-bridge/benches/serialization.rs index d398d9fdbfc..e2a0d18e9d3 100644 --- a/litellm/rust/crates/python-bridge/benches/serialization.rs +++ b/litellm/rust/crates/python-bridge/benches/serialization.rs @@ -1,8 +1,11 @@ use std::{hint::black_box, time::Duration}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use litellm_host_python::{from_py, to_py}; -use pyo3::{prelude::*, types::PyDict}; +use litellm_host_python::{from_py, json_to_py, to_py}; +use pyo3::{ + prelude::*, + types::{PyDict, PyString}, +}; use serde_json::{Value, json}; const PAYLOAD_SIZES: &[(&str, usize)] = &[ @@ -89,12 +92,81 @@ fn bridge_serialization(c: &mut Criterion) { }); } +const MESSAGES_INPUT_TEXTS: &[(&str, &str)] = &[ + ( + "ascii", + "Here is some ordinary English text about the code.\n", + ), + ( + "mixed", + "Here is some ordinary English text about the code. Done \u{2705} caf\u{e9} \u{1f600}\n", + ), + ( + "unicode", + "\u{65e5}\u{672c}\u{8a9e}\u{306e}\u{30c6}\u{30ad}\u{30b9}\u{30c8} \u{1f600}\u{1f680} caf\u{e9}\n", + ), +]; + +fn messages_body(text: &str) -> Value { + let turn = text.repeat(6); + let messages: Vec = (0..) + .map(|index| { + json!([ + {"role": "user", "content": [{"type": "text", "text": turn}]}, + {"role": "assistant", "content": [ + {"type": "text", "text": turn}, + {"type": "tool_use", "id": format!("toolu_{index:06}"), "name": "bash", "input": {"command": "ls"}} + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": format!("toolu_{index:06}"), "content": turn}]} + ]) + }) + .scan(0, |size, turns| { + *size += turns.to_string().len(); + (*size < 210 * 1024).then_some(turns) + }) + .flat_map(|turns| turns.as_array().cloned().unwrap_or_default()) + .collect(); + json!({"model": "claude-sonnet-5", "max_tokens": 4096, "messages": messages, "stream": true}) +} + +fn messages_input(c: &mut Criterion) { + Python::initialize(); + Python::attach(|py| { + let json = py.import("json").expect("Python json module should import"); + for &(label, text) in MESSAGES_INPUT_TEXTS { + let body = messages_body(text); + let dict = json_to_py(py, &body).expect("body should convert"); + c.bench_with_input( + BenchmarkId::new("messages_input_python_json_dumps", label), + &dict, + |b, dict| { + b.iter(|| { + json.call_method1("dumps", (black_box(dict),)) + .expect("body should serialize") + }) + }, + ); + c.bench_with_input( + BenchmarkId::new("messages_input_rust_raw_json", label), + &body, + |b, body| { + b.iter(|| { + let bytes = litellm_host::raw_json::to_vec(black_box(body)) + .expect("body should serialize"); + PyString::new(py, std::str::from_utf8(&bytes).expect("JSON is UTF-8")) + }) + }, + ); + } + }); +} + criterion_group! { name = benches; config = Criterion::default() .sample_size(20) .warm_up_time(Duration::from_secs(1)) .measurement_time(Duration::from_secs(4)); - targets = bridge_serialization + targets = bridge_serialization, messages_input } criterion_main!(benches); diff --git a/litellm/rust/crates/python-bridge/src/routes/messages/transport.rs b/litellm/rust/crates/python-bridge/src/routes/messages/transport.rs index 12e782f61aa..8baed9b8a21 100644 --- a/litellm/rust/crates/python-bridge/src/routes/messages/transport.rs +++ b/litellm/rust/crates/python-bridge/src/routes/messages/transport.rs @@ -6,9 +6,9 @@ use litellm_host::{ machine::HostServices, protocol::{Protocol, Reply}, }; -use litellm_host_python::{json_to_py, to_py}; +use litellm_host_python::to_py; use litellm_inference_messages::{ - Delivered, Error, Transport, TransportBody, TransportRequest, TransportResponse, + Delivered, Error, Transport, TransportRequest, TransportResponse, }; use pyo3::{ exceptions::{PyBaseException, PyException, PyStopAsyncIteration}, @@ -160,10 +160,6 @@ impl HostedTransport { .handler .as_ref() .ok_or(litellm_host_python::Error::MissingState)?; - let body = match &sent.body { - TransportBody::Json(body) => json_to_py(py, body)?.unbind(), - TransportBody::Signed(bytes) => PyBytes::new(py, bytes).into_any().unbind(), - }; let awaitable = py .import(ROUTE_HOST_MODULE)? .getattr("transport_send")? @@ -172,7 +168,7 @@ impl HostedTransport { request, &sent.url, to_py(py, &sent.headers)?, - body, + PyBytes::new(py, &sent.body), sent.stream, ))?; self.pending = Some(Pending::Send(reply)); diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 19124fcffc0..44cac1a0eb2 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -9,8 +9,8 @@ from __future__ import annotations import asyncio import contextvars import datetime +import functools import inspect -import json import sys import traceback import uuid @@ -266,22 +266,18 @@ def pre_call( logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) -def messages_input(body: Mapping[str, object], kwargs: Mapping[str, object]) -> Sequence[Mapping[str, str]]: +@functools.cache +def _messages_handler_parameters() -> tuple[str, ...]: from litellm.messages import anthropic_messages_handler + return tuple(inspect.signature(anthropic_messages_handler).parameters) + + +def messages_input_order(body: Mapping[str, object], kwargs: Mapping[str, object]) -> tuple[str, ...]: # Python's Messages handler builds its body as messages, max_tokens, model, then its own parameters in # signature order and the caller's other keywords in their order, and logs that body as serialized - order: Final = ( - "messages", - "max_tokens", - "model", - *inspect.signature(anthropic_messages_handler).parameters, - *kwargs, - *body, - ) - return [ - {"role": "user", "content": json.dumps({name: body[name] for name in dict.fromkeys(order) if name in body})} - ] + order: Final = ("messages", "max_tokens", "model", *_messages_handler_parameters(), *kwargs, *body) + return tuple(name for name in dict.fromkeys(order) if name in body) def post_call( @@ -303,9 +299,9 @@ def wire_headers(headers: Sequence[tuple[str, str]]) -> httpx.Headers: return httpx.Headers([(name.encode(), value.encode()) for name, value in headers]) -def _sent_request(sent: tuple[str, Mapping[str, str], Mapping[str, object]]) -> httpx.Request: +def _sent_request(sent: tuple[str, Mapping[str, str], bytes]) -> httpx.Request: url, headers, body = sent - return httpx.Request("POST", url, headers=list(headers.items()), content=json.dumps(body, default=str).encode()) + return httpx.Request("POST", url, headers=list(headers.items()), content=body) def record_response( @@ -313,11 +309,11 @@ def record_response( status: int, headers: Sequence[tuple[str, str]], body: str, - sent: tuple[str, Mapping[str, str], Mapping[str, object]] | None, + streamed: bool, + sent: tuple[str, Mapping[str, str], bytes] | None, ) -> None: """Stores the provider response the way Python's Messages handler does, for logging to replay.""" request: Final = None if sent is None else _sent_request(sent) - streamed: Final = sent is not None and sent[2].get("stream") is True response: Final = ( httpx.Response(status_code=status, stream=_UnreadBody(), request=request) if streamed diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index b4113bf7efa..93b4f362569 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -6,7 +6,6 @@ import errno import functools import inspect import ipaddress -import json import os import socket from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence @@ -1263,14 +1262,14 @@ async def transport_send( request: Mapping[str, object], url: str, headers: Sequence[tuple[str, str]], - body: object, + body: bytes, stream: bool, ) -> Sent: try: response: Final = await handler.post( # pyright: ignore[reportUnknownMemberType] # the handler's untyped dict params url=_sent_url(request, url), headers=dict(headers), - data=body if isinstance(body, bytes) else json.dumps(body), + data=body, stream=stream, logging_obj=cast( # cast-ok: the call's logger, handed on as Python's handler hands it "LiteLLMLoggingObj | None", request.get("litellm_logging_obj") diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index 0550b530913..c2cd3c5dd71 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -1,9 +1,11 @@ import asyncio -from collections.abc import Iterator +import json +from collections.abc import Iterator, Mapping, Sequence from contextvars import ContextVar from datetime import datetime from typing import Final +import httpx import pytest from pydantic import TypeAdapter @@ -13,11 +15,20 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.integrations.argilla import ArgillaItem from litellm.types.utils import ModelResponse, StandardLoggingPayload -from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, REQUEST, call +from tests.test_litellm_rust.messages.data import ( + BAD_REQUEST, + MESSAGES, + MESSAGES_EVENTS, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + RATE_LIMITED, + REQUEST, + call, +) from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed -from tests.test_litellm_rust.support.results import response_id, with_id +from tests.test_litellm_rust.support.results import consume, response_id, with_id from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -510,3 +521,152 @@ async def test_sync_call_raises_like_python_without_logging(server: RecordingSer assert not server.requests assert not recorder.names + + +_HTTPX_RESPONSE: Final = TypeAdapter(httpx.Response, config={"arbitrary_types_allowed": True}) +_STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) +_WIDE: Final = 2**70 +_TOOLS: Final = ( + { + "name": "get_weather", + "input_schema": {"type": "object", "properties": {"celsius": {"type": "number", "minimum": 1e-7}}}, + }, +) +_INVALID_SIGNATURE: Final = { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "messages.1.content.0: Invalid `signature` in `thinking` block", + }, +} +_THINKING_REPLAY: Final = ( + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Look it up.", "signature": "EqQBCkYIARgCKkB"}, + {"type": "tool_use", "id": "toolu_01Wide", "name": "get_weather", "input": {"big": _WIDE}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01Wide", "content": "18C"}]}, +) + + +def _messages_logger() -> Logging: + return Logging( + model=MESSAGES_MODEL, + messages=[], + stream=False, + start_time=datetime.now(), + call_type="anthropic_messages", + litellm_call_id=CALL_ID, + function_id=CALL_ID, + ) + + +_TOOL_TURN: Final = ( + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01Wide", "name": "get_weather", "input": {"big": _WIDE}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01Wide", "content": "18C"}]}, +) +_LOGGED_INPUT: Final = TypeAdapter(tuple[dict[str, str]]) +_DECODER: Final = json.JSONDecoder() +_DECODED_AT: Final = TypeAdapter(tuple[object, int]) +_KEY_AT: Final = TypeAdapter(tuple[str, int]) + + +def _members(text: str, start: int = 1) -> tuple[tuple[str, str], ...]: + if text[start] == "}": + return () + key, colon = _KEY_AT.validate_python(_DECODER.raw_decode(text, start)) + _, end = _DECODED_AT.validate_python(_DECODER.raw_decode(text, colon + 1)) + return ((key, text[colon + 1 : end]), *_members(text, end + (text[end] == ","))) + + +async def _logged_exchange( + server: RecordingServer, + stream: bool, + messages: Sequence[Mapping[str, object]] = MESSAGES, + thinking: Mapping[str, object] | None = None, +) -> tuple[httpx.Request, str]: + logging_obj: Final = _messages_logger() + if stream: + server.enqueue(_STREAM) + response: Final = await call( + server, + litellm_logging_obj=logging_obj, + stream=stream, + messages=list(messages), + thinking=thinking, + system="café \U0001f600", + tools=list(_TOOLS), + ) + if stream: + _ = await consume(response) + (logged,) = _LOGGED_INPUT.validate_python(logging_obj.model_call_details["input"]) + assert logged["role"] == "user" + return ( + _HTTPX_RESPONSE.validate_python(logging_obj.model_call_details["httpx_response"]).request, + logged["content"], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("unary", "stream")) +async def test_the_logged_provider_request_is_the_bytes_the_provider_received( + server: RecordingServer, stream: bool +) -> None: + sent, _ = await _logged_exchange(server, stream) + + assert_served_natively(server) + assert sent.content == server.requests[0].raw_body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("unary", "stream")) +async def test_the_logged_provider_request_of_a_retried_call_is_the_retried_bytes( + server: RecordingServer, stream: bool +) -> None: + server.responses.insert(0, ResponseSpec(body=_INVALID_SIGNATURE, status=400)) + + sent, _ = await _logged_exchange(server, stream, _THINKING_REPLAY, {"type": "enabled", "budget_tokens": 2048}) + + assert_served_natively(server, count=2) + assert sent.content == server.requests[1].raw_body + assert json.loads(sent.content)["messages"][1]["content"][-1]["input"] == {"big": _WIDE} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("unary", "stream")) +async def test_the_logged_input_holds_the_bytes_the_provider_received_in_the_handler_key_order( + server: RecordingServer, stream: bool +) -> None: + sent, logged = await _logged_exchange(server, stream, _TOOL_TURN) + + assert_served_natively(server) + members: Final = _members(logged) + assert dict(members) == dict(_members(server.requests[0].raw_body.decode())) + assert [name for name, _ in members][:3] == ["messages", "max_tokens", "model"] + assert sent.content == server.requests[0].raw_body + assert json.loads(logged)["messages"][1]["content"][0]["input"] == {"big": _WIDE} + + +@pytest.mark.asyncio +async def test_a_pre_api_call_edit_reaches_the_wire_but_not_the_logged_input(server: RecordingServer) -> None: + class Edit(CustomLogger): + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + request_body(kwargs)["system"] = "edited" + + recorder: Final = RecordingLogger() + + response: Final = await call(server, system="original", callbacks=[Edit(), recorder]) + + assert_served_natively(server, response) + assert _OBJECT.validate_python(server.requests[0].body)["system"] == "edited" + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + (logged,) = _LOGGED_INPUT.validate_python(pre_call["input"]) + assert json.loads(logged["content"])["system"] == "original" diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index 99b6101da4e..c8cc16a932f 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -1,5 +1,4 @@ import asyncio -import json from collections.abc import AsyncIterator, Iterator from typing import Final @@ -108,23 +107,14 @@ async def _streamed_pre_call_input(server: RecordingServer) -> object: @pytest.mark.asyncio -async def test_stream_logs_the_request_body_as_input_like_python(server: RecordingServer) -> None: +async def test_stream_logs_the_request_body_as_input_in_the_python_handler_key_order(server: RecordingServer) -> None: logged: Final = await _streamed_pre_call_input(server) assert_served_natively(server) assert logged == [ { "role": "user", - "content": json.dumps( - { - "messages": [{"role": "user", "content": "Hello"}], - "max_tokens": 64, - "model": "claude-sonnet-5", - "metadata": {"user_id": "user-1"}, - "stop_sequences": ["END"], - "stream": True, - "system": "Be brief", - } - ), + "content": '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":64,"model":"claude-sonnet-5",' + '"metadata":{"user_id":"user-1"},"stop_sequences":["END"],"stream":true,"system":"Be brief"}', } ] diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index ec1b1282a7e..271a00696cf 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -1,6 +1,5 @@ import datetime import functools -import json from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType from typing import Final @@ -1009,8 +1008,9 @@ def test_a_positive_timeout_keeps_the_native_client(timeout: object) -> None: @pytest.mark.asyncio -async def test_transport_send_posts_what_the_python_handler_posts() -> None: +async def test_transport_send_posts_the_native_wire_bytes_unchanged() -> None: logger: Final = object() + wire: Final = '{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"h\u00e9llo"}]}'.encode() timeout: Final = httpx.Timeout(30.0, connect=2.0) handler: Final = _Handler( httpx.Response(200, json={"id": "msg_1"}, headers={"request-id": "req_1"}, request=_REQUEST) @@ -1021,7 +1021,7 @@ async def test_transport_send_posts_what_the_python_handler_posts() -> None: {"litellm_logging_obj": logger, "timeout": timeout}, _URL, [("x-api-key", "sk-ant"), ("anthropic-version", "2023-06-01")], - {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "héllo"}]}, + wire, False, ) @@ -1029,7 +1029,7 @@ async def test_transport_send_posts_what_the_python_handler_posts() -> None: { "url": _URL, "headers": {"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, - "data": json.dumps({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "héllo"}]}), + "data": wire, "stream": False, "logging_obj": logger, "timeout": timeout, @@ -1055,7 +1055,7 @@ async def test_transport_send_leaves_a_stream_for_the_route_to_read() -> None: httpx.Response(200, stream=httpx.ByteStream(b"event: message_stop\n\n"), request=_REQUEST) ) - sent: Final = await messages_host.transport_send(handler, {}, _URL, [], {}, True) + sent: Final = await messages_host.transport_send(handler, {}, _URL, [], b"", True) assert sent.chunks is not None assert sent.text is None @@ -1069,7 +1069,7 @@ async def test_transport_send_reports_a_rejection_with_the_status_error(checked: error: Final = httpx.HTTPStatusError("rate limited", request=_REQUEST, response=rejected) handler: Final = _Handler(error if checked else rejected) - sent: Final = await messages_host.transport_send(handler, {}, _URL, [], {}, False) + sent: Final = await messages_host.transport_send(handler, {}, _URL, [], b"", False) assert (sent.status, sent.text) == (429, "slow down") assert ("retry-after", "7") in sent.headers @@ -1082,7 +1082,7 @@ async def test_a_rejection_python_cannot_read_reaches_the_caller_as_raised() -> unread: Final = httpx.Response(500, stream=httpx.ByteStream(b"unread"), request=_REQUEST) with pytest.raises(messages_host.RawTransportFailure) as raised: - await messages_host.transport_send(_Handler(unread), {}, _URL, [], {}, True) + await messages_host.transport_send(_Handler(unread), {}, _URL, [], b"", True) assert isinstance(raised.value.error, httpx.ResponseNotRead) assert isinstance(raised.value.error.__context__, httpx.HTTPStatusError) @@ -1093,7 +1093,7 @@ async def test_a_handler_failure_is_raised_for_the_route_to_map() -> None: failure: Final = httpx.ConnectError("refused") with pytest.raises(httpx.ConnectError) as raised: - await messages_host.transport_send(_Handler(failure), {}, _URL, [], {}, False) + await messages_host.transport_send(_Handler(failure), {}, _URL, [], b"", False) assert raised.value is failure diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index b8d55c8e9d6..7143a550f76 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -47,12 +47,8 @@ def test_native_stream_headers_reach_spend_callbacks() -> None: def test_recorded_messages_stream_keeps_the_provider_header_names() -> None: logger: Final = _supplied_logger() - sent: Final[tuple[str, Mapping[str, str], Mapping[str, object]]] = ( - "http://upstream/v1/messages", - {}, - {"stream": True}, - ) - legacy.record_response(logger, 200, [("request-id", "req_native")], "", sent) + sent: Final[tuple[str, Mapping[str, str], bytes]] = ("http://upstream/v1/messages", {}, b'{"stream":true}') + legacy.record_response(logger, 200, [("request-id", "req_native")], "", True, sent) legacy.stream_opened(logger, {"additional_headers": {"llm_provider-request-id": "req_native"}}) assert logger.stream is True @@ -278,7 +274,7 @@ def test_the_rust_contract_matches_the_shim_signatures() -> None: assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} -def test_messages_input_is_the_body_as_the_python_handler_builds_and_serializes_it() -> None: +def test_messages_input_order_is_the_key_order_the_python_handler_builds_its_body_in() -> None: body: Final = { "model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], @@ -289,17 +285,19 @@ def test_messages_input_is_the_body_as_the_python_handler_builds_and_serializes_ "output_config": {"effort": "low"}, } - logged: Final = legacy.messages_input( - body, {"output_config": {"effort": "low"}, "context_management": {"edits": []}} + order: Final = legacy.messages_input_order( + body, {"output_config": {"effort": "low"}, "context_management": {"edits": []}, "absent": 1} ) - assert logged == [ - { - "role": "user", - "content": '{"messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "model": "claude-sonnet-4-5", ' - '"stream": true, "system": "brief", "output_config": {"effort": "low"}, "context_management": {"edits": []}}', - } - ] + assert order == ( + "messages", + "max_tokens", + "model", + "stream", + "system", + "output_config", + "context_management", + ) def _messages_logger() -> Logging: @@ -317,7 +315,7 @@ def _messages_logger() -> Logging: def test_native_provider_response_is_recorded_for_logging_replay() -> None: logger: Final = _messages_logger() - legacy.record_response(logger, 200, (("request-id", "req_native"),), '{"id": "msg_native"}', None) + legacy.record_response(logger, 200, (("request-id", "req_native"),), '{"id": "msg_native"}', False, None) recorded: Final[object] = logger.model_call_details["httpx_response"] assert isinstance(recorded, httpx.Response) @@ -329,7 +327,7 @@ def test_native_provider_response_is_recorded_for_logging_replay() -> None: def test_a_non_ascii_provider_header_is_recorded_as_httpx_decodes_it() -> None: logger: Final = _messages_logger() - legacy.record_response(logger, 200, (("x-fz-unicode", "caf\u00e9"),), '{"id": "msg_native"}', None) + legacy.record_response(logger, 200, (("x-fz-unicode", "caf\u00e9"),), '{"id": "msg_native"}', False, None) assert logger.model_call_details["response_headers"]["x-fz-unicode"] == "caf\u00e9" @@ -340,17 +338,17 @@ _SENT_HEADERS: Final = MappingProxyType({"x-api-key": "sk-ant", "anthropic-beta" def test_native_provider_response_carries_the_request_that_was_sent() -> None: logger: Final = _messages_logger() - body: Final = {"model": "claude-sonnet-4-5", "thread": {"type": "create"}} + wire: Final = '{"model":"claude-sonnet-4-5","system":"caf\u00e9","temperature":1e-7}'.encode() legacy.record_response( - logger, 200, (), '{"id": "msg_native"}', ("https://api.anthropic.com/v1/messages", _SENT_HEADERS, body) + logger, 200, (), '{"id": "msg_native"}', False, ("https://api.anthropic.com/v1/messages", _SENT_HEADERS, wire) ) recorded: Final = _HTTPX_RESPONSE.validate_python(logger.model_call_details["httpx_response"]) assert recorded.request.method == "POST" assert str(recorded.request.url) == "https://api.anthropic.com/v1/messages" assert recorded.request.headers["anthropic-beta"] == "message-threads-2026-08-12" - assert recorded.request.content == b'{"model": "claude-sonnet-4-5", "thread": {"type": "create"}}' + assert recorded.request.content == wire assert recorded.content == b'{"id": "msg_native"}' @@ -362,7 +360,8 @@ def test_native_stream_response_is_recorded_unread_like_the_python_stream() -> N 200, (("request-id", "req_native"),), "", - ("https://api.anthropic.com/v1/messages", _SENT_HEADERS, {"stream": True}), + True, + ("https://api.anthropic.com/v1/messages", _SENT_HEADERS, b'{"stream":true}'), ) recorded: Final = _HTTPX_RESPONSE.validate_python(logger.model_call_details["httpx_response"])