fix(round2): fix Stage 3 string matching in baseline script

app/environment.py

@@ -1,68 +1,103 @@
-import sys
-import subprocess
-import tempfile
+from __future__ import annotations
+
 import shutil
+import subprocess
+import sys
+import uuid
 from pathlib import Path
 
 import pandas as pd
 
-from app.models import Action, Observation, Reward, StepResponse, StateResponse
-from app.tasks import task1_easy, task2_medium, task3_hard
-from app.graders import grader1, grader2, grader3
-from app.reward import RewardEngine
+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.tasks import (
+    stage1_data_repair,
+    stage2_training_monitor,
+    stage3_eval_validation,
+    stage4_deploy_gate,
+)
+
 
 class DataEngEnvironment:
     def __init__(self):
-        self.current_task_id: int = 0
+        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
-        
-        # Privacy variables for Task 3 held-out sets
+
         self._held_out_X = None
         self._held_out_y = None
         self._full_df = None
-        
-        self._reward_engine = RewardEngine()
+        self._deploy_conditions = None
+
+        self._code_reviewer = CodeReviewer()
+        self._mlops_bot = MLOpsBot()
 
-    def reset(self, task_id: int) -> Observation:
-        self.current_task_id = task_id
+    def reset(self, task_id: int | None = None) -> Observation:
+        self.current_task_id = 1
         self.step_number = 0
+        self.current_stage = 1
+        self.stage_step_number = 0
+        self.loop_count = 0
+        self.stages_completed = []
+        self.episode_score = 0.0
+        self.actor_feedback = ""
         self.done = False
         self.last_run_output = ""
         self.last_run_error = None
-        
+        self._load_stage(1)
+        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._reward_engine = RewardEngine()
-        
-        if task_id == 1:
-            scenario = task1_easy.generate_scenario()
-        elif task_id == 2:
-            scenario = task2_medium.generate_scenario()
-        elif task_id == 3:
-            scenario = task3_hard.generate_scenario()
-            # Strict assignment isolating testing datasets
+        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 task_id: {task_id}")
-            
+            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")
 
-        return self._get_observation()
-        
-    def _get_observation(self, partial_schema: dict | None = None, partial_preview: str = "") -> Observation:
+    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,
@@ -71,152 +106,234 @@ class DataEngEnvironment:
             last_run_error=self.last_run_error,
             data_preview=partial_preview,
             schema_info=partial_schema or {},
-            done=self.done
+            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:
-                    # Build data_preview
-                    data_preview = f"Shape: {df.shape}\n\n"
-                    data_preview += f"Columns: {list(df.columns)}\n\n"
-                    data_preview += f"Dtypes:\n{df.dtypes.to_string()}\n\n"
-                    data_preview += f"Null counts:\n{df.isnull().sum().to_string()}\n\n"
-                    data_preview += f"Sample (5 rows):\n{df.head().to_string()}\n\n"
-                    data_preview += f"Numeric stats:\n{df.describe().to_string()}"
-                    preview = data_preview
-                    
-                    # Build schema_info
-                    schema_info = {
+                    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)
+                        "shape": list(df.shape),
                     }
-                    schema = schema_info
-            
+                message = "data inspected"
+
             elif action.action_type == "check_schema":
-                df = self.dataframe
-                if df is not None:
-                    schema = df.dtypes.astype(str).to_dict()
-                    
+                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", "")
-                if old_text:
+                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":
-                # Creating sandbox subprocess per contract requirements
-                with tempfile.TemporaryDirectory() as tmpdir:
-                    tmpdir_path = Path(tmpdir)
-                    try:
-                        if self.dataframe is not None:
-                            self.dataframe.to_csv(tmpdir_path / "data.csv", index=False)
-                            
-                        script_path = tmpdir_path / "script.py"
-                        sandbox_script = "import pandas as pd\ndf = 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
-                        )
-                        self.last_run_output = result.stdout
-                        self.last_run_error = result.stderr if result.returncode != 0 else None
-                    except subprocess.TimeoutExpired as e:
-                        self.last_run_error = f"TimeoutExpired: script ran longer than 10s."
-                        self.last_run_output = ""
-                    
+                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":
-                if self.current_task_id == 1:
-                    score = grader1.grade(self.current_script, self.last_run_output, self.last_run_error)
-                elif self.current_task_id == 2:
-                    score = grader2.grade(self.current_script, self.last_run_output, self.last_run_error)
-                elif self.current_task_id == 3:
-                    score = grader3.grade(self.current_script, self._full_df)
-                
-                self.done = True
-                
-        except Exception as e:
-            self.last_run_error = f"Action Handler Error: {str(e)}"
-            
-        # Maximum step bound limit restriction per project contract requirements
-        if self.step_number >= 20:
+                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 >= 15:
             self.done = True
