From b7225947928da16f72e5a7fff4b5a3f7889291f2 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 13 Nov 2023 17:24:13 -0800
Subject: [PATCH] test(utils.py): add logging and fix azure streaming

---
 litellm/llms/azure.py | 2 +-
 litellm/utils.py      | 5 +++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py
index d617d977c6a..d20d27c7c71 100644
--- a/litellm/llms/azure.py
+++ b/litellm/llms/azure.py
@@ -181,7 +181,7 @@ class AzureChatCompletion(BaseLLM):
                         raise AzureOpenAIError(status_code=response.status_code, message=response.text)
                     
                     completion_stream = response.iter_lines()
-                    streamwrapper = CustomStreamWrapper(completion_stream=completion_stream, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
+                    streamwrapper = CustomStreamWrapper(completion_stream=completion_stream, model=model, custom_llm_provider="azure",logging_obj=logging_obj)
                     for transformed_chunk in streamwrapper:
                         yield transformed_chunk
 
diff --git a/litellm/utils.py b/litellm/utils.py
index e2593b4a222..f9b71567364 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -4296,6 +4296,7 @@ class CustomStreamWrapper:
         is_finished = False
         finish_reason = ""
         text = ""
+        print_verbose(f"chunk: {chunk}")
         if "data: [DONE]" in chunk:
             text = ""
             is_finished = True
@@ -4308,6 +4309,7 @@ class CustomStreamWrapper:
                 if data_json["choices"][0].get("finish_reason", None): 
                     is_finished = True
                     finish_reason = data_json["choices"][0]["finish_reason"]
+                print_verbose(f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}")
                 return {"text": text, "is_finished": is_finished, "finish_reason": finish_reason}
             except:
                 raise ValueError(f"Unable to parse response. Original response: {chunk}")
@@ -4370,6 +4372,7 @@ class CustomStreamWrapper:
             text = "" 
             is_finished = False
             finish_reason = None
+            print_verbose(f"str_line: {str_line}")
             if "data: [DONE]" in str_line:
                 text = ""
                 is_finished = True
@@ -4489,6 +4492,8 @@ class CustomStreamWrapper:
             elif self.custom_llm_provider and self.custom_llm_provider == "azure": 
                 response_obj = self.handle_azure_chunk(chunk)
                 completion_obj["content"] = response_obj["text"]
+                print_verbose(f"completion obj content: {completion_obj['content']}")
+                print_verbose(f"len(completion_obj['content']: {len(completion_obj['content'])}")
                 if response_obj["is_finished"]: 
                     model_response.choices[0].finish_reason = response_obj["finish_reason"]
             elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":