""" Model workflow with session control It is from `rdagent/app/qlib_rd_loop/model.py` and try to replace `rdagent/app/qlib_rd_loop/RDAgent.py` """ import asyncio import json from multiprocessing import Queue from pathlib import Path from typing import Any from rdagent.components.workflow.conf import BasePropSetting from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.developer import Developer from rdagent.core.proposal import ( Experiment2Feedback, ExperimentPlan, Hypothesis, Hypothesis2Experiment, HypothesisFeedback, HypothesisGen, Trace, ) from rdagent.core.scenario import Scenario from rdagent.core.utils import import_class from rdagent.log import rdagent_logger as logger from rdagent.utils.qlib import ALPHA20, validate_qlib_features from rdagent.utils.workflow import LoopBase, LoopMeta class RDLoop(LoopBase, metaclass=LoopMeta): def __init__(self, PROP_SETTING: BasePropSetting): scen: Scenario = import_class(PROP_SETTING.scen)() logger.log_object(scen, tag="scenario") logger.log_object(PROP_SETTING.model_dump(), tag="RDLOOP_SETTINGS") logger.log_object(RD_AGENT_SETTINGS.model_dump(), tag="RD_AGENT_SETTINGS") self.hypothesis_gen: HypothesisGen = ( import_class(PROP_SETTING.hypothesis_gen)(scen) if hasattr(PROP_SETTING, "hypothesis_gen") and PROP_SETTING.hypothesis_gen else None ) self.plan: ExperimentPlan = { "features": ALPHA20, "feature_codes": {}, } # for user interaction self.hypothesis2experiment: Hypothesis2Experiment = ( import_class(PROP_SETTING.hypothesis2experiment)() if hasattr(PROP_SETTING, "hypothesis2experiment") and PROP_SETTING.hypothesis2experiment else None ) self.coder: Developer = ( import_class(PROP_SETTING.coder)(scen) if hasattr(PROP_SETTING, "coder") and PROP_SETTING.coder else None ) self.runner: Developer = ( import_class(PROP_SETTING.runner)(scen) if hasattr(PROP_SETTING, "runner") and PROP_SETTING.runner else None ) self.summarizer: Experiment2Feedback = ( import_class(PROP_SETTING.summarizer)(scen) if hasattr(PROP_SETTING, "summarizer") and PROP_SETTING.summarizer else None ) self.trace = Trace(scen=scen) super().__init__() # excluded steps def _set_interactor(self, user_request_q: Queue, user_response_q: Queue): self.user_request_q = user_request_q self.user_response_q = user_response_q def _init_base_features(self, base_features_path: str | None): if base_features_path is not None: try: base_dir = Path(base_features_path) base_factors_file = base_dir / "base_factors.json" feature_codes: dict[str, str] = {} for py_file in sorted(base_dir.glob("*.py")): feature_codes[py_file.name] = py_file.read_text() self.plan["feature_codes"] = feature_codes if not base_factors_file.exists(): logger.info(f"No base_factors.json found under {base_dir}. Keeping default base features.") logger.info(f"{len(feature_codes)} feature code files loaded from {base_dir}.") else: with base_factors_file.open("r") as f: features = json.load(f) if not isinstance(features, dict): raise ValueError( "`base_factors.json` must contain a JSON object of feature_name -> expression." ) if validate_qlib_features(list(features.values())): self.plan["features"] = features logger.info( f"Loaded base features from {base_factors_file}. {len(features)} features loaded and {len(feature_codes)} feature code files loaded." ) else: logger.warning( f"Base feature validation failed for features loaded from {base_factors_file}. Using default features." ) except Exception as e: logger.warning(f"Failed to load base features from {base_features_path}: {e}. Using default features.") else: logger.info("No base features path provided. Using default features.") def _interact_init_params(self) -> None: if not (hasattr(self, "user_request_q") and hasattr(self, "user_response_q")): return logger.info("Waiting for user interaction on initial parameters...") try: self.user_request_q.put( { "user_instruction": None, } ) res_dict = self.user_response_q.get() logger.info("Received user instruction response.") self.plan.update(res_dict) if "feature_codes" not in self.plan: self.plan[ "user_instruction" ] += f"\n\n{str(list(self.plan['feature_codes'].keys()))} has been configured as the base factor; do not generate duplicate factors." fea_valid_msg = "" while True: logger.info("Requesting base feature configuration from user.") self.user_request_q.put( { "features": self.plan["features"], "feature_validation_msg": fea_valid_msg, } ) self.plan["features"] = self.user_response_q.get() logger.info("Received base feature configuration response.") if validate_qlib_features(list(self.plan["features"].values())): logger.info(f"Base feature validation passed. {len(self.plan['features'])} features selected.") break else: logger.info("Base feature validation failed. Asking user to revise.") fea_valid_msg = "Some features are invalid, please revise." except (EOFError, OSError): logger.info("User interaction failed, using default initial parameters.") return logger.info("Received user interaction on initial parameters.") def _interact_hypo(self, hypo: Hypothesis) -> Hypothesis: if not (hasattr(self, "user_request_q") and hasattr(self, "user_response_q")): return hypo logger.info("Waiting for user interaction on hypothesis...") try: self.user_request_q.put(hypo.__dict__) res_dict = self.user_response_q.get() modified_hypo = type(hypo)(**res_dict) except (EOFError, OSError, TypeError): logger.info("User interaction failed, using original hypothesis.") return hypo logger.info("Received user interaction on hypothesis.") return modified_hypo def _interact_feedback(self, feedback: HypothesisFeedback) -> HypothesisFeedback: if not (hasattr(self, "user_request_q") and hasattr(self, "user_response_q")): return feedback logger.info("Waiting for user interaction on feedback...") try: self.user_request_q.put(feedback.__dict__) res_dict = self.user_response_q.get() modified_feedback = HypothesisFeedback(**res_dict) except (EOFError, OSError, TypeError): logger.info("User interaction failed, using original feedback.") return feedback logger.info("Received user interaction on feedback.") return modified_feedback def _propose(self): hypothesis = self.hypothesis_gen.gen(self.trace, self.plan) # user can change the hypothesis here hypothesis = self._interact_hypo(hypothesis) logger.log_object(hypothesis, tag="hypothesis generation") return hypothesis def _exp_gen(self, hypothesis: Hypothesis): exp = self.hypothesis2experiment.convert(hypothesis, self.trace) logger.log_object(exp.sub_tasks, tag="experiment generation") return exp # included steps async def direct_exp_gen(self, prev_out: dict[str, Any]): while True: if self.get_unfinished_loop_cnt(self.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel(): hypo = self._propose() exp = self._exp_gen(hypo) exp.base_features = self.plan["features"] exp.base_feature_codes = self.plan["feature_codes"] if exp.based_experiments: exp.based_experiments[-1].base_features = self.plan["features"] exp.based_experiments[-1].base_feature_codes = self.plan["feature_codes"] return {"propose": hypo, "exp_gen": exp} await asyncio.sleep(1) def coding(self, prev_out: dict[str, Any]): exp = self.coder.develop(prev_out["direct_exp_gen"]["exp_gen"]) logger.log_object(exp.sub_workspace_list, tag="coder result") return exp def running(self, prev_out: dict[str, Any]): exp = self.runner.develop(prev_out["coding"]) logger.log_object(exp, tag="runner result") return exp def feedback(self, prev_out: dict[str, Any]): # TODO: the logic branch of exception should be moved to summarizer e = prev_out.get(self.EXCEPTION_KEY, None) if e is not None: feedback = HypothesisFeedback( reason=str(e), decision=False, code_change_summary="", acceptable=False, ) else: feedback = self.summarizer.generate_feedback(prev_out["running"], self.trace) feedback = self._interact_feedback(feedback) logger.log_object(feedback, tag="feedback") return feedback def record(self, prev_out: dict[str, Any]): feedback = prev_out["feedback"] exp = prev_out.get("running") or prev_out.get("coding") or prev_out.get("direct_exp_gen", {}).get("exp_gen") self.trace.sync_dag_parent_and_hist((exp, feedback), prev_out[self.LOOP_IDX_KEY])