ai-town/engine-python/app/models.py

from __future__ import annotations
from pydantic import BaseModel, Field
from typing import List, Dict, Optional
from enum import Enum


class Emotion(str, Enum):
    CALM = "calm"
    HAPPY = "happy"
    ANXIOUS = "anxious"


class Weather(str, Enum):
    SUNNY = "sunny"
    RAINY = "rainy"


class GlobalMeter(BaseModel):
    """全体能量条"""
    value: int = 0
    threshold: int = 100
    cooldown: int = 0


class WorldEffect(BaseModel):
    """持续影响效果"""
    type: str
    name: str
    intensity: int = 1
    remaining_ticks: int = 5
    mood_modifier: int = 0


class Opinion(BaseModel):
    """角色对事件的观点"""
    about: str  # 事件类型 (effect_type)
    text: str   # 观点内容
    tick: int   # 生成时的 tick


class Stance(BaseModel):
    """角色立场"""
    optimism: float = Field(default=0.5, ge=0.0, le=1.0)
    fear: float = Field(default=0.5, ge=0.0, le=1.0)


class FactionData(BaseModel):
    """单个派系的数据"""
    power: int = 0
    threshold: int = 10
    skill: str = ""
    members: List[str] = Field(default_factory=list)


class Factions(BaseModel):
    """派系分布（带 power 和 members）"""
    optimists: FactionData = Field(default_factory=FactionData)
    fearful: FactionData = Field(default_factory=FactionData)


class Votes(BaseModel):
    """投票统计"""
    optimists: int = 0
    fearful: int = 0
    # 记录本 tick 已投票的用户（用于去重）
    voted_users: List[str] = Field(default_factory=list)


class FactionSkill(BaseModel):
    """单个技能"""
    unlocked: bool = False
    cost: int = 10
    effect: str = ""
    requires: List[str] = Field(default_factory=list)


class FactionSkillTree(BaseModel):
    """单个阵营的技能树"""
    skills: Dict[str, FactionSkill] = Field(default_factory=dict)


class FactionSkills(BaseModel):
    """所有阵营的技能树"""
    optimists: FactionSkillTree = Field(default_factory=FactionSkillTree)
    fearful: FactionSkillTree = Field(default_factory=FactionSkillTree)


class FactionEventResult(BaseModel):
    """阵营事件触发结果"""
    type: Optional[str] = None  # "festival" | "panic" | None
    source_faction: Optional[str] = None  # "optimists" | "fearful"


class StoryArc(BaseModel):
    """剧情线"""
    progress: float = Field(default=0.0, ge=0.0, le=1.0)
    threshold: float = 1.0
    active: bool = True
    stage: int = 1  # 当前阶段
    description: str = ""


class StoryEventResult(BaseModel):
    """剧情事件触发结果"""
    triggered: bool = False
    arc_id: Optional[str] = None
    event_name: Optional[str] = None
    description: Optional[str] = None


class ActionFeedback(BaseModel):
    """行动反馈"""
    success: bool = True
    reason: str = ""
    remaining_ap: int = 0
    user: str = ""


class AgentState(BaseModel):
    emotion: Emotion = Emotion.CALM
    goal: str = ""
    memory: List[str] = Field(default_factory=list)
    opinion: Optional[Opinion] = None
    # 记录已评论过的事件，防止重复生成
    commented_effects: List[str] = Field(default_factory=list)
    # 角色立场
    stance: Stance = Field(default_factory=Stance)
    # 所属阵营
    faction: str = "neutral"  # "optimists" | "fearful" | "neutral"
    # 行动点系统
    action_points: int = Field(default=3, ge=0)
    max_action_points: int = Field(default=3, ge=1)
    last_action_tick: int = 0


class WorldState(BaseModel):
    tick: int = 0
    weather: Weather = Weather.SUNNY
    town_mood: int = Field(default=0, ge=-10, le=10)
    agents: Dict[str, AgentState] = Field(default_factory=dict)
    events: List[str] = Field(default_factory=list)
    global_meter: GlobalMeter = Field(default_factory=GlobalMeter)
    world_effects: List[WorldEffect] = Field(default_factory=list)
    factions: Factions = Field(default_factory=Factions)
    story_arcs: Dict[str, StoryArc] = Field(default_factory=dict)
    votes: Votes = Field(default_factory=Votes)
    faction_skills: FactionSkills = Field(default_factory=FactionSkills)


class Event(BaseModel):
    type: str
    text: str
    user: str
    ts: float
    faction: Optional[str] = None  # 用于投票事件: "optimists" | "fearful"


class StepRequest(BaseModel):
    events: List[Event] = Field(default_factory=list)


class Action(BaseModel):
    agent_id: str
    say: str
    do: str


class GlobalEventInfo(BaseModel):
    """世界级事件信息"""
    name: str
    description: str


class GlobalEventResult(BaseModel):
    """世界级事件触发结果"""
    triggered: bool = False
    event: Optional[GlobalEventInfo] = None


class StepResponse(BaseModel):
    world_state: WorldState
    actions: List[Action]
    global_event: GlobalEventResult = Field(default_factory=GlobalEventResult)
    triggered_faction_event: FactionEventResult = Field(default_factory=FactionEventResult)
    story_event: StoryEventResult = Field(default_factory=StoryEventResult)
    action_feedbacks: List[ActionFeedback] = Field(default_factory=list)