phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(test) aembedding on custom logger

Browse source
This commit is contained in:
ishaan-jaff 2023-12-06 19:09:04 -08:00
parent 7fcd17cbbe
commit f3469fdef8
1 changed files with 24 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

25
litellm/tests/test_custom_logger.py
View file

@ -13,10 +13,13 @@ class MyCustomHandler(CustomLogger):
self.success: bool = False self.success: bool = False
self.failure: bool = False self.failure: bool = False
self.async_success: bool = False self.async_success: bool = False
self.async_success_embedding: bool = False
self.async_failure: bool = False self.async_failure: bool = False
self.async_failure_embedding: bool = False
self.async_completion_kwargs = None # test if the kwargs are available, for async_succcess for completion/embedding self.async_completion_kwargs = None # test if the kwargs are available, for async_succcess for completion/embedding
self.async_embedding_kwargs = None self.async_embedding_kwargs = None
self.async_embedding_response = None
self.async_completion_kwargs_fail = None self.async_completion_kwargs_fail = None
self.async_embedding_kwargs_fail = None self.async_embedding_kwargs_fail = None
@ -43,6 +46,10 @@ class MyCustomHandler(CustomLogger):
self.async_success = True self.async_success = True
print("Value of async success: ", self.async_success) print("Value of async success: ", self.async_success)
print("\n kwargs: ", kwargs) print("\n kwargs: ", kwargs)
if kwargs.get("model") == "text-embedding-ada-002":
self.async_success_embedding = True
self.async_embedding_kwargs = kwargs
self.async_embedding_response = response_obj
self.async_completion_kwargs = kwargs self.async_completion_kwargs = kwargs
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
@ -107,6 +114,7 @@ def test_async_custom_handler():
try: try:
customHandler2 = MyCustomHandler() customHandler2 = MyCustomHandler()
litellm.callbacks = [customHandler2] litellm.callbacks = [customHandler2]
litellm.set_verbose = True
messages = [ messages = [
{"role": "system", "content": "You are a helpful assistant."}, {"role": "system", "content": "You are a helpful assistant."},
{ {
@ -134,7 +142,10 @@ def test_async_custom_handler():
async def test_2(): async def test_2():
response = await litellm.acompletion( response = await litellm.acompletion(
model="gpt-3.5-turbo", model="gpt-3.5-turbo",
messages=messages messages=[{
"role": "user",
"content": "hello from litellm test",
}]
) )
print("\n response", response) print("\n response", response)
assert customHandler2.async_success == False assert customHandler2.async_success == False
@ -142,6 +153,18 @@ def test_async_custom_handler():
assert customHandler2.async_success == True, "async success is not set to True even after success" assert customHandler2.async_success == True, "async success is not set to True even after success"
assert customHandler2.async_completion_kwargs.get("model") == "gpt-3.5-turbo" assert customHandler2.async_completion_kwargs.get("model") == "gpt-3.5-turbo"
async def test_3():
response = await litellm.aembedding(
model="text-embedding-ada-002",
input = ["hello world"],
)
print("\n response", response)
assert customHandler2.async_success_embedding == False
asyncio.run(test_3())
assert customHandler2.async_success_embedding == True, "async_success_embedding is not set to True even after success"
assert customHandler2.async_embedding_kwargs.get("model") == "text-embedding-ada-002"
assert customHandler2.async_embedding_response["usage"]["prompt_tokens"] ==2
print("Passed setting async success") print("Passed setting async success")
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")