116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据模型
|
||
"""
|
||
from datetime import datetime
|
||
from enum import StrEnum
|
||
from typing import Dict, List, Optional
|
||
|
||
from pydantic import BaseModel, Field
|
||
from pydantic_ai._uuid import uuid7
|
||
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
|
||
from sqlmodel import Field as SqlField, SQLModel
|
||
|
||
|
||
# 消息历史数据表模型
|
||
# 重新初始化：先手动删除 alembic 相关配置和文件夹，再使用 reflex db init 初始化数据库表
|
||
class MessageHistory(SQLModel, table=True):
|
||
id: str = SqlField(
|
||
default_factory=lambda: str(uuid7()),
|
||
primary_key=True,
|
||
description="消息唯一标识",
|
||
)
|
||
conversation_id: str = SqlField(index=True, description="对话唯一标识")
|
||
run_id: str = SqlField(index=True, description="运行唯一标识")
|
||
new_message: str = SqlField(description="新增消息")
|
||
|
||
@staticmethod
|
||
def adapt(
|
||
conversation_id: str, run_id: str, new_message: List[ModelMessage]
|
||
) -> "MessageHistory":
|
||
"""
|
||
适配新增消息为消息历史记录
|
||
:param conversation_id: 对话唯一标识
|
||
:param run_id: 运行唯一标识
|
||
:param new_message: 新增消息
|
||
:return: 消息历史记录
|
||
"""
|
||
return MessageHistory(
|
||
conversation_id=conversation_id,
|
||
run_id=run_id,
|
||
new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode(
|
||
"utf-8"
|
||
), # 序列化为 JSON 字符串
|
||
)
|
||
|
||
|
||
class EventKind(StrEnum):
|
||
"""事件类型"""
|
||
|
||
PART_START = "part_start"
|
||
PART_DELTA = "part_delta"
|
||
PART_END = "part_end"
|
||
FUNCTION_TOOL_CALL = "function_tool_call"
|
||
FUNCTION_TOOL_RESULT = "function_tool_result"
|
||
RUN_START = "run_start"
|
||
RUN_END = "run_end"
|
||
|
||
|
||
class PartKind(StrEnum):
|
||
"""分片类型"""
|
||
|
||
THINKING = "thinking"
|
||
TOOL_SEARCH = "tool-search"
|
||
CAPABILITY_LOAD = "capability-load"
|
||
TOOL_CALL = "tool-call"
|
||
TEXT = "text"
|
||
TOOL_RETURN = "tool-return"
|
||
RETRY_PROMPT = "retry-prompt"
|
||
RUN_RETURN = "run-return"
|
||
|
||
|
||
class Kind(StrEnum):
|
||
"""分段种类"""
|
||
|
||
THINKING = "thinking"
|
||
TOOL_SEARCH = "tool-search"
|
||
LOAD_CAPABILITY = "load-capability"
|
||
TOOL_CALL = "tool-call"
|
||
TEXT = "text"
|
||
TOOL_RETURN = "tool-return"
|
||
RETRY_PROMPT = "retry-prompt"
|
||
RUN_RETURN = "run-return"
|
||
|
||
|
||
class Reasoning(BaseModel):
|
||
"""推理类"""
|
||
|
||
kind: Kind = Field(..., description="推理类型")
|
||
content: str = Field(default="", description="推理内容")
|
||
|
||
|
||
class Run(BaseModel):
|
||
"""运行类"""
|
||
|
||
user_prompt: str = Field(..., description="用户提示词")
|
||
reasonings: Dict[int, Reasoning] = Field(
|
||
default_factory=dict, description="推理字典"
|
||
)
|
||
is_reasoning: bool = Field(
|
||
default=False,
|
||
description="推理状态，True 表示正在推理，False 表示非正在推理",
|
||
)
|
||
assistant_content: str = Field(default="", description="回复正文")
|
||
is_running: bool = Field(
|
||
default=False,
|
||
description="运行状态，True 表示正在运行，False 表示非正在运行",
|
||
)
|
||
|
||
|
||
class Conversation(BaseModel):
|
||
"""对话类"""
|
||
|
||
description: str = Field(default="新对话", description="描述")
|
||
runs: Dict[str, Run] = Field(default_factory=dict, description="运行字典")
|
||
created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
|