Merge pull request #1572 from BerriAI/litellm_use_same_response_id_across_chunks

[Fix] Streaming - Use same `response_id` across chunks

litellm/tests/test_completion.py

@@ -1408,9 +1408,15 @@ def test_completion_sagemaker_stream():
         )
 
         complete_streaming_response = ""
-
-        for chunk in response:
+        first_chunk_id, chunk_id = None, None
+        for i, chunk in enumerate(response):
             print(chunk)
+            chunk_id = chunk.id
+            print(chunk_id)
+            if i == 0:
+                first_chunk_id = chunk_id
+            else:
+                assert chunk_id == first_chunk_id
             complete_streaming_response += chunk.choices[0].delta.content or ""
         # Add any assertions here to check the response
         # print(response)

litellm/tests/test_streaming.py

@@ -733,8 +733,15 @@ def test_completion_bedrock_claude_stream():
         complete_response = ""
         has_finish_reason = False
         # Add any assertions here to check the response
+        first_chunk_id = None
         for idx, chunk in enumerate(response):
             # print
+            if idx == 0:
+                first_chunk_id = chunk.id
+            else:
+                assert (
+                    chunk.id == first_chunk_id
+                ), f"chunk ids do not match: {chunk.id} != first chunk id{first_chunk_id}"
             chunk, finished = streaming_format_tests(idx, chunk)
             has_finish_reason = finished
             complete_response += chunk

litellm/utils.py

@@ -7041,6 +7041,7 @@ class CustomStreamWrapper:
         self._hidden_params = {
             "model_id": (_model_info.get("id", None))
         }  # returned as x-litellm-model-id response header in proxy
+        self.response_id = None
 
     def __iter__(self):
         return self
@@ -7613,6 +7614,10 @@ class CustomStreamWrapper:
 
     def chunk_creator(self, chunk):
         model_response = ModelResponse(stream=True, model=self.model)
+        if self.response_id is not None:
+            model_response.id = self.response_id
+        else:
+            self.response_id = model_response.id
         model_response._hidden_params["custom_llm_provider"] = self.custom_llm_provider
         model_response.choices = [StreamingChoices()]
         model_response.choices[0].finish_reason = None