models.py

"""
Data models for the TalkingHeadBench environment.

TalkingHeadBench evaluates diagnostic reasoning across 3 coupled sub-environments
for talking-head LoRA pipelines. Each episode has three agent-facing decision points:

  Step 1 — Image Diagnostician    (Node 1): ImageDiagnosticsAction
  Step 2 — Param Anomaly Detector (Node 2): ParamAnomalyAction
  Step 3 — Phoneme Risk Assessor  (Node 8): PhonemeRiskAction

All grading is deterministic and rule-based. No LLM judge required.
Final reward: 0.25 * subenv1 + 0.35 * subenv2 + 0.40 * subenv3
"""

from __future__ import annotations

from typing import Literal, Optional

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


# ---------------------------------------------------------------------------
# Shared sub-types (inlined to keep models.py self-contained)
# ---------------------------------------------------------------------------


# ---------------------------------------------------------------------------
# Step 1 — Image Diagnostician Action (Node 1)
# ---------------------------------------------------------------------------


class ImageDiagnosticsAction(Action):
    """
    Structured diagnosis of a reference image and prompt produced by
    the Image Diagnostician agent (Node 1).
    """

    regime_classification: Literal[
        "frontal_simple",
        "non_frontal",
        "complex_background",
        "occluded",
        "low_quality",
    ] = Field(..., description="Regime classification of the reference image.")

    identified_risk_factors: list[str] = Field(
        ..., description="Specific image or prompt risk factors detected."
    )
    prompt_issues: list[str] = Field(
        default_factory=list,
        description="Conflicting or weakly-anchored terms in the prompt.",
    )
    recommended_prompt_modifications: list[str] = Field(
        default_factory=list,
        description="Actionable suggestions to improve the prompt.",
    )
    image_usability_score: float = Field(
        ..., ge=0.0, le=1.0, description="Overall image usability in [0, 1]."
    )
    reasoning: str = Field(
        ..., description="Brief rationale for the regime and primary risk factor."
    )


# ---------------------------------------------------------------------------
# Step 2 — Parameter Anomaly Detector Action (Node 2)
# ---------------------------------------------------------------------------


class ParameterAnomaly(Action):
    """A single flagged parameter anomaly."""

    parameter: str = Field(..., description='Parameter name, e.g. "cfg" or "eta".')
    issue: str = Field(..., description="Description of the anomaly.")
    severity: Literal["minor", "moderate", "severe"] = Field(..., description="Severity level.")
    linked_failure_mode: str = Field(..., description="Associated generation failure mode.")


class DirectionalFix(Action):
    """A directional (non-prescriptive) fix recommendation."""

    target: str = Field(..., description="Parameter to fix.")
    direction: Literal[
        "increase", "decrease", "enable", "disable", "reconsider"
    ] = Field(..., description="Direction of the fix.")
    rationale: str = Field(..., description="Why this fix is recommended.")
    priority: Literal["critical", "recommended", "optional"] = Field(
        ..., description="Fix priority."
    )


class ParamAnomalyAction(Action):
    """
    Structured output of the Parameter Anomaly Detector agent (Node 2).
    """

    config_risk_level: Literal["safe", "marginal", "risky", "dangerous"] = Field(
        ..., description="Overall risk level of the proposed configuration."
    )
    anomalies: list[ParameterAnomaly] = Field(
        default_factory=list, description="Flagged parameter anomalies."
    )
    predicted_failure_modes: list[
        Literal[
            "identity_collapse",
            "reference_token_dropout",
            "temporal_jitter",
            "background_bleed",
            "lip_sync_desync",
            "pose_instability",
            "overexposure_artifacts",
        ]
    ] = Field(default_factory=list, description="Predicted generation failure modes.")
    directional_fixes: list[DirectionalFix] = Field(
        default_factory=list, description="Directional fix recommendations."
    )
    summary: str = Field(..., description="One-sentence risk summary.")


