From cd53291b6270d19c8854c68fd8a754f77e76ce60 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 29 Mar 2024 07:56:10 -0700
Subject: [PATCH] fix(utils.py): support bedrock mistral streaming

---
 litellm/tests/test_streaming.py | 37 +++++++++++++++++++++++++++++++++
 litellm/utils.py                | 11 ++++++++++
 2 files changed, 48 insertions(+)

diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 80086dbe6..76bc610f2 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -974,6 +974,43 @@ def test_completion_bedrock_ai21_stream():
 # test_completion_bedrock_ai21_stream()
 
 
+def test_completion_bedrock_mistral_stream():
+    try:
+        litellm.set_verbose = False
+        response = completion(
+            model="bedrock/mistral.mixtral-8x7b-instruct-v0:1",
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Be as verbose as possible and give as many details as possible, how does a court case get to the Supreme Court?",
+                }
+            ],
+            temperature=1,
+            max_tokens=20,
+            stream=True,
+        )
+        print(response)
+        complete_response = ""
+        has_finish_reason = False
+        # Add any assertions here to check the response
+        for idx, chunk in enumerate(response):
+            # print
+            chunk, finished = streaming_format_tests(idx, chunk)
+            has_finish_reason = finished
+            complete_response += chunk
+            if finished:
+                break
+        if has_finish_reason is False:
+            raise Exception("finish reason not set for last chunk")
+        if complete_response.strip() == "":
+            raise Exception("Empty response received")
+        print(f"completion_response: {complete_response}")
+    except RateLimitError:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_sagemaker_weird_response():
     """
     When the stream ends, flush any remaining holding chunks.
diff --git a/litellm/utils.py b/litellm/utils.py
index fa15b73c2..f42bdeecc 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -9093,6 +9093,17 @@ class CustomStreamWrapper:
                 if stop_reason != None:
                     is_finished = True
                     finish_reason = stop_reason
+            ######## bedrock.mistral mappings ###############
+            elif "outputs" in chunk_data:
+                if (
+                    len(chunk_data["outputs"]) == 1
+                    and chunk_data["outputs"][0].get("text", None) is not None
+                ):
+                    text = chunk_data["outputs"][0]["text"]
+                stop_reason = chunk_data.get("stop_reason", None)
+                if stop_reason != None:
+                    is_finished = True
+                    finish_reason = stop_reason
             ######## bedrock.cohere mappings ###############
             # meta mapping
             elif "generation" in chunk_data: