From bbe6a92eb999bdfd97149286d59d1c37f72178e8 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 25 Jan 2024 14:51:08 -0800
Subject: [PATCH] fix(main.py): fix order of assembly for streaming chunks

---
 litellm/main.py                     | 6 ++++++
 litellm/tests/test_custom_logger.py | 2 +-
 litellm/utils.py                    | 1 +
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/litellm/main.py b/litellm/main.py
index fca3bd2b25..6b40354739 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -3343,6 +3343,12 @@ def stream_chunk_builder(
     chunks: list, messages: Optional[list] = None, start_time=None, end_time=None
 ):
     model_response = litellm.ModelResponse()
+    ### SORT CHUNKS BASED ON CREATED ORDER ##
+    if chunks[0]._hidden_params.get("created_at", None):
+        # Sort chunks based on created_at in ascending order
+        chunks = sorted(
+            chunks, key=lambda x: x._hidden_params.get("created_at", float("inf"))
+        )
     # set hidden params from chunk to model_response
     if model_response is not None and hasattr(model_response, "_hidden_params"):
         model_response._hidden_params = chunks[0].get("_hidden_params", {})
diff --git a/litellm/tests/test_custom_logger.py b/litellm/tests/test_custom_logger.py
index 565df5b25a..e403c3afe4 100644
--- a/litellm/tests/test_custom_logger.py
+++ b/litellm/tests/test_custom_logger.py
@@ -211,7 +211,7 @@ def test_azure_completion_stream():
             {"role": "system", "content": "You are a helpful assistant."},
             {
                 "role": "user",
-                "content": "write 1 sentence about litellm being amazing",
+                "content": f"write 1 sentence about litellm being amazing {time.time()}",
             },
         ]
         complete_streaming_response = ""
diff --git a/litellm/utils.py b/litellm/utils.py
index fb3210b1df..02ac83d065 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -7714,6 +7714,7 @@ class CustomStreamWrapper:
         else:
             self.response_id = model_response.id
         model_response._hidden_params["custom_llm_provider"] = self.custom_llm_provider
+        model_response._hidden_params["created_at"] = time.time()
         model_response.choices = [StreamingChoices()]
         model_response.choices[0].finish_reason = None
         response_obj = {}