From 1d7cc35bae7cbb54025bd69da9b2900bad7c1f11 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 6 Dec 2023 17:16:18 -0800
Subject: [PATCH] (fix) proxy edit custom logger

---
 litellm/proxy/custom_callbacks.py | 47 +++++--------------------------
 1 file changed, 7 insertions(+), 40 deletions(-)

diff --git a/litellm/proxy/custom_callbacks.py b/litellm/proxy/custom_callbacks.py
index c5a4a00c5b..18aea8a97f 100644
--- a/litellm/proxy/custom_callbacks.py
+++ b/litellm/proxy/custom_callbacks.py
@@ -31,19 +31,21 @@ class MyCustomHandler(CustomLogger):
         print(f"On Stream")
         
     def log_success_event(self, kwargs, response_obj, start_time, end_time): 
+        print("On Success!")
+    
+    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()
+        # 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`
+        # 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() #######################
+        # Calculate cost using  litellm.completion_cost()
         cost = litellm.completion_cost(completion_response=response_obj)
         response = response_obj
         # tokens used in response 
@@ -99,41 +101,6 @@ class MyCustomHandler(CustomLogger):
             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()