Skip to content

fix: recover streamed responses completed output #30933

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
emsi wants to merge 1 commit into from
Closed

fix: recover streamed responses completed output #30933

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions litellm/responses/streaming_iterator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@
)
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.sse_output_recovery import (
record_output_item_chunk,
record_output_text_chunk,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.utils import CallTypes
Expand Down Expand Up @@ -78,6 +82,8 @@ def __init__(
self._completed_response_cache_hit: Optional[bool] = None
self._persist_completed_response_before_logging = True
self._stream_created_time: float = time.time()
self._streamed_output_items: dict[int, dict[str, Any]] = {}
self._streamed_text_only_output_items: dict[int, dict[str, Any]] = {}

# track request context for hooks
self.litellm_metadata = litellm_metadata
Expand Down Expand Up @@ -138,6 +144,9 @@ def _process_chunk(self, chunk) -> Optional[Any]:

# Format as ResponsesAPIStreamingResponse
if isinstance(parsed_chunk, dict):
self._record_streamed_output_chunk(parsed_chunk)
parsed_chunk = self._backfill_completed_response_output(parsed_chunk)

if self.responses_api_provider_config is None:
raise ValueError(
"responses_api_provider_config is required to process live streaming chunks"
Expand Down Expand Up @@ -292,6 +301,54 @@ def _process_chunk(self, chunk) -> Optional[Any]:
self._handle_failure(e)
raise

def _record_streamed_output_chunk(self, parsed_chunk: dict[str, Any]) -> None:
event_type = parsed_chunk.get("type")
if event_type == ResponsesAPIStreamEvents.RESPONSE_CREATED:
self._streamed_output_items.clear()
self._streamed_text_only_output_items.clear()
return
if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE:
record_output_item_chunk(
parsed_chunk=parsed_chunk,
output_items=self._streamed_output_items,
)
return
if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE:
record_output_text_chunk(
parsed_chunk=parsed_chunk,
output_items=self._streamed_output_items,
text_only_items=self._streamed_text_only_output_items,
)

def _backfill_completed_response_output(
self, parsed_chunk: dict[str, Any]
) -> dict[str, Any]:
if parsed_chunk.get("type") != ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
return parsed_chunk

response_payload = parsed_chunk.get("response")
if not isinstance(response_payload, dict):
return parsed_chunk
if response_payload.get("output"):
return parsed_chunk

recovered_output = self._recovered_streamed_output_items()
if not recovered_output:
return parsed_chunk

completed_chunk = dict(parsed_chunk)
completed_response = dict(response_payload)
completed_response["output"] = recovered_output
completed_chunk["response"] = completed_response
return completed_chunk

def _recovered_streamed_output_items(self) -> list[dict[str, Any]]:
output_items: dict[int, dict[str, Any]] = {
**self._streamed_text_only_output_items
}
output_items.update(self._streamed_output_items)
return [dict(item) for _, item in sorted(output_items.items())]

def _log_completed_response(self, *, is_async: bool) -> None:
if self._completed_response_logged:
return
Expand Down
162 changes: 161 additions & 1 deletion tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
import json
import os
import sys
from unittest.mock import MagicMock, patch
from typing import cast
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest
Expand All @@ -19,6 +20,16 @@
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
from litellm.responses.streaming_iterator import SyncResponsesAPIStreamingIterator


def _output_text(output_item: object) -> str:
if hasattr(output_item, "model_dump"):
serialized = cast("dict[str, object]", output_item.model_dump())
else:
serialized = cast("dict[str, object]", output_item)
content = cast("list[dict[str, object]]", serialized["content"])
return cast(str, content[0]["text"])


class TestChatGPTResponsesAPITransformation:
Expand Down Expand Up @@ -248,6 +259,155 @@ def test_chatgpt_non_stream_sse_response_recovers_output_items(

assert parsed.output_text == "Hello from stream!"

def test_chatgpt_streaming_response_completed_recovers_output_item_done(self):
config = ChatGPTResponsesAPIConfig()
logging_obj = MagicMock()
logging_obj.model_call_details = {"litellm_params": {}}
logging_obj.async_success_handler = AsyncMock()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="gpt-5.5",
responses_api_provider_config=config,
logging_obj=logging_obj,
custom_llm_provider=LlmProviders.CHATGPT,
)
streamed_output_item = {
"id": "msg_from_item",
"type": "message",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "OK my lord",
"annotations": [],
"logprobs": [],
}
],
}
completed_response = {
"id": "resp_test",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "gpt-5.5",
"output": [],
}

iterator._process_chunk(
json.dumps(
{
"type": "response.output_item.done",
"output_index": 0,
"item": streamed_output_item,
}
)
)
completed_event = iterator._process_chunk(
json.dumps(
{
"type": "response.completed",
"response": completed_response,
}
)
)

assert completed_event is not None
assert completed_event.type == "response.completed"
assert _output_text(completed_event.response.output[0]) == "OK my lord"

def test_chatgpt_streaming_response_completed_keeps_authoritative_output(self):
config = ChatGPTResponsesAPIConfig()
logging_obj = MagicMock()
logging_obj.model_call_details = {"litellm_params": {}}
logging_obj.async_success_handler = AsyncMock()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="gpt-5.5",
responses_api_provider_config=config,
logging_obj=logging_obj,
custom_llm_provider=LlmProviders.CHATGPT,
)

iterator._process_chunk(
json.dumps(
{
"type": "response.output_text.done",
"output_index": 0,
"content_index": 0,
"item_id": "msg_from_stream",
"text": "Earlier stream text",
}
)
)
completed_event = iterator._process_chunk(
json.dumps(
{
"type": "response.completed",
"response": {
"id": "resp_test",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "gpt-5.5",
"output": [
{
"type": "message",
"role": "assistant",
"status": "completed",
"content": [
{
"type": "output_text",
"text": "Authoritative completed text",
"annotations": [],
}
],
}
],
},
}
)
)

assert completed_event is not None
assert (
_output_text(completed_event.response.output[0])
== "Authoritative completed text"
)

def test_chatgpt_streaming_recovered_output_items_are_copied(self):
config = ChatGPTResponsesAPIConfig()
logging_obj = MagicMock()
logging_obj.model_call_details = {"litellm_params": {}}
logging_obj.async_success_handler = AsyncMock()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="gpt-5.5",
responses_api_provider_config=config,
logging_obj=logging_obj,
custom_llm_provider=LlmProviders.CHATGPT,
)
streamed_output_item = {
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Recovered"}],
}

iterator._process_chunk(
json.dumps(
{
"type": "response.output_item.done",
"output_index": 0,
"item": streamed_output_item,
}
)
)
recovered_output = iterator._recovered_streamed_output_items()

assert recovered_output[0] == streamed_output_item
assert recovered_output[0] is not streamed_output_item

def test_chatgpt_non_stream_sse_recovers_whitespace_padded_chunks(self):
"""Chunks with leading whitespace before `data:` must still parse.

Expand Down
Loading