diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 37f3d6e49e0..99d4f65d45a 100644 --- a/litellm/proxy/proxy_server.py +++ b/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: diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 699606b5277..7577117dbe0 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/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)