from __future__ import annotations import shutil import subprocess import sys import uuid from pathlib import Path import pandas as pd from app.actors.code_reviewer import CodeReviewer from app.actors.mlops_bot import MLOpsBot from app.graders import grader_stage1, grader_stage2, grader_stage3, grader_stage4 from app.models import Action, Observation, Reward, StateResponse, StepResponse from app.reward import RewardEngine from app.tasks import ( stage1_data_repair, stage2_training_monitor, stage3_eval_validation, stage4_deploy_gate, ) class DataEngEnvironment: def __init__(self): self.current_task_id: int = 1 self.step_number: int = 0 self.current_stage: int = 1 self.stage_step_number: int = 0 self.loop_count: int = 0 self.stages_completed: list[int] = [] self.episode_score: float = 0.0 self.actor_feedback: str = "" self.current_script: str = "" self.dataframe: pd.DataFrame | None = None self.task_description: str = "" self.done: bool = False self.last_run_output: str = "" self.last_run_error: str | None = None self._held_out_X = None self._held_out_y = None self._full_df = None self._deploy_conditions = None self._code_reviewer = CodeReviewer() self._mlops_bot = MLOpsBot() self._reward_engine = RewardEngine() def reset(self, task_id: int | None = None) -> Observation: start_stage = task_id if task_id in {1, 2, 3, 4} else 1 self.current_task_id = start_stage self.step_number = 0 self.current_stage = start_stage self.stage_step_number = 0 self.loop_count = 0 self.stages_completed = [] self.episode_score = 0.0 self.actor_feedback = "" self._reward_engine = RewardEngine() self.done = False self.last_run_output = "" self.last_run_error = None self._load_stage(start_stage) return self._get_observation() def _load_stage(self, stage: int) -> None: self.current_stage = stage self.current_task_id = stage self.stage_step_number = 0 self.actor_feedback = "" self.last_run_output = "" self._held_out_X = None self._held_out_y = None self._full_df = None self._deploy_conditions = None if stage == 1: scenario = stage1_data_repair.generate_scenario() elif stage == 2: scenario = stage2_training_monitor.generate_scenario() elif stage == 3: scenario = stage3_eval_validation.generate_scenario() self._held_out_X = scenario.get("held_out_X") self._held_out_y = scenario.get("held_out_y") self._full_df = scenario.get("full_df") elif stage == 4: scenario = stage4_deploy_gate.generate_scenario() self._deploy_conditions = scenario.get("deploy_conditions") else: raise ValueError(f"Unknown stage: {stage}") self.current_script = scenario.get("broken_script", "") self.dataframe = scenario.get("dataframe") self.task_description = scenario.get("task_description", "") self.last_run_error = scenario.get("initial_error_log") def _get_observation( self, partial_schema: dict | None = None, partial_preview: str = "", ) -> Observation: return Observation( task_id=self.current_task_id, step_number=self.step_number, script_content=self.current_script, last_run_output=self.last_run_output, last_run_error=self.last_run_error, data_preview=partial_preview, schema_info=partial_schema or {}, done=self.done, current_stage=self.current_stage, stage_step_number=self.stage_step_number, actor_feedback=self.actor_feedback, loop_count=self.loop_count, stages_completed=list(self.stages_completed), ) def _make_reward(self, score: float, message: str, is_terminal: bool) -> Reward: return Reward( score=float(max(0.0, min(1.0, score))), partial_rewards={ "episode_score": float(max(0.0, min(1.0, self.episode_score))), "current_stage": float(self.current_stage), }, message=message, is_terminal=is_terminal, ) def _run_current_script(self) -> None: runs_dir = Path.cwd() / ".env_runs" runs_dir.mkdir(exist_ok=True) tmpdir = runs_dir / f"stage_{self.current_stage}_{uuid.uuid4().hex}" tmpdir.mkdir() def cleanup() -> None: shutil.rmtree(tmpdir, ignore_errors=True) try: if self.dataframe is not None: self.dataframe.to_csv(tmpdir / "data.csv", index=False) script_path = tmpdir / "script.py" sandbox_script = ( "import pandas as pd\n" "df = pd.read_csv('data.csv')\n" + self.current_script ) script_path.write_text(sandbox_script, encoding="utf-8") result = subprocess.run( [sys.executable, str(script_path)], capture_output=True, text=True, timeout=10, cwd=tmpdir, check=False, ) self.last_run_output = result.stdout self.last_run_error = result.stderr if result.returncode != 0 else None except subprocess.TimeoutExpired: self.last_run_error = "TimeoutExpired: script ran longer than 10s." self.last_run_output = "" finally: cleanup() def _grade_current_stage(self) -> float: if self.current_stage == 1: return grader_stage1.grade( self.current_script, self.last_run_output, self.last_run_error, ) if self.current_stage == 2: return grader_stage2.grade( self.current_script, self.last_run_output, self.last_run_error, ) if self.current_stage == 3: return grader_stage3.grade( self.current_script, self._full_df, self._held_out_X, self._held_out_y, ) if self.current_stage == 4: return grader_stage4.grade( self.current_script, self.dataframe, self._deploy_conditions, ) return 0.0 def _check_stage_transition(self, score: float) -> str: stage = self.current_stage if score >= 0.7: if stage not in self.stages_completed: self.stages_completed.append(stage) self.episode_score = max(0.0, min(1.0, self.episode_score + score * 0.25)) if stage == 4: self.done = True return "Stage 4 complete. Deployment approved." self._load_stage(stage + 1) return f"Stage {stage} complete. Advanced to Stage {self.current_stage}." if stage == 4: self.loop_count += 1 self.episode_score = max(0.0, self.episode_score - 0.3) if self.loop_count >= 2: self.done = True return "Stage 4 failed twice. Episode terminated." self._load_stage(1) return "Stage 4 failed. Looping back to Stage 1." return f"Stage {stage} score below threshold. Continue fixing this stage." def _handle_query_actor(self) -> None: if self.current_stage == 1: self.actor_feedback = "No reviewer assigned to this stage." elif self.current_stage in {2, 3}: self.actor_feedback = self._code_reviewer.review( self.current_script, stage=self.current_stage, ) elif self.current_stage == 4: self.actor_feedback = self._mlops_bot.check_deployment( self.current_script, self.dataframe, self._deploy_conditions, ) def step(self, action: Action) -> StepResponse: if self.done: obs = self._get_observation() return StepResponse( observation=obs, reward=self._make_reward(self.episode_score, "episode already complete", True), ) self.step_number += 1 self.stage_step_number += 1 preview = "" schema = {} score = 0.0 message = "no reward" try: if action.action_type == "inspect_data": df = self.dataframe if df is not None: preview = ( f"Shape: {df.shape}\n\n" f"Columns: {list(df.columns)}\n\n" f"Dtypes:\n{df.dtypes.to_string()}\n\n" f"Null counts:\n{df.isnull().sum().to_string()}\n\n" f"Sample (5 rows):\n{df.head().to_string()}\n\n" f"Numeric stats:\n{df.describe().to_string()}" ) schema = { "columns": list(df.columns), "dtypes": {col: str(df[col].dtype) for col in df.columns}, "null_counts": df.isnull().sum().to_dict(), "shape": list(df.shape), } message = "data inspected" elif action.action_type == "check_schema": if self.dataframe is not None: schema = self.dataframe.dtypes.astype(str).to_dict() message = "schema checked" elif action.action_type == "edit_script": old_text = action.payload.get("old", "") new_text = action.payload.get("new", "") replacement_script = action.payload.get("script") if isinstance(replacement_script, str): self.current_script = replacement_script self.last_run_error = None elif old_text: if old_text not in self.current_script: self.last_run_error = "edit_script failed: the text to replace was not found in the script" else: self.current_script = self.current_script.replace(old_text, new_text) self.last_run_error = None message = "script edited" elif action.action_type == "run_script": self._run_current_script() message = "script ran" if not self.last_run_error else "script failed" elif action.action_type == "query_actor": self._handle_query_actor() message = "actor queried" elif action.action_type == "submit": score = self._grade_current_stage() message = self._check_stage_transition(score) except Exception as exc: self.last_run_error = f"Action Handler Error: {exc}" message = self.last_run_error if self.step_number >= 60 or self.stage_step_number >= 20: self.done = True message = f"{message} Stage or episode step limit reached." obs = self._get_observation(partial_schema=schema, partial_preview=preview) # Use RewardEngine for dense partial rewards on every action shaped_reward = self._reward_engine.compute_reward( task_id=self.current_stage, action_type=action.action_type, action_payload=action.payload or {}, observation={ "last_run_error": self.last_run_error, "data_preview": preview, }, grader_score=score, step_number=self.step_number, current_script=self.current_script, ) # For submit, override score but keep partial breakdown if action.action_type == "submit": final_score = self.episode_score if self.done else score shaped_reward = Reward( score=float(max(0.0, min(1.0, final_score))), partial_rewards={ **shaped_reward.partial_rewards, "episode_score": float(max(0.0, min(1.0, self.episode_score))), "current_stage": float(self.current_stage), }, message=message, is_terminal=self.done, ) return StepResponse( observation=obs, reward=shaped_reward, ) def submit_grader(self, task_id: int | None = None) -> Reward: score = self._grade_current_stage() message = self._check_stage_transition(score) return self._make_reward(self.episode_score if self.done else score, message, self.done) def state(self) -> StateResponse: return StateResponse( task_id=self.current_task_id, step_number=self.step_number, max_steps=60, current_script=self.current_script, episode_done=self.done, task_description=self.task_description, current_stage=self.current_stage, stage_step_number=self.stage_step_number, loop_count=self.loop_count, stages_completed=list(self.stages_completed), episode_score=float(max(0.0, min(1.0, self.episode_score))), actor_feedback=self.actor_feedback, )