fix(streaming): preserve served models in OpenRouter streams (#49498)

litellm/litellm_core_utils/streaming_handler.py

@@ -905,13 +905,21 @@ class CustomStreamWrapper:
         except Exception as e:
             raise e
 
+    def _served_openrouter_model(self) -> str | None:
+        if self.custom_llm_provider == LlmProviders.OPENROUTER.value or (
+            self.custom_llm_provider == "cached_response"
+            and self.logging_obj.model_call_details.get("custom_llm_provider") == LlmProviders.OPENROUTER.value
+        ):
+            return self._provider_response_model
+        return None
+
     def model_response_creator(
         self,
         chunk: dict | None = None,
         hidden_params: Mapping[str, object] | None = None,
         delta: Delta | None = None,
     ) -> ModelResponseStream:
-        _model: Final = self._cached_model_name
+        _model: Final = self._served_openrouter_model() or self._cached_model_name
         _logging_obj_llm_provider: Final = self._cached_logging_llm_provider
 
         if chunk is None:
@@ -1699,7 +1707,7 @@ class CustomStreamWrapper:
                 return dispatch_result.value
             response_obj = dispatch_result.response_obj
 
-            model_response.model = self.model
+            model_response.model = self._served_openrouter_model() or self.model
             ## FUNCTION CALL PARSING
             original_chunk: Final[_UpstreamChunkLike | None] = (
                 response_obj.get("original_chunk") if response_obj is not None else None
@@ -2561,7 +2569,7 @@ class CustomStreamWrapper:
         if partial is None or not isinstance(usage, Usage):
             return None
         if self.model:
-            partial.model = self.model
+            partial.model = self._served_openrouter_model() or self.model
         backfill_missing_cache_usage_fields(usage)
         return partial
 

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -7176,6 +7176,80 @@ def test_rebuilding_a_priced_stream_bills_the_margin_once(monkeypatch: pytest.Mo
     )
 
 
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("has_served_model", [True, False])
+@pytest.mark.parametrize(
+    ("wrapper_provider", "logging_provider"),
+    [("openrouter", "openrouter"), ("cached_response", "openrouter"),
+     ("openai", "openai"), ("cached_response", "openai")],
+)
+@pytest.mark.asyncio
+async def test_openrouter_served_model_survives_blank_chunks_usage_and_cache_replay(
+    sync_mode: bool, has_served_model: bool, wrapper_provider: str, logging_provider: str
+) -> None:
+    selected_model: Final = "selected-model" if has_served_model else ""
+    expected_model: Final = selected_model if logging_provider == "openrouter" and selected_model else "auto"
+    logging_obj: Final = Logging(
+        model="auto",
+        messages=[{"role": "user", "content": "Say hi"}],
+        stream=True,
+        call_type="completion" if sync_mode else "acompletion",
+        start_time=1.0,
+        litellm_call_id="openrouter-model-identity",
+        function_id="test",
+    )
+    logging_obj.stream_options = {"include_usage": True}
+    logging_obj.model_call_details["custom_llm_provider"] = (
+        None if wrapper_provider == "cached_response" else logging_provider
+    )
+    chunks: Final = [
+        ModelResponseStream(
+            id="served-stream", created=1, model=None,
+            choices=[StreamingChoices(index=0, delta=Delta(role="assistant"))],
+        ),
+        ModelResponseStream(
+            id="served-stream", created=1, model=selected_model,
+            choices=[StreamingChoices(index=0, delta=Delta(content="hello"))],
+            usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5),
+        ),
+        ModelResponseStream(
+            id="served-stream", created=1, model=None,
+            choices=[StreamingChoices(index=0, delta=Delta(content=" world"))],
+        ),
+        ModelResponseStream(
+            id="served-stream", created=1, model="",
+            choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")],
+        ),
+        ModelResponseStream(
+            id="served-stream", created=1, model=None, choices=[],
+            usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5),
+        ),
+    ]
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=ModelResponseListIterator(model_responses=chunks),
+        model="auto", logging_obj=logging_obj, custom_llm_provider=wrapper_provider,
+        stream_options={"include_usage": True},
+    )
+    logging_obj.update_environment_variables(
+        litellm_params={}, optional_params={}, model="auto", custom_llm_provider=logging_provider,
+    )
+    first: Final = next(wrapper) if sync_mode else await wrapper.__anext__()
+    second: Final = next(wrapper) if sync_mode else await wrapper.__anext__()
+    partial: Final = wrapper.partial_response()
+    emitted: Final = [first, second, *(await _collect(wrapper, sync_mode))]
+    assembled: Final = litellm.stream_chunk_builder(chunks=emitted)
+
+    assert first.model == "auto"
+    assert {chunk.model for chunk in emitted[1:]} == {expected_model}
+    assert partial is not None and partial.model == expected_model
+    assert assembled is not None and assembled.model == expected_model
+    assert assembled.choices[0].message.content == "hello world"
+    assert emitted[-1].choices == []
+    assert emitted[-1].usage is not None and emitted[-1].usage.total_tokens == 5
+    assert second._hidden_params.get("provider_response_model") == (selected_model or None)
+    assert wrapper.model == logging_obj.model == logging_obj.model_call_details["model"] == "auto"
+
+
 def test_chunk_creator_keeps_provider_model_private_across_stream():
     from litellm.litellm_core_utils.add_retry_fallback_headers import (
         get_hidden_params_dict,

tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py

@@ -11,7 +11,7 @@ import pytest
 import respx
 
 import litellm
-from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
 from litellm.llms.openrouter.chat.transformation import (
     OpenRouterChatCompletionStreamingHandler,
@@ -742,7 +742,11 @@ def test_openrouter_completion_with_non_object_usage_raises_api_connection_error
 
 
 @respx.mock
-def test_openrouter_streaming_completion_assembles_text_and_finish_reason() -> None:
+@pytest.mark.parametrize("async_stream", [False, True])
+@pytest.mark.parametrize("requested_model", ["openrouter/openai/gpt-4o-mini", "openrouter/auto"])
+async def test_openrouter_streaming_completion_assembles_text_and_finish_reason(
+    async_stream: bool, requested_model: str
+) -> None:
     chunks: Final = (
         {
             "id": "gen-stream-model",
@@ -785,14 +789,20 @@ def test_openrouter_streaming_completion_assembles_text_and_finish_reason() -> N
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": "How does a case reach the Supreme Court?"},
     ]
-    stream: Final = litellm.completion(
-        model="openrouter/openai/gpt-4o-mini",
+    kwargs: Final = dict(
+        model=requested_model,
         messages=messages,
         stream=True,
         max_tokens=50,
         api_key="test-openrouter-key",
     )
-    stream_chunks: Final = tuple(stream)
+    client: Final = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport())
+    try:
+        stream: Final = await litellm.acompletion(**kwargs, client=client) if async_stream else litellm.completion(**kwargs)
+        stream_chunks: Final = tuple([chunk async for chunk in stream]) if async_stream else tuple(stream)
+    finally:
+        await client.client.aclose()
+    assert {chunk.model for chunk in stream_chunks} == {chunks[0]["model"]}
     assert stream_chunks[0].choices[0].delta.role == "assistant"
     assert tuple(
         chunk.choices[0].finish_reason for chunk in stream_chunks if chunk.choices[0].finish_reason is not None
@@ -802,6 +812,7 @@ def test_openrouter_streaming_completion_assembles_text_and_finish_reason() -> N
     ) == "The Supreme Court hears federal cases."
     assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages)
     assert assembled is not None
+    assert assembled.model == chunks[0]["model"]
     assert assembled.choices[0].message.content == "The Supreme Court hears federal cases."
     assert assembled.choices[0].finish_reason == "stop"