phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(openai.py): fix linting issue

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-22 18:20:15 -08:00
parent e917d0eee6
commit 3e8c8ef507
8 changed files with 166 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

63
litellm/tests/test_custom_callback_input.py
View file

@ -74,6 +74,7 @@ class CompletionCustomHandler(
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
try:
print(f"kwargs: {kwargs}")
self.states.append("post_api_call")
## START TIME
assert isinstance(start_time, datetime)
@ -149,7 +150,14 @@ class CompletionCustomHandler(
## END TIME
assert isinstance(end_time, datetime)
## RESPONSE OBJECT
assert isinstance(response_obj, litellm.ModelResponse)
assert isinstance(
response_obj,
Union[
litellm.ModelResponse,
litellm.EmbeddingResponse,
litellm.ImageResponse,
],
)
## KWARGS
assert isinstance(kwargs["model"], str)
assert isinstance(kwargs["messages"], list) and isinstance(
@ -177,6 +185,7 @@ class CompletionCustomHandler(
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
print(f"kwargs: {kwargs}")
self.states.append("sync_failure")
## START TIME
assert isinstance(start_time, datetime)
@ -766,6 +775,52 @@ async def test_async_embedding_azure_caching():
assert len(customHandler_caching.states) == 4 # pre, post, success, success
# asyncio.run(
# test_async_embedding_azure_caching()
# )
# Image Generation
# ## Test OpenAI + Sync
# def test_image_generation_openai():
# try:
# customHandler_success = CompletionCustomHandler()
# customHandler_failure = CompletionCustomHandler()
# litellm.callbacks = [customHandler_success]
# litellm.set_verbose = True
# response = litellm.image_generation(
# prompt="A cute baby sea otter", model="dall-e-3"
# )
# print(f"response: {response}")
# assert len(response.data) > 0
# print(f"customHandler_success.errors: {customHandler_success.errors}")
# print(f"customHandler_success.states: {customHandler_success.states}")
# assert len(customHandler_success.errors) == 0
# assert len(customHandler_success.states) == 3 # pre, post, success
# # test failure callback
# litellm.callbacks = [customHandler_failure]
# try:
# response = litellm.image_generation(
# prompt="A cute baby sea otter", model="dall-e-4"
# )
# except:
# pass
# print(f"customHandler_failure.errors: {customHandler_failure.errors}")
# print(f"customHandler_failure.states: {customHandler_failure.states}")
# assert len(customHandler_failure.errors) == 0
# assert len(customHandler_failure.states) == 3 # pre, post, failure
# except litellm.RateLimitError as e:
# pass
# except litellm.ContentPolicyViolationError:
# pass # OpenAI randomly raises these errors - skip when they occur
# except Exception as e:
# pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_openai()
## Test OpenAI + Async
## Test Azure + Sync
## Test Azure + Async