src/agent/models.py

"""Pydantic models representing game state and LLM outputs."""

from typing import Dict, List, Optional, Set

from pydantic import BaseModel, Field


class Milestone(BaseModel):
    """Milestone that can be achieved during the story."""

    id: str
    description: str


class Ending(BaseModel):
    """Possible game ending."""

    id: str
    type: str  # "good" or "bad"
    condition: str
    description: Optional[str] = None


class StoryFrame(BaseModel):
    """Overall plot information generated by the LLM."""

    lore: str
    goal: str
    milestones: List[Milestone]
    endings: List[Ending]
    setting: str
    character: Dict[str, str]
    genre: str


class StoryFrameLLM(BaseModel):
    """Structure returned by the LLM for story frame generation."""

    lore: str
    goal: str
    milestones: List[Milestone]
    endings: List[Ending]


class SceneChoice(BaseModel):
    """User choice leading to another scene."""

    text: str
    next_scene_short_desc: str


class PlayerOption(BaseModel):
    """Option presented to the player in a scene."""

    option_description: str = Field(
        description=(
            "Description of the option, e.g. '[Say] Hello!' or "
            "'Go to the forest'"
        )
    )


class Scene(BaseModel):
    """Game scene with choices and optional assets."""

    scene_id: str
    description: str
    choices: List[SceneChoice]
    image: Optional[str] = None
    music: Optional[str] = None


class SceneLLM(BaseModel):
    """Structure expected from the LLM when generating a scene."""

    description: str
    choices: List[SceneChoice]


class EndingCheckResult(BaseModel):
    """Result returned from the LLM when checking for an ending."""

    ending_reached: bool = Field(default=False)
    ending: Optional[Ending] = None


class UserChoice(BaseModel):
    """Single player choice recorded in the history."""

    scene_id: str
    choice_text: str
    timestamp: Optional[str] = None


class UserState(BaseModel):
    """State stored for each user."""

    story_frame: Optional[StoryFrame] = None
    current_scene_id: Optional[str] = None
    scenes: Dict[str, Scene] = Field(default_factory=dict)
    milestones_achieved: Set[str] = Field(default_factory=set)
    user_choices: List[UserChoice] = Field(default_factory=list)
    ending: Optional[Ending] = None
    assets: Dict[str, str] = Field(default_factory=dict)