models.py

"""
Module: models.py
Purpose: Define all typed Pydantic models and enums for the TriageFlow environment.
Part of: Medical Triage Assistant — OpenEnv Round 1
Author: Team Squirrel

Overview:
    Redesigned model layer. The agent now receives one patient at a time and
    outputs a classification (IMMEDIATE/URGENT/LESS_URGENT/NON_URGENT/ESCALATE)
    plus the full reordered priority queue.

Dependencies:
    - openenv.core.env_server: Action, Observation, State base classes
    - pydantic: Field for model metadata
    - enum: Enum support for typed constants
    - typing: Type annotations

Usage:
    from models import TriageAction, TriageObservation, TriageState, PriorityLevel
"""

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

from pydantic import Field

from openenv.core.env_server import Action, Observation, State


# ============================================================================
# Enums
# ============================================================================

class PriorityLevel(str, Enum):
    """
    Patient urgency priority levels following standard triage categories,
    plus ESCALATE for patients with incomplete data.

    Attributes:
        IMMEDIATE: Life-threatening, requires immediate intervention.
        URGENT: Serious condition, needs attention within 15 minutes.
        LESS_URGENT: Stable condition, can wait 30-60 minutes.
        NON_URGENT: Minor issue, routine care.
        ESCALATE: Cannot safely classify — send to human for data completion.
    """
    IMMEDIATE = "immediate"
    URGENT = "urgent"
    LESS_URGENT = "less_urgent"
    NON_URGENT = "non_urgent"
    ESCALATE = "escalate"


# ============================================================================
# OpenEnv Action Model
# ============================================================================

class TriageAction(Action):
    """
    Action model for the redesigned TriageFlow environment.

    Each step the agent outputs:
    1. A classification for the incoming patient (priority level or ESCALATE)
    2. The full reordered priority queue after inserting this patient

    Attributes:
        classification (str): Priority classification or "escalate".
        reordered_queue (List[str]): Full queue of patient IDs in priority order.

    Notes:
        For ESCALATE actions, the patient should NOT appear in reordered_queue.
        The queue should list all non-escalated patients seen so far, ordered
        by priority (IMMEDIATE first, then URGENT, LESS_URGENT, NON_URGENT).
    """
    classification: str = Field(
        ..., description="Priority classification: immediate/urgent/less_urgent/non_urgent/escalate"
    )
    reordered_queue: List[str] = Field(
        default_factory=list,
        description="Full queue of patient IDs in agent's chosen priority order"
    )


# ============================================================================
# OpenEnv Observation Model
# ============================================================================

class TriageObservation(Observation):
    """
    Observation model returned to the agent after each step.

    The agent sees: a new incoming patient + the current queue state.

    Attributes:
        incoming_patient (Optional[dict]): The patient who just arrived.
        current_queue (List[str]): Patient IDs in current priority order.
        step_number (int): Current step in the episode.
        total_expected_patients (int): Total patients the agent will see.
        previous_feedback (Optional[str]): Feedback on last classification.

    Notes:
        done and reward are inherited from the Observation base class.
    """
    # done: bool and reward: Optional[float] inherited from Observation
    incoming_patient: Optional[Dict[str, Any]] = Field(
        None, description="The patient who just arrived for triage"
    )
    current_queue: List[str] = Field(
        default_factory=list,
        description="Patient IDs in current priority order (agent's queue so far)"
    )
    step_number: int = Field(0, description="Current step in episode")
    total_expected_patients: int = Field(0, description="Total patients expected")
    previous_feedback: Optional[str] = Field(
        None, description="Feedback on last classification"
    )
    task_name: str = Field("", description="Name of the current task")


# ============================================================================
# OpenEnv State Model
# ============================================================================

class TriageState(State):
    """
    Full internal state of the environment, used by graders and validators.

    Contains everything needed to deterministically score the episode,
    including per-step answer keys and the agent's full trajectory.

    Attributes:
        task_name (str): Which task is active.
        patients (List[dict]): All patient data (including ground truth).
        answer_key (List[dict]): Per-step expected classification and queue.
        agent_classifications (List[dict]): Agent's classification per step.
        agent_queues (List[List[str]]): Agent's queue output per step.
        current_queue (List[str]): Current priority queue state.
        escalated_patients (List[str]): Patient IDs that were escalated.
        patient_index (int): Which patient we're presenting next.
        max_steps (int): Maximum steps allowed.
        all_patients_seen (bool): Whether all patients have been presented.

    Notes:
        episode_id and step_count are inherited from the State base class.
    """
    # episode_id: Optional[str] and step_count: int are inherited from State
    task_name: str = Field("", description="Active task name")
    patients: List[Dict[str, Any]] = Field(
        default_factory=list, description="All patient data with ground truth"
    )
    answer_key: List[Dict[str, Any]] = Field(
        default_factory=list, description="Per-step expected classification and queue"
    )
    agent_classifications: List[Dict[str, Any]] = Field(
        default_factory=list, description="Agent's classification per step"
    )
    agent_queues: List[List[str]] = Field(
        default_factory=list, description="Agent's queue output per step"
    )
    current_queue: List[str] = Field(
        default_factory=list, description="Current priority queue"
    )
    escalated_patients: List[str] = Field(
        default_factory=list, description="Patient IDs escalated to human"
    )
    patient_index: int = Field(0, description="Next patient to present")
    max_steps: int = Field(20, description="Maximum steps allowed")
    all_patients_seen: bool = Field(False, description="All patients presented")