_get_assembled_streaming_response

This commit is contained in:
Ishaan Jaff 2025-03-12 07:21:03 -07:00
parent a3d1e39164
commit c1b9c4cc7b

View file

@ -1117,7 +1117,7 @@ class Logging(LiteLLMLoggingBaseClass):
## BUILD COMPLETE STREAMED RESPONSE ## BUILD COMPLETE STREAMED RESPONSE
complete_streaming_response: Optional[ complete_streaming_response: Optional[
Union[ModelResponse, TextCompletionResponse] Union[ModelResponse, TextCompletionResponse, ResponseCompletedEvent]
] = None ] = None
if "complete_streaming_response" in self.model_call_details: if "complete_streaming_response" in self.model_call_details:
return # break out of this. return # break out of this.
@ -1639,7 +1639,7 @@ class Logging(LiteLLMLoggingBaseClass):
if "async_complete_streaming_response" in self.model_call_details: if "async_complete_streaming_response" in self.model_call_details:
return # break out of this. return # break out of this.
complete_streaming_response: Optional[ complete_streaming_response: Optional[
Union[ModelResponse, TextCompletionResponse] Union[ModelResponse, TextCompletionResponse, ResponseCompletedEvent]
] = self._get_assembled_streaming_response( ] = self._get_assembled_streaming_response(
result=result, result=result,
start_time=start_time, start_time=start_time,
@ -2349,16 +2349,24 @@ class Logging(LiteLLMLoggingBaseClass):
def _get_assembled_streaming_response( def _get_assembled_streaming_response(
self, self,
result: Union[ModelResponse, TextCompletionResponse, ModelResponseStream, Any], result: Union[
ModelResponse,
TextCompletionResponse,
ModelResponseStream,
ResponseCompletedEvent,
Any,
],
start_time: datetime.datetime, start_time: datetime.datetime,
end_time: datetime.datetime, end_time: datetime.datetime,
is_async: bool, is_async: bool,
streaming_chunks: List[Any], streaming_chunks: List[Any],
) -> Optional[Union[ModelResponse, TextCompletionResponse]]: ) -> Optional[Union[ModelResponse, TextCompletionResponse, ResponseCompletedEvent]]:
if isinstance(result, ModelResponse): if isinstance(result, ModelResponse):
return result return result
elif isinstance(result, TextCompletionResponse): elif isinstance(result, TextCompletionResponse):
return result return result
elif isinstance(result, ResponseCompletedEvent):
return result
elif isinstance(result, ModelResponseStream): elif isinstance(result, ModelResponseStream):
complete_streaming_response: Optional[ complete_streaming_response: Optional[
Union[ModelResponse, TextCompletionResponse] Union[ModelResponse, TextCompletionResponse]