diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index afc71f681dd..0c74748ec7d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -365,10 +365,10 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] -def _is_json_content(response: httpx.Response) -> bool: - content_types: Final = httpx.Headers(response.headers).get_list("content-type") - media_types: Final = (value.split(";")[0].strip().lower() for value in content_types) - return "application/json" in media_types +def _is_sized_json_body(response: httpx.Response) -> bool: + headers: Final = httpx.Headers(response.headers) + media_types: Final = (value.split(";")[0].strip().lower() for value in headers.get_list("content-type")) + return "content-length" in headers and "application/json" in media_types def _is_json_document(content: bytes) -> bool: @@ -3035,7 +3035,7 @@ class BaseLLMHTTPHandler: **body_kwargs, ) if caller_requested_stream: - if fake_stream is True or (_is_json_content(response) and _is_json_document(response.read())): + if fake_stream is True or (_is_sized_json_body(response) and _is_json_document(response.read())): return MockResponsesAPIStreamingIterator( response=response, model=model, @@ -3239,7 +3239,7 @@ class BaseLLMHTTPHandler: if caller_requested_stream: if fake_stream is True or ( - _is_json_content(response) and _is_json_document(await response.aread()) + _is_sized_json_body(response) and _is_json_document(await response.aread()) ): return MockResponsesAPIStreamingIterator( response=response, diff --git a/tests/unit/llms/bigmodel/test_bigmodel_provider.py b/tests/unit/llms/bigmodel/test_bigmodel_provider.py index bb48827f535..5c485ef0291 100644 --- a/tests/unit/llms/bigmodel/test_bigmodel_provider.py +++ b/tests/unit/llms/bigmodel/test_bigmodel_provider.py @@ -1,3 +1,4 @@ +import asyncio import contextlib import gzip import io @@ -469,6 +470,51 @@ async def test_responses_stream_answered_with_one_json_response_replays_it_as_ev assert chunks[-1].response.usage.total_tokens == 30 +class _GatedStream(httpx.AsyncByteStream): + def __init__(self, first: bytes, rest: bytes, gate: asyncio.Event) -> None: + self._first = first + self._rest = rest + self._gate = gate + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self._first + await self._gate.wait() + yield self._rest + + +@pytest.mark.asyncio +@respx.mock +async def test_responses_stream_labeled_json_without_a_length_yields_events_before_it_ends() -> None: + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": {**response_payload(), "status": "in_progress", "output": [], "usage": None}, + } + completed: Final = {"type": "response.completed", "sequence_number": 1, "response": response_payload()} + gate: Final = asyncio.Event() + respx.post(f"{BASE}/api/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "application/json"}, + stream=_GatedStream( + f"event: response.created\ndata: {json.dumps(created)}\n\n".encode(), + f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(), + gate, + ), + ) + ) + + stream: Final = await asyncio.wait_for( + litellm.aresponses(model=f"bigmodel/{MODEL}", input="Hello", stream=True), timeout=10 + ) + first: Final = await asyncio.wait_for(anext(stream), timeout=10) + gate.set() + rest: Final = tuple([chunk async for chunk in stream]) + assert first.type == "response.created" + assert rest[-1].type == "response.completed" + assert rest[-1].response.output_text == "Hello" + + @pytest.mark.parametrize("event_type", ("response.failed", "error")) @pytest.mark.asyncio @respx.mock