From fc33bb607e6af36ac7083420233fed5f428605a3 Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Sat, 8 Feb 2025 20:13:38 +0800
Subject: [PATCH] feat: diff mode fix (#569)

* feat: Add last runnable experiment feedback retrieval and improve diff generation

* lint
---
 rdagent/core/proposal.py                             |  2 ++
 rdagent/scenarios/data_science/proposal/exp_gen.py   | 12 +++++++++++-
 rdagent/scenarios/data_science/proposal/prompts.yaml | 11 ++++++++---
 rdagent/utils/repo/diff.py                           | 12 +++---------
 4 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/rdagent/core/proposal.py b/rdagent/core/proposal.py
index d102d9d4..2c341640 100644
--- a/rdagent/core/proposal.py
+++ b/rdagent/core/proposal.py
@@ -64,6 +64,8 @@ class ExperimentFeedback(Feedback):
     ) -> None:
         self.decision = decision
         self.reason = reason
+        # Exception is not None means failing to generate runnable experiments due to exception.
+        # Runable reuslts are not always good.
         self.exception: Exception | None = (
             exception  # if the experiment raises exception, it will be integrated into part of the feedback.
         )
diff --git a/rdagent/scenarios/data_science/proposal/exp_gen.py b/rdagent/scenarios/data_science/proposal/exp_gen.py
index e79329af..5d149c04 100644
--- a/rdagent/scenarios/data_science/proposal/exp_gen.py
+++ b/rdagent/scenarios/data_science/proposal/exp_gen.py
@@ -122,6 +122,15 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
                 return exp
         return None
 
+    def last_runnable_exp_fb(self) -> tuple[DSExperiment, ExperimentFeedback] | None:
+        """
+        Access the last runnable experiment (no exception, usually not all task failed) and feedback
+        """
+        for exp, ef in self.hist[::-1]:
+            if ef.exception is None:
+                return exp, ef
+        return None
+
 
 class DSExpGen(ExpGen):
     """Data Science Task Generator."""
@@ -266,7 +275,8 @@ class DSExpGen(ExpGen):
             # - Extra RAG
             sota_exp = trace.sota_experiment()
             assert sota_exp is not None, "SOTA experiment is not provided."
-            exp_and_feedback = trace.hist[-1]
+            exp_and_feedback = trace.last_runnable_exp_fb()
+            assert exp_and_feedback is not None, "Last runnable experiment is not provided."
             last_exp = exp_and_feedback[0]
 
             # Step 1: Generate component
diff --git a/rdagent/scenarios/data_science/proposal/prompts.yaml b/rdagent/scenarios/data_science/proposal/prompts.yaml
index 5ce57176..7453828e 100644
--- a/rdagent/scenarios/data_science/proposal/prompts.yaml
+++ b/rdagent/scenarios/data_science/proposal/prompts.yaml
@@ -221,8 +221,12 @@ direct_exp_gen:
     The former hypothesis and the corresponding feedbacks are as follows (focus on the last one & the new hypothesis that it provides and reasoning to see if you agree):
     {{recent_trace_desc}}
 
-    # The difference from the best experiments to the last one
-    {{last_exp_diff}}
+    {% if last_exp_diff %}
+    # Here are the differences between the latest version of implementation and the current best version of implementation
+    It is presented in diff format, highlighting changes from the best version to the latest version.
+    {{ last_exp_diff }}
+    {% endif %}
+
     {% endif %}
 
 
@@ -254,7 +258,8 @@ component_gen:
     {{sota_exp_desc}}
 
     {% if last_exp_diff %}
-    # Here is the latest version of implementation different from the sota_exp_desc
+    # Here are the differences between the latest version of implementation and the current best version of implementation
+    It is presented in diff format, highlighting changes from the best version to the latest version.
     {{ last_exp_diff }}
     {% endif %}
 
diff --git a/rdagent/utils/repo/diff.py b/rdagent/utils/repo/diff.py
index d3acf2e0..4f271fcf 100644
--- a/rdagent/utils/repo/diff.py
+++ b/rdagent/utils/repo/diff.py
@@ -29,26 +29,20 @@ def generate_diff(dir1: str, dir2: str) -> List[str]:
 
         if file1.exists() and file2.exists():
             with file1.open() as f1, file2.open() as f2:
-                diff = list(
-                    difflib.unified_diff(f1.readlines(), f2.readlines(), fromfile=str(file1), tofile=str(file2))
-                )
+                diff = list(difflib.unified_diff(f1.readlines(), f2.readlines(), fromfile=str(file), tofile=str(file)))
                 if diff:
                     diff_files.extend(diff)
         else:
             if file1.exists():
                 with file1.open() as f1:
                     diff = list(
-                        difflib.unified_diff(
-                            f1.readlines(), [], fromfile=str(file1), tofile=str(file2) + " (empty file)"
-                        )
+                        difflib.unified_diff(f1.readlines(), [], fromfile=str(file), tofile=str(file) + " (empty file)")
                     )
                     diff_files.extend(diff)
             elif file2.exists():
                 with file2.open() as f2:
                     diff = list(
-                        difflib.unified_diff(
-                            [], f2.readlines(), fromfile=str(file1) + " (empty file)", tofile=str(file2)
-                        )
+                        difflib.unified_diff([], f2.readlines(), fromfile=str(file) + " (empty file)", tofile=str(file))
                     )
                     diff_files.extend(diff)