env_models.py

"""
models.py — Typed Pydantic models for IT Support Triage OpenEnv environment.
All inputs/outputs are strongly typed for OpenEnv spec compliance.
"""

from enum import Enum
from typing import Optional, List
from pydantic import BaseModel, Field


# ─── Enumerations ────────────────────────────────────────────────────────────

class TicketCategory(str, Enum):
    HARDWARE        = "hardware"
    SOFTWARE        = "software"
    NETWORK         = "network"
    SECURITY        = "security"
    ACCESS          = "access"
    EMAIL           = "email"
    PRINTER         = "printer"
    OTHER           = "other"


class TicketPriority(str, Enum):
    CRITICAL  = "critical"   # P1 — system down, security breach, data loss
    HIGH      = "high"       # P2 — major function impaired
    MEDIUM    = "medium"     # P3 — degraded but workaround exists
    LOW       = "low"        # P4 — cosmetic / minor inconvenience


class EscalationDecision(str, Enum):
    ESCALATE     = "escalate"
    SELF_RESOLVE = "self_resolve"
    MONITOR      = "monitor"


class Department(str, Enum):
    TIER1_HELPDESK  = "tier1_helpdesk"
    TIER2_SUPPORT   = "tier2_support"
    SECURITY_TEAM   = "security_team"
    NETWORK_OPS     = "network_ops"
    SYSADMIN        = "sysadmin"
    VENDOR_SUPPORT  = "vendor_support"


# ─── Observation (what the agent sees) ───────────────────────────────────────

class TicketObservation(BaseModel):
    """The full context provided to the agent for each ticket."""
    ticket_id:        str                   = Field(..., description="Unique ticket identifier")
    subject:          str                   = Field(..., description="Ticket subject line")
    body:             str                   = Field(..., description="Full ticket description from user")
    reporter_name:    str                   = Field(..., description="Name of the person filing the ticket")
    reporter_role:    str                   = Field(..., description="Job role of the reporter")
    system_info:      Optional[str]         = Field(None, description="OS/software version if provided")
    timestamp:        str                   = Field(..., description="ISO 8601 submission timestamp")
    previous_tickets: int                   = Field(0, description="Number of prior tickets from this user")
    task_instruction: str                   = Field(..., description="What the agent must do for this task")
    valid_categories: List[str]             = Field(..., description="Allowed category values")
    valid_priorities: List[str]             = Field(..., description="Allowed priority values")
    valid_departments: List[str]            = Field(..., description="Allowed department routing values")


# ─── Action (what the agent does) ────────────────────────────────────────────

class TriageAction(BaseModel):
    """The agent's triage decision for a ticket."""
    category:    TicketCategory           = Field(..., description="Ticket category classification")
    priority:    TicketPriority           = Field(..., description="Assigned priority level")
    department:  Department               = Field(..., description="Routing destination")
    escalate:    EscalationDecision       = Field(..., description="Escalation decision")
    response:    str                      = Field(..., description="Response message to send to the user (max 500 chars)")
    reasoning:   str                      = Field(..., description="Internal reasoning for the triage decision")

    class Config:
        use_enum_values = True


# ─── Step result ─────────────────────────────────────────────────────────────

class StepResult(BaseModel):
    """Returned after each step() call."""
    observation:  Optional[TicketObservation] = Field(None, description="Next observation (None if episode done)")
    reward:       float                        = Field(..., description="Reward for this step (0.0–1.0)")
    done:         bool                         = Field(..., description="Whether the episode has ended")
    info:         dict                         = Field(default_factory=dict, description="Grader breakdown and metadata")


# ─── Environment state ────────────────────────────────────────────────────────

class EnvState(BaseModel):
    """Full environment state snapshot."""
    task_id:          str
    task_name:        str
    task_description: str
    current_step:     int
    max_steps:        int
    total_reward:     float
    done:             bool
    current_ticket:   Optional[TicketObservation]
    history:          List[dict]