albert-einstein-09/codedark

1"""2CodeDark Environment3 4OpenEnv-compatible environment for multi-turn data analytics tasks.5Agents analyze CSV data using Python/Pandas tools and submit answers.6"""7 8import json9import uuid10from pathlib import Path11from typing import Any, Dict, List, Optional12 13import pandas as pd14 15from ..models import CodeDarkAction, CodeDarkObservation, CodeDarkState16from .tools import (17    run_python,18    read_notes,19    save_note,20    clarify,21    submit_answer,22    parse_tool_call,23)24from .scoring import compute_reward25 26 27class CodeDarkEnvironment:28    """CodeDark environment for multi-turn data analytics.29 30    Features:31    - Multi-turn agent evaluation32    - 5 tools: run_python, read_notes, save_note, clarify, submit_answer33    - Shaped rewards: correctness (80%) + efficiency (10%) + token cost (10%)34    - Supports bank and road datasets35    """36 37    def __init__(38        self,39        data_dir: Optional[str] = None,40        tasks_path: Optional[str] = None,41        max_turns: int = 10,42        max_clarifications: int = 2,43    ):44        """Initialize CodeDark environment.45 46        Args:47            data_dir: Path to directory containing CSV files48            tasks_path: Path to tasks.jsonl file49            max_turns: Maximum turns per episode (default: 10)50            max_clarifications: Maximum clarifications per episode (default: 2)51        """52        self.max_turns = max_turns53        self.max_clarifications = max_clarifications54 55        # Resolve paths56        if data_dir:57            self.data_dir = Path(data_dir)58        else:59            # Default to data/ relative to this file's parent60            self.data_dir = Path(__file__).parent.parent / "data"61 62        if tasks_path:63            self.tasks_path = Path(tasks_path)64        else:65            self.tasks_path = self.data_dir / "tasks" / "final_25_tasks.jsonl"66 67        # Load tasks68        self.tasks = self._load_tasks()69        self._tasks_by_id = {t["id"]: t for t in self.tasks}70        self._task_index = 071 72        # Current episode state73        self._state: Optional[CodeDarkState] = None74        self._df: Optional[pd.DataFrame] = None75        self._current_task: Optional[Dict] = None76 77    def _load_tasks(self) -> List[Dict]:78        """Load tasks from JSONL file."""79        if not self.tasks_path.exists():80            return []81 82        tasks = []83        with open(self.tasks_path) as f:84            for line in f:85                if line.strip():86                    tasks.append(json.loads(line))87        return tasks88 89    def _load_data_for_task(self, task: Dict) -> Optional[pd.DataFrame]:90        """Load the appropriate CSV for a task.91 92        Args:93            task: Task dictionary with 'dataset' field94 95        Returns:96            DataFrame or None if not found97        """98        dataset = task.get("dataset", "bank")99        csv_path = self.data_dir / f"{dataset}.csv"100 101        if csv_path.exists():102            return pd.read_csv(csv_path)103        return None104 105    @property106    def state(self) -> CodeDarkState:107        """Return current environment state."""108        if self._state is None:109            self._state = CodeDarkState()110        return self._state111 112    def reset(113        self, task_id: Optional[str] = None, seed: Optional[int] = None114    ) -> CodeDarkObservation:115        """Reset environment for a new episode.116 117        Args:118            task_id: Specific task to load (optional)119            seed: Random seed for task selection (optional)120 121        Returns:122            Initial observation with task question123        """124        # Select task125        if task_id and task_id in self._tasks_by_id:126            task = self._tasks_by_id[task_id]127        elif self.tasks:128            if seed is not None:129                import random130 131                random.seed(seed)132                task = random.choice(self.tasks)133            else:134                # Round-robin through tasks135                task = self.tasks[self._task_index % len(self.tasks)]136                self._task_index += 1137        else:138            # No tasks loaded - return error observation139            return CodeDarkObservation(140                stderr="Error: No tasks loaded",141                exit_code=1,142                done=True,143            )144 145        self._current_task = task146 147        # Load data for this task148        self._df = self._load_data_for_task(task)149        if self._df is None:150            return CodeDarkObservation(151                stderr=f"Error: Could not load data for dataset '{task.get('dataset', 'bank')}'",152                exit_code=1,153                done=True,154            )155 156        # Initialize state157        self._state = CodeDarkState(158            episode_id=str(uuid.uuid4()),159            step_count=0,160            task_id=task["id"],161            dataset=task.get("dataset", "bank"),162            notes=[],163            turn_count=0,164            error_count=0,165            clarify_count=0,166            submitted=False,167            submitted_answer=None,168            expected_answer=task["golden"]["answer_value"],169            tolerance=task["golden"].get("tolerance", 0.01),170        )171 172        # Return initial observation173        return CodeDarkObservation(174            stdout=f"Task loaded. DataFrame shape: {self._df.shape}",175            turn=0,176            max_turns=self.max_turns,177            notes=[],178            task_id=task["id"],179            question=task["goal"],180            difficulty=task.get("level", "L5"),181            dataset=task.get("dataset", "bank"),182            done=False,183            submitted=False,184        )185 186    def step(self, action: CodeDarkAction) -> CodeDarkObservation:187        """Execute an action and return observation.188 189        Args:190            action: CodeDarkAction with tool name and args191 192        Returns:193            CodeDarkObservation with results194        """195        if self._state is None or self._current_task is None:196            return CodeDarkObservation(197                stderr="Error: Environment not reset. Call reset() first.",198                exit_code=1,199                done=True,200            )201 202        if self._state.submitted:203            return self._make_final_observation()204 205        # Increment turn206        self._state.turn_count += 1207        self._state.step_count += 1208 209        # Check turn limit210        if self._state.turn_count > self.max_turns:211            self._state.submitted = True212            return self._make_final_observation()213 214        # Parse tool-specific args215        parsed_content, parse_error = parse_tool_call(action.args, action.tool)216 217        if parse_error:218            self._state.error_count += 1219            return CodeDarkObservation(220                stderr=f"{action.tool} Error: {parse_error}",221                exit_code=1,222                turn=self._state.turn_count,223                max_turns=self.max_turns,224                notes=self._state.notes.copy(),225                task_id=self._state.task_id,226                question=self._current_task["goal"],227                difficulty=self._current_task.get("level", "L5"),228                dataset=self._state.dataset,229                done=False,230                submitted=False,231            )232 233        # Execute tool234        stdout, stderr, exit_code = "", "", 0235 236        if action.tool == "run_python":237            stdout, stderr, exit_code = run_python(parsed_content, self._df)238 239        elif action.tool == "read_notes":240            stdout, stderr, exit_code = read_notes(self._state.notes)241 242        elif action.tool == "save_note":243            stdout, stderr, exit_code = save_note(parsed_content, self._state.notes)244 245        elif action.tool == "clarify":246            stdout, stderr, exit_code, new_count = clarify(247                question=parsed_content,248                clarify_count=self._state.clarify_count,249                max_clarifications=self.max_clarifications,250                ambiguities=self._current_task.get("ambiguities", []),251                answer_type=self._current_task.get("golden", {}).get(252                    "answer_type", "scalar"253                ),254            )255            self._state.clarify_count = new_count256 257        elif action.tool == "submit_answer":258            stdout, stderr, exit_code, answer = submit_answer(parsed_content)259            if exit_code == 0:260                self._state.submitted = True261                self._state.submitted_answer = answer262                return self._make_final_observation()263 264        # Track errors265        if exit_code != 0:266            self._state.error_count += 1267 268        return CodeDarkObservation(269            stdout=stdout,270            stderr=stderr,271            exit_code=exit_code,272            turn=self._state.turn_count,273            max_turns=self.max_turns,274            notes=self._state.notes.copy(),275            task_id=self._state.task_id,276            question=self._current_task["goal"],277            difficulty=self._current_task.get("level", "L5"),278            dataset=self._state.dataset,279            done=False,280            submitted=False,281        )282 283    def _make_final_observation(self) -> CodeDarkObservation:284        """Create final observation with reward computation."""285        if self._state is None or self._current_task is None:286            return CodeDarkObservation(done=True)287 288        # Compute reward289        reward, correctness, efficiency, token_cost = compute_reward(290            submitted=self._state.submitted_answer,291            expected=self._state.expected_answer,292            tolerance=self._state.tolerance,293            turns=self._state.turn_count,294            max_turns=self.max_turns,295        )296 297        return CodeDarkObservation(298            stdout="[EPISODE COMPLETE]",299            turn=self._state.turn_count,300            max_turns=self.max_turns,301            notes=self._state.notes.copy(),302            task_id=self._state.task_id,303            question=self._current_task["goal"],304            difficulty=self._current_task.get("level", "L5"),305            dataset=self._state.dataset,306            done=True,307            submitted=self._state.submitted,308            reward=reward,309            correctness=correctness,310            efficiency=efficiency,311            metadata={312                "submitted_answer": self._state.submitted_answer,313                "expected_answer": self._state.expected_answer,314                "tolerance": self._state.tolerance,315                "error_count": self._state.error_count,316                "clarify_count": self._state.clarify_count,317                "token_cost_usd": token_cost,318            },319        )320