models.py

"""
PII-Scrub-Assistant-v1 -- Pydantic Models
==========================================
OpenEnv 2026 Sync Protocol schemas.
"""
from __future__ import annotations
from typing import Optional
from pydantic import BaseModel, Field


class ScrubActionPayload(BaseModel):
    """Wrapped-int action payload."""
    action_id: int = Field(..., description="Protocol-required int. 1=submit, 0=skip.")
    redacted_text: str = Field(..., description="Scrubbed text with PII replaced by [REDACTED].")


class ScrubAction(BaseModel):
    """Top-level wrapped action envelope."""
    action: ScrubActionPayload


class ScrubObservation(BaseModel):
    """Returned after reset() or step()."""
    task_id: str = Field(..., description="Current task identifier.")
    original_text: str = Field(..., description="Raw email to redact.")
    instruction: str = Field(..., description="Task instruction.")
    score: Optional[float] = Field(None, ge=0.0, le=1.0)
    reward: Optional[float] = None
    done: bool = False
    feedback: Optional[str] = None


class ScrubState(BaseModel):
    """Full environment state snapshot."""
    current_task_id: Optional[str] = None
    player_id: Optional[str] = None
    session_id: Optional[str] = None
    step_count: int = 0
    done: bool = False
    last_score: Optional[float] = None
    last_reward: Optional[float] = None
    original_text: Optional[str] = None


class ResetRequest(BaseModel):
    """POST /reset body."""
    player_id: str = Field(..., description="Agent ID for turn validation.")
    session_id: str = Field(..., description="Session token for dedup.")
    task_id: Optional[str] = Field(None, description="task_1|task_2|task_3")


class StepRequest(BaseModel):
    """POST /step body."""
    player_id: str = Field(..., description="Must match /reset player_id.")
    session_id: str = Field(..., description="Must match /reset session_id.")
    action: ScrubActionPayload