_handle_logging_completed_response

This commit is contained in:
Ishaan Jaff 2025-03-12 09:42:48 -07:00
parent d6a49f6b66
commit 35e9bba154

View file

@ -6,7 +6,9 @@ from typing import Any, AsyncIterator, Dict, Optional, Union
import httpx import httpx
from litellm.constants import STREAM_SSE_DONE_STRING from litellm.constants import STREAM_SSE_DONE_STRING
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.types.llms.openai import ( from litellm.types.llms.openai import (
ResponsesAPIStreamEvents, ResponsesAPIStreamEvents,
@ -72,7 +74,7 @@ class BaseResponsesAPIStreamingIterator:
== ResponsesAPIStreamEvents.RESPONSE_COMPLETED == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
): ):
self.completed_response = openai_responses_api_chunk self.completed_response = openai_responses_api_chunk
self._handle_completed_response() self._handle_logging_completed_response()
return openai_responses_api_chunk return openai_responses_api_chunk
@ -81,7 +83,7 @@ class BaseResponsesAPIStreamingIterator:
# If we can't parse the chunk, continue # If we can't parse the chunk, continue
return None return None
def _handle_completed_response(self): def _handle_logging_completed_response(self):
"""Base implementation - should be overridden by subclasses""" """Base implementation - should be overridden by subclasses"""
pass pass
@ -128,7 +130,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
self.finished = True self.finished = True
raise e raise e
def _handle_completed_response(self): def _handle_logging_completed_response(self):
"""Handle logging for completed responses in async context""" """Handle logging for completed responses in async context"""
asyncio.create_task( asyncio.create_task(
self.logging_obj.async_success_handler( self.logging_obj.async_success_handler(
@ -139,6 +141,14 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
) )
) )
executor.submit(
self.logging_obj.success_handler,
result=self.completed_response,
cache_hit=None,
start_time=self.start_time,
end_time=datetime.now(),
)
class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
""" """
@ -181,11 +191,20 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
self.finished = True self.finished = True
raise e raise e
def _handle_completed_response(self): def _handle_logging_completed_response(self):
"""Handle logging for completed responses in sync context""" """Handle logging for completed responses in sync context"""
self.logging_obj.success_handler( run_async_function(
async_function=self.logging_obj.async_success_handler,
result=self.completed_response, result=self.completed_response,
start_time=self.start_time, start_time=self.start_time,
end_time=datetime.now(), end_time=datetime.now(),
cache_hit=None, cache_hit=None,
) )
executor.submit(
self.logging_obj.success_handler,
result=self.completed_response,
cache_hit=None,
start_time=self.start_time,
end_time=datetime.now(),
)