""" Email Triage OpenEnv Environment. Implements a realistic email management simulation where AI agents learn to triage emails into appropriate folders. Supports deterministic reset with seed, multi-step episodes, and multi-component reward shaping. Usage: >>> env = EmailTriageEnv(task_id="basic_triage", seed=42) >>> obs = env.reset() >>> action = Action(action_type="move", email_id=0, target_folder="work") >>> obs, reward, done, info = env.step(action) """ from __future__ import annotations import logging from typing import Any, Dict, List, Optional, Tuple from src.data_generator import generate_realistic_emails from src.models import ( AVAILABLE_FOLDERS, MAX_STEPS_PER_EPISODE, Action, Email, Observation, Reward, StepRecord, ) from src.reward_shaper import compute_reward logger = logging.getLogger(__name__) # ── Task Configuration ─────────────────────────────────────────────────────── TASK_CONFIG: Dict[str, Dict[str, Any]] = { "basic_triage": { "email_count": 5, "difficulty": "easy", "max_steps": 15, "description": "Sort 5 emails into work vs. spam", }, "multi_folder_triage": { "email_count": 15, "difficulty": "medium", "max_steps": 30, "description": "Sort 15 emails into 4 folders", }, "advanced_triage_with_urgency": { "email_count": 30, "difficulty": "hard", "max_steps": 50, "description": "Sort 30 emails with VIP and urgency handling", }, } class EmailTriageEnv: """Email triage environment following the OpenEnv interface. The agent observes an inbox of emails and must sort them into the correct folders. Each action processes one email. The episode ends when the inbox is empty or the step limit is reached. Attributes: task_id: Which task configuration to use. seed: Random seed for deterministic episode generation. current_step: Steps taken so far in this episode. max_steps: Maximum steps allowed for this task. done: Whether the current episode has ended. history: List of StepRecords for grading. """ def __init__( self, task_id: str = "basic_triage", seed: Optional[int] = None, ) -> None: if task_id not in TASK_CONFIG: raise ValueError( f"Unknown task_id '{task_id}'. " f"Available: {list(TASK_CONFIG.keys())}" ) self.task_id = task_id self.seed = seed self._config = TASK_CONFIG[task_id] # Episode state — populated on reset() self._inbox: List[Email] = [] self._folders: Dict[str, List[Email]] = {f: [] for f in AVAILABLE_FOLDERS} self._flagged_emails: set = set() self._action_history: List[Action] = [] self.history: List[StepRecord] = [] self.current_step: int = 0 self.max_steps: int = self._config["max_steps"] self.done: bool = False self._episode_num: int = 0 self._email_lookup: Dict[int, Email] = {} # ── OpenEnv Interface ──────────────────────────────────────────────── def reset(self, seed: Optional[int] = None) -> Observation: """Reset the environment to a fresh episode. Args: seed: Optional override for the random seed. Returns: Initial observation with a full inbox. """ if seed is not None: self.seed = seed self._episode_num += 1 self.current_step = 0 self.done = False self._action_history = [] self.history = [] self._flagged_emails = set() self._folders = {f: [] for f in AVAILABLE_FOLDERS} # Generate emails deterministically effective_seed = ( self.seed if self.seed is not None else self._episode_num * 1000 + hash(self.task_id) % 10000 ) self._inbox = generate_realistic_emails( count=self._config["email_count"], difficulty=self._config["difficulty"], seed=effective_seed, ) self._email_lookup = {e.id: e for e in self._inbox} logger.info( "Reset environment: task=%s, emails=%d, seed=%s", self.task_id, len(self._inbox), effective_seed, ) return self._build_observation() def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]: """Execute one action in the environment. Args: action: The agent's chosen action. Returns: Tuple of (observation, reward, done, info). Raises: RuntimeError: If the episode has already ended. ValueError: If the target email doesn't exist in the inbox. """ if self.done: raise RuntimeError( "Episode is done. Call reset() to start a new episode." ) self.current_step += 1 info: Dict[str, Any] = {"task_id": self.task_id} # Validate email exists email = self._email_lookup.get(action.email_id) if email is None: # Email not found — return zero reward and continue reward = Reward( value=0.0, components={"correctness": 0.0, "efficiency": 0.0}, reason=f"Email ID {action.email_id} not found in inbox", ) info["error"] = f"Invalid email_id: {action.email_id}" self._check_done() obs = self._build_observation() self._record_step(action, email or Email( id=action.email_id, subject="(unknown)", sender="unknown@unknown.com", sender_domain="unknown.com", timestamp="2024-01-01T00:00:00", ), reward) return obs, reward, self.done, info # Compute reward before processing the action reward = compute_reward( action=action, email=email, current_step=self.current_step, max_steps=self.max_steps, action_history=self._action_history, flagged_emails=self._flagged_emails, ) # Process the action self._process_action(action, email) self._action_history.append(action) # Record for grading self._record_step(action, email, reward) # Check termination self._check_done() obs = self._build_observation() info["reward_components"] = reward.components info["emails_remaining"] = len(self._inbox) return obs, reward, self.done, info def state(self) -> Observation: """Return the current observation without advancing the environment.""" return self._build_observation() # ── Internal Methods ───────────────────────────────────────────────── def _build_observation(self) -> Observation: """Construct the current observation from internal state.""" return Observation( inbox_emails=list(self._inbox), available_folders=list(AVAILABLE_FOLDERS), current_step=self.current_step, max_steps=self.max_steps, episode_num=self._episode_num, done=self.done, info={ "task_id": self.task_id, "emails_processed": len(self.history), "folders_used": { f: len(emails) for f, emails in self._folders.items() if emails }, }, ) def _process_action(self, action: Action, email: Email) -> None: """Update internal state based on the agent's action.""" if action.action_type == "move": self._move_email(email, action.target_folder) elif action.action_type == "delete": self._remove_from_inbox(email) elif action.action_type == "mark_spam": self._move_email(email, "spam") elif action.action_type == "flag": self._flagged_emails.add(email.id) # Flagging doesn't remove from inbox elif action.action_type == "snooze": # Snoozing temporarily removes from inbox self._remove_from_inbox(email) def _move_email(self, email: Email, folder: str) -> None: """Move an email from the inbox to a target folder.""" self._remove_from_inbox(email) if folder in self._folders: self._folders[folder].append(email) else: logger.warning("Unknown folder '%s', defaulting to inbox", folder) self._folders["inbox"].append(email) def _remove_from_inbox(self, email: Email) -> None: """Remove an email from the inbox list.""" self._inbox = [e for e in self._inbox if e.id != email.id] def _check_done(self) -> None: """Determine whether the episode should end.""" if self.current_step >= self.max_steps: self.done = True logger.info("Episode done: max steps (%d) reached", self.max_steps) elif not self._inbox: self.done = True logger.info("Episode done: inbox empty") def _record_step(self, action: Action, email: Email, reward: Reward) -> None: """Append a step record for later grading.""" self.history.append( StepRecord( step_num=self.current_step, action=action, email=email, reward=reward, done=self.done, ) )