Blame for models.py

"""
Pydantic models for the InformationIntegrity OpenEnv environment.
All models are strict and versioned for deterministic serialisation.
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field, model_validator


# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------

class Stance(str, Enum):
    SUPPORTING = "supporting"
    CONTRADICTING = "contradicting"
    NEUTRAL = "neutral"


class VerificationStatus(str, Enum):
    UNVERIFIED = "unverified"
    TRUE = "true"
    FALSE = "false"
    PARTIALLY_TRUE = "partially_true"
    UNKNOWN = "unknown"


class DecisionType(str, Enum):
    ALLOW = "allow"
    FLAG = "flag"
    REMOVE = "remove"


class ActionType(str, Enum):
    VERIFY_CLAIM = "verify_claim"
    CHECK_SOURCE = "check_source"
    DECIDE = "decide"


class TaskDifficulty(str, Enum):
    EASY = "easy"
    MEDIUM = "medium"
    HARD = "hard"


# ---------------------------------------------------------------------------
# Core domain models
# ---------------------------------------------------------------------------

class SourceInfo(BaseModel):
    """A single evidence source presented to the agent."""

    id: str = Field(..., description="Unique source identifier, e.g. 'src_001'")
    title: str = Field(..., description="Headline or document title")
    text: str = Field(..., description="Excerpt or full text snippet")
    url: str = Field(..., description="Canonical URL of the source")
    stance: Stance = Field(..., description="Relationship of source to the primary claim")
    reliability: float = Field(..., ge=0.0, le=1.0,
                               description="0=unreliable, 1=highly authoritative")
    source_type: str = Field(default="unknown",
                             description="official | news | blog | social | academic")

    class Config:
        use_enum_values = True


class Claim(BaseModel):
    """A single verifiable claim extracted from the content."""

    id: str = Field(..., description="Unique claim identifier, e.g. 'claim_001'")
    text: str = Field(..., description="Verbatim claim text")
    verification_status: VerificationStatus = Field(
        default=VerificationStatus.UNVERIFIED,
        description="Agent-updated verification result"
    )
    confidence: float = Field(
        default=0.0, ge=0.0, le=1.0,
        description="Agent confidence in the verification (0–1)"
    )
    verified_by_sources: List[str] = Field(
        default_factory=list,
        description="List of source IDs used to verify this claim"
    )

    class Config:
        use_enum_values = True


class StatusFlags(BaseModel):
    """Bookkeeping flags that track episode-level constraints."""

    decision_made: bool = False
    sufficient_evidence: bool = False
    premature_decision: bool = False
    redundant_actions: int = Field(default=0, ge=0)
    max_steps_reached: bool = False


class InteractionRecord(BaseModel):
    """One entry in the interaction history."""

    step: int
    action_type: str
    action_payload: Dict[str, Any]
    result: Dict[str, Any]
    reward_delta: float = 0.0


class Observation(BaseModel):
    """Full observable state returned to the agent after each step."""

    content: str = Field(..., description="The piece of content under review")
    claims: List[Claim] = Field(..., description="Extracted claims (may be partially verified)")
    available_sources: List[SourceInfo] = Field(..., description="Evidence pool for this scenario")
    interaction_history: List[InteractionRecord] = Field(default_factory=list)
    status_flags: StatusFlags = Field(default_factory=StatusFlags)
    step_count: int = Field(default=0, ge=0)
    episode_id: str = Field(default="", description="Unique identifier for this episode")
    task_difficulty: str = Field(default="easy")

    class Config:
        use_enum_values = True


# ---------------------------------------------------------------------------
# Action model
# ---------------------------------------------------------------------------

class Action(BaseModel):
    """Agent action submitted to the environment."""

    type: ActionType = Field(..., description="Which action to perform")
    claim_id: Optional[str] = Field(
        default=None,
        description="Required when type=verify_claim"
    )
    source_id: Optional[str] = Field(
        default=None,
        description="Required when type=check_source"
    )
    decision: Optional[DecisionType] = Field(
        default=None,
        description="Required when type=decide"
    )
    reasoning: Optional[str] = Field(
        default=None,
        description="Agent's free-text reasoning (logged, not evaluated directly)"
    )

    @model_validator(mode="after")
    def validate_action_payload(self) -> "Action":
        if self.type == ActionType.VERIFY_CLAIM and not self.claim_id:
            raise ValueError("claim_id is required for verify_claim actions")
        if self.type == ActionType.CHECK_SOURCE and not self.source_id:
            raise ValueError("source_id is required for check_source actions")
        if self.type == ActionType.DECIDE and not self.decision:
            raise ValueError("decision is required for decide actions")
        return self

    class Config:
        use_enum_values = True


# ---------------------------------------------------------------------------
# Reward models
# ---------------------------------------------------------------------------

class RewardBreakdown(BaseModel):
    """Itemised reward components for transparency."""

    verify_claim_reward: float = 0.0
    check_source_reward: float = 0.0
    multi_source_bonus: float = 0.0
    decision_reward: float = 0.0
    premature_penalty: float = 0.0
    redundant_penalty: float = 0.0
    inefficiency_penalty: float = 0.0
    aggression_penalty: float = 0.0

    @property
    def total(self) -> float:
        return (
            self.verify_claim_reward
            + self.check_source_reward
            + self.multi_source_bonus
            + self.decision_reward
            + self.premature_penalty
            + self.redundant_penalty
            + self.inefficiency_penalty
            + self.aggression_penalty
        )


class Reward(BaseModel):
    """Reward signal returned after each step."""

    score: float = Field(..., description="Clipped total reward for this step")
    breakdown: RewardBreakdown
    message: str = Field(default="", description="Human-readable reward explanation")
    is_terminal: bool = Field(default=False)

    class Config:
        use_enum_values = True


# ---------------------------------------------------------------------------
# Episode result (returned at episode end)
# ---------------------------------------------------------------------------

class EpisodeResult(BaseModel):
    """Accumulated statistics at the end of an episode."""

    episode_id: str
    task_difficulty: str
    scenario_index: int
    total_reward: float
    decision: Optional[str] = None
    ground_truth: str
    correct: bool
    steps_taken: int
    claims_verified: int
    sources_checked: int
    grader_score: float = Field(ge=0.0, le=1.0)
    reward_breakdown: RewardBreakdown
    explanation: Optional[str] = None

    class Config:
        use_enum_values = True


# ---------------------------------------------------------------------------
# Scenario (loaded from dataset JSON)
# ---------------------------------------------------------------------------

class ScenarioSource(BaseModel):
    """Source as stored in the dataset file."""

    id: str
    title: str
    text: str
    url: str
    stance: str
    reliability: float = Field(ge=0.0, le=1.0)
    source_type: str = "unknown"
    # Hidden from agent — used only by environment internals
    verification_weight: float = Field(default=1.0, ge=0.0, le=1.0)


class ScenarioClaim(BaseModel):
    """Claim as stored in the dataset file."""

    id: str
    text: str
    # Hidden ground truth for the verifier
    ground_truth_status: str
    ground_truth_confidence: float = Field(default=1.0, ge=0.0, le=1.0)


class Scenario(BaseModel):
    """A complete evaluation scenario loaded from a dataset file."""

    id: str
    difficulty: str
    content: str
    claims: List[ScenarioClaim]
    sources: List[ScenarioSource]
    ground_truth: str = Field(..., description="allow | flag | remove")
    topic: str = ""
    tags: List[str] = Field(default_factory=list)
    explanation_hint: str = Field(
        default="",
        description="Template hint for explanation generation"
    )