(fix) proxy use async logging

litellm/proxy/custom_callbacks.py

@@ -1,9 +1,26 @@
 from litellm.integrations.custom_logger import CustomLogger
 import litellm
+import inspect
 
 # This file includes the custom callbacks for LiteLLM Proxy
 # Once defined, these can be passed in proxy_config.yaml
 class MyCustomHandler(CustomLogger):
+    def __init__(self):
+        blue_color_code = "\033[94m"
+        reset_color_code = "\033[0m"
+        print(f"{blue_color_code}Initialized LiteLLM custom logger")
+        try:
+            print(f"Logger Initialized with following methods:")
+            methods = [method for method in dir(self) if inspect.ismethod(getattr(self, method))]
+            
+            # Pretty print the methods
+            for method in methods:
+                print(f" - {method}")
+            print(f"{reset_color_code}")
+        except:
+            pass
+        
+
     def log_pre_api_call(self, model, messages, kwargs): 
         print(f"Pre-API Call")
     
@@ -45,81 +62,81 @@ class MyCustomHandler(CustomLogger):
         )
         return
 
-    def log_failure_event(self, kwargs, response_obj, start_time, end_time): 
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): 
-        print(f"On Failure")
+        try:
+            print(f"On Async Failure !")
+            print("\nkwargs", kwargs)
+            # Access kwargs passed to litellm.completion()
+            model = kwargs.get("model", None)
+            messages = kwargs.get("messages", None)
+            user = kwargs.get("user", None)
+
+            # Access litellm_params passed to litellm.completion(), example access `metadata`
+            litellm_params = kwargs.get("litellm_params", {})
+            metadata = litellm_params.get("metadata", {})   # headers passed to LiteLLM proxy, can be found here
+
+            # Acess Exceptions & Traceback
+            exception_event = kwargs.get("exception", None)
+            traceback_event = kwargs.get("traceback_exception", None)
+
+            # Calculate cost using  litellm.completion_cost()
+            cost = litellm.completion_cost(completion_response=response_obj)
+            print("now checking response obj")
+            
+            print(
+                f"""
+                    Model: {model},
+                    Messages: {messages},
+                    User: {user},
+                    Cost: {cost},
+                    Response: {response_obj}
+                    Proxy Metadata: {metadata}
+                    Exception: {exception_event}
+                    Traceback: {traceback_event}
+                """
+            )
+        except Exception as e:
+            print(f"Exception: {e}")
+
+
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        print(f"On Async Success!")
+        # log: key, user, model, prompt, response, tokens, cost
+        print("\nOn Success")
+        # Access kwargs passed to litellm.completion()
+        model = kwargs.get("model", None)
+        messages = kwargs.get("messages", None)
+        user = kwargs.get("user", None)
+
+        # Access litellm_params passed to litellm.completion(), example access `metadata`
+        litellm_params = kwargs.get("litellm_params", {})
+        metadata = litellm_params.get("metadata", {})   # headers passed to LiteLLM proxy, can be found here
+
+        # Calculate cost using  litellm.completion_cost()
+        cost = litellm.completion_cost(completion_response=response_obj)
+        response = response_obj
+        # tokens used in response 
+        usage = response_obj["usage"]
+
+        print(
+            f"""
+                Model: {model},
+                Messages: {messages},
+                User: {user},
+                Usage: {usage},
+                Cost: {cost},
+                Response: {response}
+                Proxy Metadata: {metadata}
+            """
+        )
+        return
+
+        
 
 proxy_handler_instance = MyCustomHandler()
 
+
 # need to set litellm.callbacks = [customHandler] # on the proxy
 
-## setting only one function 
-async def async_on_succes_logger(kwargs, response_obj, start_time, end_time):
-    print(f"On Async Success!")
-    # log: key, user, model, prompt, response, tokens, cost
-    print("\nOn Success")
-    # Access kwargs passed to litellm.completion()
-    model = kwargs.get("model", None)
-    messages = kwargs.get("messages", None)
-    user = kwargs.get("user", None)
-
-    # Access litellm_params passed to litellm.completion(), example access `metadata`
-    litellm_params = kwargs.get("litellm_params", {})
-    metadata = litellm_params.get("metadata", {})   # headers passed to LiteLLM proxy, can be found here
-
-    # Calculate cost using  litellm.completion_cost()
-    cost = litellm.completion_cost(completion_response=response_obj)
-    response = response_obj
-    # tokens used in response 
-    usage = response_obj["usage"]
-
-    print(
-        f"""
-            Model: {model},
-            Messages: {messages},
-            User: {user},
-            Usage: {usage},
-            Cost: {cost},
-            Response: {response}
-            Proxy Metadata: {metadata}
-        """
-    )
-    return
-
-
-async def async_on_fail_logger(kwargs, response_obj, start_time, end_time):
-    print(f"On Async Failure!")
-
-    # Access kwargs passed to litellm.completion()
-    model = kwargs.get("model", None)
-    messages = kwargs.get("messages", None)
-    user = kwargs.get("user", None)
-
-    # Access litellm_params passed to litellm.completion(), example access `metadata`
-    litellm_params = kwargs.get("litellm_params", {})
-    metadata = litellm_params.get("metadata", {})   # headers passed to LiteLLM proxy, can be found here
-
-    # Acess Exceptions & Traceback
-    exception_event = kwargs.get("exception", None)
-    traceback_event = kwargs.get("traceback_exception", None)
-
-    # Calculate cost using  litellm.completion_cost()
-    cost = litellm.completion_cost(completion_response=response_obj)
-    response = response_obj
-    # tokens used in response 
-    usage = response_obj.get("usage", {})
-
-    print(
-        f"""
-            Model: {model},
-            Messages: {messages},
-            User: {user},
-            Usage: {usage},
-            Cost: {cost},
-            Response: {response}
-            Proxy Metadata: {metadata}
-            Exception: {exception_event}
-            Traceback: {traceback_event}
-        """
-    )
-
 # litellm.success_callback = [async_on_succes_logger]