[Feat] Add streaming support for using bedrock invoke models with /v1/messages (#10710)

* add basic bedrock transform

* test_anthropic_messages_streaming_bedrock_invoke

* fix: typing ant

* fix: get async response iterator

* fix: code quality check

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,

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",
+        )

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")

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)
 
 

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

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,

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():
     """