(feat) fix custom handler bug

litellm/proxy/custom_callbacks.py

@@ -0,0 +1,66 @@
+from litellm.integrations.custom_logger import CustomLogger
+import litellm
+
+
+# This file includes the custom callbacks for LiteLLM Proxy
+# Once defined, these can be passed in proxy_config.yaml
+class MyCustomHandler(CustomLogger):
+    def log_pre_api_call(self, model, messages, kwargs):
+        print(f"Pre-API Call")  # noqa
+
+    def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
+        print(f"Post-API Call")  # noqa
+
+    def log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        print(f"On Stream")  # noqa
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        print("On Success")  # noqa
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        print(f"On Failure")  # noqa
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        print(f"ishaan async_log_success_event")  # noqa
+        # log: key, user, model, prompt, response, tokens, cost
+        # 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
+
+        return
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        try:
+            print(f"On Async Failure !")  # noqa
+            print("\nkwargs", kwargs)  # noqa
+            # 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()
+        except Exception as e:
+            print(f"Exception: {e}")  # noqa
+
+
+proxy_handler_instance = MyCustomHandler()
+
+# Set litellm.callbacks = [proxy_handler_instance] on the proxy
+# need to set litellm.callbacks = [proxy_handler_instance] # on the proxy