diff --git a/litellm/constants.py b/litellm/constants.py index 37d44f64c5..148cb9847c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -141,6 +141,7 @@ DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2 #### Networking settings #### request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds STREAM_SSE_DONE_STRING: str = "[DONE]" +STREAM_SSE_DATA_PREFIX: str = "data: " ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND = float( os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 78078b93f8..2395df3faf 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -12,6 +12,8 @@ from fastapi.responses import StreamingResponse import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import STREAM_SSE_DATA_PREFIX +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -22,6 +24,24 @@ from litellm.proxy.utils import ProxyLogging router = APIRouter() +def return_anthropic_chunk(chunk: str | dict) -> str: + """ + Helper function to format streaming chunks for Anthropic API format + + Args: + chunk: A string or dictionary to be returned in SSE format + + Returns: + str: A properly formatted SSE chunk string + """ + if isinstance(chunk, dict): + # Use safe_dumps for proper JSON serialization with circular reference detection + chunk_str = safe_dumps(chunk) + return f"{STREAM_SSE_DATA_PREFIX}{chunk_str}\n\n" + else: + return chunk + + async def async_data_generator_anthropic( response, user_api_key_dict: UserAPIKeyAuth, @@ -40,7 +60,8 @@ async def async_data_generator_anthropic( user_api_key_dict=user_api_key_dict, response=chunk ) - yield chunk + # Format chunk using helper function + yield return_anthropic_chunk(chunk) except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format( @@ -69,7 +90,7 @@ async def async_data_generator_anthropic( code=getattr(e, "status_code", 500), ) error_returned = json.dumps({"error": proxy_exception.to_dict()}) - yield f"data: {error_returned}\n\n" + yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n" @router.post( diff --git a/tests/litellm/proxy/anthropic_endpoints/__init__.py b/tests/litellm/proxy/anthropic_endpoints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/litellm/proxy/anthropic_endpoints/test_endpoints.py new file mode 100644 index 0000000000..2dbdf34504 --- /dev/null +++ b/tests/litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -0,0 +1,61 @@ +""" +Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objects in streaming responses +""" + +import json +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.anthropic_endpoints.endpoints import async_data_generator_anthropic + + +class TestAnthropicEndpoints(unittest.TestCase): + @patch("litellm.litellm_core_utils.safe_json_dumps.safe_dumps") + @pytest.mark.asyncio + async def test_async_data_generator_anthropic_dict_handling(self, mock_safe_dumps): + """Test async_data_generator_anthropic handles dictionary chunks properly""" + # Setup + mock_response = AsyncMock() + mock_response.__aiter__.return_value = [ + {"type": "message_start", "message": {"id": "msg_123"}}, + "text chunk data", + {"type": "content_block_delta", "delta": {"text": "more data"}}, + "text chunk data again", + ] + + mock_user_api_key_dict = MagicMock() + mock_request_data = {} + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + + # Configure safe_dumps to return a properly formatted JSON string + mock_safe_dumps.side_effect = lambda chunk: json.dumps(chunk) + + # Execute + result = [chunk async for chunk in async_data_generator_anthropic( + response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + request_data=mock_request_data, + proxy_logging_obj=mock_proxy_logging_obj, + )] + + # Verify + expected_result = [ + 'data: {"type": "message_start", "message": {"id": "msg_123"}}\n\n', + 'text chunk data', + 'data: {"type": "content_block_delta", "delta": {"text": "more data"}}\n\n', + 'text chunk data again', + ] + + self.assertEqual(result, expected_result) + + # Assert safe_dumps was called for dictionary objects + mock_safe_dumps.assert_any_call({"type": "message_start", "message": {"id": "msg_123"}}) + mock_safe_dumps.assert_any_call({"type": "content_block_delta", "delta": {"text": "more data"}}) + assert mock_safe_dumps.call_count == 2 # Called twice, once for each dict object + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file