From 65baa839287e09a9169f5b360885097ba15824e9 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 10 Jun 2024 14:47:49 -0700
Subject: [PATCH] fix(bedrock_httpx.py): returning correct finish reason on
 streaming completion

---
 litellm/llms/bedrock_httpx.py   | 8 +++++++-
 litellm/tests/test_streaming.py | 3 +--
 litellm/utils.py                | 5 ++---
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py
index 336cbd3bbb1..b011d951295 100644
--- a/litellm/llms/bedrock_httpx.py
+++ b/litellm/llms/bedrock_httpx.py
@@ -1898,8 +1898,10 @@ class AWSEventStreamDecoder:
                     }
             elif "stopReason" in chunk_data:
                 finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
+                is_finished = True
             elif "usage" in chunk_data:
                 usage = ConverseTokenUsageBlock(**chunk_data["usage"])  # type: ignore
+
             response = GenericStreamingChunk(
                 text=text,
                 tool_use=tool_use,
@@ -1924,7 +1926,11 @@ class AWSEventStreamDecoder:
             is_finished = True
             finish_reason = "stop"
         ######## bedrock.anthropic mappings ###############
-        elif "contentBlockIndex" in chunk_data:
+        elif (
+            "contentBlockIndex" in chunk_data
+            or "stopReason" in chunk_data
+            or "metrics" in chunk_data
+        ):
             return self.converse_chunk_parser(chunk_data=chunk_data)
         ######## bedrock.mistral mappings ###############
         elif "outputs" in chunk_data:
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index ee9f1cec68e..ce69c0bf995 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -2573,9 +2573,8 @@ def test_streaming_and_function_calling(model):
         # Add any assertions here to check the response
         for idx, chunk in enumerate(response):
             # continue
-            # print("\n{}\n".format(chunk))
+            print("\n{}\n".format(chunk))
             if idx == 0:
-                print(chunk)
                 assert (
                     chunk.choices[0].delta.tool_calls[0].function.arguments is not None
                 )
diff --git a/litellm/utils.py b/litellm/utils.py
index fc6ac9fecd0..6a5afabe77b 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -11458,7 +11458,6 @@ class CustomStreamWrapper:
                     raise StopIteration
                 response_obj: GenericStreamingChunk = chunk
                 completion_obj["content"] = response_obj["text"]
-
                 if response_obj["is_finished"]:
                     self.received_finish_reason = response_obj["finish_reason"]
 
@@ -11591,7 +11590,7 @@ class CustomStreamWrapper:
                         # for azure, we need to pass the model from the orignal chunk
                         self.model = chunk.model
                 response_obj = self.handle_openai_chat_completion_chunk(chunk)
-                if response_obj == None:
+                if response_obj is None:
                     return
                 completion_obj["content"] = response_obj["text"]
                 print_verbose(f"completion obj content: {completion_obj['content']}")
@@ -11821,7 +11820,7 @@ class CustomStreamWrapper:
                 else:
                     return
             elif self.received_finish_reason is not None:
-                if self.sent_last_chunk == True:
+                if self.sent_last_chunk is True:
                     raise StopIteration
                 # flush any remaining holding chunk
                 if len(self.holding_chunk) > 0: