custom_logger for litellm - callback_func

litellm/integrations/custom_logger.py

@@ -25,7 +25,6 @@ class CustomLogger:
                 response_obj,
                 start_time,
                 end_time,
-                print_verbose
             )
             print_verbose(
                 f"Custom Logger - final response object: {response_obj}"

litellm/tests/test_custom_logger.py

@@ -0,0 +1,43 @@
+### What this tests ####
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath('../..'))
+
+from litellm import completion, embedding
+import litellm
+
+def custom_callback(
+        kwargs,
+        completion_response,
+        start_time,
+        end_time,
+):
+    print(
+        "in custom callback func"
+    )
+    print("kwargs", kwargs)
+    print(completion_response)
+    print(start_time)
+    print(end_time)
+litellm.success_callback = [custom_callback]
+
+
+litellm.set_verbose = True
+
+
+# def test_chat_openai():
+#     try:
+#         response = completion(model="gpt-3.5-turbo",
+#                               messages=[{
+#                                   "role": "user",
+#                                   "content": "Hi 👋 - i'm openai"
+#                               }])
+
+#         print(response)
+
+#     except Exception as e:
+#         print(e)
+
+
+# test_chat_openai()

litellm/utils.py

@@ -1078,7 +1078,7 @@ def set_callbacks(callback_list, function_id=None):
                 else:
                     liteDebuggerClient = LiteDebugger(email=str(uuid.uuid4()))
             elif callable(callback):
-                customLogger = CustomLogger()
+                customLogger = CustomLogger(callback_func=callback)
     except Exception as e:
         raise e