models.py

from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Dict, Any

class Email(BaseModel):
    id: str = Field(..., description="Unique identifier for the email")
    sender: str = Field(..., description="Email address of the sender")
    subject: str = Field(..., description="Subject line of the email")
    body: str = Field(..., description="Content of the email")
    is_read: bool = Field(False, description="Whether the email has been read")
    folder: str = Field("inbox", description="The folder the email is currently in")
    is_flagged: bool = Field(False, description="Whether the email is flagged as urgent/important")

class Observation(BaseModel):
    inbox_summary: List[Dict[str, Any]] = Field(
        ..., description="List of emails in the current folder (abbreviated view: id, sender, subject, is_read, is_flagged)"
    )
    current_folder: str = Field(..., description="The folder currently being viewed")
    opened_email: Optional[Email] = Field(None, description="The full email currently open, or null if none is open")
    last_action_status: str = Field("Environment initialized.", description="Feedback from the last action taken")

class Action(BaseModel):
    action_type: Literal[
        "open_email", "close_email", "move_email", 
        "reply", "delete_email", "flag_email", "submit"
    ] = Field(..., description="The type of action to perform.")
    email_id: Optional[str] = Field(None, description="ID of the email to act on (required for open, move, reply, delete, flag)")
    folder_name: Optional[str] = Field(None, description="Folder to move the email to (required for move_email)")
    reply_body: Optional[str] = Field(None, description="Body of the reply (required for reply)")

class RewardInfo(BaseModel):
    score: float = Field(..., description="Reward score from 0.0 to 1.0")
    is_done: bool = Field(..., description="Whether the episode has terminated")
    metrics: Dict[str, Any] = Field(default_factory=dict, description="Detailed metrics about performance")