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"])