fix(proxy): stream completed model responses

litellm/proxy/proxy_server.py

@@ -64,6 +64,7 @@ from litellm.litellm_core_utils.litellm_logging import (
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
+from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
 from litellm.proxy._types import (
     UI_TEAM_ID,
     CallbackDelete,
@@ -7135,6 +7136,7 @@ async def async_data_generator(
         model_mismatch_logged = False
         fallback_metadata_event_sent = False
         include_fallback_errors = _should_include_fallback_errors(request_data)
+        streaming_response = MockResponseIterator(response) if isinstance(response, ModelResponse) else response
         # Use a running string instead of list + join to avoid O(n^2) overhead.
         # Previously "".join(str_so_far_parts) was called every chunk, re-joining
         # the entire accumulated response. String += is O(n) amortized total.
@@ -7156,11 +7158,11 @@ async def async_data_generator(
         if needs_iterator_wrap:
             stream_iterator = proxy_logging_obj.async_post_call_streaming_iterator_hook(
                 user_api_key_dict=user_api_key_dict,
-                response=response,
+                response=streaming_response,
                 request_data=request_data,
             )
         else:
-            stream_iterator = response
+            stream_iterator = streaming_response
 
         async for chunk in stream_iterator:
             if needs_per_chunk_hook:

tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py

@@ -36,7 +36,13 @@ from litellm.proxy.proxy_server import (
     data_generator,
     select_data_generator,
 )
-from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
+from litellm.types.utils import (
+    Delta,
+    ModelResponse,
+    ModelResponseStream,
+    StreamingChoices,
+    Usage,
+)
 
 from .conftest import normalize
 
@@ -667,6 +673,77 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
     }
 
 
+@pytest.mark.asyncio
+async def test_async_data_generator_adapts_completed_model_response(monkeypatch):
+    _patch_logging_flags(monkeypatch, needs_wrap=True)
+
+    completed_response = ModelResponse(
+        id="chatcmpl-completed",
+        model="gpt-5.6",
+        choices=[
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "OK"},
+                "finish_reason": "stop",
+            }
+        ],
+        usage={
+            "prompt_tokens": 36,
+            "completion_tokens": 22,
+            "total_tokens": 58,
+        },
+    )
+    iterator_hook_responses = []
+
+    def iterator_hook(*, response, **_kwargs):
+        iterator_hook_responses.append(response)
+        return response
+
+    async def failure_hook(**kwargs):
+        return None
+
+    monkeypatch.setattr(
+        ps.proxy_logging_obj,
+        "async_post_call_streaming_iterator_hook",
+        iterator_hook,
+    )
+    monkeypatch.setattr(
+        ps.proxy_logging_obj,
+        "post_call_failure_hook",
+        failure_hook,
+    )
+
+    output = [
+        chunk
+        async for chunk in async_data_generator(
+            response=completed_response,
+            user_api_key_dict=_user_auth(),
+            request_data={"model": "gpt-5.6"},
+        )
+    ]
+
+    assert len(iterator_hook_responses) == 1
+    assert iterator_hook_responses[0] is not completed_response
+    assert hasattr(iterator_hook_responses[0], "__aiter__")
+    assert output[-1] == "data: [DONE]\n\n"
+    first_chunk = output[0]
+    assert isinstance(first_chunk, str)
+    payload = json.loads(first_chunk.removeprefix("data: ").removesuffix("\n\n"))
+    assert payload["id"] == completed_response.id
+    assert payload["model"] == completed_response.model
+    assert payload["object"] == "chat.completion.chunk"
+    assert payload["choices"][0]["delta"] == {
+        "role": "assistant",
+        "content": "OK",
+    }
+    assert payload["choices"][0]["finish_reason"] == "stop"
+    assert payload["usage"] == {
+        "prompt_tokens": 36,
+        "completion_tokens": 22,
+        "total_tokens": 58,
+    }
+
+
 @pytest.mark.asyncio
 async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch):
     _patch_logging_flags(monkeypatch)