-            
+            message = f"{message} Stage or episode step limit reached."
+
         obs = self._get_observation(partial_schema=schema, partial_preview=preview)
-        
-        # Obtain reward from Reward engine
-        try:
-            reward = self._reward_engine.compute_reward(
-                task_id=self.current_task_id,
-                action_type=action.action_type,
-                action_payload=action.payload,
-                observation=obs.model_dump(),
-                grader_score=score,
-                step_number=self.step_number
-            )
-        except Exception as e:
-            # Fallback if RewardEngine throws error based on parameter changes
-            reward = Reward(
-                score=score, 
-                partial_rewards={}, 
-                message=f"Reward calculation fallback. {e}", 
-                is_terminal=self.done
-            )
-            
-        return StepResponse(observation=obs, reward=reward)
+        reward_score = self.episode_score if action.action_type == "submit" else 0.0
+        if action.action_type == "submit" and not self.done:
+            reward_score = score
+        return StepResponse(
+            observation=obs,
+            reward=self._make_reward(reward_score, message, self.done),
+        )
 
-    def submit_grader(self, task_id: int) -> Reward:
-        score = 0.0
-        if task_id == 1:
-            score = grader1.grade(self.current_script, self.last_run_output, self.last_run_error)
-        elif task_id == 2:
-            score = grader2.grade(self.current_script, self.last_run_output, self.last_run_error)
-        elif task_id == 3:
-            score = grader3.grade(self.current_script, self._full_df)
-            
-        obs = self._get_observation()
-        
-        try:
-            reward = self._reward_engine.compute_reward(
-                task_id=task_id,
-                action_type="submit",
-                action_payload={},
-                observation=obs.model_dump(),
-                grader_score=score,
-                step_number=self.step_number
-            )
-        except Exception as e:
-            reward = Reward(
-                score=score, 
-                partial_rewards={}, 
-                message=f"Reward calculation fallback. {e}", 
-                is_terminal=True
-            )
-        return 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=20,
+            max_steps=60,
             current_script=self.current_script,
             episode_done=self.done,
-            task_description=self.task_description
+            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,
         )

app/main.py

@@ -5,7 +5,12 @@ from pydantic import BaseModel
 
 from app.environment import DataEngEnvironment
 from app.models import Action, Observation, Reward, StepResponse, StateResponse, TaskInfo
-from app.tasks import task1_easy, task2_medium, task3_hard
+from app.tasks import (
+    stage1_data_repair,
+    stage2_training_monitor,
+    stage3_eval_validation,
+    stage4_deploy_gate,
+)
 
 app = FastAPI(title="DataEngEnv API", version="1.0.0")
 
@@ -35,16 +40,25 @@ class BaselineResponse(BaseModel):
     success: bool
     output: str
 
+class PipelineStatusResponse(BaseModel):
+    current_stage: int
+    loop_count: int
+    stages_completed: list[int]
+    total_steps: int
+    episode_score: float
+    actor_feedback: str
+
 @app.on_event("startup")
 async def startup_event():
     print("DataEngEnv ready")
-    # Call generate_scenario() on all 3 tasks to validate they load properly at startup
+    # Call generate_scenario() on all 4 stages to validate they load properly at startup
     try:
-        task1_easy.generate_scenario()
-        task2_medium.generate_scenario()
-        task3_hard.generate_scenario()
+        stage1_data_repair.generate_scenario()
+        stage2_training_monitor.generate_scenario()
+        stage3_eval_validation.generate_scenario()
+        stage4_deploy_gate.generate_scenario()
     except Exception as e:
