From afdb87d7a278ca27d15cb279d4c53458bf44f4db Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Fri, 6 Nov 2026 09:56:00 +0000 Subject: [PATCH] fix(messages): log the request body as input on native messages calls Python's messages handler calls pre_call with the request body as a json user message, and a stream's spend-log cache_key is computed over that input. The native adapter logged the string Messages, so a streamed call's cache_key differed from Python's --- .../python_contract.json | 4 +++ .../callbacks-legacy-python/src/adapter.rs | 29 ++++++++++++------ .../callbacks-legacy-python/src/callbacks.rs | 4 +-- .../callbacks-legacy-python/src/python.rs | 2 ++ .../src/test_support.rs | 1 + .../host/callbacks_legacy_python.py | 24 ++++++++++++++- .../messages/test_streaming.py | 30 ++++++++++++++++++- .../host/test_callbacks_legacy_python.py | 24 +++++++++++++++ 8 files changed, 105 insertions(+), 13 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..f88b3b38d94 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -27,6 +27,10 @@ "api_key", "additional_args" ], + "messages_input": [ + "body", + "kwargs" + ], "post_call": [ "logger", "original_response", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..e6428bbcbda 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -14,7 +14,7 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::{PyDateTime, PyDict, PyList}, + types::{PyDateTime, PyDict, PyList, PyString}, }; use serde_json::Value; @@ -22,7 +22,7 @@ use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, finalize, is_internal_call, - python::Streaming, + python::{Logging, Streaming}, setup, }; @@ -158,13 +158,20 @@ impl LegacyLogging { } } - fn input_description(&self) -> &'static str { - match self.operation { + fn logged_input<'py>( + &self, + py: Python<'py>, + body: &Bound<'py, PyDict>, + ) -> PyResult> { + let description = match self.operation { LoggingOperation::Completion => "Chat completions", LoggingOperation::Responses => "Responses", - LoggingOperation::Messages => "Messages", + LoggingOperation::Messages => { + return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py))); + } LoggingOperation::Ocr => "OCR document processing", - } + }; + Ok(PyString::new(py, description).into_any()) } /// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key. @@ -469,9 +476,10 @@ impl LegacyLogging { headers: headers.clone().unbind(), context: context.clone(), }); + let input = self.logged_input(py, &body)?; self.logger()?.pre_call( py, - self.input_description(), + &input, self.logged_api_key(context), &body, &headers, @@ -1276,7 +1284,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update #[rstest] #[case::completion(crate::LoggingOperation::Completion, "Chat completions")] #[case::responses(crate::LoggingOperation::Responses, "Responses")] - #[case::messages(crate::LoggingOperation::Messages, "Messages")] + #[case::messages( + crate::LoggingOperation::Messages, + 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( #[case] operation: crate::LoggingOperation, @@ -1317,7 +1328,7 @@ prepared = {'pages': replacement} &locals, c" assert logger.pre['complete_input_dict']['pages'] is replacement -assert logger.pre_input == description +assert str(logger.pre_input) == description assert original == [0] ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 1951c9ebfb4..a700bd14dbc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -26,7 +26,7 @@ pub trait LegacyCallbacks { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, @@ -148,7 +148,7 @@ impl LegacyCallbacks for PythonLogger { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6d910d35f53..a77d2ef280c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -42,6 +42,8 @@ pub(crate) enum Logging { Update, #[strum(serialize = "pre_call")] PreCall, + #[strum(serialize = "messages_input")] + MessagesInput, #[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 8acfa7a898b..b774bdc1049 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -62,6 +62,7 @@ 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)}], 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92590531cdc..4a4911d2396 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -9,6 +9,8 @@ from __future__ import annotations import asyncio import contextvars import datetime +import inspect +import json import traceback import uuid import weakref @@ -234,10 +236,30 @@ def update_logging( ) -def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: +def pre_call( + logger: LoggingSurface, input: str | list[dict[str, str]], api_key: str | None, additional_args: dict[str, object] +) -> None: logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) +def messages_input(body: Mapping[str, object], kwargs: Mapping[str, object]) -> list[dict[str, str]]: + from litellm.llms.anthropic.messages.native.handler import anthropic_messages_handler + + # 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})} + ] + + def post_call( logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] ) -> None: diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index e57631f48b9..a3e2fe2bed8 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -3,10 +3,14 @@ from typing import Final import pytest +import litellm +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +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.routing import assert_served_natively +from tests.test_litellm_rust.support.routing import assert_served_natively, routed pytestmark = pytest.mark.requires_rust_extension @@ -68,3 +72,27 @@ async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> assert len(await recorder.wait_for_async("async_log_success_event")) == 1 with pytest.raises(StopAsyncIteration): await anext(stream) + + +async def _streamed_pre_call_input(server: RecordingServer) -> object: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + with rebound(litellm, "callbacks", [recorder]): + stream: Final = await call( + server, stream=True, system="Be brief", metadata={"user_id": "user-1"}, stop_sequences=["END"] + ) + assert isinstance(stream, AsyncIterator) + _ = [chunk async for chunk in stream] + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + return pre_call["input"] + + +@pytest.mark.asyncio +async def test_stream_logs_the_request_body_as_input_like_python(server: RecordingServer) -> None: + native: Final = await _streamed_pre_call_input(server) + assert_served_natively(server) + with routed(RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY)): + python: Final = await _streamed_pre_call_input(server) + + assert native == python 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 1ba0188e76d..8fcfbfbd462 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -247,6 +247,30 @@ 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: + body: Final = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "system": "brief", + "stream": True, + "context_management": {"edits": []}, + "output_config": {"effort": "low"}, + } + + logged: Final = legacy.messages_input( + body, {"output_config": {"effort": "low"}, "context_management": {"edits": []}} + ) + + 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": []}}', + } + ] + + def _messages_logger() -> Logging: return Logging( model="claude-sonnet-4-5",