From 7d0086d742ed6666ea7b8251a3dcef49dbafe069 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 17 Apr 2024 17:43:41 -0700
Subject: [PATCH] fix(utils.py): ensure streaming output parsing only applied
 for hf / sagemaker models

selectively applies the <s>
</s> checking
---
 litellm/tests/test_streaming.py | 14 ++++++++++++++
 litellm/utils.py                |  9 +++++++++
 2 files changed, 23 insertions(+)

diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index a18da18192f..aa2a91b9f32 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -220,6 +220,20 @@ tools_schema = [
 # test_completion_cohere_stream()
 
 
+def test_completion_azure_stream_special_char():
+    messages = [
+        {"role": "user", "content": "Respond with the '<' sign and nothing else."}
+    ]
+    response = completion(model="azure/chatgpt-v-2", messages=messages, stream=True)
+    response_str = ""
+    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():
     try:
         litellm.cache = None
diff --git a/litellm/utils.py b/litellm/utils.py
index 2ae1467d076..bea24c02fef 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8856,7 +8856,16 @@ class CustomStreamWrapper:
             raise e
 
     def check_special_tokens(self, chunk: str, finish_reason: Optional[str]):
+        """
+        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 finish_reason:
             for token in self.special_tokens:
                 if token in chunk: