fix(streaming_iterator.py): return done chunks in the order expected by responses api

ensures responses api sdk's (e.g. openai ruby) work when calling non-openai models
This commit is contained in:
Krrish Dholakia 2025-10-10 15:30:10 -07:00
parent 15b5e6f5d9
commit c74eb6403d
2 changed files with 185 additions and 36 deletions

View File

@ -1,6 +1,6 @@
import time
import uuid
from typing import List, Optional, Union
from typing import List, Optional, Union, cast
import litellm
from litellm.main import stream_chunk_builder
@ -12,8 +12,13 @@ from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
ContentPartAddedEvent,
ContentPartDoneEvent,
ContentPartDonePartOutputText,
ContentPartDonePartReasoningText,
OutputItemAddedEvent,
OutputItemDoneEvent,
OutputTextDeltaEvent,
OutputTextDoneEvent,
ReasoningSummaryTextDeltaEvent,
ResponseCompletedEvent,
ResponseCreatedEvent,
@ -64,6 +69,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.sent_response_in_progress_event: bool = False
self.sent_output_item_added_event: bool = False
self.sent_content_part_added_event: bool = False
self.sent_output_text_done_event: bool = False
self.sent_output_content_part_done_event: bool = False
self.sent_output_item_done_event: bool = False
self.litellm_model_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = None
self.final_text: str = ""
def _default_response_created_event_data(self) -> dict:
response_created_event_data = {
@ -161,6 +173,97 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
),
)
def create_litellm_model_response(
self,
) -> Optional[ModelResponse]:
return cast(
Optional[ModelResponse],
stream_chunk_builder(
chunks=self.collected_chat_completion_chunks,
logging_obj=self.litellm_logging_obj,
),
)
def create_output_text_done_event(
self, litellm_complete_object: ModelResponse
) -> OutputTextDoneEvent:
return OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=f"msg_{str(uuid.uuid4())}",
output_index=0,
content_index=0,
text=getattr(litellm_complete_object.choices[0].message, "content", "") # type: ignore
or "",
)
def create_output_content_part_done_event(
self, litellm_complete_object: ModelResponse
) -> ContentPartDoneEvent:
text = getattr(litellm_complete_object.choices[0].message, "content", "") or "" # type: ignore
reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" # type: ignore
if reasoning_content:
part = ContentPartDonePartReasoningText(
type="reasoning_text",
reasoning=reasoning_content,
)
else:
part = ContentPartDonePartOutputText(
type="output_text",
text=text,
annotations=[],
logprobs=None,
)
return ContentPartDoneEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=f"msg_{str(uuid.uuid4())}",
output_index=0,
content_index=0,
part=part,
)
def create_output_item_done_event(
self, litellm_complete_object: ModelResponse
) -> OutputItemDoneEvent:
text = self.litellm_model_response.choices[0].message.content or "" # type: ignore
return OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
sequence_number=1,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": f"msg_{str(uuid.uuid4())}",
"status": "completed",
"type": "message",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": text,
"annotations": [],
}
],
}
),
)
def return_default_done_events(
self, litellm_complete_object: ModelResponse
) -> Optional[BaseLiteLLMOpenAIResponseObject]:
if self.sent_output_text_done_event is False:
self.sent_output_text_done_event = True
return self.create_output_text_done_event(litellm_complete_object)
if self.sent_output_content_part_done_event is False:
self.sent_output_content_part_done_event = True
return self.create_output_content_part_done_event(litellm_complete_object)
if self.sent_output_item_done_event is False:
self.sent_output_item_done_event = True
return self.create_output_item_done_event(litellm_complete_object)
return None
def return_default_initial_events(
self,
) -> Optional[BaseLiteLLMOpenAIResponseObject]:
@ -178,6 +281,44 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return self.create_content_part_added_event()
return None
def is_stream_finished(self) -> bool:
if (
self.sent_output_text_done_event is True
and self.sent_output_content_part_done_event is True
and self.sent_output_item_done_event is True
):
return True
return False
def common_done_event_logic(
self, sync_mode: bool = True
) -> BaseLiteLLMOpenAIResponseObject:
if not self.litellm_model_response or isinstance(
self.litellm_model_response, TextCompletionResponse
):
self.litellm_model_response = self.create_litellm_model_response()
if self.litellm_model_response:
done_event = self.return_default_done_events(self.litellm_model_response)
if done_event:
return done_event
else:
if sync_mode:
raise StopIteration
else:
raise StopAsyncIteration
self.finished = self.is_stream_finished()
response_completed_event = self._emit_response_completed_event(
self.litellm_model_response
)
if response_completed_event:
return response_completed_event
else:
if sync_mode:
raise StopIteration
else:
raise StopAsyncIteration
async def __anext__(
self,
) -> Union[
@ -205,12 +346,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if response_api_chunk:
return response_api_chunk
except StopAsyncIteration:
self.finished = True
response_completed_event = self._emit_response_completed_event()
if response_completed_event:
return response_completed_event
else:
raise StopAsyncIteration
return self.common_done_event_logic(sync_mode=False)
except Exception as e:
# Handle HTTP errors
@ -247,13 +383,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if response_api_chunk:
return response_api_chunk
except StopIteration:
self.finished = True
response_completed_event = self._emit_response_completed_event()
if response_completed_event:
return response_completed_event
else:
raise StopIteration
return self.common_done_event_logic(sync_mode=True)
except Exception as e:
# Handle HTTP errors
self.finished = True
@ -308,14 +438,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
chat_completion_delta: ChatCompletionDelta = choice.delta
return chat_completion_delta.content or ""
def _emit_response_completed_event(self) -> Optional[ResponseCompletedEvent]:
litellm_model_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = stream_chunk_builder(
chunks=self.collected_chat_completion_chunks,
logging_obj=self.litellm_logging_obj,
)
if litellm_model_response and isinstance(litellm_model_response, ModelResponse):
def _emit_response_completed_event(
self, litellm_model_response: ModelResponse
) -> Optional[ResponseCompletedEvent]:
if litellm_model_response:
# Add cost to usage object if include_cost_in_streaming_usage is True
if (
litellm.include_cost_in_streaming_usage

View File

@ -1189,9 +1189,23 @@ class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject):
class OutputItemDoneEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE]
output_index: int
sequence_number: int = 1
item: BaseLiteLLMOpenAIResponseObject
class OpenAIChatCompletionLogprobsContentTopLogprobs(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
class OpenAIChatCompletionLogprobsContent(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
top_logprobs: List[OpenAIChatCompletionLogprobsContentTopLogprobs]
class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_ADDED]
item_id: str
@ -1200,12 +1214,33 @@ class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject):
part: BaseLiteLLMOpenAIResponseObject
class ContentPartDonePartOutputText(BaseLiteLLMOpenAIResponseObject):
type: Literal["output_text"]
text: str
annotations: List[BaseLiteLLMOpenAIResponseObject]
logprobs: Optional[List[OpenAIChatCompletionLogprobsContent]]
class ContentPartDonePartRefusal(BaseLiteLLMOpenAIResponseObject):
type: Literal["refusal"]
refusal: str
class ContentPartDonePartReasoningText(BaseLiteLLMOpenAIResponseObject):
type: Literal["reasoning_text"]
reasoning: str
class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_DONE]
item_id: str
output_index: int
content_index: int
part: BaseLiteLLMOpenAIResponseObject
part: Union[
ContentPartDonePartOutputText,
ContentPartDonePartRefusal,
ContentPartDonePartReasoningText,
]
class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
@ -1735,19 +1770,6 @@ class OpenAIModerationResponse(BaseLiteLLMOpenAIResponseObject):
_hidden_params: dict = PrivateAttr(default_factory=dict)
class OpenAIChatCompletionLogprobsContentTopLogprobs(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
class OpenAIChatCompletionLogprobsContent(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
top_logprobs: List[OpenAIChatCompletionLogprobsContentTopLogprobs]
class OpenAIChatCompletionLogprobs(TypedDict, total=False):
content: List[OpenAIChatCompletionLogprobsContent]
refusal: List[OpenAIChatCompletionLogprobsContent]