diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 53e8605011a..dbd3ff889c8 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 5aa14b0442a..642770efeb3 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,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 a132df62b08..3fa488d0051 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/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 6d7d62e47aa..47c5b6619fe 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -11,13 +11,18 @@ 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, 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 @@ -2367,3 +2372,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..200967db11a 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,5 +1,6 @@ import copy import json +from queue import Queue from typing import Final from unittest.mock import MagicMock, patch @@ -17,6 +18,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.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.pass_through.messages.transformation import ( @@ -26,6 +28,7 @@ from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig 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, ) @@ -6798,3 +6801,86 @@ 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.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 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 bfcdd8da3d8..9010f708348 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")