79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据模型
|
||
"""
|
||
from enum import StrEnum
|
||
from time import time_ns
|
||
from typing import List
|
||
from typing import Annotated
|
||
from uuid import uuid4
|
||
|
||
from pydantic import BaseModel, Field
|
||
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
|
||
from sqlmodel import SQLModel, Field as SQLField
|
||
|
||
|
||
# 数据库表模型
|
||
class HistoryMessage(SQLModel, table=True):
|
||
id: int = SQLField(default_factory=int, primary_key=True)
|
||
chat_id: str
|
||
new_message: str
|
||
timestamp: int
|
||
|
||
@staticmethod
|
||
def adapt(chat_id: str, new_message: List[ModelMessage]) -> "HistoryMessage":
|
||
return HistoryMessage(
|
||
chat_id=chat_id,
|
||
new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode("utf-8"),
|
||
timestamp=time_ns() // 1000, # 微秒级时间戳
|
||
)
|
||
|
||
|
||
"""
|
||
聊天、对话和消息关系：
|
||
一次聊天包含若干轮对话，每轮对话包含一条输入消息（input_message）和若干条输出消息（output_messages）。其中，输入消息和输出消息合称消息（message）。
|
||
"""
|
||
|
||
|
||
class Type_(StrEnum):
|
||
"""消息类型类"""
|
||
|
||
THINKING = "thinking"
|
||
TEXT = "text"
|
||
CALL = "call"
|
||
TOOL_ARGS = "tool_args"
|
||
TOOL_RETURN = "tool_return"
|
||
RESULT = "result"
|
||
ERROR = "error"
|
||
|
||
|
||
# 前缀映射表（动态生成）
|
||
PREFIX_MAPING = {f"{i:02d}:": t for i, t in enumerate(Type_)}
|
||
|
||
|
||
class Message(BaseModel):
|
||
"""消息类"""
|
||
|
||
id: str = Field(default_factory=lambda: uuid4().hex, description="消息唯一标识")
|
||
type_: Type_ = Field(..., description="消息类型")
|
||
content: str = Field(default="", description="消息内容")
|
||
|
||
|
||
class Dialog(BaseModel):
|
||
"""对话类"""
|
||
|
||
id: str = Field(default_factory=lambda: uuid4().hex, description="对话唯一标识")
|
||
input_: str = Field(..., description="输入消息")
|
||
output: List[Message] = Field(default_factory=list, description="输出消息")
|
||
|
||
|
||
class Chat(BaseModel):
|
||
"""聊天类"""
|
||
|
||
description: str = Field(default="新聊天", description="描述")
|
||
is_streaming: bool = Field(
|
||
default=False,
|
||
description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
|
||
)
|
||
dialogs: List[Dialog] = Field(default_factory=list, description="对话列表")
|