From 5fc61f6d7a51d5feba077104733aeb582ec49cbf Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 11 Dec 2023 12:39:28 -0800 Subject: [PATCH] (test) custom logger --- litellm/tests/test_aamazing_proxy_custom_logger.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_aamazing_proxy_custom_logger.py b/litellm/tests/test_aamazing_proxy_custom_logger.py index a90be7d6bda..b92f5aa5a30 100644 --- a/litellm/tests/test_aamazing_proxy_custom_logger.py +++ b/litellm/tests/test_aamazing_proxy_custom_logger.py @@ -62,8 +62,12 @@ def test_embedding(client): print(litellm.callbacks) # assert len(litellm.callbacks) == 1 # assert litellm is initialized with 1 callback - + my_custom_logger = litellm.callbacks[0] + for callback in litellm.callbacks: + if "MyCustomHandler" in str(callback): + my_custom_logger = callback + break assert my_custom_logger.async_success_embedding == False test_data = { @@ -101,6 +105,10 @@ def test_chat_completion(client): # assert len(litellm.callbacks) == 1 # assert litellm is initialized with 1 callback my_custom_logger = litellm.callbacks[0] + for callback in litellm.callbacks: + if "MyCustomHandler" in str(callback): + my_custom_logger = callback + break assert my_custom_logger.async_success == False test_data = { @@ -148,6 +156,10 @@ def test_chat_completion_stream(client): # assert len(litellm.callbacks) == 1 # assert litellm is initialized with 1 callback my_custom_logger = litellm.callbacks[0] + for callback in litellm.callbacks: + if "MyCustomHandler" in str(callback): + my_custom_logger = callback + break assert my_custom_logger.streaming_response_obj == None # no streaming response obj is set pre call