(test) catch litellm.ContentPolicyViolationError

This commit is contained in:
ishaan-jaff 2024-01-09 17:04:04 +05:30
parent 186fc4614d
commit c0b56b6575

View file

@ -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)}")
@ -38,15 +40,20 @@ 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()
@ -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)}")
@ -79,9 +88,12 @@ 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())
@ -94,5 +106,7 @@ async def test_async_image_generation_azure():
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)}")