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) 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 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