fixes streaming iterator

This commit is contained in:
Ishaan Jaff 2025-04-18 18:07:43 -07:00
parent 0b2f68ddd2
commit e91bdff9f4

View file

@ -5,12 +5,11 @@ from litellm.main import stream_chunk_builder
from litellm.responses.litellm_completion_transformation.transformation import ( from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig, LiteLLMCompletionResponsesConfig,
) )
from litellm.responses.streaming_iterator import ( from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
ResponsesAPIStreamingIterator,
SyncResponsesAPIStreamingIterator,
)
from litellm.types.llms.openai import ( from litellm.types.llms.openai import (
ResponseCompletedEvent, ResponseCompletedEvent,
ResponseInputParam,
ResponsesAPIOptionalRequestParams,
ResponsesAPIStreamEvents, ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse, ResponsesAPIStreamingResponse,
) )
@ -29,10 +28,16 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def __init__( def __init__(
self, self,
litellm_custom_stream_wrapper: litellm.CustomStreamWrapper, litellm_custom_stream_wrapper: litellm.CustomStreamWrapper,
request_input: Union[str, ResponseInputParam],
responses_api_request: ResponsesAPIOptionalRequestParams,
): ):
self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = ( self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = (
litellm_custom_stream_wrapper litellm_custom_stream_wrapper
) )
self.request_input: Union[str, ResponseInputParam] = request_input
self.responses_api_request: ResponsesAPIOptionalRequestParams = (
responses_api_request
)
self.collected_chunks: List[ModelResponseStream] = [] self.collected_chunks: List[ModelResponseStream] = []
self.finished: bool = False self.finished: bool = False
@ -65,10 +70,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
Union[ModelResponse, TextCompletionResponse] Union[ModelResponse, TextCompletionResponse]
] = stream_chunk_builder(chunks=self.collected_chunks) ] = stream_chunk_builder(chunks=self.collected_chunks)
if litellm_model_response and isinstance(litellm_model_response, ModelResponse): if litellm_model_response and isinstance(litellm_model_response, ModelResponse):
return ResponseCompletedEvent( return ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( response=LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
litellm_model_response request_input=self.request_input,
chat_completion_response=litellm_model_response,
responses_api_request=self.responses_api_request,
), ),
) )
else: else: