fix: fix task return dict with wrong format (#558)

* refactor: Simplify exp_gen and enhance workflow with init_kwargs_update_func

* fix: Correct type hint for init_kwargs_udpate_func parameter

* lint

* lint

* lint

* lint

rdagent/scenarios/data_science/proposal/exp_gen.py

@@ -1,23 +1,12 @@
 import json
-import re
-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
 from rdagent.components.coder.data_science.model.exp import ModelTask
 from rdagent.components.coder.data_science.raw_data_loader.exp import DataLoaderTask
 from rdagent.components.coder.data_science.workflow.exp import WorkflowTask
-from rdagent.core.experiment import Experiment, Workspace
 from rdagent.core.knowledge_base import KnowledgeBase
-from rdagent.core.proposal import (
-    ExperimentFeedback,
-    ExpGen,
-    Hypothesis,
-    HypothesisFeedback,
-    Trace,
-)
+from rdagent.core.proposal import ExperimentFeedback, ExpGen, Hypothesis, Trace
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.data_science.experiment.experiment import COMPONENT, DSExperiment
 from rdagent.scenarios.data_science.scen import DataScienceScen
@@ -225,7 +214,10 @@ class DSExpGen(ExpGen):
         task = task_cls(
             name=component if component != "Model" else resp_dict.pop("model_name"),
             description=resp_dict.get("description", f"{component} description not provided"),
-            **resp_dict.get("extra_params", {}),
+            **{
+                k: resp_dict.get("extra_params", {}).get(k, v)
+                for k, v in COMPONENT_TASK_MAPPING[component].get("extra_params", {}).items()
+            },
         )
 
         exp = DSExperiment(pending_tasks_list=[[task]], hypothesis=DSHypothesis(component))

rdagent/utils/agent/workflow.py

@@ -1,5 +1,5 @@
 import json
-from typing import Type, TypeVar
+from typing import Any, Callable, Type, TypeVar, Union, cast
 
 from rdagent.core.exception import FormatError
 from rdagent.log import rdagent_logger as logger
@@ -9,7 +9,12 @@ T = TypeVar("T")
 
 
 def build_cls_from_json_with_retry(
-    cls: Type[T], system_prompt: str, user_prompt: str, retry_n: int = 5, **kwargs: dict
+    cls: Type[T],
+    system_prompt: str,
+    user_prompt: str,
+    retry_n: int = 5,
+    init_kwargs_update_func: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    **kwargs: dict,
 ) -> T:
     """
     Parameters
@@ -22,6 +27,10 @@ def build_cls_from_json_with_retry(
         The prompt given by the user to guide the response generation.
     retry_n : int
         The number of attempts to retry in case of failure.
+    init_kwargs_update_func : Union[Callable[[dict], dict], None]
+        A function that takes the initial keyword arguments as input and returns the updated keyword arguments.
+        This function can be used to modify the response data before it is used to instantiate the class.
+
     **kwargs
         Additional keyword arguments passed to the API call.
 
@@ -36,9 +45,11 @@ def build_cls_from_json_with_retry(
             user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True, **kwargs  # type: ignore[arg-type]
         )
         try:
-            return cls(**json.loads(resp))
+            resp_dict = json.loads(resp)
+            if init_kwargs_update_func:
+                resp_dict = init_kwargs_update_func(resp_dict)
+            return cls(**resp_dict)
         except Exception as e:
             logger.warning(f"Attempt {i + 1}: The previous attempt didn't work due to: {e}")
             user_prompt = user_prompt + f"\n\nAttempt {i + 1}: The previous attempt didn't work due to: {e}"
-    else:
-        raise FormatError("Unable to produce a JSON response that meets the specified requirements.")
+    raise FormatError("Unable to produce a JSON response that meets the specified requirements.")