From 223ac464d78509e9574b429014acf8edcff3352e Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 12 Mar 2024 09:50:43 -0700
Subject: [PATCH] (fix) support streaming for azure/instruct models

---
 litellm/llms/azure_text.py       |  4 ++--
 litellm/tests/test_completion.py | 14 ++++++++++++++
 litellm/utils.py                 | 30 ++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/azure_text.py b/litellm/llms/azure_text.py
index 690d67b8862..17cf4b6b211 100644
--- a/litellm/llms/azure_text.py
+++ b/litellm/llms/azure_text.py
@@ -441,7 +441,7 @@ class AzureTextCompletion(BaseLLM):
         streamwrapper = CustomStreamWrapper(
             completion_stream=response,
             model=model,
-            custom_llm_provider="azure",
+            custom_llm_provider="azure_text",
             logging_obj=logging_obj,
         )
         return streamwrapper
@@ -500,7 +500,7 @@ class AzureTextCompletion(BaseLLM):
             streamwrapper = CustomStreamWrapper(
                 completion_stream=response,
                 model=model,
-                custom_llm_provider="azure",
+                custom_llm_provider="azure_text",
                 logging_obj=logging_obj,
             )
             return streamwrapper  ## DO NOT make this into an async for ... loop, it will yield an async generator, which won't raise errors if the response fails
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index c4760a10a95..2b372f57a6f 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -1162,6 +1162,20 @@ def test_azure_instruct():
     print("response", response)
 
 
+@pytest.mark.asyncio
+async def test_azure_instruct_stream():
+    litellm.set_verbose = False
+    response = await litellm.acompletion(
+        model="azure_text/instruct-model",
+        messages=[{"role": "user", "content": "What is the weather like in Boston?"}],
+        max_tokens=10,
+        stream=True,
+    )
+    print("response", response)
+    async for chunk in response:
+        print(chunk)
+
+
 async def test_re_use_azure_async_client():
     try:
         print("azure gpt-3.5 ASYNC with clie nttest\n\n")
diff --git a/litellm/utils.py b/litellm/utils.py
index 3b6169770df..262935faa58 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8656,6 +8656,27 @@ class CustomStreamWrapper:
             traceback.print_exc()
             raise e
 
+    def handle_azure_text_completion_chunk(self, chunk):
+        try:
+            print_verbose(f"\nRaw OpenAI Chunk\n{chunk}\n")
+            text = ""
+            is_finished = False
+            finish_reason = None
+            choices = getattr(chunk, "choices", [])
+            if len(choices) > 0:
+                text = choices[0].text
+                if choices[0].finish_reason is not None:
+                    is_finished = True
+                    finish_reason = choices[0].finish_reason
+            return {
+                "text": text,
+                "is_finished": is_finished,
+                "finish_reason": finish_reason,
+            }
+
+        except Exception as e:
+            raise e
+
     def handle_openai_text_completion_chunk(self, chunk):
         try:
             print_verbose(f"\nRaw OpenAI Chunk\n{chunk}\n")
@@ -9129,6 +9150,14 @@ class CustomStreamWrapper:
                     model_response.choices[0].finish_reason = response_obj[
                         "finish_reason"
                     ]
+            elif self.custom_llm_provider == "azure_text":
+                response_obj = self.handle_azure_text_completion_chunk(chunk)
+                completion_obj["content"] = response_obj["text"]
+                print_verbose(f"completion obj content: {completion_obj['content']}")
+                if response_obj["is_finished"]:
+                    model_response.choices[0].finish_reason = response_obj[
+                        "finish_reason"
+                    ]
             elif self.custom_llm_provider == "cached_response":
                 response_obj = {
                     "text": chunk.choices[0].delta.content,
@@ -9406,6 +9435,7 @@ class CustomStreamWrapper:
                 or self.custom_llm_provider == "azure"
                 or self.custom_llm_provider == "custom_openai"
                 or self.custom_llm_provider == "text-completion-openai"
+                or self.custom_llm_provider == "azure_text"
                 or self.custom_llm_provider == "huggingface"
                 or self.custom_llm_provider == "ollama"
                 or self.custom_llm_provider == "ollama_chat"