environment/models.py

"""
OpenEnv typed models for Customer Support Inbox environment.
Implements Observation, Action, Reward per OpenEnv spec.
"""
from __future__ import annotations

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


# ─── Enums ────────────────────────────────────────────────────────────────────

class TicketCategory(str, Enum):
    BILLING = "billing"
    TECHNICAL = "technical"
    SHIPPING = "shipping"
    RETURNS = "returns"
    GENERAL = "general"
    ACCOUNT = "account"
    COMPLAINT = "complaint"


class TicketPriority(str, Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    URGENT = "urgent"


class TicketStatus(str, Enum):
    OPEN = "open"
    IN_PROGRESS = "in_progress"
    WAITING_CUSTOMER = "waiting_customer"
    ESCALATED = "escalated"
    RESOLVED = "resolved"
    CLOSED = "closed"


class ActionType(str, Enum):
    CLASSIFY = "classify"           # Assign category + priority
    RESPOND = "respond"             # Send a message to the customer
    ESCALATE = "escalate"           # Escalate ticket to senior team
    RESOLVE = "resolve"             # Mark ticket as resolved with resolution notes
    REQUEST_INFO = "request_info"   # Ask customer for more information
    TAG = "tag"                     # Add searchable tags
    ASSIGN = "assign"               # Assign to department/agent


class SentimentLabel(str, Enum):
    POSITIVE = "positive"
    NEUTRAL = "neutral"
    NEGATIVE = "negative"
    ANGRY = "angry"


# ─── Sub-models ───────────────────────────────────────────────────────────────

class Message(BaseModel):
    """A single message in a ticket conversation."""
    role: str = Field(..., description="'customer' or 'agent'")
    content: str = Field(..., description="Message text")
    timestamp: str = Field(..., description="ISO 8601 timestamp")
    metadata: Dict[str, Any] = Field(default_factory=dict)


class CustomerProfile(BaseModel):
    """Basic customer profile information."""
    customer_id: str
    name: str
    email: str
    account_tier: str = Field(default="standard", description="standard | premium | enterprise")
    total_orders: int = Field(default=0)
    previous_tickets: int = Field(default=0)
    sentiment: SentimentLabel = Field(default=SentimentLabel.NEUTRAL)


class TicketState(BaseModel):
    """Full internal state of a support ticket."""
    ticket_id: str
    subject: str
    body: str
    category: Optional[TicketCategory] = None
    priority: Optional[TicketPriority] = None
    status: TicketStatus = TicketStatus.OPEN
    tags: List[str] = Field(default_factory=list)
    assigned_to: Optional[str] = None
    conversation: List[Message] = Field(default_factory=list)
    customer: CustomerProfile = Field(...)
    created_at: str = Field(...)
    updated_at: str = Field(...)
    sla_deadline: str = Field(...)
    resolution_notes: Optional[str] = None
    escalation_reason: Optional[str] = None
    turn_count: int = Field(default=0, description="Number of agent actions taken")
    max_turns: int = Field(default=10, description="Maximum allowed turns before episode ends")
    task_id: str = Field(default="task1")
    task_metadata: Dict[str, Any] = Field(default_factory=dict)


# ─── OpenEnv Core Models ───────────────────────────────────────────────────────

class Observation(BaseModel):
    """
    OpenEnv Observation — what the agent sees each step.
    Contains all information needed to decide the next action.
    """
    # Ticket basics
    ticket_id: str = Field(..., description="Unique ticket identifier")
    subject: str = Field(..., description="Ticket subject line")
    body: str = Field(..., description="Original customer message")
    status: TicketStatus = Field(..., description="Current ticket status")
    category: Optional[TicketCategory] = Field(None, description="Assigned category (None if unclassified)")
    priority: Optional[TicketPriority] = Field(None, description="Assigned priority (None if unset)")
    tags: List[str] = Field(default_factory=list, description="Applied tags")

    # Customer context
    customer: CustomerProfile = Field(..., description="Customer profile info")

    # Conversation history
    conversation: List[Message] = Field(default_factory=list, description="Full message thread")

    # Episode state
    turn_count: int = Field(..., description="Actions taken so far this episode")
    turns_remaining: int = Field(..., description="Actions left before forced termination")
    time_to_sla: str = Field(..., description="Human-readable time remaining on SLA")

    # Task framing
    task_id: str = Field(..., description="Current task identifier")
    task_description: str = Field(..., description="What the agent must accomplish")
    task_objectives: List[str] = Field(default_factory=list, description="Specific success criteria")

    # Available actions
    available_actions: List[str] = Field(..., description="Valid ActionTypes for current state")

    # Feedback from last action
    last_action_result: Optional[str] = Field(None, description="Outcome message from previous step")
    last_reward: float = Field(default=0.0, description="Reward earned on last step")

    # Episode meta
    episode_done: bool = Field(default=False, description="Whether the episode has ended")
    info: Dict[str, Any] = Field(default_factory=dict, description="Extra diagnostic info")


class Action(BaseModel):
    """
    OpenEnv Action — what the agent can do each step.
    Not all fields are required; depends on action_type.
    """
    action_type: ActionType = Field(..., description="The type of action to perform")

    # For CLASSIFY
    category: Optional[TicketCategory] = Field(None, description="Category to assign (for CLASSIFY)")
    priority: Optional[TicketPriority] = Field(None, description="Priority to assign (for CLASSIFY)")

    # For RESPOND / REQUEST_INFO
    response_text: Optional[str] = Field(None, description="Message to send to customer (for RESPOND/REQUEST_INFO)")

    # For ESCALATE
    escalation_reason: Optional[str] = Field(None, description="Reason for escalation (for ESCALATE)")
    escalation_team: Optional[str] = Field(None, description="Team to escalate to: 'tier2' | 'billing_specialist' | 'engineering'")

    # For RESOLVE
    resolution_notes: Optional[str] = Field(None, description="Resolution summary (for RESOLVE)")
    resolution_category: Optional[str] = Field(None, description="How was it resolved: 'fixed' | 'refunded' | 'explained' | 'workaround'")

    # For TAG
    tags: Optional[List[str]] = Field(None, description="Tags to add (for TAG)")

    # For ASSIGN
    assigned_to: Optional[str] = Field(None, description="Department/agent to assign (for ASSIGN)")


class RewardBreakdown(BaseModel):
    """Detailed breakdown of reward components."""
    classification_accuracy: float = Field(default=0.0, ge=0.0, le=1.0)
    response_quality: float = Field(default=0.0, ge=0.0, le=1.0)
    resolution_completeness: float = Field(default=0.0, ge=0.0, le=1.0)
    sla_compliance: float = Field(default=0.0, ge=0.0, le=1.0)
    efficiency: float = Field(default=0.0, ge=0.0, le=1.0)
    customer_sentiment_improvement: float = Field(default=0.0, ge=0.0, le=1.0)
    escalation_appropriateness: float = Field(default=0.0, ge=0.0, le=1.0)
    step_penalty: float = Field(default=0.0, ge=-1.0, le=0.0)


class Reward(BaseModel):
    """
    OpenEnv Reward — returned with each step.
    Provides fine-grained partial credit signals.
    """
    score: float = Field(..., ge=0.0, le=1.0, description="Overall step reward (0.0–1.0)")
    cumulative_score: float = Field(default=0.0, ge=0.0, le=1.0, description="Accumulated reward this episode")
    done: bool = Field(..., description="Whether this action ended the episode")
    breakdown: RewardBreakdown = Field(default_factory=RewardBreakdown, description="Per-component reward details")
    feedback: str = Field(default="", description="Human-readable reward explanation")
    task_complete: bool = Field(default=False, description="Whether the task objective was fully met")