dnd-helpers/src/llm/models.py

from typing import List, Optional

from pydantic import BaseModel, Field


class LoreUpdate(BaseModel):
    category: str = Field(
        ...,
        description="Category of lore: 'NPC', 'Location', 'WorldBuilding', or 'Plot'",
    )
    entity_name: Optional[str] = Field(
        None, description="The name of the NPC, Location, or entity being updated"
    )
    content: str = Field(..., description="The actual lore fact or update")
    context: Optional[str] = Field(
        None, description="Brief context from the conversation that led to this update"
    )


class InventoryChange(BaseModel):
    item: str = Field(..., description="Name of the item")
    quantity: int = Field(1, description="Quantity of the item")
    action: str = Field(..., description="Either 'added' or 'removed'")


class CharacterStateUpdate(BaseModel):
    character_name: str = Field(
        ..., description="Name of the character whose state is changing"
    )
    hp_change: Optional[int] = Field(
        None, description="Change in HP (negative for damage, positive for healing)"
    )
    status_effects_added: List[str] = Field(
        default_factory=list,
        description="List of status effects applied to the character",
    )
    status_effects_removed: List[str] = Field(
        default_factory=list,
        description="List of status effects removed from the character",
    )
    inventory_changes: List[InventoryChange] = Field(
        default_factory=list,
        description="List of items added or removed from inventory",
    )


class ContextUpdate(BaseModel):
    query: str = Field(..., description="The search query used to retrieve the context")
    snippet: str = Field(
        ..., description="The relevant text snippet retrieved from the source"
    )
    source: str = Field(
        ..., description="The source of the snippet (e.g., 'PHB p. 12')"
    )


class FilterResult(BaseModel):
    contextual_info: str = Field(
        ...,
        description="Information interesting to the user but not useful for structured extraction",
    )
    filtered_text: str = Field(
        ..., description="Cleaned transcript used for structured data extraction"
    )


class ExtractionResult(BaseModel):
    lore_updates: List[LoreUpdate] = Field(
        default_factory=list, description="List of discovered lore facts", alias="lore"
    )
    character_updates: List[CharacterStateUpdate] = Field(
        default_factory=list,
        description="List of character state changes",
        alias="character_state",
    )
    significant_events: List[str] = Field(
        default_factory=list,
        description="List of significant plot points or events",
        alias="events",
    )
    context_updates: List[ContextUpdate] = Field(
        default_factory=list,
        description="List of context updates",
        alias="context",
    )

    class Config:
        populate_by_name = True