diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 1f936ad27f..49815e5dc8 100644 --- a/litellm/integrations/custom_logger.py +++ b/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}" diff --git a/litellm/tests/test_custom_logger.py b/litellm/tests/test_custom_logger.py new file mode 100644 index 0000000000..50231bb83d --- /dev/null +++ b/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() diff --git a/litellm/utils.py b/litellm/utils.py index 614d824fed..507e130f55 100644 --- a/litellm/utils.py +++ b/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