Tolan/backend/models.py

"""
数据模型定义
"""
from typing import Literal, Optional, Dict
from pydantic import BaseModel


# 消息类型
MessageType = Literal[
    "audio_chunk",
    "audio_end",
    "cancel",
    "text_result",
    "emotion",
    "audio_response",
    "error",
    "ping",
    "pong",
    "energy_stone_spawn",
    "energy_stone_event",
]

# 情感类型
EmotionType = Literal[
    # 基础情感
    "neutral",      # 中性/平静
    "happy",        # 开心/愉悦
    "sad",          # 难过/悲伤
    "excited",      # 兴奋/激动
    "angry",        # 生气/愤怒
    "surprised",    # 惊讶/意外
    "thoughtful",   # 深思/沉思
    "encouraging",  # 鼓励/支持
    # 扩展情感
    "confused",     # 困惑/疑惑
    "worried",      # 担忧/焦虑
    "satisfied",    # 满意/满足
    "friendly",     # 友好/亲切
    "sorry",        # 抱歉/歉意
]

# 错误代码
ErrorCode = Literal[
    "STT_ERROR",
    "LLM_ERROR",
    "TTS_ERROR",
    "NETWORK_ERROR",
    "RATE_LIMIT",
    "INVALID_AUDIO",
    "CANCELLED",
]


class BaseMessage(BaseModel):
    """基础消息"""
    type: MessageType
    requestId: str
    timestamp: int


class AudioChunkMessage(BaseMessage):
    """音频数据块消息"""
    type: Literal["audio_chunk"] = "audio_chunk"
    data: str
    sampleRate: int = 16000
    channels: int = 1
    bitDepth: int = 16
    chunkIndex: int


class AudioEndMessage(BaseMessage):
    """音频结束消息"""
    type: Literal["audio_end"] = "audio_end"
    totalChunks: int
    speechDuration: Optional[int] = None  # 说话时长（毫秒）


class CancelMessage(BaseMessage):
    """取消消息"""
    type: Literal["cancel"] = "cancel"


class TextResultMessage(BaseMessage):
    """文本结果消息"""
    type: Literal["text_result"] = "text_result"
    text: str
    isFinal: bool


class EmotionMessage(BaseMessage):
    """情感消息"""
    type: Literal["emotion"] = "emotion"
    emotion: EmotionType
    intensity: Optional[float] = None


class AudioResponseMessage(BaseMessage):
    """音频响应消息"""
    type: Literal["audio_response"] = "audio_response"
    data: str
    sampleRate: int = 24000  # 流式 TTS 使用 24kHz
    channels: int = 1
    bitDepth: int = 16
    format: Literal["pcm", "wav"] = "pcm"  # 音频格式：pcm（流式）或 wav（完整）
    chunkIndex: int
    isLast: bool


class ErrorMessage(BaseMessage):
    """错误消息"""
    type: Literal["error"] = "error"
    code: ErrorCode
    message: str


class PingMessage(BaseMessage):
    """Ping消息"""
    type: Literal["ping"] = "ping"


class PongMessage(BaseMessage):
    """Pong消息"""
    type: Literal["pong"] = "pong"


StoneType = Literal["mission", "conversation"]
EnergyStoneEvent = Literal["click", "timeout"]


class EnergyStoneSpawnMessage(BaseMessage):
    """能量石生成消息"""
    type: Literal["energy_stone_spawn"] = "energy_stone_spawn"
    stoneId: str
    stoneType: StoneType
    ttlSeconds: int
    appearAt: int
    metadata: Optional[Dict] = None


class EnergyStoneEventMessage(BaseMessage):
    """能量石交互事件"""
    type: Literal["energy_stone_event"] = "energy_stone_event"
    stoneId: str
    event: EnergyStoneEvent
    clientTs: Optional[int] = None
    position: Optional[Dict] = None