cace_env/models.py

"""
cace_env/models.py
Pydantic models for CACE — Action, Observation, State.
All inherit from openenv.core base classes as required by OpenEnv spec.
"""

from typing import Any, Dict, List, Optional
from openenv.core import Action, Observation, State


# ── Action ────────────────────────────────────────────────────────────────────

class CACEAction(Action):
    """
    The Decision Agent's moderation action.

    action_int: int in [0, 4]
        0 = ALLOW                  — content is legitimate
        1 = REMOVE                 — content violates policy
        2 = ALLOW_WITH_LABEL       — allow with warning label
        3 = ESCALATE               — refer to human reviewer
        4 = RESTRICT_DISTRIBUTION  — limit reach without removal

    V2 only: selected_indices
        Which posts to review (from 20 seeded posts, pick 8).
        Only used when mode="v2" and prioritisation is active.
    """
    action_int: int
    selected_indices: Optional[List[int]] = None  # V2 prioritisation

    @property
    def action_str(self) -> str:
        return {
            0: "ALLOW",
            1: "REMOVE",
            2: "ALLOW_WITH_LABEL",
            3: "ESCALATE",
            4: "RESTRICT_DISTRIBUTION",
        }.get(self.action_int, "ESCALATE")


# ── Observation ───────────────────────────────────────────────────────────────

class CACEObservation(Observation):
    """
    What the Decision Agent sees after reset() or step().
    Inherits: done, reward, metadata from openenv.core.Observation

    observation: The full enriched case prompt (post + 4-agent deliberation)
    case_id:     Unique Oversight Board case identifier
    language:    Detected language of the post
    region:      Geographic region
    complexity:  low | medium | high
    mode:        v1 (single case) | v2 (network batch)

    V2 only:
    batch_posts: List of 20 post summaries with spread signals
    """
    observation: str
    case_id: str
    language: str = "Unknown"
    region: str = "Unknown"
    complexity: str = "medium"
    culture_flag: bool = False
    mode: str = "unified"  # always unified

    # Network fields
    batch_posts: Optional[List[Dict[str, Any]]] = None
    network_step: Optional[int] = None

    # Reward breakdown (populated after step())
    reward_breakdown: Optional[Dict[str, Any]] = None


# ── State ─────────────────────────────────────────────────────────────────────

class CACEState(State):
    """
    Internal environment state — visible via state() endpoint.
    Inherits: episode_id, step_count from openenv.core.State

    Used for debugging, demo rendering, and training metrics.
    """
    case_id: str = ""
    post_text: str = ""
    language: str = "Unknown"
    region: str = "Unknown"
    policy_clause: str = "Unknown"
    cultural_brief: str = ""
    challenge_brief: str = ""
    policy_anchor: str = ""
    ground_truth: str = ""
    complexity: str = "medium"
    mode: str = "unified"
    total_episodes: int = 0
    correct_decisions: int = 0
    accuracy: float = 0.0
    avg_reward_last_50: float = 0.0

    # V2 network state
    network_nodes: Optional[int] = None
    network_edges: Optional[int] = None
    posts_in_batch: Optional[int] = None
    posts_selected: Optional[int] = None