Spaces:
Sleeping
Sleeping
| import gymnasium as gym | |
| from gymnasium import spaces | |
| import numpy as np | |
| from typing import List, Dict, Optional | |
| from models import Email, CalendarEvent, Observation, Action | |
| class EmailEnv(gym.Env): | |
| """ | |
| Email Triage & Scheduling Assistant: A real-world human-task environment. | |
| Simulates inbox management, spam filtering, and meeting coordination. | |
| """ | |
| def __init__(self): | |
| super(EmailEnv, self).__init__() | |
| self.action_space = spaces.Discrete(10) # Placeholder for discrete actions if needed | |
| self._setup_inbox() | |
| self.max_steps = 30 | |
| self.reset() | |
| def _setup_inbox(self): | |
| # Sample structured data for tasks | |
| self.sample_emails = [ | |
| Email(id=1, sender="spam@bott.io", subject="CASH NOW!!", body="Claim your 1M dollars", folder="Inbox", priority=3), | |
| Email(id=2, sender="boss@corp.com", subject="Urgent: Project Update", body="Send the report by 5 PM.", folder="Inbox", priority=1), | |
| Email(id=3, sender="calendar@corp.com", subject="Meeting Request: Sync", body="Let's sync at 2 PM.", folder="Inbox", priority=2), | |
| Email(id=4, sender="news@daily.com", subject="Daily Briefing", body="Top stories of the day.", folder="Inbox", priority=3), | |
| Email(id=5, sender="friend@web.com", subject="Coffee?", body="Are you free tomorrow at 10 AM?", folder="Inbox", priority=3) | |
| ] | |
| self.sample_calendar = [ | |
| CalendarEvent(title="Sprint Review", start_time="10:00", end_time="11:00"), | |
| CalendarEvent(title="Lunch", start_time="12:00", end_time="13:00") | |
| ] | |
| def reset(self, seed=None, options=None): | |
| super().reset(seed=seed) | |
| self.current_level = options.get("level", 1) if options else 1 | |
| self.inbox = [e.model_copy() for e in self.sample_emails] | |
| self.calendar = [c.model_copy() for c in self.sample_calendar] | |
| self.steps = 0 | |
| self.completed_tasks = 0 | |
| return self._get_observation(), {} | |
| def _get_observation(self) -> Dict: | |
| # Pydantic conversion to dict for Gym-compatible step/reset returns | |
| obs = Observation( | |
| inbox_count=len([e for e in self.inbox if e.folder == "Inbox"]), | |
| current_email=self.inbox[0] if self.inbox else None, | |
| calendar=self.calendar | |
| ) | |
| return obs.model_dump() | |
| def step(self, action_dict: Dict): | |
| """ | |
| Receives an Action model mapping (dict) and applies it to the state. | |
| Returns: observation, reward, terminated, truncated, info | |
| """ | |
| self.steps += 1 | |
| action = Action(**action_dict) | |
| reward = 0.0 | |
| terminated = False | |
| info = {"is_success": False} | |
| # Logic for Task 1: Deleting Spam (Easy) | |
| if self.current_level == 1: | |
| if action.type == "MOVE" and action.email_id == 1 and action.target_folder == "Spam": | |
| reward = 1.0 # Solved Task 1 | |
| self.inbox[0].folder = "Spam" | |
| terminated = True | |
| info["is_success"] = True | |
| else: | |
| reward = -0.1 # Logical error penalty | |
| # Logic for Task 2: Categorization (Medium) | |
| elif self.current_level == 2: | |
| target_ids = [2, 4] # Boss to Work, News to Archive | |
| if action.type == "MOVE": | |
| email = next((e for e in self.inbox if e.id == action.email_id), None) | |
| if email: | |
| if email.id == 2 and action.target_folder == "Work": | |
| reward += 0.4 | |
| email.folder = "Work" | |
| elif email.id == 4 and action.target_folder == "Archive": | |
| reward += 0.4 | |
| email.folder = "Archive" | |
| else: | |
| reward -= 0.1 | |
| # Check for completion | |
| if all(e.folder != "Inbox" for e in self.inbox if e.id in target_ids): | |
| reward += 0.2 # Completion bonus | |
| terminated = True | |
| info["is_success"] = True | |
| # Logic for Task 3: Scheduling (Hard) | |
| elif self.current_level == 3: | |
| # Task: Schedule a meeting at 2 PM (No conflict) vs Avoiding 10 AM (Conflict) | |
| if action.type == "SCHEDULE": | |
| if "2 PM" in (action.reply_text or ""): | |
| reward = 1.0 | |
| terminated = True | |
| info["is_success"] = True | |
| elif "10 AM" in (action.reply_text or ""): | |
| reward = -0.5 # Fail: Calendar conflict! | |
| terminated = True | |
| else: | |
| reward = -0.1 | |
| truncated = self.steps >= self.max_steps | |
| return self._get_observation(), reward, terminated, truncated, info | |
| def state(self) -> Dict: | |
| """Required by OpenEnv for full state snapshot.""" | |
| return { | |
| "inbox_snapshot": [e.model_dump() for e in self.inbox], | |
| "calendar_snapshot": [c.model_dump() for c in self.calendar], | |
| "level": self.current_level | |
| } | |