From 15ae7a8314fec6bb1e7d77933efd442de39eac33 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 17 Apr 2024 18:03:40 -0700
Subject: [PATCH] fix(utils.py): fix streaming special character flushing logic

---
 litellm/tests/test_streaming.py |  3 +--
 litellm/utils.py                | 13 +++++++------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index aa2a91b9f3..3297697823 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -221,6 +221,7 @@ tools_schema = [
 
 
 def test_completion_azure_stream_special_char():
+    litellm.set_verbose = True
     messages = [
         {"role": "user", "content": "Respond with the '<' sign and nothing else."}
     ]
@@ -229,9 +230,7 @@ def test_completion_azure_stream_special_char():
     for part in response:
         response_str += part.choices[0].delta.content or ""
 
-    print(f"response_str: {response_str}")
     assert len(response_str) > 0
-    raise Exception("it worked")
 
 
 def test_completion_cohere_stream_bad_key():
diff --git a/litellm/utils.py b/litellm/utils.py
index bea24c02fe..3e31874bd0 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8860,11 +8860,11 @@ class CustomStreamWrapper:
         Output parse <s> / </s> special tokens for sagemaker + hf streaming.
         """
         hold = False
-        if (
-            self.custom_llm_provider != "huggingface"
-            and self.custom_llm_provider != "sagemaker"
-        ):
-            return hold, chunk
+        # if (
+        #     self.custom_llm_provider != "huggingface"
+        #     and self.custom_llm_provider != "sagemaker"
+        # ):
+        #     return hold, chunk
 
         if finish_reason:
             for token in self.special_tokens:
@@ -8881,6 +8881,7 @@ class CustomStreamWrapper:
         for token in self.special_tokens:
             if len(curr_chunk) < len(token) and curr_chunk in token:
                 hold = True
+                self.holding_chunk = curr_chunk
             elif len(curr_chunk) >= len(token):
                 if token in curr_chunk:
                     self.holding_chunk = curr_chunk.replace(token, "")
@@ -9962,6 +9963,7 @@ class CustomStreamWrapper:
                 f"model_response.choices[0].delta: {model_response.choices[0].delta}; completion_obj: {completion_obj}"
             )
             print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}")
+
             ## RETURN ARG
             if (
                 "content" in completion_obj
@@ -10034,7 +10036,6 @@ class CustomStreamWrapper:
             elif self.received_finish_reason is not None:
                 if self.sent_last_chunk == True:
                     raise StopIteration
-
                 # flush any remaining holding chunk
                 if len(self.holding_chunk) > 0:
                     if model_response.choices[0].delta.content is None: