pass bedrock_invoke_provider to make_call

Author: ishaan-jaff

litellm/llms/bedrock/chat/invoke_handler.py

@@ -40,6 +40,9 @@
     parse_xml_params,
     prompt_factory,
 )
+from litellm.llms.anthropic.chat.handler import (
+    ModelResponseIterator as AnthropicModelResponseIterator,
+)
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -177,6 +180,7 @@ async def make_call(
     logging_obj: Logging,
     fake_stream: bool = False,
     json_mode: Optional[bool] = False,
+    bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None,
 ):
     try:
         if client is None:
@@ -214,6 +218,14 @@ async def make_call(
             completion_stream: Any = MockResponseIterator(
                 model_response=model_response, json_mode=json_mode
             )
+        elif bedrock_invoke_provider == "anthropic":
+            decoder = AmazonAnthropicClaudeStreamDecoder(
+                model=model,
+                sync_stream=False,
+            )
+            completion_stream = decoder.aiter_bytes(
+                response.aiter_bytes(chunk_size=1024)
+            )
         else:
             decoder = AWSEventStreamDecoder(model=model)
             completion_stream = decoder.aiter_bytes(
@@ -248,6 +260,7 @@ def make_sync_call(
     logging_obj: Logging,
     fake_stream: bool = False,
     json_mode: Optional[bool] = False,
+    bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None,
 ):
     try:
         if client is None:
@@ -283,6 +296,12 @@ def make_sync_call(
             completion_stream: Any = MockResponseIterator(
                 model_response=model_response, json_mode=json_mode
             )
+        elif bedrock_invoke_provider == "anthropic":
+            decoder = AmazonAnthropicClaudeStreamDecoder(
+                model=model,
+                sync_stream=True,
+            )
+            completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024))
         else:
             decoder = AWSEventStreamDecoder(model=model)
             completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024))
@@ -1323,7 +1342,7 @@ def _chunk_parser(self, chunk_data: dict) -> GChunk:
             text = chunk_data.get("completions")[0].get("data").get("text")  # type: ignore
             is_finished = True
             finish_reason = "stop"
-        ######## bedrock.anthropic mappings ###############
+        ######## converse bedrock.anthropic mappings ###############
         elif (
             "contentBlockIndex" in chunk_data
             or "stopReason" in chunk_data
@@ -1429,6 +1448,22 @@ def _parse_message_from_event(self, event) -> Optional[str]:
             return chunk.decode()  # type: ignore[no-any-return]
 
 
+class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
+    def __init__(
+        self,
+        model: str,
+        sync_stream: bool,
+    ) -> None:
+        super().__init__(model=model)
+        self.anthropic_model_response_iterator = AnthropicModelResponseIterator(
+            streaming_response=None,
+            sync_stream=sync_stream,
+        )
+
+    def _chunk_parser(self, chunk_data: dict) -> GChunk:
+        return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
+
+
 class MockResponseIterator:  # for returning ai21 streaming responses
     def __init__(self, model_response, json_mode: Optional[bool] = False):
         self.model_response = model_response

litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py

@@ -83,19 +83,3 @@ def transform_response(
             api_key=api_key,
             json_mode=json_mode,
         )
-
-
-class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
-    def __init__(
-        self,
-        model: str,
-        sync_stream: bool,
-    ) -> None:
-        super().__init__(model=model)
-        self.anthropic_model_response_iterator = AnthropicModelResponseIterator(
-            streaming_response=None,
-            sync_stream=sync_stream,
-        )
-
-    def _chunk_parser(self, chunk_data: dict) -> GChunk:
-        return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)

litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py

@@ -440,6 +440,7 @@ def get_async_custom_stream_wrapper(
                 messages=messages,
                 logging_obj=logging_obj,
                 fake_stream=True if "ai21" in api_base else False,
+                bedrock_invoke_provider=self.get_bedrock_invoke_provider(model),
             ),
             model=model,
             custom_llm_provider="bedrock",
@@ -473,6 +474,7 @@ def get_sync_custom_stream_wrapper(
                 messages=messages,
                 logging_obj=logging_obj,
                 fake_stream=True if "ai21" in api_base else False,
+                bedrock_invoke_provider=self.get_bedrock_invoke_provider(model),
             ),
             model=model,
             custom_llm_provider="bedrock",