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:
parent
15b5e6f5d9
commit
c74eb6403d
@ -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
|
||||
|
||||
@ -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]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user