From b3f039627ea132f4bd19040f4b9c7730501adfbe Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 6 Dec 2023 14:41:40 -0800
Subject: [PATCH] (feat) litellm - add _async_failure_callback

---
 litellm/__init__.py                   |  1 +
 litellm/integrations/custom_logger.py | 20 +++++++++++++
 litellm/utils.py                      | 42 +++++++++++++++++++++++++++
 3 files changed, 63 insertions(+)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index 1ef9b9af207..9e4a859c797 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -10,6 +10,7 @@ success_callback: List[Union[str, Callable]] = []
 failure_callback: List[Union[str, Callable]] = []
 callbacks: List[Callable] = []
 _async_success_callback: List[Callable] = [] # internal variable - async custom callbacks are routed here. 
+_async_failure_callback: List[Callable] = [] # internal variable - async custom callbacks are routed here. 
 pre_call_rules: List[Callable] = []
 post_call_rules: List[Callable] = []
 email: Optional[
diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py
index e502439a95f..bcfef0fc2b6 100644
--- a/litellm/integrations/custom_logger.py
+++ b/litellm/integrations/custom_logger.py
@@ -81,3 +81,23 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback
             # traceback.print_exc()
             print_verbose(f"Custom Logger Error - {traceback.format_exc()}")
             pass
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time, print_verbose, callback_func):
+        # Method definition
+        try:
+            kwargs["log_event_type"] = "post_api_call"
+            await callback_func(
+                kwargs, # kwargs to func
+                response_obj,
+                start_time,
+                end_time,
+            )
+            print_verbose(
+                f"Custom Logger - final response object: {response_obj}"
+            )
+        except:
+            # traceback.print_exc()
+            print_verbose(f"Custom Logger Error - {traceback.format_exc()}")
+            pass
+    
+    
diff --git a/litellm/utils.py b/litellm/utils.py
index 9e93f6b6472..08c1066686e 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1113,6 +1113,36 @@ class Logging:
                 f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {traceback.format_exc()}"
             )
             pass
+    async def async_failure_handler(self, exception, traceback_exception, start_time=None, end_time=None):
+        """
+        Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
+        """
+        # on some exceptions, model_call_details is not always initialized, this ensures that we still log those exceptions
+        if not hasattr(self, "model_call_details"):
+            self.model_call_details = {}
+
+        self.model_call_details["log_event_type"] = "failed_api_call"
+        self.model_call_details["exception"] = exception
+        self.model_call_details["traceback_exception"] = traceback_exception
+        self.model_call_details["end_time"] = end_time
+        result = {} # result sent to all loggers, init this to None incase it's not created 
+
+        for callback in litellm._async_failure_callback:
+            try: 
+                if callable(callback): # custom logger functions
+                    await customLogger.async_log_failure_event(
+                        kwargs=self.model_call_details,
+                        response_obj=result,
+                        start_time=start_time,
+                        end_time=end_time,
+                        print_verbose=print_verbose,
+                        callback_func=callback
+                    )
+            except: 
+                print_verbose(
+                    f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {traceback.format_exc()}"
+                )
+
 
 
 def exception_logging(
@@ -1236,6 +1266,17 @@ def client(original_function):
                 # Pop the async items from success_callback in reverse order to avoid index issues
                 for index in reversed(removed_async_items):
                     litellm.success_callback.pop(index)
+            
+            if len(litellm.failure_callback) > 0: 
+                removed_async_items = []
+                for index, callback in enumerate(litellm.failure_callback): 
+                    if inspect.iscoroutinefunction(callback): 
+                        litellm._async_failure_callback.append(callback)
+                        removed_async_items.append(index)
+
+                # Pop the async items from success_callback in reverse order to avoid index issues
+                for index in reversed(removed_async_items):
+                    litellm.success_callback.pop(index)
             if add_breadcrumb:
                 add_breadcrumb(
                     category="litellm.llm_call",
@@ -1513,6 +1554,7 @@ def client(original_function):
             end_time = datetime.datetime.now()
             if logging_obj:
                 logging_obj.failure_handler(e, traceback_exception, start_time, end_time) # DO NOT MAKE THREADED - router retry fallback relies on this!
+                asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time))
             raise e
 
     is_coroutine = inspect.iscoroutinefunction(original_function)