From 5c7ceec37fa36d5d46cb7a32d44656caeef35495 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Thu, 6 Feb 2025 19:24:14 +0800 Subject: [PATCH] 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 --- .../data_science/proposal/exp_gen.py | 18 +++++----------- rdagent/utils/agent/workflow.py | 21 ++++++++++++++----- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/rdagent/scenarios/data_science/proposal/exp_gen.py b/rdagent/scenarios/data_science/proposal/exp_gen.py index 2c1e09df..e79329af 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen.py +++ b/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)) diff --git a/rdagent/utils/agent/workflow.py b/rdagent/utils/agent/workflow.py index 9d76731e..33a4362d 100644 --- a/rdagent/utils/agent/workflow.py +++ b/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.")