From 669862643b82cbfe06be621fd3b63db94105c028 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 12 Dec 2023 16:13:31 -0800
Subject: [PATCH] fix(utils.py): flush holding chunk for streaming, on stream
 end

---
 litellm/tests/test_streaming.py | 22 ++++++++++++++++++++--
 litellm/utils.py                |  7 +++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index b718d00e0b1..33199ea323f 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -245,7 +245,6 @@ def test_completion_azure_stream():
         complete_response = ""
         # Add any assertions here to check the response
         for idx, init_chunk in enumerate(response):
-            print(f"azure chunk: {init_chunk}")
             chunk, finished = streaming_format_tests(idx, init_chunk)
             complete_response += chunk
             if finished:
@@ -255,7 +254,7 @@ def test_completion_azure_stream():
             raise Exception("Empty response received")
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
-test_completion_azure_stream() 
+# test_completion_azure_stream() 
 
 def test_completion_azure_function_calling_stream():
     try:
@@ -636,6 +635,25 @@ def test_completion_bedrock_ai21_stream():
 
 # test_completion_bedrock_ai21_stream()
 
+def test_sagemaker_weird_response(): 
+    """
+    When the stream ends, flush any remaining holding chunks.
+    """
+    try: 
+        chunk = """<s>[INST] Hey, how's it going? [/INST]
+
+    I'm doing well, thanks for asking! How about you? Is there anything you'd like to chat about or ask? I'm here to help with any questions you might have."""
+
+        logging_obj = litellm.Logging(model="berri-benchmarking-Llama-2-70b-chat-hf-4", messages=messages, stream=True, litellm_call_id="1234", function_id="function_id", call_type="acompletion", start_time=time.time())
+        response = litellm.CustomStreamWrapper(completion_stream=chunk, model="berri-benchmarking-Llama-2-70b-chat-hf-4", custom_llm_provider="sagemaker", logging_obj=logging_obj)
+        complete_response = ""
+        for chunk in response:
+            complete_response += chunk["choices"][0]["delta"]["content"]
+        assert len(complete_response) > 0
+    except Exception as e: 
+        pytest.fail(f"An exception occurred - {str(e)}")
+# test_sagemaker_weird_response()
+
 # def test_completion_sagemaker_stream():
 #     try:
 #         response = completion(
diff --git a/litellm/utils.py b/litellm/utils.py
index 3d45bf1e0a6..900fc45edca 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -5682,6 +5682,13 @@ class CustomStreamWrapper:
                 else: 
                     return 
             elif model_response.choices[0].finish_reason:
+                # flush any remaining holding chunk 
+                if len(self.holding_chunk) > 0:
+                    if model_response.choices[0].delta.content is None:
+                        model_response.choices[0].delta.content = self.holding_chunk
+                    else:
+                        model_response.choices[0].delta.content = self.holding_chunk + model_response.choices[0].delta.content
+                    self.holding_chunk = "" 
                 model_response.choices[0].finish_reason = map_finish_reason(model_response.choices[0].finish_reason) # ensure consistent output to openai
                 return model_response
             elif response_obj is not None and response_obj.get("original_chunk", None) is not None: # function / tool calling branch - only set for openai/azure compatible endpoints