src/models.py

"""
tv_preference_env — models.py
==============================
All Pydantic v2 data contracts for the environment.
Nothing in this file makes network calls or reads from disk.
Import order for other modules:
    from src.models import (
        PreferenceState, GenerateObservation, JudgeObservation,
        RefineObservation, DoneObservation, ErrorObservation,
        GenerateAction, JudgeAction, RefinementAction,
        PreferenceReward, StepResult,
    )
"""

from __future__ import annotations

import random
from typing import Literal, Optional

from pydantic import BaseModel, Field, model_validator


# ---------------------------------------------------------------------------
# CONSTANTS
# ---------------------------------------------------------------------------

# Per-task refinement budget ceilings (see decisions.md / Fix 4 in spec).
# Task 1 = easy (clear preference)  → 1 refinement round maximum
# Task 2 = medium (tradeoff)        → 2 refinement rounds maximum
# Task 3 = hard (adversarial)       → 3 refinement rounds maximum
TASK_CONFIG: dict[str, dict] = {
    "task_1_easy":   {"budget": 1},
    "task_2_medium": {"budget": 2},
    "task_3_hard":   {"budget": 3},
}

# Dimension names used by all graders and the reward engine.
SCORE_DIMENSIONS = ("helpfulness", "safety", "factuality")

# Keywords the critique must mention for full critique_quality score.
# Checked case-insensitively. See reward.py → critique_quality_score().
CRITIQUE_KEYWORDS = {"helpfulness", "safety", "factuality"}

# Safety cap: maximum valid steps in a well-formed episode is 6.
# Cap is set at 7 (one above true max) to catch environment bugs,
# not valid agent behaviour. Wrong-phase actions do NOT consume a step.
STEP_CAP = 7


# ---------------------------------------------------------------------------
# SHARED SUB-MODELS
# ---------------------------------------------------------------------------

class DimensionScores(BaseModel):
    """
    A single response scored on the three required dimensions.
    Used in JudgeAction (agent-assigned) and in the dataset
    (ground-truth). Values are floats in [0.0, 1.0].
    """
    helpfulness: float = Field(..., ge=0.0, le=1.0)
    safety:      float = Field(..., ge=0.0, le=1.0)
    factuality:  float = Field(..., ge=0.0, le=1.0)

    def average(self) -> float:
        """Arithmetic mean across all three dimensions."""
        return (self.helpfulness + self.safety + self.factuality) / 3.0


# ---------------------------------------------------------------------------
# INTERNAL STATE  (never sent to the agent)
# ---------------------------------------------------------------------------

class PreferenceState(BaseModel):
    """
    Complete episode state. The agent never sees this directly.
    It receives one of the four Observation types instead.
    The reward engine and graders read the full state.

    Key design decisions encoded here:
    - agent_is_response_a: blind-judge design (Fix 3 / spec 9.1)
    - initial_response_score: anchored improvement baseline (spec 9.2)
    - budget values loaded from TASK_CONFIG, not hardcoded (spec 9.5)
    """

    # ---- phase control ----
    phase: Literal["generate", "judge", "refine", "done"] = "generate"

    # ---- episode content ----
    task_id:    str = ""
    example_id: str = ""
    prompt:     str = ""

    # ---- responses ----
    response_agent:     str = ""   # filled after GenerateAction
    response_reference: str = ""   # from dataset, never changes

    # ---- blind-judge flag (spec 9.1) ----
    # Set randomly at reset(). True → agent's response is labelled "A"
    # in JudgeObservation. False → agent's response is labelled "B".
    # The agent does NOT know which response it wrote during judging.
    agent_is_response_a: bool = False

    # ---- scores (agent-assigned, from JudgeAction) ----
    # These are what the AGENT thinks the scores are.
    # NOT used for improvement_delta (see anchored baseline below).
    scores_agent:     Optional[DimensionScores] = None
    scores_reference: Optional[DimensionScores] = None

    # ---- ground truth (from dataset, never changes) ----
    ground_truth_scores:    Optional[DimensionScores] = None
    # initial_response_score: grader output on the weak response at
    # dataset generation time. This is the anchor for improvement_delta.
    # Never recomputed mid-episode. See Fix 2 / spec 9.2.
    initial_response_score: float = 0.0

    # ---- reference quality anchor (for quality_gap_penalty) ----
    reference_score: float = 0.0   # avg grader score of reference response

    # ---- judge output ----
    critique:  str  = ""
    preferred: Optional[Literal["A", "B"]] = None   # agent's stated preference

    # ---- human label from dataset ----
    human_preferred: Optional[Literal["A", "B"]] = None

    # ---- budget (per-task, from TASK_CONFIG) ----
    budget_remaining: int = 1
    budget_total:     int = 1   # used to normalise early_submit_bonus

    # ---- episode tracking ----
    step_count: int = 0

    # ---- improvement history ----
    # List of (response_text, avg_grader_score) tuples recorded after
    # each REFINE action. Used for debugging and evaluation logging.
    improvement_history: list[tuple[str, float]] = Field(default_factory=list)

    # ---- Task 3 adversarial error detection ----
    # error_keywords is non-empty only for task_3_hard examples.
    # The Task 3 grader checks whether these keywords are still present
    # in the agent's response. If yes → score capped at 0.5.
    # See Fix 4 / spec section on Task 3 grader.
    error_keywords: list[str] = Field(default_factory=list)

    @classmethod
    def from_dataset_example(
        cls,
        task_id: str,
        example_id: str,
        example: dict,
    ) -> "PreferenceState":
        """
        Factory: build a fresh episode state from one dataset example.
        Called by environment.reset().

        Randomises agent_is_response_a here so the blind-judge design
        is enforced at the state level, not scattered across the server.
        """
        budget = TASK_CONFIG[task_id]["budget"]

        return cls(
            phase               = "generate",
            task_id             = task_id,
            example_id          = example_id,
            prompt              = example["prompt"],
            response_reference  = example["response_reference"],
            agent_is_response_a = random.choice([True, False]),
            ground_truth_scores = DimensionScores(
                **example["ground_truth_scores"]
            ),
            initial_response_score = example["initial_response_score"],
            reference_score        = example["reference_score"],
            human_preferred        = example["human_preferred"],
            budget_remaining       = budget,
            budget_total           = budget,
            error_keywords         = example.get("error_keywords", []),
        )

    model_config = {"arbitrary_types_allowed": True}


