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

fix(openai.py): fix image gen logging

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-26 21:04:06 -08:00
parent dad578f96a
commit 618d216c11
1 changed files with 15 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

28
litellm/tests/test_custom_callback_input.py
View file

@ -824,26 +824,28 @@ def test_image_generation_openai():
try: try:
customHandler_success = CompletionCustomHandler() customHandler_success = CompletionCustomHandler()
customHandler_failure = CompletionCustomHandler() customHandler_failure = CompletionCustomHandler()
litellm.callbacks = [customHandler_success] # litellm.callbacks = [customHandler_success]
litellm.set_verbose = True # litellm.set_verbose = True
response = litellm.image_generation( # response = litellm.image_generation(
prompt="A cute baby sea otter", model="dall-e-3" # prompt="A cute baby sea otter", model="dall-e-3"
) # )
print(f"response: {response}") # print(f"response: {response}")
assert len(response.data) > 0 # assert len(response.data) > 0
print(f"customHandler_success.errors: {customHandler_success.errors}") # print(f"customHandler_success.errors: {customHandler_success.errors}")
print(f"customHandler_success.states: {customHandler_success.states}") # print(f"customHandler_success.states: {customHandler_success.states}")
assert len(customHandler_success.errors) == 0 # assert len(customHandler_success.errors) == 0
assert len(customHandler_success.states) == 3 # pre, post, success # assert len(customHandler_success.states) == 3 # pre, post, success
# test failure callback # test failure callback
litellm.callbacks = [customHandler_failure] litellm.callbacks = [customHandler_failure]
try: try:
response = litellm.image_generation( response = litellm.image_generation(
prompt="A cute baby sea otter", model="dall-e-4" prompt="A cute baby sea otter",
model="dall-e-2",
api_key="my-bad-api-key",
) )
except: except:
pass pass
@ -859,7 +861,7 @@ def test_image_generation_openai():
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_openai() test_image_generation_openai()
## Test OpenAI + Async ## Test OpenAI + Async
## Test Azure + Sync ## Test Azure + Sync