fix: replace API call with build_cls_from_json_with_retry function (#548

)

* refactor: Replace API call with build_cls_from_json_with_retry function

* fix lint error

* fix lint errors

* lint

* trigger

rdagent/components/coder/data_science/ensemble/eval.py

@@ -12,6 +12,7 @@
 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 @@ def evaluate(
             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
+        )

rdagent/components/coder/data_science/feature/eval.py

@@ -10,6 +10,7 @@
 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 @@ def evaluate(
             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)

rdagent/components/coder/data_science/model/eval.py

@@ -17,6 +17,7 @@
 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 @@ def evaluate(
             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)

rdagent/components/coder/data_science/raw_data_loader/eval.py

@@ -14,6 +14,7 @@
 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 @@ def evaluate(
             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
+        )

rdagent/components/coder/data_science/workflow/eval.py

@@ -9,12 +9,12 @@
     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 @@ def evaluate(
         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 @@ def evaluate(
             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
+        )

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

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"]

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.")