From ab58e32e90af3793e048dca20524a75a3006ec95 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 15 Sep 2023 17:46:35 -0700
Subject: [PATCH] streaming fixes

---
 litellm/utils.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/litellm/utils.py b/litellm/utils.py
index 220ec82781..6ee25ac486 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -2454,7 +2454,6 @@ class CustomStreamWrapper:
     def __next__(self):
         try:
             # return this for all models
-            model_response = ModelResponse(stream=True)
             completion_obj = {"content": ""} # default to role being assistant
             if self.model in litellm.anthropic_models:
                 chunk = next(self.completion_stream)
@@ -2506,8 +2505,11 @@ class CustomStreamWrapper:
 
             # LOGGING
             threading.Thread(target=self.logging_obj.success_handler, args=(completion_obj,)).start()
+            model_response = ModelResponse(stream=True)
+            model_response.choices[0].delta = completion_obj
+            model_response.model = self.model
 
-            if model_response.choices[0].delta.content == "<special_litellm_token>":
+            if model_response.choices[0].delta['content'] == "<special_litellm_token>":
                 model_response.choices[0].delta = {
                     "content": completion_obj["content"],
                 }