Merge 0f114d399c into 7bc80994b7

litellm/litellm_core_utils/streaming_handler.py

@@ -1454,8 +1454,9 @@ class CustomStreamWrapper:
                 LlmProviders.AZURE_AI.value,
             ]:
                 if isinstance(chunk, BaseModel) and hasattr(chunk, "model"):
-                    # for azure, we need to pass the model from the original chunk
-                    self.model = getattr(chunk, "model", self.model)
+                    chunk_model: Final = getattr(chunk, "model", None)
+                    if chunk_model is not None:
+                        self.model = chunk_model
             response_obj = self.handle_openai_chat_completion_chunk(chunk)
             if response_obj is None:
                 return _ProviderChunkEarlyReturn(None)

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -3203,6 +3203,35 @@ def test_chunk_creator_drops_empty_finish_chunk(
     assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
 
 
+@pytest.mark.parametrize(
+    ("chunk_model", "expected_model"),
+    [(None, "azure/request-model"), ("azure/chunk-model", "azure/chunk-model")],
+)
+def test_azure_chunk_model_preserves_or_updates_request_model(
+    chunk_model: str | None, expected_model: str
+):
+    wrapper = CustomStreamWrapper(
+        completion_stream=None,
+        model="azure/request-model",
+        logging_obj=MagicMock(),
+        custom_llm_provider="azure",
+    )
+    chunk = ModelResponseStream(
+        model=chunk_model,
+        choices=[
+            StreamingChoices(
+                finish_reason=None,
+                index=0,
+                delta=Delta(content="hello", role="assistant"),
+            )
+        ],
+    )
+
+    wrapper.chunk_creator(chunk=chunk)
+
+    assert wrapper.model == expected_model
+
+
 def test_chunk_creator_stops_iteration_on_trailing_chunk(
     initialized_custom_stream_wrapper: CustomStreamWrapper,
 ):