From 00a789440d6ead260c28f5769e515346b36232bb Mon Sep 17 00:00:00 2001
From: daarko10 <34120174+daarko10@users.noreply.github.com>
Date: Wed, 21 May 2025 17:27:17 +0300
Subject: [PATCH] Improve response_id propagation logic and add tests for
 valid/empty ID handling in streaming. (#11006)

---
 .../litellm_core_utils/streaming_handler.py   |  7 +++--
 .../test_streaming_handler.py                 | 27 +++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 785186b8ab..a0767c09ee 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -695,10 +695,13 @@ class CustomStreamWrapper:
 
         Ensure model id is always the same across all chunks.
 
-        If first chunk sent + id set, use that id for all chunks.
+        If a valid ID is received in any chunk, use it for the response.
         """
-        if self.response_id is None:
+        if id and isinstance(id, str) and id.strip():
             self.response_id = id
+        elif self.response_id is None:
+            self.response_id = id
+
         if self.response_id is not None and isinstance(self.response_id, str):
             model_response.id = self.response_id
         return model_response
diff --git a/tests/litellm/litellm_core_utils/test_streaming_handler.py b/tests/litellm/litellm_core_utils/test_streaming_handler.py
index 81bde88f39..c38709735f 100644
--- a/tests/litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/litellm/litellm_core_utils/test_streaming_handler.py
@@ -612,6 +612,33 @@ def test_streaming_handler_with_stop_chunk(
     assert returned_chunk is None
 
 
+def test_set_response_id_propagation_empty_to_valid(initialized_custom_stream_wrapper: CustomStreamWrapper):
+    """Test that response_id is properly set when first chunk has empty ID and second chunk has valid ID"""
+
+    model_response1 = ModelResponseStream(id="", created=1742056047, model=None)
+    model_response1 = initialized_custom_stream_wrapper.set_model_id(model_response1.id, model_response1)
+    assert model_response1.id == ""
+
+    model_response2 = ModelResponseStream(id="valid-id-123", created=1742056048, model=None)
+    model_response2 = initialized_custom_stream_wrapper.set_model_id("valid-id-123", model_response2)
+    assert model_response2.id == "valid-id-123"
+    assert initialized_custom_stream_wrapper.response_id == "valid-id-123"
+
+
+def test_set_response_id_propagation_valid_to_invalid(initialized_custom_stream_wrapper: CustomStreamWrapper):
+    """Test that response_id is maintained when first chunk has valid ID and second chunk has invalid ID"""
+
+    model_response1 = ModelResponseStream(id="first-valid-id", created=1742056049, model=None)
+    model_response1 = initialized_custom_stream_wrapper.set_model_id("first-valid-id", model_response1)
+    assert model_response1.id == "first-valid-id"
+    assert initialized_custom_stream_wrapper.response_id == "first-valid-id"
+
+    model_response2 = ModelResponseStream(id="", created=1742056050, model=None)
+    model_response2 = initialized_custom_stream_wrapper.set_model_id("", model_response2)
+    assert model_response2.id == "first-valid-id"
+    assert initialized_custom_stream_wrapper.response_id == "first-valid-id"
+
+
 @pytest.mark.asyncio
 async def test_streaming_completion_start_time(logging_obj: Logging):
     """Test that the start time is set correctly"""