AIChatRoom/backend/models/chatroom.py

"""
聊天室数据模型
定义讨论聊天室的配置结构
"""
from datetime import datetime
from typing import Optional, Dict, Any, List
from enum import Enum
from pydantic import Field
from beanie import Document


class ChatRoomStatus(str, Enum):
    """聊天室状态枚举"""
    IDLE = "idle"               # 空闲，等待开始
    ACTIVE = "active"           # 讨论进行中
    PAUSED = "paused"           # 暂停
    COMPLETED = "completed"     # 已完成
    ERROR = "error"             # 出错


class ChatRoomConfig:
    """聊天室配置"""
    max_rounds: int = 50                  # 最大轮数（备用终止条件）
    message_history_size: int = 20        # 上下文消息数
    consensus_threshold: float = 0.8      # 共识阈值
    round_interval: float = 1.0           # 轮次间隔(秒)
    allow_user_interrupt: bool = True     # 允许用户中断
    
    def __init__(
        self,
        max_rounds: int = 50,
        message_history_size: int = 20,
        consensus_threshold: float = 0.8,
        round_interval: float = 1.0,
        allow_user_interrupt: bool = True
    ):
        self.max_rounds = max_rounds
        self.message_history_size = message_history_size
        self.consensus_threshold = consensus_threshold
        self.round_interval = round_interval
        self.allow_user_interrupt = allow_user_interrupt
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "max_rounds": self.max_rounds,
            "message_history_size": self.message_history_size,
            "consensus_threshold": self.consensus_threshold,
            "round_interval": self.round_interval,
            "allow_user_interrupt": self.allow_user_interrupt
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatRoomConfig":
        """从字典创建"""
        if not data:
            return cls()
        return cls(
            max_rounds=data.get("max_rounds", 50),
            message_history_size=data.get("message_history_size", 20),
            consensus_threshold=data.get("consensus_threshold", 0.8),
            round_interval=data.get("round_interval", 1.0),
            allow_user_interrupt=data.get("allow_user_interrupt", True)
        )


class ChatRoom(Document):
    """
    聊天室文档模型
    存储讨论聊天室的配置信息
    """
    room_id: str = Field(..., description="唯一标识")
    name: str = Field(..., description="聊天室名称")
    description: str = Field(default="", description="描述")
    objective: str = Field(default="", description="当前讨论目标")
    
    # 参与者
    agents: List[str] = Field(default_factory=list, description="Agent ID列表")
    moderator_agent_id: Optional[str] = Field(default=None, description="共识判断Agent ID")
    
    # 配置
    config: Dict[str, Any] = Field(
        default_factory=lambda: {
            "max_rounds": 50,
            "message_history_size": 20,
            "consensus_threshold": 0.8,
            "round_interval": 1.0,
            "allow_user_interrupt": True
        },
        description="聊天室配置"
    )
    
    # 状态
    status: str = Field(default=ChatRoomStatus.IDLE.value, description="当前状态")
    current_round: int = Field(default=0, description="当前轮次")
    current_discussion_id: Optional[str] = Field(default=None, description="当前讨论ID")
    
    # 元数据
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    completed_at: Optional[datetime] = Field(default=None, description="完成时间")
    
    class Settings:
        name = "chatrooms"
    
    def get_config(self) -> ChatRoomConfig:
        """获取配置对象"""
        return ChatRoomConfig.from_dict(self.config)
    
    def is_active(self) -> bool:
        """检查聊天室是否处于活动状态"""
        return self.status == ChatRoomStatus.ACTIVE.value
    
    class Config:
        json_schema_extra = {
            "example": {
                "room_id": "product-design-room",
                "name": "产品设计讨论室",
                "description": "用于讨论新产品功能设计",
                "objective": "设计一个用户友好的登录系统",
                "agents": ["product-manager", "designer", "developer"],
                "moderator_agent_id": "moderator",
                "config": {
                    "max_rounds": 50,
                    "message_history_size": 20,
                    "consensus_threshold": 0.8
                },
                "status": "idle",
                "current_round": 0
            }
        }