refine component choosing part (#630)

rdagent/scenarios/data_science/proposal/exp_gen.py

@@ -1,4 +1,7 @@
 import json
+from typing import Literal
+
+import pandas as pd
 
 from rdagent.components.coder.data_science.ensemble.exp import EnsembleTask
 from rdagent.components.coder.data_science.feature.exp import FeatureTask
@@ -110,14 +113,33 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
         return False
 
     def experiment_and_feedback_list_after_init(
-        self, return_success=True
+        self, return_type: Literal["sota", "failed", "all"]
     ) -> list[tuple[DSExperiment, ExperimentFeedback]]:
+        """
+        Retrieve a list of experiments and feedbacks based on the return_type.
+
+        Parameters
+        ----------
+        return_type : str
+            One of "sota", "failed", "all".
+
+        Returns
+        -------
+        list[tuple[DSExperiment, ExperimentFeedback]]
+            List of experiments and feedbacks.
+        """
+
         final_component = self.COMPLETE_ORDER[-1]
         has_final_component = False
         exp_and_feedback_list = []
         for exp, fb in self.hist:
-            if has_final_component and bool(fb) == return_success:
-                exp_and_feedback_list.append((exp, fb))
+            if has_final_component:
+                if return_type == "all":
+                    exp_and_feedback_list.append((exp, fb))
+                elif return_type == "failed" and not fb.decision:
+                    exp_and_feedback_list.append((exp, fb))
+                elif return_type == "sota" and fb.decision:
+                    exp_and_feedback_list.append((exp, fb))
             if exp.hypothesis.component == final_component and fb:
                 has_final_component = True
         return exp_and_feedback_list
@@ -307,10 +329,18 @@ class DSExpGen(ExpGen):
                 )
             )  # we use file_dict for hitting the cache when replicate the experiment in another machine.
 
-            sota_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_success=True)
-            failed_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_success=False)[
+            sota_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_type="sota")
+            failed_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_type="failed")[
                 -self.max_trace_hist :
             ]
+            all_exp_feedback_list = trace.experiment_and_feedback_list_after_init(return_type="all")
+            trace_component_to_feedback_df = pd.DataFrame(columns=["component", "hypothesis", "decision"])
+            for index, (exp, fb) in enumerate(all_exp_feedback_list):
+                trace_component_to_feedback_df.loc[f"trial {index + 1}"] = [
+                    exp.hypothesis.component,
+                    exp.hypothesis.hypothesis,
+                    fb.decision,
+                ]
 
             sota_exp_feedback_list_desc = T("scenarios.data_science.share:describe.trace").r(
                 exp_and_feedback_list=sota_exp_feedback_list,
@@ -332,6 +362,11 @@ class DSExpGen(ExpGen):
             component_user_prompt = T(".prompts:component_gen.user").r(
                 sota_exp_and_feedback_list_desc=sota_exp_feedback_list_desc,
                 failed_exp_and_feedback_list_desc=failed_exp_feedback_list_desc,
+                component_and_feedback_df=(
+                    trace_component_to_feedback_df.to_string()
+                    if len(trace_component_to_feedback_df) > 0
+                    else "No experiment and feedback provided"
+                ),
             )
 
             resp_dict_component: dict = json.loads(
@@ -341,6 +376,7 @@ class DSExpGen(ExpGen):
             )
 
             component = resp_dict_component.get("component", "Component not provided")
+            component_reason = resp_dict_component.get("reason", "Reason not provided")
             sota_exp_model_file_count = len(
                 [
                     k
@@ -396,7 +432,7 @@ class DSExpGen(ExpGen):
                     hypothesis = DSHypothesis(
                         component=component,
                         hypothesis=hypothesis_proposal.get("hypothesis", ""),
-                        reason=hypothesis_proposal.get("reason", ""),
+                        reason=component_reason + "\n" + hypothesis_proposal.get("reason", ""),
                         concise_reason=hypothesis_proposal.get("concise_reason", ""),
                         concise_observation=hypothesis_proposal.get("concise_observation", ""),
                         concise_justification=hypothesis_proposal.get("concise_justification", ""),

rdagent/scenarios/data_science/proposal/prompts.yaml

@@ -277,6 +277,11 @@ component_gen:
     Also, here's the former failed experiments and their feedbacks:
     {{ failed_exp_and_feedback_list_desc }}
 
+    All former trials and their feedbacks are provided in pandas DataFrame format. The user has already made several hypothesis on this scenario and did several evaluation on them:
+    {{ component_and_feedback_df }}
+    
+    Please choose the most proper component to focus on based on the information above. Please balance the exploration and exploitation.
+
 
 exp_and_feedback: |-
   {% for experiment, feedback in trace.hist[-10:] %}
@@ -298,6 +303,7 @@ hypothesis_specification: |-
 output_format:
   component: |-
     {
+      "reason": "The reason why you choose this component. Based on the current status and former trials, why this component is the most promising one to focus on.",
       "component": "The component you suggest to focus on. It must be one of ['DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow']."
     }
   hypothesis: |-