test: mark flaky tests

litellm/litellm_core_utils/exception_mapping_utils.py

@@ -311,6 +311,9 @@ def exception_type(  # type: ignore  # noqa: PLR0915
                 elif (
                     "invalid_request_error" in error_str
                     and "content_policy_violation" in error_str
+                ) or (
+                    "Invalid prompt" in error_str
+                    and "violating our usage policy" in error_str
                 ):
                     exception_mapping_worked = True
                     raise ContentPolicyViolationError(

tests/llm_translation/base_llm_unit_tests.py

@@ -811,6 +811,7 @@ class BaseLLMChatTest(ABC):
 
         return url
     
+    @pytest.mark.flaky(retries=3, delay=1)
     def test_empty_tools(self):
         """
         Related Issue: https://github.com/BerriAI/litellm/issues/9080
@@ -833,6 +834,8 @@ class BaseLLMChatTest(ABC):
             response = completion(**base_completion_call_args, messages=[{"role": "user", "content": "Hello, how are you?"}], tools=[]) # just make sure call doesn't fail
             print("response: ", response)
             assert response is not None
+        except litellm.ContentPolicyViolationError:
+            pass
         except litellm.InternalServerError:
             pytest.skip("Model is overloaded")
         except litellm.RateLimitError:
@@ -840,6 +843,7 @@ class BaseLLMChatTest(ABC):
         except Exception as e:
             pytest.fail(f"Error occurred: {e}")
 
+    @pytest.mark.flaky(retries=3, delay=1)
     def test_basic_tool_calling(self):
         try:
             from litellm import completion, ModelResponse