102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据模型
|
||
"""
|
||
from enum import StrEnum
|
||
from time import time_ns
|
||
from typing import List, Dict
|
||
from uuid import uuid4
|
||
|
||
from pydantic import BaseModel, Field
|
||
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
|
||
from sqlmodel import Field as SqlField, SQLModel
|
||
|
||
|
||
# 消息历史数据表模型
|
||
# 需使用 reflex db init 初始化数据库表，若重新初始化需手动删除 alembic 相关配置和文件夹
|
||
class MessageHistory(SQLModel, table=True):
|
||
id: str = SqlField(
|
||
default_factory=lambda: uuid4().hex,
|
||
primary_key=True,
|
||
description="消息唯一标识",
|
||
)
|
||
chat_id: str = SqlField(index=True, description="聊天唯一标识")
|
||
new_message: str = SqlField(description="新消息")
|
||
create_at: int = SqlField(
|
||
index=True,
|
||
description="创建时间戳（毫秒级）",
|
||
)
|
||
|
||
@staticmethod
|
||
def adapt(chat_id: str, new_message: List[ModelMessage]) -> "MessageHistory":
|
||
return MessageHistory(
|
||
chat_id=chat_id,
|
||
new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode(
|
||
"utf-8"
|
||
), # 序列化为 JSON 字符串
|
||
create_at=time_ns() // 1_000_000, # 毫秒级时间戳
|
||
)
|
||
|
||
|
||
"""
|
||
聊天、对话和消息关系：
|
||
一次聊天包含若干轮对话，每轮对话包含用户提示词（user_prompt）和输出（output）
|
||
其中，
|
||
输出包含若干片段（part）
|
||
"""
|
||
|
||
|
||
class PartType(StrEnum):
|
||
"""片段类型（适配前端渲染）"""
|
||
|
||
THINKING = "thinking"
|
||
TEXT = "text"
|
||
TOOL_NAME = "tool_name"
|
||
TOOL_ARGS = "tool_args"
|
||
TOOL_RETURN = "tool_return"
|
||
FINISHED = "finished"
|
||
ERROR = "error"
|
||
|
||
|
||
# 动态生成前缀和片段类型映射表
|
||
PREFIX_MAPING = {f"{i:02d}": t for i, t in enumerate(PartType)}
|
||
|
||
|
||
class Part(BaseModel):
|
||
"""片段类（仅就输出消息）"""
|
||
|
||
part_type: PartType = Field(..., description="片段类型")
|
||
content: str = Field(default="", description="片段内容")
|
||
|
||
is_streaming: bool = Field(
|
||
default=False,
|
||
description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
|
||
)
|
||
|
||
is_open: bool = Field(
|
||
default=False,
|
||
description="折叠面板打开状态，True表示打开，False表示关闭",
|
||
)
|
||
|
||
|
||
class Dialog(BaseModel):
|
||
"""对话类"""
|
||
|
||
user_prompt: str = Field(..., description="用户提示词")
|
||
output: Dict[str, Part] = Field(
|
||
default_factory=dict, description="输出，键为片段唯一标识，值为片段对象"
|
||
)
|
||
|
||
|
||
class Chat(BaseModel):
|
||
"""聊天类"""
|
||
|
||
description: str = Field(default="新聊天", description="聊天描述")
|
||
is_streaming: bool = Field(
|
||
default=False,
|
||
description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
|
||
)
|
||
dialogs: Dict[str, Dialog] = Field(
|
||
default_factory=dict, description="对话列表，键为对话唯一标识，值为对话对象"
|
||
)
|