# ---------------------------------------------------------------------------
# Step 3 — Phoneme Risk Assessor Action (Node 8)
# ---------------------------------------------------------------------------


class PhonemeRiskEntry(Action):
    """Risk profile for a single phoneme."""

    phoneme: str = Field(..., description="Phoneme string, e.g. 'AH' or 'S'.")
    risk_score: float = Field(..., ge=0.0, le=1.0, description="Risk score in [0, 1].")
    risk_type: Literal[
        "identity_trigger",
        "expression_trigger",
        "motion_trigger",
        "artifact_trigger",
        "unknown_anomaly",
    ] = Field(..., description="Category of behavioral risk.")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence in [0, 1].")
    evidence: str = Field(..., description="Evidence justifying the risk score.")


class BehaviorTriggerPrediction(Action):
    """Predicted phoneme → behavior association."""

    trigger_phoneme: str = Field(..., description="Phoneme that triggers the behavior.")
    triggered_behavior: str = Field(
        ..., description='Behavior label, e.g. "smile" or "blink".'
    )
    association_strength: float = Field(..., ge=0.0, le=1.0)
    is_intended: bool = Field(..., description="Whether the association is intentional.")
    concern_level: Literal["none", "low", "medium", "high"] = Field(..., description="Concern level.")


class PhonemeCluster(Action):
    """A cluster of phonemes sharing a common risk pattern."""

    phonemes: list[str] = Field(..., description="Phonemes in the cluster.")
    cluster_risk_type: str = Field(..., description="Description of the shared risk pattern.")
    combined_risk_score: float = Field(..., ge=0.0, le=1.0)
    interaction_description: str = Field(..., description="How the phonemes interact.")


class MitigationRecommendation(Action):
    """A directional mitigation for a phoneme-level behavioral risk."""

    target: str = Field(..., description="Phoneme or layer to target.")
    action: Literal[
        "retrain_with_more_data",
        "remove_from_dataset",
        "add_counter_examples",
        "reduce_lora_rank",
        "apply_weight_regularization",
        "flag_for_manual_review",
    ] = Field(..., description="Recommended mitigation action.")
    rationale: str = Field(..., description="Why this mitigation is recommended.")
    priority: Literal["critical", "recommended", "optional"] = Field(..., description="Priority level.")


class PhonemeRiskAction(Action):
    """
    Structured behavioral risk profile produced by the Phoneme Risk Assessor (Node 8).
    """

    phoneme_risk_ranking: list[PhonemeRiskEntry] = Field(
        default_factory=list, description="Ranked list of at-risk phonemes."
    )
    predicted_behavior_triggers: list[BehaviorTriggerPrediction] = Field(
        default_factory=list, description="Predicted phoneme-to-behavior associations."
    )
    risky_phoneme_clusters: list[PhonemeCluster] = Field(
        default_factory=list, description="Clusters of phonemes with shared risk patterns."
    )
    model_behavioral_safety: Literal[
        "safe", "minor_concerns", "moderate_risk", "high_risk", "unsafe"
    ] = Field(..., description="Overall behavioral safety rating.")
    mitigation_recommendations: list[MitigationRecommendation] = Field(
        default_factory=list, description="Recommended mitigations."
    )
    summary: str = Field(..., description="One-paragraph summary of behavioral risks.")


# ---------------------------------------------------------------------------
# Shared Observation type
# ---------------------------------------------------------------------------


class TalkingHeadObservation(Observation):
    """
    Observation emitted by TalkingHeadBench at each step.

    The ``metadata`` dict carries the full observation payload:
      - ``node``:                  which node produced this observation
      - ``observation``:           the typed signal dict for the agent
      - ``expected_action_schema``: name of the Action class to return
      - ``instruction``:           natural-language task description
      - ``step``:                  current episode step index (0–3)
      - ``scores`` (on done=True): per-sub-env and final score breakdown
    """

    pass