edbddf855d
- 实现Agent管理，支持AI辅助生成系统提示词 - 支持多个AI提供商（OpenRouter、智谱、MiniMax等） - 实现聊天室和讨论引擎 - WebSocket实时消息推送 - 前端使用React + Ant Design - 后端使用FastAPI + MongoDB Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
"""
|
|
消息数据模型
|
|
定义聊天消息的结构
|
|
"""
|
|
from datetime import datetime
|
|
from typing import Optional, Dict, Any, List
|
|
from enum import Enum
|
|
from pydantic import Field
|
|
from beanie import Document
|
|
|
|
|
|
class MessageType(str, Enum):
|
|
"""消息类型枚举"""
|
|
TEXT = "text" # 纯文本
|
|
IMAGE = "image" # 图片
|
|
FILE = "file" # 文件
|
|
SYSTEM = "system" # 系统消息
|
|
ACTION = "action" # 动作消息（如调用工具）
|
|
|
|
|
|
class MessageAttachment:
|
|
"""消息附件"""
|
|
attachment_type: str # 附件类型: image, file
|
|
url: str # 资源URL
|
|
name: str # 文件名
|
|
size: int = 0 # 文件大小(字节)
|
|
mime_type: str = "" # MIME类型
|
|
|
|
def __init__(
|
|
self,
|
|
attachment_type: str,
|
|
url: str,
|
|
name: str,
|
|
size: int = 0,
|
|
mime_type: str = ""
|
|
):
|
|
self.attachment_type = attachment_type
|
|
self.url = url
|
|
self.name = name
|
|
self.size = size
|
|
self.mime_type = mime_type
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""转换为字典"""
|
|
return {
|
|
"attachment_type": self.attachment_type,
|
|
"url": self.url,
|
|
"name": self.name,
|
|
"size": self.size,
|
|
"mime_type": self.mime_type
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> "MessageAttachment":
|
|
"""从字典创建"""
|
|
return cls(
|
|
attachment_type=data.get("attachment_type", ""),
|
|
url=data.get("url", ""),
|
|
name=data.get("name", ""),
|
|
size=data.get("size", 0),
|
|
mime_type=data.get("mime_type", "")
|
|
)
|
|
|
|
|
|
class Message(Document):
|
|
"""
|
|
消息文档模型
|
|
存储聊天消息
|
|
"""
|
|
message_id: str = Field(..., description="唯一标识")
|
|
room_id: str = Field(..., description="聊天室ID")
|
|
discussion_id: str = Field(..., description="讨论ID")
|
|
agent_id: Optional[str] = Field(default=None, description="发送Agent ID(系统消息为空)")
|
|
|
|
# 消息内容
|
|
content: str = Field(..., description="消息内容")
|
|
message_type: str = Field(default=MessageType.TEXT.value, description="消息类型")
|
|
attachments: List[Dict[str, Any]] = Field(default_factory=list, description="附件列表")
|
|
|
|
# 元数据
|
|
round: int = Field(default=0, description="所属轮次")
|
|
token_count: int = Field(default=0, description="token数量")
|
|
|
|
# 工具调用相关
|
|
tool_calls: List[Dict[str, Any]] = Field(default_factory=list, description="工具调用记录")
|
|
tool_results: List[Dict[str, Any]] = Field(default_factory=list, description="工具调用结果")
|
|
|
|
# 时间戳
|
|
created_at: datetime = Field(default_factory=datetime.utcnow)
|
|
|
|
class Settings:
|
|
name = "messages"
|
|
indexes = [
|
|
[("room_id", 1), ("created_at", 1)],
|
|
[("discussion_id", 1)],
|
|
[("agent_id", 1)],
|
|
]
|
|
|
|
def get_attachments(self) -> List[MessageAttachment]:
|
|
"""获取附件对象列表"""
|
|
return [MessageAttachment.from_dict(a) for a in self.attachments]
|
|
|
|
def is_from_agent(self, agent_id: str) -> bool:
|
|
"""检查消息是否来自指定Agent"""
|
|
return self.agent_id == agent_id
|
|
|
|
def is_system_message(self) -> bool:
|
|
"""检查是否为系统消息"""
|
|
return self.message_type == MessageType.SYSTEM.value
|
|
|
|
class Config:
|
|
json_schema_extra = {
|
|
"example": {
|
|
"message_id": "msg-001",
|
|
"room_id": "product-design-room",
|
|
"discussion_id": "disc-001",
|
|
"agent_id": "product-manager",
|
|
"content": "我认为登录系统应该支持多种认证方式...",
|
|
"message_type": "text",
|
|
"round": 1,
|
|
"token_count": 150
|
|
}
|
|
}
|