fix: add wait_retry to exp_gen v2 (#783)

RolandMinrui · Xu · WinstonLiyt · web-flow · commit b9fb7cf4e307 · 2025-04-14T11:02:45.000+08:00
* add wait retry to v2

* format

* fix a bug

---------

Co-authored-by: Xu &lt;v-xuminrui@microsoft.com&gt;
Co-authored-by: yuanteli &lt;1957922024@qq.com&gt;
diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py
@@ -242,13 +242,7 @@ def identify_scenario_problem(self, scenario_desc: str, competition_desc: str, s
         )
         return json.loads(response)
 
-    def identify_feedback_problem(
-        self,
-        scenario_desc: str,
-        exp_feedback_list_desc: str,
-        sota_exp_desc: str,
-        pipeline: bool,
-    ) -> Dict:
+    def identify_feedback_problem(self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str) -> Dict:
         sys_prompt = T(".prompts_v2:scenario_problem.system").r(
             problem_spec=T(".prompts_v2:specification.problem").r(),
             problem_output_format=T(".prompts_v2:output_format.problem").r(),
@@ -266,6 +260,13 @@ def identify_feedback_problem(
         )
         return json.loads(response)
 
+    def _append_retry(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
+        # Only modify the user_prompt on retries (i > 0)
+        user_prompt = args[0]
+        user_prompt += "\n\nretrying..."
+        return (user_prompt,), kwargs
+
+    @wait_retry(retry_n=5, transform_args_fn=_append_retry)
     def hypothesis_gen(
         self,
         component_desc: str,
@@ -293,7 +294,13 @@ def hypothesis_gen(
             json_mode=True,
             json_target_type=Dict[str, Dict[str, str | Dict[str, str | int]]],
         )
-        return json.loads(response)
+        resp_dict = json.loads(response)
+        for key, value in resp_dict.items():
+            assert "reason" in value, "Reason not provided."
+            assert "component" in value, "Component not provided."
+            assert "hypothesis" in value, "Hypothesis not provided."
+            assert "evaluation" in value, "Evaluation not provided."
+        return resp_dict
 
     def hypothesis_rank(self, hypothesis_dict: dict, problem_dict: dict, pipeline: bool) -> DSHypothesis:
         weights = {
@@ -361,7 +368,6 @@ def task_gen(
             component_desc=component_desc,
             workflow_check=not pipeline and hypothesis.component != "Workflow",
         )
-
         user_prompt = T(".prompts_v2:task_gen.user").r(
             scenario_desc=scenario_desc,
             sota_exp_desc=sota_exp_desc,
@@ -394,7 +400,6 @@ def task_gen(
         # exp.experiment_workspace.inject_code_from_folder(sota_exp.experiment_workspace.workspace_path)
         if sota_exp is not None:
             exp.experiment_workspace.inject_code_from_file_dict(sota_exp.experiment_workspace)
-
         if not pipeline and new_workflow_desc != "No update needed":
             workflow_task = WorkflowTask(
                 name="Workflow",
@@ -442,7 +447,6 @@ def gen(self, trace: DSTrace, pipeline: bool = False) -> DSExperiment:
             scenario_desc=scenario_desc,
             exp_feedback_list_desc=exp_feedback_list_desc,
             sota_exp_desc=sota_exp_desc,
-            pipeline=pipeline,
         )
         all_problems = {**scen_problems, **fb_problems}
 
