From 5b47c0ab5c1abf5ead154362d0e43dc5874cd87f Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Thu, 22 May 2025 16:54:34 +0800
Subject: [PATCH] feat: raise policy violation (#894)

---
 rdagent/oai/backend/base.py | 17 +++++++++++++++++
 rdagent/oai/llm_conf.py     |  1 +
 2 files changed, 18 insertions(+)

diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py
index 357f0f86..e9d37c71 100644
--- a/rdagent/oai/backend/base.py
+++ b/rdagent/oai/backend/base.py
@@ -22,6 +22,7 @@ from rdagent.oai.llm_conf import LLM_SETTINGS
 from rdagent.utils import md5_hash
 
 try:
+    import litellm
     import openai
 
     openai_imported = True
@@ -332,6 +333,7 @@ class APIBackend(ABC):
         assert not (chat_completion and embedding), "chat_completion and embedding cannot be True at the same time"
         max_retry = LLM_SETTINGS.max_retry if LLM_SETTINGS.max_retry is not None else max_retry
         timeout_count = 0
+        violation_count = 0
         for i in range(max_retry):
             API_start_time = datetime.now()
             try:
@@ -352,6 +354,21 @@ class APIBackend(ABC):
                 else:
                     RD_Agent_TIMER_wrapper.api_fail_count += 1
                     RD_Agent_TIMER_wrapper.latest_api_fail_time = datetime.now(pytz.timezone("Asia/Shanghai"))
+
+                    if (
+                        openai_imported
+                        and isinstance(e, litellm.BadRequestError)
+                        and (
+                            isinstance(e.__cause__, litellm.ContentPolicyViolationError)
+                            or "The response was filtered due to the prompt triggering Azure OpenAI's content management policy"
+                            in str(e)
+                        )
+                    ):
+                        violation_count += 1
+                        if violation_count >= LLM_SETTINGS.violation_fail_limit:
+                            logger.warning("Content policy violation detected.")
+                            raise e
+
                     if (
                         openai_imported
                         and isinstance(e, openai.APITimeoutError)
diff --git a/rdagent/oai/llm_conf.py b/rdagent/oai/llm_conf.py
index 8e36249a..1446bf54 100644
--- a/rdagent/oai/llm_conf.py
+++ b/rdagent/oai/llm_conf.py
@@ -38,6 +38,7 @@ class LLMSettings(ExtendedBaseSettings):
     prompt_cache_path: str = str(Path.cwd() / "prompt_cache.db")
     max_past_message_include: int = 10
     timeout_fail_limit: int = 10
+    violation_fail_limit: int = 1
 
     # Behavior of returning answers to the same question when caching is enabled
     use_auto_chat_cache_seed_gen: bool = False