From e3638be2e0b53c5eeea901e7033e000e3e564ea1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:05:58 +0800 Subject: [PATCH 1/2] fix(vertex-ai): preserve provider metadata in Claude responses --- litellm/llms/anthropic/chat/handler.py | 6 +- litellm/llms/anthropic/chat/transformation.py | 11 +-- .../chat/test_anthropic_chat_handler.py | 75 ++++++++++++++++++- .../test_anthropic_chat_transformation.py | 31 ++++++++ 4 files changed, 115 insertions(+), 8 deletions(-) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index ef0f45d8f8b..82af602996c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -206,6 +206,7 @@ class AnthropicChatCompletion(BaseLLM): model: str, messages: list, api_base: str, + custom_llm_provider: str, custom_prompt_dict: dict, model_response: ModelResponse, print_verbose: Callable, @@ -245,7 +246,7 @@ class AnthropicChatCompletion(BaseLLM): streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, model=model, - custom_llm_provider="anthropic", + custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, _response_headers=process_anthropic_headers(headers), ) @@ -413,6 +414,7 @@ class AnthropicChatCompletion(BaseLLM): messages=messages, data=data, api_base=api_base, + custom_llm_provider=custom_llm_provider, custom_prompt_dict=custom_prompt_dict, model_response=model_response, print_verbose=print_verbose, @@ -489,7 +491,7 @@ class AnthropicChatCompletion(BaseLLM): return CustomStreamWrapper( completion_stream=completion_stream, model=model, - custom_llm_provider="anthropic", + custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, _response_headers=process_anthropic_headers(headers), ) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 3bffee48d6a..cd021caa988 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2593,8 +2593,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, ): - _hidden_params: Final[dict] = {} - _hidden_params["additional_headers"] = process_anthropic_headers(dict(raw_response.headers)) if "error" in completion_response: response_headers: Final = getattr(raw_response, "headers", None) raise AnthropicError( @@ -2668,7 +2666,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _message = json_mode_message model_response.choices[0].message = _message - model_response._hidden_params["original_response"] = completion_response["content"] model_response.choices[0].finish_reason = cast( OpenAIChatCompletionFinishReason, map_finish_reason(completion_response["stop_reason"]), @@ -2685,8 +2682,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model_response.created = int(time.time()) model_response.model = completion_response["model"] - _hidden_params["provider_specific_fields"] = provider_specific_fields - model_response._hidden_params = _hidden_params + model_response._hidden_params = { + **model_response._hidden_params, + "additional_headers": process_anthropic_headers(dict(raw_response.headers)), + "original_response": completion_response["content"], + "provider_specific_fields": provider_specific_fields, + } return model_response def get_prefix_prompt(self, messages: list[AllMessageValues]) -> str | None: diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 1e0d2e55373..9828d760f0e 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -10,13 +10,18 @@ import pytest import litellm from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME -from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call +from litellm.llms.anthropic.chat.handler import ( + AnthropicChatCompletion, + ModelResponseIterator, + make_call, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ) from litellm.types.responses.main import OutputCodeInterpreterCall +from litellm.types.utils import ModelResponse @pytest.mark.asyncio @@ -2452,3 +2457,71 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): 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 + + +def test_streaming_completion_preserves_custom_llm_provider(): + response = MagicMock() + response.status_code = 200 + response.headers = {} + response.iter_lines.return_value = iter(()) + client = MagicMock(spec=HTTPHandler) + client.post.return_value = response + logging_obj = MagicMock() + logging_obj.model_call_details = { + "custom_llm_provider": "vertex_ai", + "litellm_params": {}, + } + + result = AnthropicChatCompletion().completion( + model="claude-sonnet-4-5@20250929", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://example.com/v1/messages", + custom_llm_provider="vertex_ai", + custom_prompt_dict={}, + model_response=ModelResponse(), + print_verbose=MagicMock(), + encoding=None, + api_key="test-key", + logging_obj=logging_obj, + optional_params={"stream": True, "max_tokens": 16, "is_vertex_request": True}, + timeout=60.0, + litellm_params={}, + client=client, + ) + + assert result.custom_llm_provider == "vertex_ai" + + +@pytest.mark.asyncio +async def test_async_streaming_completion_preserves_custom_llm_provider(): + response = MagicMock() + response.headers = {} + response.aiter_lines.return_value = iter(()) + client = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=response) + logging_obj = MagicMock() + logging_obj.model_call_details = { + "custom_llm_provider": "vertex_ai", + "litellm_params": {}, + } + + completion = AnthropicChatCompletion().completion( + model="claude-sonnet-4-5@20250929", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://example.com/v1/messages", + custom_llm_provider="vertex_ai", + custom_prompt_dict={}, + model_response=ModelResponse(), + print_verbose=MagicMock(), + encoding=None, + api_key="test-key", + logging_obj=logging_obj, + optional_params={"stream": True, "max_tokens": 16, "is_vertex_request": True}, + timeout=60.0, + litellm_params={}, + acompletion=True, + client=client, + ) + result = await completion + + assert result.custom_llm_provider == "vertex_ai" diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 332153b4c7d..8219df9ae91 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6798,3 +6798,34 @@ def test_chat_dummy_tool_result_for_an_orphaned_tool_call_replays_a_byte_identic _assert_prefix_stable(requests) assert [m["role"] for m in requests[0]["messages"]] == ["user", "assistant", "user"] assert requests[0]["messages"][2]["content"][0]["type"] == "tool_result" + + +def test_transform_parsed_response_preserves_existing_hidden_params(): + from litellm.types.utils import ModelResponse + + config = AnthropicConfig() + raw_response = MagicMock() + raw_response.headers = {"request-id": "req_vertex"} + raw_response.status_code = 200 + completion_response = { + "id": "msg_vertex", + "model": "claude-sonnet-4-5@20250929", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + "content": [{"type": "text", "text": "Hello"}], + } + model_response = ModelResponse() + model_response._hidden_params = { + "custom_llm_provider": "vertex_ai", + "region_name": "us-east5", + } + + result = config.transform_parsed_response( + completion_response=completion_response, + raw_response=raw_response, + model_response=model_response, + ) + + assert result._hidden_params["custom_llm_provider"] == "vertex_ai" + assert result._hidden_params["region_name"] == "us-east5" + assert result._hidden_params["original_response"] == completion_response["content"] From b53fb1f445bf2f02b4b38e227d0ff7a199d1b07f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:28:33 +0800 Subject: [PATCH 2/2] fix(azure-ai): pass provider to async Claude streams --- litellm/llms/azure_ai/anthropic/handler.py | 1 + .../claude/test_azure_anthropic_handler.py | 58 ++++++++++++++++++- 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure_ai/anthropic/handler.py b/litellm/llms/azure_ai/anthropic/handler.py index 24ee76b31d0..c0454cfd6b4 100644 --- a/litellm/llms/azure_ai/anthropic/handler.py +++ b/litellm/llms/azure_ai/anthropic/handler.py @@ -104,6 +104,7 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion): messages=messages, data=data, api_base=api_base, + custom_llm_provider=custom_llm_provider, custom_prompt_dict=custom_prompt_dict, model_response=model_response, print_verbose=print_verbose, diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py index dcb6aec8091..586490193d8 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py @@ -6,11 +6,12 @@ sys.path.insert( ) import json -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import ModelResponse @@ -64,7 +65,10 @@ class TestAzureAnthropicChatCompletion: headers = {} with patch.object( - handler, "acompletion_function", return_value=ModelResponse() + handler, + "acompletion_function", + new_callable=MagicMock, + return_value=ModelResponse(), ) as mock_acompletion: handler.completion( model=model, @@ -87,6 +91,7 @@ class TestAzureAnthropicChatCompletion: # Verify AzureAnthropicConfig was used mock_azure_config.assert_called_once() mock_config_instance.validate_environment.assert_called_once() + mock_acompletion.assert_called_once() @patch("litellm.llms.anthropic.chat.handler.make_sync_call") @patch("litellm.utils.ProviderConfigManager") @@ -158,6 +163,55 @@ class TestAzureAnthropicChatCompletion: mock_make_sync_call.assert_called_once() assert result is not None + @pytest.mark.asyncio + @patch("litellm.llms.azure_ai.anthropic.handler.AzureAnthropicConfig") + async def test_completion_async_streaming_preserves_provider(self, mock_azure_config): + handler = AzureAnthropicChatCompletion() + mock_config = MagicMock() + mock_config.validate_environment.return_value = { + "api-key": "test-api-key", + "anthropic-version": "2023-06-01", + } + mock_config.transform_request.return_value = { + "model": "claude-sonnet-4-5", + "messages": [], + "stream": True, + } + mock_azure_config.return_value = mock_config + + response = MagicMock() + response.headers = {} + response.aiter_lines.return_value = iter(()) + client = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=response) + logging_obj = MagicMock() + logging_obj.model_call_details = { + "custom_llm_provider": "azure_ai", + "litellm_params": {}, + } + + completion = handler.completion( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://test.services.ai.azure.com/anthropic/v1/messages", + custom_llm_provider="azure_ai", + custom_prompt_dict={}, + model_response=ModelResponse(), + print_verbose=MagicMock(), + encoding=None, + api_key="test-api-key", + logging_obj=logging_obj, + optional_params={"stream": True}, + timeout=60.0, + litellm_params={"api_key": "test-api-key"}, + acompletion=True, + client=client, + ) + result = await completion + + client.post.assert_awaited_once() + assert result.custom_llm_provider == "azure_ai" + @patch("litellm.llms.custom_httpx.http_handler._get_httpx_client") @patch("litellm.utils.ProviderConfigManager") @patch("litellm.llms.azure_ai.anthropic.handler.AzureAnthropicConfig")