diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 13ceea3e31b..4b2f8ed79e9 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -333,6 +333,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, @@ -375,7 +376,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), ) @@ -543,6 +544,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, @@ -622,7 +624,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 9a63afd5cc0..277d6b5a4a3 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2979,8 +2979,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): tool_name_reverse_map: dict[str, str] | None = None, response_format_schema: ResponseFormatSchema | 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( @@ -3063,7 +3061,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"]), @@ -3080,8 +3077,11 @@ 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)), + "provider_specific_fields": provider_specific_fields, + } return model_response def get_prefix_prompt(self, messages: list[AllMessageValues]) -> str | None: diff --git a/litellm/llms/azure_ai/anthropic/handler.py b/litellm/llms/azure_ai/anthropic/handler.py index ca4dad5ee3e..84066deb3dc 100644 --- a/litellm/llms/azure_ai/anthropic/handler.py +++ b/litellm/llms/azure_ai/anthropic/handler.py @@ -105,6 +105,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/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 1f4a5aaa256..ddad246b651 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -15,7 +15,7 @@ import litellm from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call +from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion, ModelResponseIterator, _envelope_verdict, make_call from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( @@ -23,7 +23,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ) from litellm.types.responses.main import OutputCodeInterpreterCall -from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream +from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponse, ModelResponseStream from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @@ -2796,6 +2796,74 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): 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" + + _STREAMED_RISK_SCHEMA: Final = { "type": "object", "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, 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 32f18b952ac..fbf4472d3ee 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,6 +1,7 @@ import copy import json from itertools import chain +from queue import Queue from typing import Final from unittest.mock import MagicMock, patch @@ -20,6 +21,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, RESPONSE_FORMAT_TOOL_NAME, ) +from litellm.integrations.datadog.datadog import DataDogLogger from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt from litellm.llms.anthropic.chat import ModelResponseIterator @@ -34,6 +36,7 @@ from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePla from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) @@ -7198,6 +7201,89 @@ def test_chat_dummy_tool_result_for_an_orphaned_tool_call_replays_a_byte_identic 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.choices[0].message.content == "Hello" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turn_off_message_logging", [True, False]) +async def test_anthropic_completion_respects_message_redaction_in_datadog_v1( + monkeypatch: pytest.MonkeyPatch, turn_off_message_logging: bool +) -> None: + monkeypatch.setattr(litellm, "datadog_use_v1", True) + private_output: Final = "private-claude-output-for-redaction-regression" + logged_requests: Final[Queue[httpx.Request]] = Queue() + + def provider_response(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_redaction", + "type": "message", + "role": "assistant", + "model": "claude-fable-5-1", + "content": [{"type": "text", "text": private_output}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + def datadog_response(request: httpx.Request) -> httpx.Response: + logged_requests.put(request) + return httpx.Response(202) + + logger: Final = DataDogLogger(dd_api_key="test-key", dd_site="datadoghq.com", allow_env_credentials=False) + provider_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(provider_response))) + datadog_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(datadog_response))) + logger.sync_client = datadog_client + try: + response: Final = litellm.completion( + model="anthropic/claude-fable-5-1", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-key", + client=provider_client, + success_callback=[logger], + turn_off_message_logging=turn_off_message_logging, + ) + request: Final = logged_requests.get(timeout=10) + finally: + provider_client.close() + datadog_client.close() + + assert response.choices[0].message.content == private_output + assert (private_output in request.content.decode()) is (not turn_off_message_logging) + payload: Final = json.loads(json.loads(request.content)["message"]) + assert payload["usage"]["prompt_tokens"] == 10 + assert payload["usage"]["completion_tokens"] == 5 + + anthropic_chunk_list: Final = [ { "type": "content_block_start", 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 052c4c422f2..ba44ba89adb 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,7 +6,7 @@ sys.path.insert( ) import json -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import respx @@ -14,6 +14,7 @@ from httpx import Response import litellm from litellm.llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import ModelResponse @@ -67,7 +68,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, @@ -90,6 +94,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") @@ -161,6 +166,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")