From a7c6fca8b452e5c1432b7df12bbc237a62f3390d Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Thu, 29 May 2025 17:23:06 +0800
Subject: [PATCH] chore: withdraw for policy error (#907)

* chore: withdraw for policy error

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>

* Apply suggestions from code review
---
 rdagent/core/exception.py              |  6 ++++++
 rdagent/oai/backend/base.py            |  3 ++-
 rdagent/scenarios/data_science/loop.py |  3 ++-
 rdagent/utils/workflow.py              | 29 ++++++++++++++++++++++++++
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/rdagent/core/exception.py b/rdagent/core/exception.py
index 12cbe504..014a9174 100644
--- a/rdagent/core/exception.py
+++ b/rdagent/core/exception.py
@@ -57,3 +57,9 @@ class KaggleError(Exception):
     """
     Exceptions raised when calling Kaggle API
     """
+
+
+class PolicyError(Exception):
+    """
+    Exceptions raised due to content management policy
+    """
diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py
index 846773ca..b8ee8ac5 100644
--- a/rdagent/oai/backend/base.py
+++ b/rdagent/oai/backend/base.py
@@ -14,6 +14,7 @@ from typing import Any, Optional, cast
 import pytz
 from pydantic import TypeAdapter
 
+from rdagent.core.exception import PolicyError
 from rdagent.core.utils import LLM_CACHE_SEED_GEN, SingletonBaseClass
 from rdagent.log import LogColors
 from rdagent.log import rdagent_logger as logger
@@ -368,7 +369,7 @@ class APIBackend(ABC):
                         violation_count += 1
                         if violation_count >= LLM_SETTINGS.violation_fail_limit:
                             logger.warning("Content policy violation detected.")
-                            raise e
+                            raise PolicyError(e)
 
                     if (
                         openai_imported
diff --git a/rdagent/scenarios/data_science/loop.py b/rdagent/scenarios/data_science/loop.py
index 44929e85..0f46f4c6 100644
--- a/rdagent/scenarios/data_science/loop.py
+++ b/rdagent/scenarios/data_science/loop.py
@@ -21,7 +21,7 @@ from rdagent.components.coder.data_science.workflow.exp import WorkflowTask
 from rdagent.components.workflow.conf import BasePropSetting
 from rdagent.components.workflow.rd_loop import RDLoop
 from rdagent.core.conf import RD_AGENT_SETTINGS
-from rdagent.core.exception import CoderError, RunnerError
+from rdagent.core.exception import CoderError, PolicyError, RunnerError
 from rdagent.core.proposal import ExperimentFeedback, ExpGen
 from rdagent.core.scenario import Scenario
 from rdagent.core.utils import import_class
@@ -36,6 +36,7 @@ from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledg
 class DataScienceRDLoop(RDLoop):
     # NOTE: we move the DataScienceRDLoop here to be easier to be imported
     skip_loop_error = (CoderError, RunnerError)
+    withdraw_loop_error = (PolicyError,)
 
     @staticmethod
     def _get_exp_gen(class_uri: str, scen: Scenario):
diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py
index f68e1c57..1081d9ab 100644
--- a/rdagent/utils/workflow.py
+++ b/rdagent/utils/workflow.py
@@ -88,6 +88,9 @@ class LoopBase:
     loop_trace: dict[int, list[LoopTrace]]
 
     skip_loop_error: tuple[type[BaseException], ...] = ()  # you can define a list of error that will skip current loop
+    withdraw_loop_error: tuple[
+        type[BaseException], ...
+    ] = ()  # you can define a list of error that will withdraw current loop
 
     EXCEPTION_KEY = "_EXCEPTION"
 
@@ -183,6 +186,11 @@ class LoopBase:
                             self.step_idx = len(self.steps) - 1  # directly jump to the last step.
                             self.loop_prev_out[self.EXCEPTION_KEY] = e
                             continue
+                        elif isinstance(e, self.withdraw_loop_error):
+                            logger.warning(f"Withdraw loop {li} due to {e}")
+                            # Back to previous loop
+                            self.step_backward(li - 1)
+                            continue
                         else:
                             raise
                     finally:
@@ -207,6 +215,27 @@ class LoopBase:
 
                 self.dump(self.session_folder / f"{li}" / f"{si}_{name}")  # save a snapshot after the session
 
+    def step_backward(self, li: int) -> None:
+        prev_session_dir = self.session_folder / str(li)
+        prev_path = min(
+            (p for p in prev_session_dir.glob("*_*") if p.is_file()),
+            key=lambda item: int(item.name.split("_", 1)[0]),
+            default=None,
+        )
+        if prev_path:
+            loaded = type(self).load(
+                prev_path,
+                output_path=self.session_folder.parent,
+                do_truncate=False,
+                replace_timer=True,
+            )
+            logger.info(f"Load previous session from {prev_path}")
+            # Overwrite current instance state
+            self.__dict__ = loaded.__dict__
+        else:
+            logger.error(f"No previous dump found at {prev_session_dir}, cannot withdraw loop {li}")
+            raise
+
     def dump(self, path: str | Path) -> None:
         if RD_Agent_TIMER_wrapper.timer.started:
             RD_Agent_TIMER_wrapper.timer.update_remain_time()