From 39b4f19bd8acedc77cae67b881f1dde4149fc9fc Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 23 Jan 2024 12:56:03 -0800
Subject: [PATCH 1/3] (fix) same response_id across chunk

---
 litellm/utils.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/litellm/utils.py b/litellm/utils.py
index 85d160334ea..a400a899e8a 100644
--- a/litellm/utils.py
+++ b/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

From d455833dfb5477818f1269e72fd5a067f61aedaf Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 23 Jan 2024 12:57:04 -0800
Subject: [PATCH 2/3] (test) same response id across chunks

---
 litellm/tests/test_completion.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 4215802530a..b2c69804ccc 100644
--- a/litellm/tests/test_completion.py
+++ b/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)

From 64a387d09bf5e3e6e7e8ec3faa18414a79b33fb0 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 23 Jan 2024 13:06:05 -0800
Subject: [PATCH 3/3] (test) test chunk_ids match across chunks for bedrock

---
 litellm/tests/test_streaming.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 959e63d5945..14b1a721030 100644
--- a/litellm/tests/test_streaming.py
+++ b/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