From d9cee55063b88cdc65d8f7f23ce707fb06a874bb Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Wed, 16 Apr 2025 19:01:46 +0800
Subject: [PATCH] fix retry when hypothesis gen (#796)

---
 .../data_science/proposal/exp_gen/prompts_v2.yaml         | 4 +---
 .../scenarios/data_science/proposal/exp_gen/proposal.py   | 8 +-------
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml b/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml
index 50bbf2d1..a33e401b 100644
--- a/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml
+++ b/rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml
@@ -201,9 +201,7 @@ output_format:
     {
       "problem name 1 (Should be exactly same as the problem name provided)": {
         "reason": "Provide a clear, logical progression from problem identification to hypothesis formulation, grounded in evidence (e.g., trace history, domain principles, or competition constraints). Refer to the Hypothesis Guidelines for better understanding. Reason should be short with no more than two sentences.",
-        {% if not pipeline %}"component": "The component name that the hypothesis focus on. Must be one of ('DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow').",
-        {% else %}"component": "The component name that the hypothesis mainly focus on. Must be one of ('DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow').",
-        {% endif %}
+        "component": "The component name that the hypothesis {% if pipeline %}mainly {% endif %}focuses on. Must be one of ('DataLoadSpec', 'FeatureEng', 'Model', 'Ensemble', 'Workflow').",
         "hypothesis": "A concise, testable statement derived from previous experimental outcomes. Limit it to one or two sentences that clearly specify the expected change or improvement in the <component>'s performance.",
         "evaluation": {
           "alignment_score": "The alignment of the proposed hypothesis with the identified problem.",
diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py
index ddf52ad2..8325a138 100644
--- a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py
+++ b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py
@@ -259,13 +259,7 @@ class DSProposalV2ExpGen(ExpGen):
         )
         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)
+    @wait_retry(retry_n=5)
     def hypothesis_gen(
         self,
         component_desc: str,