From b1bd799be86cafdb008ae4360fe886e242708879 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 4 Dec 2023 08:26:06 -0800
Subject: [PATCH] (feat) proxy: use custom_logger

---
 litellm/proxy/custom_logger.py  | 19 +++++++++++++++++++
 litellm/proxy/proxy_config.yaml |  3 +++
 2 files changed, 22 insertions(+)
 create mode 100644 litellm/proxy/custom_logger.py

diff --git a/litellm/proxy/custom_logger.py b/litellm/proxy/custom_logger.py
new file mode 100644
index 00000000000..8a1a824ace3
--- /dev/null
+++ b/litellm/proxy/custom_logger.py
@@ -0,0 +1,19 @@
+from litellm.integrations.custom_logger import CustomLogger
+class MyCustomHandler(CustomLogger):
+    def log_pre_api_call(self, model, messages, kwargs): 
+        print(f"Pre-API Call")
+    
+    def log_post_api_call(self, kwargs, response_obj, start_time, end_time): 
+        # log: key, user, model, prompt, response, tokens, cost
+        print(f"Post-API Call")
+    
+    def log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        print(f"On Stream")
+        
+    def log_success_event(self, kwargs, response_obj, start_time, end_time): 
+        print(f"On Success")
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time): 
+        print(f"On Failure")
+
+customHandler = MyCustomHandler()
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index e7848f0dc50..2d45dd91b0c 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -3,6 +3,9 @@ model_list:
     litellm_params:
       model: gpt-3.5-turbo
 
+litellm_settings:
+  callbacks: [custom_logger.customHandler] # sets litellm.callbacks = [module.module_variable]
+
 general_settings: 
   # otel: True          # OpenTelemetry Logger
   # master_key: sk-1234 # [OPTIONAL] Only use this if you to require all calls to contain this key (Authorization: Bearer sk-1234)