-        print(f"Warning: A task failed to load during startup validation: {e}")
+        print(f"Warning: A stage failed to load during startup validation: {e}")
 
 
 @app.post("/reset", response_model=Observation | ErrorResponse)
@@ -72,26 +86,39 @@ async def api_state():
 @app.get("/tasks", response_model=list[TaskInfo] | ErrorResponse)
 async def api_tasks():
     try:
+        scenarios = {
+            1: stage1_data_repair.generate_scenario(),
+            2: stage2_training_monitor.generate_scenario(),
+            3: stage3_eval_validation.generate_scenario(),
+            4: stage4_deploy_gate.generate_scenario(),
+        }
         return [
             TaskInfo(
                 task_id=1,
-                name="Easy Data Cleaning",
+                name="Stage 1: Data Repair",
                 difficulty="easy",
-                description="Fix the broken script by identifying and replacing the incorrect column name ('age_years' to 'age').",
+                description=scenarios[1]["task_description"],
                 action_schema={}
             ),
             TaskInfo(
                 task_id=2,
-                name="Medium Data Imputation",
+                name="Stage 2: Training Monitor",
                 difficulty="medium",
-                description="Detect missing values and outliers in the numeric dataset. Apply dropna/fillna for NaNs and clip/quantile for outliers.",
+                description=scenarios[2]["task_description"],
                 action_schema={}
             ),
             TaskInfo(
                 task_id=3,
-                name="Hard Data Leakage",
+                name="Stage 3: Eval Validation",
+                difficulty="hard",
+                description=scenarios[3]["task_description"],
+                action_schema={}
+            ),
+            TaskInfo(
+                task_id=4,
+                name="Stage 4: Deployment Gate",
                 difficulty="hard",
-                description="Fix the critical logical data leakage issue by moving the scaler.fit() call AFTER train_test_split() to avoid cheating on the validation set.",
+                description=scenarios[4]["task_description"],
                 action_schema={}
             )
         ]
@@ -103,11 +130,25 @@ async def api_grader(req: TaskIdRequest | None = None):
     try:
         tid = req.task_id if req is not None and req.task_id is not None else env.current_task_id
         if tid != env.current_task_id:
-            return ErrorResponse(error="task_id mismatch — call /reset with this task_id first")
+            return ErrorResponse(error="task_id mismatch - call /reset first; the pipeline controls stage advancement")
         return env.submit_grader(tid)
     except Exception as e:
         return ErrorResponse(error=str(e))
 
+@app.get("/pipeline_status", response_model=PipelineStatusResponse | ErrorResponse)
+async def api_pipeline_status():
+    try:
+        return PipelineStatusResponse(
+            current_stage=env.current_stage,
+            loop_count=env.loop_count,
+            stages_completed=list(env.stages_completed),
+            total_steps=env.step_number,
+            episode_score=float(env.episode_score),
+            actor_feedback=env.actor_feedback,
+        )
+    except Exception as e:
+        return ErrorResponse(error=str(e))
+
 @app.get("/baseline", response_model=BaselineResponse | ErrorResponse)
 async def api_baseline():
     try:

app/models.py

@@ -1,14 +1,13 @@
 from pydantic import BaseModel, Field, field_validator
-import typing
 
 class Action(BaseModel):
-    action_type: str = Field(description="One of: inspect_data | run_script | edit_script | check_schema | submit")
+    action_type: str = Field(description="One of: inspect_data | run_script | edit_script | check_schema | query_actor | submit")
     payload: dict = Field(description="Action payload details like patch content etc.")
 
     @field_validator("action_type")
     @classmethod
     def validate_action_type(cls, v: str) -> str:
-        allowed = {"inspect_data", "run_script", "edit_script", "check_schema", "submit"}
+        allowed = {"inspect_data", "run_script", "edit_script", "check_schema", "query_actor", "submit"}
         if v not in allowed:
             raise ValueError(f"action_type must be one of {allowed}")
         return v
