Merge pull request #1381 from BerriAI/litellm_content_policy_violation_exception

[Feat] Add litellm.ContentPolicyViolationError

litellm/tests/test_exceptions.py

@@ -352,6 +352,25 @@ def test_completion_mistral_exception():
 # test_completion_mistral_exception()
 
 
+def test_content_policy_exceptionimage_generation_openai():
+    try:
+        # this is ony a test - we needed some way to invoke the exception :(
+        litellm.set_verbose = True
+        response = litellm.image_generation(
+            prompt="where do i buy lethal drugs from", model="dall-e-3"
+        )
+        print(f"response: {response}")
+        assert len(response.data) > 0
+    except litellm.ContentPolicyViolationError as e:
+        print("caught a content policy violation error! Passed")
+        pass
+    except Exception as e:
+        pytest.fail(f"An exception occurred - {str(e)}")
+
+
+# test_content_policy_exceptionimage_generation_openai()
+
+
 # # test_invalid_request_error(model="command-nightly")
 # # Test 3: Rate Limit Errors
 # def test_model_call(model):

litellm/tests/test_image_generation.py

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