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

(test) catch litellm.ContentPolicyViolationError

Browse source
This commit is contained in:
ishaan-jaff 2024-01-09 17:04:04 +05:30
parent 186fc4614d
commit c0b56b6575
1 changed files with 20 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

26
litellm/tests/test_image_generation.py
View file

@ -19,7 +19,7 @@ import litellm
def test_image_generation_openai(): def test_image_generation_openai():
try: try:
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"
@ -28,6 +28,8 @@ def test_image_generation_openai():
assert len(response.data) > 0 assert len(response.data) > 0
except litellm.RateLimitError as e: except litellm.RateLimitError as e:
pass pass
except litellm.ContentPolicyViolationError:
pass # OpenAI randomly raises these errors - skip when they occur
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
@ -36,22 +38,27 @@ def test_image_generation_openai():
def test_image_generation_azure(): def test_image_generation_azure():
try: try:
response = litellm.image_generation( response = litellm.image_generation(
prompt="A cute baby sea otter", model="azure/", api_version="2023-06-01-preview" prompt="A cute baby sea otter",
model="azure/",
api_version="2023-06-01-preview",
) )
print(f"response: {response}") print(f"response: {response}")
assert len(response.data) > 0 assert len(response.data) > 0
except litellm.RateLimitError as e: except litellm.RateLimitError as e:
pass pass
except litellm.ContentPolicyViolationError:
pass # Azure randomly raises these errors - skip when they occur
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_azure() # test_image_generation_azure()
def test_image_generation_azure_dall_e_3(): def test_image_generation_azure_dall_e_3():
try: try:
litellm.set_verbose = True litellm.set_verbose = True
response = litellm.image_generation( response = litellm.image_generation(
prompt="A cute baby sea otter", prompt="A cute baby sea otter",
@ -64,6 +71,8 @@ def test_image_generation_azure_dall_e_3():
assert len(response.data) > 0 assert len(response.data) > 0
except litellm.RateLimitError as e: except litellm.RateLimitError as e:
pass pass
except litellm.ContentPolicyViolationError:
pass # OpenAI randomly raises these errors - skip when they occur
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
@ -71,7 +80,7 @@ def test_image_generation_azure_dall_e_3():
# test_image_generation_azure_dall_e_3() # test_image_generation_azure_dall_e_3()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_image_generation_openai(): async def test_async_image_generation_openai():
try: try:
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"
) )
@ -79,20 +88,25 @@ async def test_async_image_generation_openai():
assert len(response.data) > 0 assert len(response.data) > 0
except litellm.RateLimitError as e: except litellm.RateLimitError as e:
pass pass
except litellm.ContentPolicyViolationError:
pass # openai randomly raises these errors - skip when they occur
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
# asyncio.run(test_async_image_generation_openai()) # asyncio.run(test_async_image_generation_openai())
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_image_generation_azure(): async def test_async_image_generation_azure():
try: try:
response = await litellm.aimage_generation( response = await litellm.aimage_generation(
prompt="A cute baby sea otter", model="azure/dall-e-3-test" prompt="A cute baby sea otter", model="azure/dall-e-3-test"
) )
print(f"response: {response}") print(f"response: {response}")
except litellm.RateLimitError as e: except litellm.RateLimitError as e:
pass pass
except litellm.ContentPolicyViolationError:
pass # Azure randomly raises these errors - skip when they occur
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")