agents/master/schema.py

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Annotated, Literal, TypeAlias

from pydantic import BaseModel, ConfigDict, Field

from agents.shared.openenv_compat import Action, Observation, State


class StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class WorldMeta(StrictModel):
    title: str
    difficulty_target: float
    start_node_id: str
    win_condition: "WinCondition"


class WinCondition(StrictModel):
    type: Literal["deduce"]
    target_npc_id: str
    answer_string: str


class BaseNode(StrictModel):
    id: str
    label: str
    description: str


class LocationNode(BaseNode):
    type: Literal["location"]


class JunctionNode(BaseNode):
    type: Literal["junction"]


class ContainerNode(BaseNode):
    type: Literal["container"]
    parent_id: str
    open: bool = False
    locked: bool = False
    lock_key_id: str | None = None


class DoorNode(BaseNode):
    type: Literal["door"]
    open: bool = False
    locked: bool = False
    lock_key_id: str | None = None


class ReadableNode(BaseNode):
    type: Literal["readable"]
    parent_id: str
    clue_id: str
    requires_item_id: str | None = None
    consumes_item: bool = False
    text_content: str


class FixtureNode(BaseNode):
    type: Literal["fixture"]
    parent_id: str
    requires_item_id: str
    reveals_item_id: str | None = None
    reveals_readable_id: str | None = None
    consumes_item: bool = False


class NpcNode(BaseNode):
    type: Literal["npc"]
    parent_id: str
    requires_item_id: str | None = None
    gives_item_id: str | None = None
    gives_clue_id: str | None = None


WorldNode: TypeAlias = Annotated[
    LocationNode | JunctionNode | ContainerNode | DoorNode | ReadableNode | FixtureNode | NpcNode,
    Field(discriminator="type"),
]


class Edge(StrictModel):
    id: str
    from_node_id: str
    to_node_id: str
    direction: Literal["north", "south", "east", "west", "up", "down", "in", "out"]
    type: Literal["passage", "locked_passage"]
    required_item_id: str | None = None
    door_node_id: str | None = None


class Item(StrictModel):
    id: str
    label: str
    description: str
    subtype: Literal["key", "puzzle"]
    start_node_id: str | None = None


class Clue(StrictModel):
    id: str
    text: str


class Recipe(StrictModel):
    id: str
    input_item_ids: list[str] = Field(min_length=2, max_length=2)
    output_item_id: str


class QuestStep(StrictModel):
    step_id: str
    description: str
    requires_step_ids: list[str] = Field(default_factory=list)
    action: str


class WorldDefinition(StrictModel):
    meta: WorldMeta
    nodes: list[WorldNode]
    edges: list[Edge]
    items: list[Item]
    clues: list[Clue]
    recipes: list[Recipe] = Field(default_factory=list)
    quest_chain: list[QuestStep]


class DMAction(Action):
    world_definition: WorldDefinition


class Turn(StrictModel):
    step: int
    player_action: str
    textworld_command: str
    observation: str
    game_state_delta: dict[str, object]


class DMFeedback(StrictModel):
    unreachable_nodes: list[str]
    unused_items: list[str]
    clues_missed: list[str]
    mean_steps_per_room: float
    invalid_command_count: int = 0
    wrong_submit_count: int = 0


class DMRewardBreakdown(StrictModel):
    reward_mode: Literal["gaussian_target_ratio", "compile_failure_penalty"] = "gaussian_target_ratio"
    player_won: bool
    raw_ratio: float | None = None
    clamped_ratio: float | None = None
    target_ratio: float
    target_ratio_delta: float | None = None
    efficiency_score: float | None = None
    quality_score: float = 0.0
    reward: float


class DMObservation(Observation):
    episode_transcript: list[Turn] = Field(default_factory=list)
    player_won: bool | None = None
    steps_taken: int | None = None
    min_steps: int | None = None
    ratio: float | None = None
    compile_error: str | None = None
    feedback: DMFeedback | None = None
    reward_breakdown: DMRewardBreakdown | None = None
    target_ratio_used: float | None = None


class DMState(State):
    current_world: WorldDefinition | None = None
    compile_status: Literal["valid", "invalid", "pending"] = "pending"
    episode_status: Literal["running", "complete", "failed"] = "running"
    cumulative_success_rate: float = 0.0
    target_ratio: float = 0.0
    difficulty_hint: float | None = None


@dataclass(frozen=True)
class GoAction:
    target_node_id: str


@dataclass(frozen=True)
class OpenAction:
    target_node_id: str


@dataclass(frozen=True)
class UnlockAction:
    door_id: str
    key_id: str


@dataclass(frozen=True)
class TakeAction:
    item_id: str
    source_node_id: str


@dataclass(frozen=True)
class ReadAction:
    target_node_id: str


@dataclass(frozen=True)
class UseAction:
    item_id: str
    target_node_id: str


@dataclass(frozen=True)
class CombineAction:
    item_a_id: str
    item_b_id: str


@dataclass(frozen=True)
class GiveAction:
    item_id: str
    npc_id: str


@dataclass(frozen=True)
class TalkAction:
    target_node_id: str


@dataclass(frozen=True)
class SubmitAction:
    answer_text: str


QuestAction = (
    GoAction
    | OpenAction
    | UnlockAction
    | TakeAction
    | ReadAction
    | UseAction
    | CombineAction
    | GiveAction
    | TalkAction
    | SubmitAction
)


@dataclass(frozen=True)
class NpcTrade:
    required_item_id: str
    gives_item_id: str | None
    gives_clue_id: str | None


@dataclass(frozen=True)
class UseEffect:
    required_item_id: str
    clue_id: str | None = None
    reveals_item_id: str | None = None
    reveals_readable_id: str | None = None
    consumes_item: bool = False


@dataclass
class CompiledWorld:
    episode_id: str
    world: WorldDefinition
    artifacts_dir: Path
    game_file: Path
    walkthrough_commands: list[str]
    solver_policy: list[str]
    correct_answer_normalized: str
    correct_submit_command: str
    guardian_id: str
    guardian_room_id: str
    room_name_to_id: dict[str, str]
    node_command_names: dict[str, str]
    item_command_names: dict[str, str]
    item_start_locations: dict[str, str | None]
    clue_text_by_id: dict[str, str]
    readable_clue_by_id: dict[str, str]
    npc_trade_map: dict[str, NpcTrade]
    recipe_map: dict[frozenset[str], str]
    use_effects: dict[str, UseEffect]
    produced_item_ids: set[str]
    room_edges_by_target: dict[tuple[str, str], Edge]
    room_edges_by_direction: dict[tuple[str, str], Edge]
    door_rooms: dict[str, frozenset[str]]


@dataclass
class SimulationState:
    current_room_id: str
    inventory: set[str] = field(default_factory=set)
    item_locations: dict[str, str | None] = field(default_factory=dict)
    open_nodes: set[str] = field(default_factory=set)
    locked_nodes: set[str] = field(default_factory=set)
    discovered_clues: set[str] = field(default_factory=set)
    consulted_npcs: set[str] = field(default_factory=set)
    satisfied_npcs: set[str] = field(default_factory=set)
    revealed_readables: set[str] = field(default_factory=set)
    prepared_readables: set[str] = field(default_factory=set)
    used_fixtures: set[str] = field(default_factory=set)
    produced_items: set[str] = field(default_factory=set)
    visited_nodes: set[str] = field(default_factory=set)