diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4dd0deeb62b..e3912138d8c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -632,6 +632,7 @@ class ModelResponseIterator: self.tool_name_reverse_map: dict[str, str] = tool_name_reverse_map or {} # Generate response ID once per stream to match OpenAI-compatible behavior self.response_id = _generate_id() + self.served_model: str | None = None # Track if we're currently streaming a response_format tool self.is_response_format_tool: bool = False @@ -1067,6 +1068,8 @@ class ModelResponseIterator: } """ message_start_block: Final = MessageStartBlock(**chunk) + start_message: Final = message_start_block["message"] + self.served_model = start_message["model"] if "model" in start_message else None if "usage" in message_start_block["message"]: usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"]) elif type_chunk == "error": @@ -1098,6 +1101,7 @@ class ModelResponseIterator: ], usage=usage, id=self.response_id, + model=self.served_model, ) return returned_chunk diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index c3400dc40c3..1e6639ef537 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2719,3 +2719,81 @@ class TestRustChatCompletionsHook: "model": "m", "messages": [], } + + +def _served_model_stream_chunks(model: str | None) -> list[dict]: + message: Final = { + "id": "msg_served", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + if model is not None: + message["model"] = model + return [ + {"type": "message_start", "message": message}, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ] + + +def test_message_start_model_is_carried_on_stream_chunks(): + iterator: Final = ModelResponseIterator(None, sync_stream=True) + + parsed: Final = [ + iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks("claude-served-1") + ] + + assert all(chunk.model == "claude-served-1" for chunk in parsed) + + +def test_message_start_without_model_leaves_chunk_model_unset(): + iterator: Final = ModelResponseIterator(None, sync_stream=True) + + parsed: Final = [ + iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks(None) + ] + + assert all(chunk.model is None for chunk in parsed) + + +def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + served_model: Final = "claude-served-1" + sse_lines: Final = [ + f"data: {json.dumps(chunk)}\n".encode() for chunk in _served_model_stream_chunks(served_model) + ] + iterator: Final = ModelResponseIterator(iter(sse_lines), sync_stream=True) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(iterator), + model="anthropic/claude-requested", + custom_llm_provider="anthropic", + logging_obj=MagicMock(), + ) + + chunks: Final = list(wrapper) + + assert len(chunks) > 1 + for chunk in chunks[1:]: + assert chunk._hidden_params["provider_response_model"] == served_model + assembled: Final = litellm.stream_chunk_builder( + chunks=list(chunks), messages=[{"role": "user", "content": "hi"}] + ) + assert assembled._hidden_params["provider_response_model"] == served_model