src/models.py

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

Defines the core data structures: Email, EmailTriageObservation, EmailTriageAction,
Reward, and StepRecord. Action and Observation extend the OpenEnv SDK base types
for full compatibility with create_app() and the OpenEnv validation toolchain.
"""

from __future__ import annotations

import logging
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional

from openenv.core.env_server.types import (
    Action as OpenEnvAction,
    Observation as OpenEnvObservation,
)
from pydantic import BaseModel, Field, model_validator

logger = logging.getLogger(__name__)

# ── Constants ────────────────────────────────────────────────────────────────

AVAILABLE_FOLDERS: List[str] = [
    "inbox", "work", "finance", "meetings", "spam", "archive",
]

ACTION_TYPES: List[str] = ["move", "delete", "snooze", "flag", "mark_spam"]

MAX_STEPS_PER_EPISODE: int = 50


# ── Email ────────────────────────────────────────────────────────────────────

class Email(BaseModel):
    """A single email in the agent's inbox."""

    id: int = Field(..., ge=0, description="Unique email identifier")
    subject: str = Field(..., min_length=1, max_length=200)
    sender: str = Field(..., min_length=3)
    sender_domain: str = Field(..., min_length=1)
    timestamp: datetime
    body_preview: str = Field("", max_length=500)
    priority_flag: bool = False
    has_attachment: bool = False
    is_vip_sender: bool = False
    category_hint: Optional[str] = None
    ground_truth_folder: str = Field(
        "inbox",
        description="Correct folder for grading (not exposed to agent)",
    )
    is_ambiguous: bool = False
    acceptable_folders: List[str] = Field(default_factory=list)

    def to_agent_view(self) -> Dict[str, Any]:
        """Return a dict without ground-truth fields (safe to show the agent)."""
        return self.model_dump(
            exclude={"ground_truth_folder", "is_ambiguous", "acceptable_folders"}
        )


# ── Observation (extends OpenEnv SDK) ────────────────────────────────────────

class EmailTriageObservation(OpenEnvObservation):
    """What the agent sees at each step.

    Extends openenv.core.env_server.types.Observation which provides
    done, reward, and metadata fields automatically.
    """

    inbox_emails: List[Dict[str, Any]] = Field(
        default_factory=list, description="Emails currently in the inbox"
    )
    available_folders: List[str] = Field(
        default_factory=lambda: list(AVAILABLE_FOLDERS),
        description="Folders the agent can move emails into",
    )
    current_step: int = Field(0, ge=0)
    max_steps: int = Field(MAX_STEPS_PER_EPISODE, ge=1)
    episode_num: int = Field(1, ge=1)
    info: Dict[str, Any] = Field(default_factory=dict)


# ── Action (extends OpenEnv SDK) ─────────────────────────────────────────────

class EmailTriageAction(OpenEnvAction):
    """An action the agent takes on an email.

    Extends openenv.core.env_server.types.Action which provides
    a metadata field automatically.

    Supported action types:
      - move: Move email to a target folder (requires target_folder).
      - delete: Permanently delete email.
      - snooze: Temporarily hide email.
      - flag: Mark email as important.
      - mark_spam: Move email to spam.
    """

    action_type: Literal["move", "delete", "snooze", "flag", "mark_spam"] = Field(
        ..., description="Type of action to perform"
    )
    email_id: int = Field(..., ge=0, description="ID of the target email")
    target_folder: Optional[str] = Field(
        None, description="Destination folder (required for 'move')"
    )
    reason: Optional[str] = Field(
        None, max_length=300, description="Optional explanation for the action"
    )

    @model_validator(mode="after")
    def validate_move_has_folder(self) -> "EmailTriageAction":
        if self.action_type == "move" and not self.target_folder:
            raise ValueError("target_folder is required for 'move' action")
        if (
            self.target_folder
            and self.target_folder not in AVAILABLE_FOLDERS
        ):
            raise ValueError(
                f"target_folder must be one of {AVAILABLE_FOLDERS}, "
                f"got '{self.target_folder}'"
            )
        return self


# ── Internal Observation (used by core environment.py) ───────────────────────

class Observation(BaseModel):
    """Internal observation model used by EmailTriageEnv (no SDK dependency)."""

    inbox_emails: List[Email] = Field(default_factory=list)
    available_folders: List[str] = Field(
        default_factory=lambda: list(AVAILABLE_FOLDERS),
    )
    current_step: int = Field(0, ge=0)
    max_steps: int = Field(MAX_STEPS_PER_EPISODE, ge=1)
    episode_num: int = Field(1, ge=1)
    done: bool = False
    info: Dict[str, Any] = Field(default_factory=dict)

    def to_agent_view(self) -> Dict[str, Any]:
        """Serialize for the agent, stripping ground-truth fields from emails."""
        data = self.model_dump()
        data["inbox_emails"] = [e.to_agent_view() for e in self.inbox_emails]
        return data


# ── Legacy Action alias (used by internal modules) ───────────────────────────
# Internal modules (reward_shaper, grader, inference) use this simpler model
# that doesn't depend on the SDK for easier testing.

class Action(BaseModel):
    """Standalone action model for internal use (graders, reward shaper)."""

    action_type: Literal["move", "delete", "snooze", "flag", "mark_spam"] = Field(
        ..., description="Type of action to perform"
    )
    email_id: int = Field(..., ge=0, description="ID of the target email")
    target_folder: Optional[str] = Field(None)
    reason: Optional[str] = Field(None, max_length=300)

    @model_validator(mode="after")
    def validate_move_has_folder(self) -> "Action":
        if self.action_type == "move" and not self.target_folder:
            raise ValueError("target_folder is required for 'move' action")
        if self.target_folder and self.target_folder not in AVAILABLE_FOLDERS:
            raise ValueError(
                f"target_folder must be one of {AVAILABLE_FOLDERS}, "
                f"got '{self.target_folder}'"
            )
        return self


# ── Reward ───────────────────────────────────────────────────────────────────

class Reward(BaseModel):
    """Graded feedback returned after each step."""

    value: float = Field(..., ge=0.0, le=1.0, description="Overall reward [0, 1]")
    components: Dict[str, float] = Field(
        default_factory=dict,
        description="Breakdown of reward by component",
    )
    reason: str = Field("", description="Human-readable explanation")
    is_terminal: bool = Field(
        False, description="Whether this reward ends the episode"
    )


# ── StepRecord ───────────────────────────────────────────────────────────────

class StepRecord(BaseModel):
    """A single step's data, stored for episode history and grading."""

    step_num: int
    action: Action
    email: Email
    reward: Reward
    done: bool