diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 50384d2247c..b7c8fb56502 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -10,8 +10,6 @@ import contextvars from functools import partial from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union -import httpx - import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -33,42 +31,6 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -class AnthropicMessagesHandler: - @staticmethod - async def _handle_anthropic_streaming( - response: httpx.Response, - request_body: dict, - litellm_logging_obj: LiteLLMLoggingObj, - ) -> AsyncIterator: - """Helper function to handle Anthropic streaming responses using the existing logging handlers""" - from datetime import datetime - - from litellm.proxy.pass_through_endpoints.streaming_handler import ( - PassThroughStreamingHandler, - ) - from litellm.proxy.pass_through_endpoints.success_handler import ( - PassThroughEndpointLogging, - ) - from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - EndpointType, - ) - - # Create success handler object - passthrough_success_handler_obj = PassThroughEndpointLogging() - - # Use the existing streaming handler for Anthropic - start_time = datetime.now() - return PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=request_body, - litellm_logging_obj=litellm_logging_obj, - endpoint_type=EndpointType.ANTHROPIC, - start_time=start_time, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route="/v1/messages", - ) - - @client async def anthropic_messages( max_tokens: int, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index c153b4fbcf1..5b5e2e6f36d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional +from typing import Any, AsyncIterator, Dict, List, Optional import httpx @@ -113,3 +113,38 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): message=raw_response.text, status_code=raw_response.status_code ) return AnthropicMessagesResponse(**raw_response_json) + + def get_async_streaming_response_iterator( + self, + model: str, + httpx_response: httpx.Response, + request_body: dict, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> AsyncIterator: + """Helper function to handle Anthropic streaming responses using the existing logging handlers""" + from datetime import datetime + + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, + ) + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + + # Create success handler object + passthrough_success_handler_obj = PassThroughEndpointLogging() + + # Use the existing streaming handler for Anthropic + start_time = datetime.now() + return PassThroughStreamingHandler.chunk_processor( + response=httpx_response, + request_body=request_body, + litellm_logging_obj=litellm_logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=start_time, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route="/v1/messages", + ) diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 29ac0cf2f28..cd74efb2025 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional import httpx @@ -96,3 +96,12 @@ class BaseAnthropicMessagesConfig(ABC): For all other providers, this is a no-op and we just return the headers """ return headers + + def get_async_streaming_response_iterator( + self, + model: str, + httpx_response: httpx.Response, + request_body: dict, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> AsyncIterator: + raise NotImplementedError("Subclasses must implement this method") diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index dfd16585434..be80cb58814 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1413,7 +1413,9 @@ class AWSEventStreamDecoder: except Exception as e: raise Exception("Received streaming error - {}".format(str(e))) - def _chunk_parser(self, chunk_data: dict) -> Union[GChunk, ModelResponseStream]: + def _chunk_parser( + self, chunk_data: dict + ) -> Union[GChunk, ModelResponseStream, dict]: text = "" is_finished = False finish_reason = "" @@ -1473,7 +1475,7 @@ class AWSEventStreamDecoder: def iter_bytes( self, iterator: Iterator[bytes] - ) -> Iterator[Union[GChunk, ModelResponseStream]]: + ) -> Iterator[Union[GChunk, ModelResponseStream, dict]]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer @@ -1489,7 +1491,7 @@ class AWSEventStreamDecoder: async def aiter_bytes( self, iterator: AsyncIterator[bytes] - ) -> AsyncIterator[Union[GChunk, ModelResponseStream]]: + ) -> AsyncIterator[Union[GChunk, ModelResponseStream, dict]]: """Given an async iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer @@ -1576,7 +1578,9 @@ class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder): sync_stream=sync_stream, ) - def _chunk_parser(self, chunk_data: dict) -> Union[GChunk, ModelResponseStream]: + def _chunk_parser( + self, chunk_data: dict + ) -> Union[GChunk, ModelResponseStream, dict]: return self.deepseek_model_response_iterator.chunk_parser(chunk=chunk_data) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 5eba423bc38..f11a3b0b85e 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Union + +import httpx from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, @@ -6,10 +8,13 @@ from litellm.llms.anthropic.experimental_pass_through.messages.transformation im from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) +from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import GenericStreamingChunk as GChunk +from litellm.types.utils import ModelResponseStream if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -120,3 +125,42 @@ class AmazonAnthropicClaude3MessagesConfig( if "model" in anthropic_messages_request: anthropic_messages_request.pop("model", None) return anthropic_messages_request + + def get_async_streaming_response_iterator( + self, + model: str, + httpx_response: httpx.Response, + request_body: dict, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> AsyncIterator: + aws_decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model=model, + ) + completion_stream = aws_decoder.aiter_bytes( + httpx_response.aiter_bytes(chunk_size=aws_decoder.DEFAULT_CHUNK_SIZE) + ) + return completion_stream + + +class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): + def __init__( + self, + model: str, + ) -> None: + """ + Iterator to return Bedrock invoke response in anthropic /messages format + """ + super().__init__(model=model) + self.DEFAULT_CHUNK_SIZE = 1024 + + def _chunk_parser( + self, chunk_data: dict + ) -> Union[GChunk, ModelResponseStream, dict]: + """ + Parse the chunk data into anthropic /messages format + + No transformation is needed for anthropic /messages format + + since bedrock invoke returns the response in the correct format + """ + return chunk_data diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 81739970bd1..000d511d031 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1034,10 +1034,6 @@ class BaseLLMHTTPHandler: stream: Optional[bool] = False, kwargs: Optional[Dict[str, Any]] = None, ) -> Union[AnthropicMessagesResponse, AsyncIterator]: - from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( - AnthropicMessagesHandler, - ) - if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders.ANTHROPIC @@ -1131,11 +1127,13 @@ class BaseLLMHTTPHandler: logging_obj.model_call_details["httpx_response"] = response if stream: - return await AnthropicMessagesHandler._handle_anthropic_streaming( - response=response, + completion_stream = anthropic_messages_provider_config.get_async_streaming_response_iterator( + model=model, + httpx_response=response, request_body=request_body, litellm_logging_obj=logging_obj, ) + return completion_stream else: return anthropic_messages_provider_config.transform_anthropic_messages_response( model=model, diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py index 739c24fc424..68a6d860041 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py @@ -157,6 +157,37 @@ async def test_anthropic_messages_streaming(): print("chunk=", chunk) +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_bedrock_invoke(): + """ + Test the anthropic_messages with streaming request + """ + # Get API key from environment + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not found in environment") + + # Set up test parameters + messages = [{"role": "user", "content": "Hello, can you tell me a short joke?"}] + + # Call the handler + async_httpx_client = AsyncHTTPHandler() + response = await litellm.anthropic.messages.acreate( + messages=messages, + api_key=api_key, + model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + max_tokens=100, + stream=True, + client=async_httpx_client, + ) + collected_chunks = [] + if isinstance(response, AsyncIterator): + async for chunk in response: + print("chunk=", chunk) + collected_chunks.append(chunk) + + print("collected_chunks=", collected_chunks) + @pytest.mark.asyncio async def test_anthropic_messages_streaming_with_bad_request(): """