fix: move next_component_required logic to DSTrace class and accurate implement (#612)

* refactor: Move next_component_required logic to DSTrace class

* lint

rdagent/app/data_science/loop.py

@@ -93,7 +93,7 @@ class DataScienceRDLoop(RDLoop):
 
     def running(self, prev_out: dict[str, Any]):
         exp: DSExperiment = prev_out["coding"]
-        if exp.next_component_required() is None:
+        if self.trace.next_component_required() is None:
             new_exp = self.runner.develop(exp)
             logger.log_object(new_exp)
             return new_exp
@@ -102,7 +102,7 @@ class DataScienceRDLoop(RDLoop):
 
     def feedback(self, prev_out: dict[str, Any]) -> ExperimentFeedback:
         exp: DSExperiment = prev_out["running"]
-        if exp.next_component_required() is None:
+        if self.trace.next_component_required() is None:
             feedback = self.summarizer.generate_feedback(exp, self.trace)
         else:
             feedback = ExperimentFeedback(
@@ -124,15 +124,11 @@ class DataScienceRDLoop(RDLoop):
                 )
             )
             if self.trace.sota_experiment() is None and len(self.trace.hist) >= DS_RD_SETTING.consecutive_errors:
-                trace_exp_next_component_list = [
-                    type(exp.pending_tasks_list[0][0])
-                    for exp, _ in self.trace.hist[-DS_RD_SETTING.consecutive_errors :]
-                ]
-                last_successful_exp = self.trace.last_successful_exp()
-                if (
-                    last_successful_exp not in [exp for exp, _ in self.trace.hist[-DS_RD_SETTING.consecutive_errors :]]
-                    and len(set(trace_exp_next_component_list)) == 1
-                ):
+                # if {in inital/drafting stage} and {tried enough times}
+                for _, fb in self.trace.hist[-DS_RD_SETTING.consecutive_errors :]:
+                    if fb:
+                        break  # any success will stop restarting.
+                else:  # otherwise restart it
                     logger.error("Consecutive errors reached the limit. Dumping trace.")
                     logger.log_object(self.trace, tag="trace before restart")
                     self.trace = DSTrace(scen=self.trace.scen, knowledge_base=self.trace.knowledge_base)

rdagent/scenarios/data_science/experiment/experiment.py

@@ -19,17 +19,3 @@ class DSExperiment(Experiment[Task, FBWorkspace, FBWorkspace]):
         self.experiment_workspace = FBWorkspace()
         self.pending_tasks_list = pending_tasks_list
         self.format_check_result = None
-
-    def next_component_required(self) -> COMPONENT | None:
-        files = list(self.experiment_workspace.file_dict.keys())
-        if "load_data.py" not in files:
-            return "DataLoadSpec"
-        if "feature.py" not in files:
-            return "FeatureEng"
-        if not any(re.match(r"model.*\.py", file) for file in files):
-            return "Model"
-        if "ensemble.py" not in files:
-            return "Ensemble"
-        if "main.py" not in files:
-            return "Workflow"
-        return None

rdagent/scenarios/data_science/proposal/exp_gen.py

@@ -90,6 +90,21 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
         self.hist: list[tuple[DSExperiment, ExperimentFeedback]] = []
         self.knowledge_base = knowledge_base
 
+    COMPLETE_ORDER = ("DataLoadSpec", "FeatureEng", "Model", "Ensemble", "Workflow")
+
+    def next_component_required(self) -> COMPONENT | None:
+        for c in self.COMPLETE_ORDER:
+            if not self.has_compponent(c):
+                return c
+        return None
+
+    def has_compponent(self, component: COMPONENT) -> bool:
+        for exp, fb in self.hist:
+            assert isinstance(exp.hypothesis, DSHypothesis), "Hypothesis should be DSHypothesis (and not None)"
+            if exp.hypothesis.component == component and fb:
+                return True
+        return False
+
     def sota_experiment(self, last_n: int = -1) -> DSExperiment | None:
         """
         Access the last experiment result.
@@ -108,7 +123,7 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
         assert last_n < 0
         for exp, ef in self.hist[::-1]:
             # the sota exp should be accepted decision and all required components are completed.
-            if ef.decision and exp.next_component_required() is None:
+            if ef.decision and self.next_component_required() is None:
                 last_n += 1
                 if last_n == 0:
                     return exp
@@ -237,10 +252,7 @@ class DSExpGen(ExpGen):
         scenario_desc = trace.scen.get_scenario_all_desc()
         last_successful_exp = trace.last_successful_exp()
 
-        if len(trace.hist) == 0 or last_successful_exp is None:
-            next_missing_component = "DataLoadSpec"
-        else:
-            next_missing_component = last_successful_exp.next_component_required()
+        next_missing_component = trace.next_component_required()
 
         init_component_config = {
             "DataLoadSpec": {"task_cls": DataLoaderTask, "spec_file": None, "component_prompt_key": "data_loader"},