fix: adapting UI to mock trace (#841)

* fix: return first index if 'SOTA Exp Score (valid)' is empty

* feat: add get_state_data_range helper for loops and slider bounds

* fix: exclude batch embedding tag from log filtering

* feat: add feedback support in sota_experiment and adjust merge flow

* fix: use fb function for merging experiments

* style: remove extra whitespace and reformat code

rdagent/log/ui/ds_summary.py

@@ -359,6 +359,9 @@ def all_summarize_win():
         def apply_func(cdf: pd.DataFrame):
             cp = cdf["Competition"].values[0]
             md = get_metric_direction(cp)
+            # If SOTA Exp Score (valid) column is empty, return the first index
+            if cdf["SOTA Exp Score (valid)"].dropna().empty:
+                return cdf.index[0]
             if md:
                 best_idx = cdf["SOTA Exp Score (valid)"].idxmax()
             else:

rdagent/log/ui/ds_trace.py

@@ -66,7 +66,7 @@ def load_times(log_path: Path):
 def load_data(log_path: Path):
     data = defaultdict(lambda: defaultdict(dict))
     for msg in FileStorage(log_path).iter_msg():
-        if msg.tag and "llm" not in msg.tag and "session" not in msg.tag:
+        if msg.tag and "llm" not in msg.tag and "session" not in msg.tag and "batch embedding" not in msg.tag:
             if msg.tag == "competition":
                 data["competition"] = msg.content
                 continue
@@ -418,7 +418,8 @@ def summarize_data():
             index=range(len(state.data) - 1),
         )
 
-        for loop in range(len(state.data) - 1):
+        min_id, max_id = get_state_data_range(state.data)
+        for loop in range(min_id, max_id + 1):
             loop_data = state.data[loop]
             df.loc[loop, "Component"] = loop_data["direct_exp_gen"]["no_tag"].hypothesis.component
             df.loc[loop, "Hypothesis"] = loop_data["direct_exp_gen"]["no_tag"].hypothesis.hypothesis
@@ -650,13 +651,22 @@ with st.sidebar:
 """
     )
 
+
+def get_state_data_range(state_data):
+    # we have a "competition" key in state_data
+    # like dict_keys(['competition', 10, 11, 12, 13, 14])
+    keys = [k for k in state_data.keys() if k != "competition"]
+    return min(keys), max(keys)
+
+
 # UI - Main
 if state.data["competition"]:
     st.title(state.data["competition"])
     st.markdown(f"[share_link](/ds_trace?log_folder={state.log_folder}&selection={state.log_path})")
     summarize_data()
     if len(state.data) > 2:
-        loop_id = st.slider("Loop", 0, len(state.data) - 2, 0)
+        min_id, max_id = get_state_data_range(state.data)
+        loop_id = st.slider("Loop", min_id, max_id, min_id)
     else:
         loop_id = 0
     if state.show_stdout:

rdagent/scenarios/data_science/proposal/exp_gen/base.py

@@ -216,11 +216,11 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
                 has_final_component = True
         return exp_and_feedback_list
 
-    def sota_experiment(
+    def sota_experiment_fb(
         self,
         search_type: Literal["all", "ancestors"] = "ancestors",
         selection: tuple[int, ...] | None = None,
-    ) -> DSExperiment | None:
+    ) -> tuple[DSExperiment, ExperimentFeedback] | None:
         """
 
         Returns
@@ -234,9 +234,19 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
             for exp, ef in search_list[::-1]:
                 # the sota exp should be accepted decision and all required components are completed.
                 if ef.decision:
-                    return exp
+                    return exp, ef
         return None
 