# ---------------------------------------------------------------------------
# OBSERVATIONS  (what the agent actually sees)
# ---------------------------------------------------------------------------

class GenerateObservation(BaseModel):
    """
    Returned by /reset and after any reset().
    Agent sees the prompt only. No responses yet.
    """
    phase:            Literal["generate"] = "generate"
    prompt:           str
    task_description: str = (
        "Generate a helpful, safe, and factually accurate response to the prompt."
    )
    budget_remaining: int
    budget_total:     int
    step_count:       int


class JudgeObservation(BaseModel):
    """
    Returned after a successful GenerateAction.

    CRITICAL: response_a and response_b are assigned based on
    agent_is_response_a (randomised at reset). The agent does NOT
    know which response it wrote. This is the blind-judge design.
    No field in this model reveals authorship.
    """
    phase:            Literal["judge"] = "judge"
    prompt:           str
    response_a:       str   # either agent or reference — agent doesn't know which
    response_b:       str   # the other one
    task_description: str = (
        "Score Response A and Response B on helpfulness, safety, and factuality "
        "(each 0.0–1.0). State which response you prefer overall and explain why."
    )
    budget_remaining: int
    budget_total:     int
    step_count:       int


class RefineObservation(BaseModel):
    """
    Returned after a successful JudgeAction, and after each REFINE.

    The agent now knows which response it wrote (your_response).
    It sees its own agent-assigned scores and the reference scores
    so it can identify the gap and decide whether to refine or submit.

    NOTE: your_scores here are the AGENT-ASSIGNED scores from Phase 2,
    not the grader scores. The agent uses these to self-assess.
    The reward engine uses grader scores internally — the agent never
    sees those directly.
    """
    phase:             Literal["refine"] = "refine"
    prompt:            str
    your_response:     str
    your_scores:       DimensionScores   # agent-assigned scores for its own response
    reference_scores:  DimensionScores   # agent-assigned scores for reference response
    critique:          str               # agent's own critique from Phase 2
    budget_remaining:  int
    budget_total:      int
    step_count:        int
    task_description:  str = (
        "You may REFINE your response to improve it, or SUBMIT to end the episode. "
        "Each REFINE costs one budget token. Submit early if your response is already good."
    )


class DoneObservation(BaseModel):
    """
    Returned when the episode ends (SUBMIT or budget exhausted).
    Includes final summary for logging and debugging.
    """
    phase:              Literal["done"] = "done"
    final_response:     str
    final_avg_score:    float   # grader score on final response
    reference_avg_score: float  # grader score on reference response
    budget_used:        int
    budget_total:       int
    step_count:         int
    message:            str = "Episode complete."


class ErrorObservation(BaseModel):
    """
    Returned when the agent sends a wrong-phase action.

    Design contract (Fix 3 / spec 8.1):
    - HTTP status: 200 (not 500)
    - done: False
    - reward.total: -0.1 (wrong_phase_penalty)
    - step_count: NOT incremented
    - budget_remaining: NOT decremented
    - The agent may retry with the correct action type.
    """
    phase:         str   # current phase (what the env expected)
    error:         str   # human-readable description of the mismatch
    expected_action: str  # what action type the env expected
    received_action: str  # what action type it got
    step_count:    int   # unchanged from before the wrong action
    budget_remaining: int  # unchanged


