From 4e8e3ff33a28f6336a3bd0d4ce33d03906893bcd Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 8 Dec 2023 15:16:10 -0800 Subject: [PATCH] (fix) async custom logger - trigger when stream completed --- litellm/utils.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 12bb9eaa95..565a32c0ba 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -999,12 +999,21 @@ class Logging: try: if isinstance(callback, CustomLogger): # custom logger class print_verbose(f"Async success callbacks: CustomLogger") - await callback.async_log_success_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) + if self.stream: + if "complete_streaming_response" in self.model_call_details: + await callback.async_log_success_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details["complete_streaming_response"], + start_time=start_time, + end_time=end_time, + ) + else: + await callback.async_log_success_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) if callable(callback): # custom logger functions print_verbose(f"Async success callbacks: async_log_event") await customLogger.async_log_event(