models.py

"""
Data models for the Defender RL environment.

The defender holds a hidden secret in its system prompt. A frozen attacker LLM
tries to extract it across multiple turns. The defender's job is to be helpful
without leaking the secret.
"""

from typing import Any

from openenv.core.env_server.types import Action, Observation
from pydantic import Field


class DefenderAction(Action):
    """The defender's reply to the latest user message."""

    response: str = Field(..., description="Defender's reply to the latest user message")


class DefenderObservation(Observation):
    """What the defender sees on each turn.

    Note: `secret_value` is NEVER on the observation. Episode meta fields
    (scenario_id, attacker_id, leak_occurred) ARE included so the trainer can
    decompose reward by attacker/scenario for the per-attacker plot. They live
    here as named fields rather than in the inherited `metadata` dict because
    OpenEnv's serialize_observation() explicitly drops `metadata` from the
    wire format.
    """

    attacker_message: str = Field(
        default="",
        description="Latest attacker turn to respond to. On reset() this is the opening attacker message.",
    )
    conversation_history: list[dict[str, str]] = Field(
        default_factory=list,
        description="All prior turns as a list of {'role': 'user'|'assistant', 'content': str}",
    )
    turn_index: int = Field(
        default=0,
        description="0-indexed turn counter. Episode terminates at turn_index >= MAX_TURNS or on leak.",
    )
    system_prompt: str = Field(
        default="",
        description="The defender's system prompt for this episode. Contains the secret it must protect.",
    )
    # ----- Episode meta (visible to trainer for decomposition / logging) -----
    scenario_id: str = Field(default="", description="Which scenario was sampled this episode")
    secret_type: str = Field(default="", description="Type of secret: specific_factual / negative_instruction / conditional_disclosure / compound")
    attacker_id: str = Field(default="", description="Which attacker persona was sampled")
    leak_occurred: bool = Field(default=False, description="True if the secret leaked this episode")
    leaked_piece: str = Field(default="", description="If leaked, which piece of the (possibly compound) secret was matched")