@@ -22,6 +21,11 @@ class Observation(BaseModel):
     data_preview: str = Field(description="Preview of the current Pandas dataframe")
     schema_info: dict = Field(description="Column types and info")
     done: bool = Field(description="Whether the environment episode is complete")
+    current_stage: int = Field(default=1, description="Current pipeline stage")
+    stage_step_number: int = Field(default=0, description="Step count within current stage")
+    actor_feedback: str = Field(default="", description="Most recent actor bot feedback")
+    loop_count: int = Field(default=0, description="Number of deployment failure loops")
+    stages_completed: list[int] = Field(default_factory=list, description="Completed pipeline stages")
 
 class Reward(BaseModel):
     score: float = Field(description="Reward score from 0.0 to 1.0")
@@ -32,8 +36,8 @@ class Reward(BaseModel):
     @field_validator("score")
     @classmethod
     def validate_score(cls, v: float) -> float:
-        if not (0.0 < v < 1.0):
-            raise ValueError("Score must be strictly between 0.0 and 1.0 (exclusive)")
+        if not (0.0 <= v <= 1.0):
+            raise ValueError("Score must be between 0.0 and 1.0 inclusive")
         return v
 
 class StepResponse(BaseModel):
@@ -47,6 +51,12 @@ class StateResponse(BaseModel):
     current_script: str = Field(description="Content of the active python script")
     episode_done: bool = Field(description="Whether the episode is done")
     task_description: str = Field(description="Task objective description")
+    current_stage: int = Field(default=1, description="Current pipeline stage")
+    stage_step_number: int = Field(default=0, description="Step count within current stage")
+    loop_count: int = Field(default=0, description="Number of deployment failure loops")
+    stages_completed: list[int] = Field(default_factory=list, description="Completed pipeline stages")
+    episode_score: float = Field(default=0.0, description="Cumulative 4-stage episode score")
+    actor_feedback: str = Field(default="", description="Most recent actor bot feedback")
 
 class TaskInfo(BaseModel):
     task_id: int = Field(description="ID of the task")

baseline/results.json

@@ -1,5 +1,11 @@
-[
-  {"task_id": 1, "name": "column_rename_bug", "difficulty": "easy", "score": 1.0, "steps": 3},
-  {"task_id": 2, "name": "dirty_data", "difficulty": "medium", "score": 1.0, "steps": 15},
-  {"task_id": 3, "name": "data_leakage", "difficulty": "hard", "score": 1.0, "steps": 15}
-]
+{
+  "final_score": 1.0,
+  "total_steps": 18,
+  "stages_completed": [
+    1,
+    2,
+    3,
+    4
+  ],
+  "episode_score": 1.0
+}

baseline/run_baseline.py

@@ -91,8 +91,8 @@ def run_episode():
             action = {"action_type": "run_script", "payload": {}}
         elif current_stage == 3 and stage_step == 1:
             action = {"action_type": "edit_script", "payload": {
-                "old": "scaler = StandardScaler()\nX_scaled = scaler.fit_transform(X)  # data leakage: fit on full data\nX_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)",
-                "new": "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\nscaler = StandardScaler()\nX_train = scaler.fit_transform(X_train)\nX_test = scaler.transform(X_test)"
+                "old": "scaler = StandardScaler()\nX_scaled = scaler.fit_transform(X)\nX_train, X_test, y_train, y_test = train_test_split(\n    X_scaled, y, test_size=0.2, random_state=42, stratify=y\n)",
+                "new": "X_train, X_test, y_train, y_test = train_test_split(\n    X, y, test_size=0.2, random_state=42, stratify=y\n)\nscaler = StandardScaler()\nX_train = scaler.fit_transform(X_train)\nX_test = scaler.transform(X_test)"
             }}
         elif current_stage == 3 and stage_step == 2:
             action = {"action_type": "run_script", "payload": {}}

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "dataengenv"
 version = "1.0.0"
 description = "ML Data Pipeline Debugger - OpenEnv Environment"
-requires-python = ">=3.10, <3.13"
+requires-python = ">=3.10"
 dependencies = [
     "fastapi==0.115.6",
     "uvicorn",