93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据模型
|
||
"""
|
||
from datetime import datetime
|
||
from enum import StrEnum
|
||
from typing import Dict, List
|
||
|
||
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
|
||
|
||
|
||
# 消息历史数据表模型
|
||
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 ReasoningKind(StrEnum):
|
||
"""推理种类"""
|
||
|
||
THINKING = "thinking"
|
||
TOOL_SEARCH = "tool-search"
|
||
CAPABILITY_LOAD = "load-capability"
|
||
TOOL_CALL = "tool-call"
|
||
TEXT = "text"
|
||
TOOL_RETURN = "tool-return"
|
||
RETRY_PROMPT = "retry-prompt"
|
||
RUN_RETURN = "run-return"
|
||
|
||
|
||
class Reasoning(BaseModel):
|
||
"""推理类"""
|
||
|
||
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_reasoning_panel_open: 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="运行字典")
|