models.py

"""Action and observation models for Subtext Arena."""
from __future__ import annotations

from typing import Any, Dict, List, Optional

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


class SubtextArenaAction(Action):
    """Single agent action: pick a tool, optionally with arguments."""

    tool: str = Field(
        ...,
        description=(
            "One of: get_transcript, get_prosody_features, "
            "get_pitch_contour, submit_belief"
        ),
    )
    tool_args: Dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Tool-specific arguments. "
            "get_prosody_features / get_pitch_contour: optional {start: float, end: float} in seconds. "
            "submit_belief: {label: 'sarcastic'|'sincere', confidence: float}."
        ),
    )


class SubtextArenaObservation(Observation):
    """Observation returned after each action."""

    clip_id: str = Field(default="", description="MUStARD clip identifier (e.g. '1_60')")
    speaker: str = Field(default="", description="Speaker name when available (Friends/BBT)")
    duration_s: float = Field(default=0.0, description="Total clip duration in seconds")
    is_pivot: bool = Field(default=False, description="True if clip is in the Prosody-Pivot Set")

    tool_used: str = Field(default="", description="Which tool the agent just invoked")
    tool_output: str = Field(default="", description="Text output of the tool call")

    step: int = Field(default=0, description="Current step in this episode (0-indexed)")
    max_steps: int = Field(default=6, description="Hard cap on tool calls before forced submission")
    audio_calls_so_far: int = Field(
        default=0,
        description="How many of the prior calls were audio-tool calls (used by reward)",
    )

    error: Optional[str] = Field(default=None, description="Error message when the agent's action was malformed")