environment/types.py

"""Pydantic models for Email Triage OpenEnv."""

from datetime import datetime
from enum import Enum
from typing import Optional, List, Dict, Any

from pydantic import BaseModel, Field

class EmailCategory(str, Enum):
    """Email classification categories"""
    SPAM = "spam"
    NORMAL = "normal"
    URGENT = "urgent"
    BILLING = "billing"

class Team(str, Enum):
    """Teams to route emails to"""
    SUPPORT = "support"
    SALES = "sales"
    BILLING = "billing"
    NONE = "none"

class Email(BaseModel):
    """Represents an email message"""
    email_id: str
    subject: str
    body: str
    sender_domain: str
    timestamp: datetime
    is_vip_sender: bool = False
    sla_hours: Optional[int] = None

class GroundTruth(BaseModel):
    """Ground truth labels for an email"""
    email_id: str
    category: EmailCategory
    team: Team
    priority: int = Field(ge=0, le=3)  # 0=low, 3=high

class Observation(BaseModel):
    """Observation returned after each step"""
    current_email: Email
    inbox_state: Dict[str, int] = Field(
        default_factory=lambda: {
            "pending": 0,
            "spam": 0,
            "urgent": 0,
            "processed": 0
        }
    )
    step_count: int = 0
    task_name: str = ""
    info: Dict[str, Any] = Field(default_factory=dict)

class Action(BaseModel):
    """Action taken by the agent"""
    classification: EmailCategory
    team: Team = Team.NONE
    priority: int = Field(ge=0, le=3, default=1)

class Reward(BaseModel):
    """Reward signal for an action"""
    value: float = Field(ge=0.0, le=1.0)
    breakdown: Dict[str, float] = Field(default_factory=dict)

class State(BaseModel):
    """Complete environment state"""
    current_observation: Observation
    current_reward: float
    done: bool
    info: Dict[str, Any] = Field(default_factory=dict)
    history: List[Dict[str, Any]] = Field(default_factory=list)