From 5492da1c8a5b0d5597a6126cf66d1aa842763fa1 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 28 Feb 2024 14:48:55 -0800
Subject: [PATCH] (feat) fix custom handler bug

---
 litellm/proxy/custom_callbacks.py | 66 +++++++++++++++++++++++++++++++
 litellm/proxy/proxy_config.yaml   |  4 +-
 2 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 litellm/proxy/custom_callbacks.py

diff --git a/litellm/proxy/custom_callbacks.py b/litellm/proxy/custom_callbacks.py
new file mode 100644
index 0000000000..40fc0d369b
--- /dev/null
+++ b/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
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 5755293e00..6b4b7a8f62 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -43,9 +43,9 @@ model_list:
       api_key: os.environ/OPENAI_API_KEY
 litellm_settings:
   fallbacks: [{"openai-gpt-3.5": ["azure-gpt-3.5"]}]
-  success_callback: ['clickhouse', 'langfuse']
+  success_callback: ['langfuse']
   # setting callback class
-  # callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance]
+  callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance]
 
 general_settings: 
   master_key: sk-1234