From c24927cf2ac363af4c55cf701181732f1849c8f7 Mon Sep 17 00:00:00 2001 From: tin Date: Tue, 4 Aug 2026 20:31:55 +0000 Subject: [PATCH] fix(proxy): report requested model on Anthropic streaming message_start Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_model_restamp.py | 79 +++++++++++++ litellm/proxy/common_request_processing.py | 25 +++- .../test_streaming_model_restamp.py | 109 ++++++++++++++++++ 3 files changed, 212 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/anthropic_endpoints/streaming_model_restamp.py create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py new file mode 100644 index 00000000000..857b6abd065 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -0,0 +1,79 @@ +""" +Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only +stream event carrying a model, so streamed responses report the requested model like +non-streaming ones do. + +Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the +provider passthrough path) or as event dicts (fake-stream and agentic paths). +""" + +import json + +from pydantic import TypeAdapter, ValidationError + +_MESSAGE_START_EVENT = "message_start" +_SSE_DATA_FIELD = "data:" + +_EVENT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +def _restamped_event(event: dict[str, object], requested_model: str) -> dict[str, object] | None: + message = event.get("message") + if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict): + return None + if message.get("model") == requested_model: + return None + return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is + + +def _restamped_data_line(line: str, requested_model: str) -> str | None: + stripped = line.strip() + if not stripped.startswith(_SSE_DATA_FIELD): + return None + payload = stripped[len(_SSE_DATA_FIELD) :].strip() + if not payload or payload == "[DONE]": + return None + try: + event = _EVENT_ADAPTER.validate_json(payload) + except ValidationError: + return None + restamped = _restamped_event(event, requested_model) + if restamped is None: + return None + return f"data: {json.dumps(restamped, separators=(',', ':'))}" + + +def _restamped_frame(frame: str, requested_model: str) -> str | None: + lines = frame.split("\n") + restamped = tuple(_restamped_data_line(line, requested_model) for line in lines) + if all(line is None for line in restamped): + return None + return "\n".join(new if new is not None else old for new, old in zip(restamped, lines)) + + +def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: + """ + Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``. + + Chunks that carry no model are returned unchanged. + """ + if isinstance(chunk, dict): + try: + event = _EVENT_ADAPTER.validate_python(chunk) + except ValidationError: + return chunk + return _restamped_event(event, requested_model) or chunk + + if isinstance(chunk, (bytes, bytearray)): + if _MESSAGE_START_EVENT.encode() not in chunk: + return chunk + restamped = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model) + return chunk if restamped is None else restamped.encode("utf-8") + + if isinstance(chunk, str): + if _MESSAGE_START_EVENT not in chunk: + return chunk + restamped = _restamped_frame(chunk, requested_model) + return chunk if restamped is None else restamped + + return chunk diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f9cad283166..fd5de0debec 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -70,6 +70,9 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + restamp_anthropic_stream_chunk_model, +) from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.types.utils import ( ModelResponse, @@ -1953,6 +1956,9 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, proxy_logging_obj=proxy_logging_obj, request=request, + restamp_model=( + None if _should_return_raw_model_name(self.data) else requested_model_from_client + ), ) return await create_response( generator=selected_data_generator, @@ -2801,6 +2807,18 @@ class ProxyBaseLLMRequestProcessing: else: return chunk + @staticmethod + def _sse_chunk_serializer(restamp_model: str | None) -> StreamChunkSerializer: + if not restamp_model: + return ProxyBaseLLMRequestProcessing.return_sse_chunk + + def serialize(chunk: object) -> str: + return ProxyBaseLLMRequestProcessing.return_sse_chunk( + restamp_anthropic_stream_chunk_model(chunk, restamp_model) + ) + + return serialize + @staticmethod async def _finalize_streaming_generator_cleanup( request: Request | None, @@ -2990,6 +3008,7 @@ class ProxyBaseLLMRequestProcessing: request_data: dict, proxy_logging_obj: ProxyLogging, request: Request | None = None, + restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events. @@ -2998,13 +3017,17 @@ class ProxyBaseLLMRequestProcessing: SSE serializers directly (rather than re-wrapping it in another ``async for: yield`` trampoline), so a streamed chunk traverses one fewer async-generator layer / coroutine resume on the hot path. + + ``restamp_model`` publishes that name on the Anthropic ``message_start`` + event in place of the provider's model, matching what the non-streaming + response reports. """ return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamp_model), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py new file mode 100644 index 00000000000..6e2c3f49445 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -0,0 +1,109 @@ +""" +Tests for restamping the public model on Anthropic Messages streaming chunks. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + restamp_anthropic_stream_chunk_model, +) +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + +def _message_start_frame(model: str) -> bytes: + payload = { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, + } + return f"event: message_start\ndata: {json.dumps(payload)}\n\n".encode() + + +def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: + async def _iterator_hook(**_kwargs): + for frame in frames: + yield frame + + proxy_logging_obj = MagicMock() + proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook + proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + return proxy_logging_obj + + +def _model_from_frame(frame: bytes | str) -> str: + text = frame.decode("utf-8") if isinstance(frame, bytes) else frame + data_line = next(line for line in text.split("\n") if line.startswith("data:")) + return json.loads(data_line[len("data:") :])["message"]["model"] + + +def test_restamps_sse_bytes_frame(): + restamped = restamp_anthropic_stream_chunk_model( + _message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1" + ) + + assert isinstance(restamped, bytes) + assert _model_from_frame(restamped) == "claude-auto-1" + assert b"event: message_start" in restamped + + +def test_restamps_event_dict(): + chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} + + restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2") + + assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}} + assert chunk["message"]["model"] == "claude-sonnet-4-6" + + +@pytest.mark.parametrize( + "chunk", + [ + b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + {"type": "content_block_delta", "delta": {"text": "hi"}}, + {"type": "message_start", "message": "not-a-dict"}, + b"event: message_start\ndata: not-json\n\n", + b"data: [DONE]\n\n", + ], +) +def test_leaves_chunks_without_a_model_untouched(chunk): + assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk + + +@pytest.mark.asyncio +async def test_sse_generator_publishes_requested_model_on_message_start(): + """The message_start event reports the requested model, not the provider's.""" + delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta_frame + + +@pytest.mark.asyncio +async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001"