diff --git a/rdagent/components/coder/data_science/ensemble/eval.py b/rdagent/components/coder/data_science/ensemble/eval.py index b27c9b5f..b94aad14 100644 --- a/rdagent/components/coder/data_science/ensemble/eval.py +++ b/rdagent/components/coder/data_science/ensemble/eval.py @@ -12,6 +12,7 @@ from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T +from rdagent.utils.agent.workflow import build_cls_from_json_with_retry from rdagent.utils.env import DockerEnv, DSDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -80,6 +81,6 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator): stdout=stdout, workflow_stdout=workflow_stdout, ) - - resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt, json_mode=True) - return EnsembleEvalFeedback(**json.loads(resp)) + return build_cls_from_json_with_retry( + EnsembleEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt + ) diff --git a/rdagent/components/coder/data_science/feature/eval.py b/rdagent/components/coder/data_science/feature/eval.py index 33648c2d..27b3e2ec 100644 --- a/rdagent/components/coder/data_science/feature/eval.py +++ b/rdagent/components/coder/data_science/feature/eval.py @@ -10,6 +10,7 @@ from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T +from rdagent.utils.agent.workflow import build_cls_from_json_with_retry from rdagent.utils.env import DockerEnv, DSDockerConf from rdagent.utils.fmt import shrink_text @@ -74,5 +75,4 @@ class FeatureCoSTEEREvaluator(CoSTEEREvaluator): workflow_stdout=workflow_stdout, ) - resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt, json_mode=True) - return FeatureEvalFeedback(**json.loads(resp)) + return build_cls_from_json_with_retry(FeatureEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt) diff --git a/rdagent/components/coder/data_science/model/eval.py b/rdagent/components/coder/data_science/model/eval.py index 1add5f12..aa10d662 100644 --- a/rdagent/components/coder/data_science/model/eval.py +++ b/rdagent/components/coder/data_science/model/eval.py @@ -17,6 +17,7 @@ from rdagent.core.exception import CoderError from rdagent.core.experiment import FBWorkspace, Task from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T +from rdagent.utils.agent.workflow import build_cls_from_json_with_retry from rdagent.utils.env import DockerEnv, DSDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -91,5 +92,4 @@ class ModelGeneralCaseSpecEvaluator(CoSTEEREvaluator): code=implementation.file_dict[f"{target_task.name}.py"], workflow_stdout=workflow_stdout, ) - resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt, json_mode=True) - return ModelSingleFeedback(**json.loads(resp)) + return build_cls_from_json_with_retry(ModelSingleFeedback, system_prompt=system_prompt, user_prompt=user_prompt) diff --git a/rdagent/components/coder/data_science/raw_data_loader/eval.py b/rdagent/components/coder/data_science/raw_data_loader/eval.py index 8adf829e..64f8362a 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/eval.py +++ b/rdagent/components/coder/data_science/raw_data_loader/eval.py @@ -14,6 +14,7 @@ from rdagent.components.coder.CoSTEER.knowledge_management import ( from rdagent.core.experiment import FBWorkspace, Task from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T +from rdagent.utils.agent.workflow import build_cls_from_json_with_retry from rdagent.utils.env import DockerEnv, DSDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -75,5 +76,6 @@ class DataLoaderCoSTEEREvaluator(CoSTEEREvaluator): workflow_stdout=workflow_stdout, ) - resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt, json_mode=True) - return DataLoaderEvalFeedback(**json.loads(resp)) + return build_cls_from_json_with_retry( + DataLoaderEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt + ) diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py index 7867f6e5..254d364f 100644 --- a/rdagent/components/coder/data_science/workflow/eval.py +++ b/rdagent/components/coder/data_science/workflow/eval.py @@ -9,12 +9,12 @@ from rdagent.components.coder.CoSTEER.evaluators import ( CoSTEEREvaluator, CoSTEERMultiFeedback, CoSTEERSingleFeedback, - CoSTEERSingleFeedbackDeprecated, ) from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T +from rdagent.utils.agent.workflow import build_cls_from_json_with_retry from rdagent.utils.env import DockerEnv, DSDockerConf, MLEBDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -39,7 +39,7 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): gt_implementation: FBWorkspace, queried_knowledge: QueriedKnowledge = None, **kwargs, - ) -> CoSTEERSingleFeedbackDeprecated: + ) -> CoSTEERSingleFeedback: target_task_information = target_task.get_task_information() if ( queried_knowledge is not None @@ -120,5 +120,6 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): stdout=stdout.strip(), code=implementation.file_dict["main.py"], ) - resp = APIBackend().build_messages_and_create_chat_completion(user_prompt, system_prompt, json_mode=True) - return WorkflowSingleFeedback(**json.loads(resp)) + return build_cls_from_json_with_retry( + WorkflowSingleFeedback, system_prompt=system_prompt, user_prompt=user_prompt + ) diff --git a/rdagent/core/exception.py b/rdagent/core/exception.py index f2415633..526b1544 100644 --- a/rdagent/core/exception.py +++ b/rdagent/core/exception.py @@ -1,4 +1,16 @@ -class CoderError(Exception): +class WorkflowError(Exception): + """ + Exception indicating an error that the current loop cannot handle, preventing further progress. + """ + + +class FormatError(WorkflowError): + """ + After multiple attempts, we are unable to obtain the answer in the correct format to proceed. + """ + + +class CoderError(WorkflowError): """ Exceptions raised when Implementing and running code. - start: FactorTask => FactorGenerator diff --git a/rdagent/utils/agent/__init__.py b/rdagent/utils/agent/__init__.py index e69de29b..863916f4 100644 --- a/rdagent/utils/agent/__init__.py +++ b/rdagent/utils/agent/__init__.py @@ -0,0 +1,3 @@ +from .workflow import build_cls_from_json_with_retry + +__all__ = ["build_cls_from_json_with_retry"] diff --git a/rdagent/utils/agent/workflow.py b/rdagent/utils/agent/workflow.py new file mode 100644 index 00000000..9d76731e --- /dev/null +++ b/rdagent/utils/agent/workflow.py @@ -0,0 +1,44 @@ +import json +from typing import Type, TypeVar + +from rdagent.core.exception import FormatError +from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_utils import APIBackend + +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 +) -> T: + """ + Parameters + ---------- + cls : Type[T] + The class type to be instantiated with the response data. + system_prompt : str + The initial prompt provided to the system for context. + user_prompt : str + The prompt given by the user to guide the response generation. + retry_n : int + The number of attempts to retry in case of failure. + **kwargs + Additional keyword arguments passed to the API call. + + Returns + ------- + T + An instance of the specified class type created from the response data. + """ + for i in range(retry_n): + # currently, it only handle exception caused by initial class + resp = APIBackend().build_messages_and_create_chat_completion( + user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True, **kwargs # type: ignore[arg-type] + ) + try: + return cls(**json.loads(resp)) + 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.")