Repository: microsoft/RD-Agent
File: rdagent/scenarios/data_science/loop.py (Python)

Code before:
from rdagent.components.coder.data_science.pipeline import PipelineCoSTEER
from rdagent.components.coder.data_science.pipeline.exp import PipelineTask
from rdagent.components.coder.data_science.raw_data_loader import DataLoaderCoSTEER
from rdagent.components.coder.data_science.raw_data_loader.exp import DataLoaderTask
from rdagent.components.coder.data_science.share.doc import DocDev
from rdagent.components.coder.data_science.workflow import WorkflowCoSTEER
from rdagent.components.coder.data_science.workflow.exp import WorkflowTask
from rdagent.components.workflow.conf import BasePropSetting
from rdagent.components.workflow.rd_loop import RDLoop
from rdagent.core.conf import RD_AGENT_SETTINGS
from rdagent.core.exception import CoderError, PolicyError, RunnerError
from rdagent.core.proposal import ExperimentFeedback, ExpGen
from rdagent.core.scenario import Scenario
from rdagent.core.utils import import_class
from rdagent.log import rdagent_logger as logger
from rdagent.scenarios.data_science.dev.feedback import DSExperiment2Feedback
from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
from rdagent.scenarios.data_science.proposal.exp_gen import DSTrace
from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledgeBase


class DataScienceRDLoop(RDLoop):
    # NOTE: we move the DataScienceRDLoop here to be easier to be imported
    skip_loop_error = (CoderError, RunnerError)
    withdraw_loop_error = PolicyError

    @staticmethod
    def _get_exp_gen(class_uri: str, scen: Scenario):
        """
        Just for compatibility with the old version of the code.
        """
        # TODO: remove me in the future. I don't have to be this complicated.
        # It is just for compatibility with the old version of the code and configuration.
        from rdagent.scenarios.data_science.proposal.exp_gen.proposal import (
            DSProposalV1ExpGen,
            DSProposalV2ExpGen,
            DSProposalV3ExpGen,
        )

        if class_uri == "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen":
            if DS_RD_SETTING.proposal_version not in ["v1", "v2", "v3"]:
                return import_class(DS_RD_SETTING.proposal_version)(scen=scen)
            if DS_RD_SETTING.proposal_version == "v3":
                return DSProposalV3ExpGen(scen=scen)
            if DS_RD_SETTING.proposal_version == "v1":
                return DSProposalV1ExpGen(scen=scen)
            if DS_RD_SETTING.proposal_version == "v2":
                return DSProposalV2ExpGen(scen=scen)
        return import_class(class_uri)(scen)


Review comment (suggestion):
```suggestion
    withdraw_loop_error = (PolicyError,)
```

Code after:
from rdagent.components.coder.data_science.pipeline import PipelineCoSTEER
from rdagent.components.coder.data_science.pipeline.exp import PipelineTask
from rdagent.components.coder.data_science.raw_data_loader import DataLoaderCoSTEER
from rdagent.components.coder.data_science.raw_data_loader.exp import DataLoaderTask
from rdagent.components.coder.data_science.share.doc import DocDev
from rdagent.components.coder.data_science.workflow import WorkflowCoSTEER
from rdagent.components.coder.data_science.workflow.exp import WorkflowTask
from rdagent.components.workflow.conf import BasePropSetting
from rdagent.components.workflow.rd_loop import RDLoop
from rdagent.core.conf import RD_AGENT_SETTINGS
from rdagent.core.exception import CoderError, PolicyError, RunnerError
from rdagent.core.proposal import ExperimentFeedback, ExpGen
from rdagent.core.scenario import Scenario
from rdagent.core.utils import import_class
from rdagent.log import rdagent_logger as logger
from rdagent.scenarios.data_science.dev.feedback import DSExperiment2Feedback
from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
from rdagent.scenarios.data_science.proposal.exp_gen import DSTrace
from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledgeBase


class DataScienceRDLoop(RDLoop):
    # NOTE: we move the DataScienceRDLoop here to be easier to be imported
    skip_loop_error = (CoderError, RunnerError)
    withdraw_loop_error = (PolicyError,)

    @staticmethod
    def _get_exp_gen(class_uri: str, scen: Scenario):
        """
        Just for compatibility with the old version of the code.
        """
        # TODO: remove me in the future. I don't have to be this complicated.
        # It is just for compatibility with the old version of the code and configuration.
        from rdagent.scenarios.data_science.proposal.exp_gen.proposal import (
            DSProposalV1ExpGen,
            DSProposalV2ExpGen,
            DSProposalV3ExpGen,
        )

        if class_uri == "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen":
            if DS_RD_SETTING.proposal_version not in ["v1", "v2", "v3"]:
                return import_class(DS_RD_SETTING.proposal_version)(scen=scen)
            if DS_RD_SETTING.proposal_version == "v3":
                return DSProposalV3ExpGen(scen=scen)
            if DS_RD_SETTING.proposal_version == "v1":
                return DSProposalV1ExpGen(scen=scen)
            if DS_RD_SETTING.proposal_version == "v2":
                return DSProposalV2ExpGen(scen=scen)
        return import_class(class_uri)(scen)