agents/master/policy.py

from __future__ import annotations

from typing import Protocol

from pydantic import Field

from agents.shared.llm_client import StructuredModelClient
from agents.shared.model_schema import StrictModel

from .schema import WorldDefinition


class DungeonMasterPolicyError(RuntimeError):
    pass


class DungeonMasterPolicy(Protocol):
    def generate_world(
        self,
        *,
        target_ratio: float,
        repair_context: "DMRepairContext | None" = None,
    ) -> WorldDefinition:
        ...


class DMRepairContext(StrictModel):
    attempt_number: int
    error_message: str
    previous_candidate_json: str | None = None


class WinConditionCandidate(StrictModel):
    type: str
    target_npc_id: str
    answer_string: str


class WorldMetaCandidate(StrictModel):
    title: str
    difficulty_target: float
    start_node_id: str
    win_condition: WinConditionCandidate


class WorldNodeCandidate(StrictModel):
    id: str
    type: str
    label: str
    description: str
    parent_id: str | None = None
    open: bool | None = None
    locked: bool | None = None
    lock_key_id: str | None = None
    clue_id: str | None = None
    requires_item_id: str | None = None
    consumes_item: bool | None = None
    text_content: str | None = None
    reveals_item_id: str | None = None
    reveals_readable_id: str | None = None
    gives_item_id: str | None = None
    gives_clue_id: str | None = None


class EdgeCandidate(StrictModel):
    id: str
    from_node_id: str
    to_node_id: str
    direction: str
    type: str
    required_item_id: str | None = None
    door_node_id: str | None = None


class ItemCandidate(StrictModel):
    id: str
    label: str
    description: str
    subtype: str
    start_node_id: str | None = None


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


class RecipeCandidate(StrictModel):
    id: str
    input_item_ids: list[str]
    output_item_id: str


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


class WorldDefinitionCandidate(StrictModel):
    meta: WorldMetaCandidate
    nodes: list[WorldNodeCandidate]
    edges: list[EdgeCandidate]
    items: list[ItemCandidate]
    clues: list[ClueCandidate]
    recipes: list[RecipeCandidate] = Field(default_factory=list)
    quest_chain: list[QuestStepCandidate]


class DungeonMasterLLMPolicy:
    def __init__(
        self,
        client: StructuredModelClient,
        *,
        model_name: str,
        temperature: float = 0.0,
        max_output_tokens: int = 8192,
    ) -> None:
        self.client = client
        self.model_name = model_name
        self.temperature = temperature
        self.max_output_tokens = max_output_tokens

    def generate_world(
        self,
        *,
        target_ratio: float,
        repair_context: DMRepairContext | None = None,
    ) -> WorldDefinition:
        from .prompt import build_dm_world_messages

        try:
            candidate = self.client.generate_structured(
                build_dm_world_messages(target_ratio=target_ratio, repair_context=repair_context),
                WorldDefinitionCandidate,
                model_name=self.model_name,
                temperature=self.temperature,
                max_output_tokens=self.max_output_tokens,
            )
            return WorldDefinition.model_validate(candidate.model_dump(mode="json", exclude_none=True))
        except Exception as exc:
            raise DungeonMasterPolicyError(self._normalize_error(exc)) from exc

    @staticmethod
    def _normalize_error(exc: Exception) -> str:
        return " ".join(str(exc).split()) or exc.__class__.__name__