""" Core environment: ExecOpsEnv Manages the CEO simulation state machine. """ from typing import Optional from .models import Action, Email, CalendarEvent, Goal, EnvState, StepResult from .tasks import get_task from .reward import compute_reward # 30-minute time slots starting at 9:00 AM TIME_SLOTS = [ "9:00 AM", "9:30 AM", "10:00 AM", "10:30 AM", "11:00 AM", "11:30 AM", "12:00 PM", "12:30 PM", "1:00 PM", "1:30 PM", "2:00 PM", "2:30 PM", "3:00 PM", "3:30 PM", "4:00 PM", "4:30 PM", "5:00 PM", ] URGENCY_DECAY = 0.15 # Urgency drops by this amount each step class ExecOpsEnv: def __init__(self, task_id: str = "easy"): self.task_id = task_id self._task_data = get_task(task_id) self._state: Optional[EnvState] = None self.reset() def reset(self) -> dict: """Reinitialize state from task definition. Returns initial observation.""" task = get_task(self.task_id) emails = [Email(**e) for e in task["emails"]] calendar = [CalendarEvent(**c) for c in task["calendar"]] goals = [Goal(**g) for g in task["goals"]] self._state = EnvState( current_step=0, max_steps=task["max_steps"], current_time=TIME_SLOTS[task.get("start_time_slot", 0)], inbox=emails, calendar=calendar, goals=goals, history=[], total_reward=0.0, done=False, ) return self._build_observation() def step(self, action: Action) -> dict: """ Execute one step of the environment. Order: validate → apply → update state → advance time → decay urgency → compute reward → check done → return StepResult """ if self._state is None: raise RuntimeError("Environment not initialized. Call reset() first.") if self._state.done: raise RuntimeError("Environment is done. Call reset() to restart.") # 1. Validate action valid_action_types = {"reply", "schedule", "delegate", "ignore"} if action.type not in valid_action_types: raise ValueError( f"Invalid action type '{action.type}'. Must be one of: {sorted(valid_action_types)}" ) if not action.email_id: raise ValueError("Action is missing required field 'email_id'.") email = self._find_email(action.email_id) if email is None: raise ValueError(f"Email '{action.email_id}' not found in inbox.") if email.handled: raise ValueError(f"Email '{action.email_id}' is already handled.") # 2. Apply action email.handled = True email.action_taken = action.type # Update calendar if scheduling if action.type == "schedule": self._add_to_calendar(email) # 3. Record action in history step_entry = { "step": self._state.current_step, "time": self._state.current_time, "action": action.type, "email_id": action.email_id, "email_subject": email.subject, "email_priority": email.priority, } # 6. Compute reward BEFORE marking goals complete (so reward fn sees pending goals) reward = compute_reward(action, email, self._state.goals, self._state) # Now mark goals completed for goal in self._state.goals: if goal.target_email_id == action.email_id and not goal.completed: if goal.required_action == action.type: goal.completed = True step_entry["reward"] = reward self._state.total_reward += reward self._state.history.append(step_entry) # 4. Advance time self._advance_time() # 5. Decay urgency on unhandled emails self._decay_urgency() # 7. Check termination self._state.current_step += 1 self._check_done() obs = self._build_observation() return { "observation": obs, "reward": reward, "done": self._state.done, "info": { "step": self._state.current_step, "total_reward": self._state.total_reward, "action_applied": action.type, "email_subject": email.subject, }, } def get_state(self) -> dict: """Return full serialized state.""" if self._state is None: raise RuntimeError("Environment not initialized.") return { "current_step": self._state.current_step, "max_steps": self._state.max_steps, "current_time": self._state.current_time, "inbox": [e.model_dump() for e in self._state.inbox], "calendar": [c.model_dump() for c in self._state.calendar], "goals": [g.model_dump() for g in self._state.goals], "history": self._state.history, "total_reward": self._state.total_reward, "done": self._state.done, "task_id": self.task_id, "task_name": self._task_data.get("task_name", ""), "task_description": self._task_data.get("description", ""), } # ------------------------- # Private helpers # ------------------------- def _build_observation(self) -> dict: """Build the observation dict that the agent sees.""" unhandled = [e for e in self._state.inbox if not e.handled] pending_goals = [g for g in self._state.goals if not g.completed] return { "time": self._state.current_time, "step": self._state.current_step, "max_steps": self._state.max_steps, "steps_remaining": self._state.max_steps - self._state.current_step, "inbox": [e.model_dump() for e in self._state.inbox], "unhandled_count": len(unhandled), "calendar": [c.model_dump() for c in self._state.calendar], "goals": [g.model_dump() for g in self._state.goals], "pending_goals": [g.model_dump() for g in pending_goals], "total_reward": self._state.total_reward, "done": self._state.done, "task_id": self.task_id, "task_name": self._task_data.get("task_name", ""), } def _find_email(self, email_id: str) -> Optional[Email]: for email in self._state.inbox: if email.id == email_id: return email return None def _advance_time(self): """Move to the next 30-minute time slot.""" try: idx = TIME_SLOTS.index(self._state.current_time) if idx + 1 < len(TIME_SLOTS): self._state.current_time = TIME_SLOTS[idx + 1] except ValueError: pass def _decay_urgency(self): """Reduce urgency of all unhandled emails.""" for email in self._state.inbox: if not email.handled: email.urgency = max(0.0, email.urgency - URGENCY_DECAY) def _check_done(self): """Mark done if max steps reached or all goals completed.""" all_goals_done = all(g.completed for g in self._state.goals) if self._state.current_step >= self._state.max_steps or all_goals_done: self._state.done = True def _add_to_calendar(self, email: Email): """Add a scheduling event to the calendar for scheduled emails.""" # Find next available slot locked_slots = {c.time_slot for c in self._state.calendar if c.locked} used_slots = {c.time_slot for c in self._state.calendar} for slot in range(len(TIME_SLOTS)): if slot not in used_slots and slot not in locked_slots: self._state.calendar.append( CalendarEvent( time_slot=slot, title=f"Scheduled: {email.subject[:40]}", attendee=email.sender.split("<")[0].strip(), locked=False, ) ) break