144 lines
4.5 KiB
Python
144 lines
4.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据模型
|
||
"""
|
||
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
|
||
|
||
|
||
# 消息历史数据表模型
|
||
# 需使用 reflex db init 初始化数据库表，若重新初始化需手动删除 alembic 相关配置和文件夹
|
||
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="运行唯一标识")
|
||
run_new_message: str = SqlField(description="运行新增消息")
|
||
|
||
@staticmethod
|
||
def adapt(
|
||
conversation_id: str, run_id: str, run_new_message: List[ModelMessage]
|
||
) -> "MessageHistory":
|
||
"""
|
||
适配运行新增消息为消息历史
|
||
:param conversation_id: 会话唯一标识
|
||
:param run_id: 运行唯一标识
|
||
:param run_new_message: 运行新增消息
|
||
:return: 消息历史
|
||
"""
|
||
return MessageHistory(
|
||
conversation_id=conversation_id,
|
||
run_id=run_id,
|
||
run_new_message=ModelMessagesTypeAdapter.dump_json(run_new_message).decode(
|
||
"utf-8"
|
||
), # 序列化为 JSON 字符串
|
||
)
|
||
|
||
|
||
"""
|
||
会话、运行和消息关系：
|
||
一次会话（conversation）包含若干次运行（run），每次运行包含用户提示词（user_prompt）和输出（output）。其中，输出包含推理和回答
|
||
"""
|
||
|
||
|
||
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 Event(BaseModel):
|
||
"""
|
||
事件类
|
||
"""
|
||
|
||
event_kind: EventKind = Field(..., description="事件类型")
|
||
event_content: str = Field(default="", description="事件内容")
|
||
part_index: Optional[int] = Field(default=None, description="分片索引")
|
||
previous_part_kind: Optional[PartKind] = Field(
|
||
default=None, description="上个分片类型"
|
||
)
|
||
part_kind: Optional[PartKind] = Field(default=None, description="分片类型")
|
||
next_part_kind: Optional[PartKind] = Field(default=None, description="下个分片类型")
|
||
tool_name: Optional[str] = Field(default=None, description="工具名称")
|
||
run_id: str = Field(..., description="运行唯一标识")
|
||
run_new_messages: List[ModelMessage] = Field(
|
||
default=[], description="运行新增消息列表"
|
||
)
|
||
|
||
|
||
class ReasoningKind(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 Reasoning(BaseModel):
|
||
"""推理类"""
|
||
|
||
reasoning_kind: ReasoningKind = 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 表示非正在推理",
|
||
)
|
||
is_expanded: bool = Field(
|
||
default=False,
|
||
description="推理折叠面板展开状态，True 表示展开，False 表示折叠",
|
||
)
|
||
answer: str = Field(default="", description="回答")
|
||
is_streaming: bool = Field(
|
||
default=False,
|
||
description="运行流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
|
||
)
|
||
|
||
|
||
class Conversation(BaseModel):
|
||
"""会话类"""
|
||
|
||
description: str = Field(default="新会话", description="会话描述")
|
||
runs: Dict[str, Run] = Field(default_factory=dict, description="运行字典")
|