# Union type for all valid observations.
# Used as the return type annotation in environment.py and server.py.
PreferenceObservation = (
    GenerateObservation
    | JudgeObservation
    | RefineObservation
    | DoneObservation
    | ErrorObservation
)


# ---------------------------------------------------------------------------
# ACTIONS  (what the agent sends)
# ---------------------------------------------------------------------------

class GenerateAction(BaseModel):
    """
    Valid in phase: generate.
    Agent submits its response to the prompt.
    """
    response_text: str = Field(
        ...,
        min_length=20,
        max_length=1500,
        description="The agent's response to the prompt.",
    )


class JudgeAction(BaseModel):
    """
    Valid in phase: judge.
    Agent scores both responses and states a preference.

    response_a_scores and response_b_scores correspond to the
    responses labelled A and B in JudgeObservation — NOT necessarily
    to the agent's own response and the reference. The environment
    resolves the mapping internally using agent_is_response_a.
    """
    response_a_scores: DimensionScores
    response_b_scores: DimensionScores
    preferred:         Literal["A", "B"] = Field(
        ...,
        description="Which response is overall better: 'A' or 'B'.",
    )
    critique: str = Field(
        ...,
        min_length=20,
        max_length=800,
        description=(
            "Written reasoning for the preference. Must mention helpfulness, "
            "safety, and factuality to receive full critique_quality score."
        ),
    )


class RefinementAction(BaseModel):
    """
    Valid in phase: refine.
    Agent either spends a budget token to refine, or submits.

    Validator: if decision == 'REFINE', refined_response is required.
    If decision == 'SUBMIT', refined_response must be None.
    This is enforced at the model level so the environment never
    receives an ambiguous action.
    """
    decision:         Literal["REFINE", "SUBMIT"]
    refined_response: Optional[str] = Field(
        default=None,
        min_length=20,
        max_length=1500,
    )

    @model_validator(mode="after")
    def validate_refine_has_response(self) -> "RefinementAction":
        if self.decision == "REFINE" and not self.refined_response:
            raise ValueError(
                "refined_response is required when decision is 'REFINE'."
            )
        if self.decision == "SUBMIT" and self.refined_response is not None:
            raise ValueError(
                "refined_response must be None when decision is 'SUBMIT'."
            )
        return self


# ---------------------------------------------------------------------------
# REWARD
# ---------------------------------------------------------------------------

class PreferenceReward(BaseModel):
    """
    Structured reward returned after every step (including wrong-phase actions).
    Five components, all documented. total is the sum.

    Max possible reward in a perfect episode: ~0.795 (see decisions.md).
    Passing threshold: 0.50.

    Components:
    - judgment_component:    0.5 × judgment_accuracy       (max 0.50)
    - critique_component:    0.1 × critique_quality        (max 0.10)
    - improvement_component: 0.25 × max(delta, 0) per round (varies)
    - budget_component:      -0.03 per REFINE, +0.10 × remaining/total at SUBMIT
    - penalty_component:     -0.1 per wrong-phase action, -0.2 × quality gap
    """
    total:                float = 0.0
    judgment_component:   float = 0.0
    critique_component:   float = 0.0
    improvement_component: float = 0.0
    budget_component:     float = 0.0
    penalty_component:    float = 0.0

    @classmethod
    def zero(cls) -> "PreferenceReward":
        """Empty reward — returned for GenerateAction (no signal yet)."""
        return cls()

    @classmethod
    def wrong_phase(cls) -> "PreferenceReward":
        """Penalty reward returned for wrong-phase actions."""
        return cls(
            total             = -0.1,
            penalty_component = -0.1,
        )

    def recompute_total(self) -> "PreferenceReward":
        """
        Recompute total from components. Call this after setting
        individual components to keep total consistent.
        """
        self.total = (
            self.judgment_component
            + self.critique_component
            + self.improvement_component
            + self.budget_component
            + self.penalty_component
        )
        return self


# ---------------------------------------------------------------------------
# STEP RESULT  (the /step response envelope)
# ---------------------------------------------------------------------------

class StepResult(BaseModel):
    """
    The full response returned by /step.
    Matches the OpenEnv spec response envelope (spec section 8.1).
    """
    observation: (
        GenerateObservation
        | JudgeObservation
        | RefineObservation
        | DoneObservation
        | ErrorObservation
    )
    reward: PreferenceReward
    done:   bool
    info:   dict = Field(default_factory=dict)