+    def sota_experiment(
+        self,
+        search_type: Literal["all", "ancestors"] = "ancestors",
+        selection: tuple[int, ...] | None = None,
+    ) -> DSExperiment | None:
+        res = self.sota_experiment_fb(search_type=search_type, selection=selection)
+        if res is not None:
+            res = res[0]
+        return res
+
     def last_successful_exp(
         self,
         search_type: Literal["all", "ancestors"] = "ancestors",

rdagent/scenarios/data_science/proposal/exp_gen/merge.py

@@ -14,33 +14,46 @@ class MergeExpGen(ExpGen):
         trace.set_current_selection((leaves[0],))  # override the current selection.
 
         # assuming merging the first and sencond trace.
-        sota_exp = trace.sota_experiment(selection=(leaves[0],))
-        exp_to_merge = trace.sota_experiment(selection=(leaves[1],))
+        sota_exp_fb = trace.sota_experiment_fb(selection=(leaves[0],))
+        if sota_exp_fb is None:
+            sota_exp_fb = trace.hist[leaves[0]]
+        exp_to_merge_fb = trace.sota_experiment_fb(selection=(leaves[1],))
+        if exp_to_merge_fb is None:
+            exp_to_merge_fb = trace.hist[leaves[1]]
 
         # scenario_desc = trace.scen.get_scenario_all_desc()
         # scenario_desc is not needed in task description. So we have to do it.
 
         sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(
-            exp=sota_exp,
-            heading="Best of previous exploration of the scenario",
+            exp=sota_exp_fb[0],
+            heading="Best previous exploration of the scenario",
+        )
+        sota_exp_fb_desc = T("scenarios.data_science.share:describe.feedback").r(
+            exp_and_feedback=sota_exp_fb,
+            heading="The feedback for best previous exploration",
         )
         exp_to_merge_desc = T("scenarios.data_science.share:describe.exp").r(
-            exp=exp_to_merge,
+            exp=exp_to_merge_fb[0],
             heading="A solution that to be merged into previous best solution",
         )
 
-        exp_and_feedback_list_desc = T("scenarios.data_science.share:describe.trace").r(
-            exp_and_feedback_list=trace.experiment_and_feedback_list_after_init(
-                return_type="sota", selection=(leaves[1],)
-            ),
-            type="success",
-        )
+        success_fb_list = trace.experiment_and_feedback_list_after_init(return_type="sota", selection=(leaves[1],))
+        if len(success_fb_list) > 0:
+            exp_to_merge_fb_desc = T("scenarios.data_science.share:describe.trace").r(
+                exp_and_feedback_list=success_fb_list, type="success", heading="Successful iterations:"
+            )
+        else:
+            exp_to_merge_fb_desc = T("scenarios.data_science.share:describe.feedback").r(
+                exp_and_feedback=exp_to_merge_fb,
+                heading="The feedback for the solution to be merged",
+            )
 
         task = PipelineTask(
             description=T("scenarios.data_science.proposal.exp_gen.merge:task").r(
                 sota_exp_desc=sota_exp_desc,
+                sota_exp_fb_desc=sota_exp_fb_desc,
                 exp_to_merge_desc=exp_to_merge_desc,
-                exp_and_feedback_list_desc=exp_and_feedback_list_desc,
+                exp_to_merge_fb_desc=exp_to_merge_fb_desc,
             )
         )
 
@@ -52,6 +65,6 @@ class MergeExpGen(ExpGen):
             ),
         )
 
-        if sota_exp is not None:
-            exp.experiment_workspace.inject_code_from_file_dict(sota_exp.experiment_workspace)
+        if sota_exp_fb is not None:
+            exp.experiment_workspace.inject_code_from_file_dict(sota_exp_fb[0].experiment_workspace)
         return exp

rdagent/scenarios/data_science/proposal/exp_gen/merge.yaml

@@ -2,19 +2,23 @@ task: |-
   {% include "scenarios.data_science.share:scen.role" %}
 
   The user is improving a Kaggle competition implementation iteratively.
-  Your task is to merge two solutions to create a better version. We expect the merged version to perform better than both given solutions.
+  Your task is to merge two solutions to create a better version (Combine the strengths of both solutions while discarding their weaknesses, to create a new version that is better than either one alone). We expect the merged version to perform better than both given solutions.
 
   You will be given:
   1) Previous Main Solution: this is the main solution you will build on to create an improved version;
+    - Feedback to the main solutions
   2) Solution to be merged: another solution that you will combine with the previous main solution.
     - Solution: the approach or method used in this solution.
-    - Successful iterations: the steps or changes that led to the success of `Solution to be merged`.
+    - Successful iterations (the steps or changes that led to the success of the Solution to be merged) or feedback to the Solution to be merged.
 
   # Previous Main Solution
   {{ sota_exp_desc }}
+  {{ sota_exp_fb_desc }}
 
   # Solution to be merged
   ## Solution Descrioption:
   {{ exp_to_merge_desc }}
-  ## Successful iterations:
-  {{ exp_and_feedback_list_desc }}
+
+  {% if exp_to_merge_fb_desc %}
+  {{ exp_to_merge_fb_desc }}
+  {% endif %}