195 lines
5.4 KiB
Python
195 lines
5.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
面向 reflex.state 的类
|
||
"""
|
||
from datetime import datetime
|
||
from enum import StrEnum
|
||
from typing import Any
|
||
|
||
from pydantic import BaseModel, Field, TypeAdapter
|
||
from pydantic_ai import RunUsage
|
||
from pydantic_ai._uuid import uuid7
|
||
|
||
|
||
class MessageType(StrEnum):
|
||
"""
|
||
消息类型枚举
|
||
"""
|
||
|
||
USER_PROMPT = "user_prompt"
|
||
THINKING = "thinking"
|
||
TOOL_CALL = "tool_call"
|
||
TEXT = "text"
|
||
|
||
|
||
class Approval(BaseModel):
|
||
"""
|
||
审批类
|
||
"""
|
||
|
||
tool_call_id: str = Field(..., description="工具调用唯一标识")
|
||
content: str = Field(default="", description="审批内容")
|
||
result: bool | None = Field(
|
||
default=None,
|
||
description="审批结果，True 表示审批通过，False 表示审批拒绝，None 表示待审批",
|
||
)
|
||
|
||
|
||
class Message(BaseModel):
|
||
"""
|
||
消息类
|
||
"""
|
||
|
||
id: str = Field(default_factory=lambda: str(uuid7()), description="消息唯一标识")
|
||
type: MessageType = Field(..., description="消息类型")
|
||
title: str = Field(default="", description="消息标题")
|
||
content: str = Field(default="", description="消息内容")
|
||
approval: Approval | None = Field(default=None, description="审批")
|
||
is_running: bool = Field(
|
||
default=False, description="正在运行，True 表示正在运行， False 表示运行结束"
|
||
)
|
||
is_shown: bool = Field(
|
||
default=False, description="展示组件，True 表示展示，False 表示隐藏"
|
||
)
|
||
|
||
|
||
class CapabilityComponentType(StrEnum):
|
||
"""
|
||
功能组件类型枚举
|
||
"""
|
||
|
||
TALK = "对话"
|
||
BOOK_FLIGHT = "预定航班"
|
||
BANK_CUSTOMER_SERVICE = "银行客户服务"
|
||
WEB_RETRIEVE = "联网检索"
|
||
KNOWLEDGE_BASE_RETRIEVE = "知识库检索"
|
||
|
||
|
||
class CapabilityComponent(BaseModel):
|
||
"""
|
||
功能组件类
|
||
"""
|
||
|
||
type: CapabilityComponentType = Field(
|
||
default=CapabilityComponentType.TALK, description="功能组件类型"
|
||
)
|
||
deps: Any = Field(default=None, description="功能组件依赖项")
|
||
tool_call_id_map_to_message_id: dict[str, str] = Field(
|
||
default_factory=dict, description="工具调用唯一标识映射为消息实例唯一标识字典"
|
||
)
|
||
|
||
|
||
class Task(BaseModel):
|
||
"""
|
||
任务类
|
||
"""
|
||
|
||
id: str = Field(..., description="任务唯一标识")
|
||
description: str = Field(..., description="任务描述")
|
||
last_run_at: datetime = Field(..., description="上一次运行日期时间")
|
||
user_prompt: str = Field(default="", description="用户提示词")
|
||
usage: RunUsage = Field(..., description="任务使用量")
|
||
messages: dict[str, Message] = Field(default_factory=dict, description="消息字典")
|
||
capability_component: CapabilityComponent = Field(
|
||
default_factory=CapabilityComponent, description="功能组件"
|
||
)
|
||
is_running: bool = Field(
|
||
default=False, description="运行状态，True 表示正在运行， False 表示运行结束"
|
||
)
|
||
awaiting_stream_output: bool = Field(
|
||
default=False,
|
||
description="等待流式输出状态，True 表示等待流式输出，False 表示已开始流式输出或已完成",
|
||
)
|
||
|
||
|
||
class TaskHistoryItem(BaseModel):
|
||
"""
|
||
任务历史项类
|
||
"""
|
||
|
||
id: str = Field(..., description="任务唯一标识")
|
||
description: str = Field(..., description="任务描述")
|
||
last_run_at: str = Field(..., description="上一次运行日期时间")
|
||
is_running: bool = Field(
|
||
default=False, description="正在运行，True 表示正在运行， False 表示运行结束"
|
||
)
|
||
|
||
|
||
class MessageHistoryItem(Message):
|
||
"""
|
||
消息历史项类
|
||
"""
|
||
|
||
has_approval: bool = Field(default=False, description="包含审批实例")
|
||
approval_content: str = Field(default="", description="审批内容")
|
||
|
||
|
||
# 功能组件适配器
|
||
CapabilityComponentAdapter = TypeAdapter(CapabilityComponent)
|
||
|
||
|
||
def capability_component_validate_json(capability_component: str) -> Any:
|
||
"""
|
||
将功能组件反序列化
|
||
:param capability_component: 功能组件
|
||
:return: 反序列化后功能组件
|
||
"""
|
||
if not capability_component:
|
||
return CapabilityComponent()
|
||
return CapabilityComponentAdapter.validate_json(capability_component)
|
||
|
||
|
||
def capability_component_dump_json(capability_component: CapabilityComponent) -> str:
|
||
"""
|
||
将功能组件序列化
|
||
:param capability_component: 功能组件
|
||
:return: 序列化后功能组件
|
||
"""
|
||
return CapabilityComponentAdapter.dump_json(capability_component).decode("utf-8")
|
||
|
||
|
||
# 使用量（RunUsage）适配器
|
||
UsageAdapter = TypeAdapter(RunUsage)
|
||
|
||
|
||
def usage_validate_json(usage: str) -> RunUsage:
|
||
"""
|
||
将使用量反序列化
|
||
"""
|
||
if not usage:
|
||
return RunUsage()
|
||
return UsageAdapter.validate_json(usage)
|
||
|
||
|
||
def usage_dump_json(usage: RunUsage) -> str:
|
||
"""
|
||
将使用量序列化
|
||
:param usage: 使用量
|
||
:return: 字符串
|
||
"""
|
||
return UsageAdapter.dump_json(usage).decode("utf-8")
|
||
|
||
|
||
# 审批适配器
|
||
ApprovalAdapter = TypeAdapter(Approval)
|
||
|
||
|
||
def approval_validate_json(approval: str) -> Approval | None:
|
||
"""
|
||
将审批反序列化
|
||
"""
|
||
if not approval:
|
||
return None
|
||
return ApprovalAdapter.validate_json(approval)
|
||
|
||
|
||
def approval_dump_json(approval: Approval | None) -> str:
|
||
"""
|
||
将审批序列化
|
||
:param approval: 审批
|
||
:return: 字符串
|
||
"""
|
||
if approval is None:
|
||
return ""
|
||
return ApprovalAdapter.dump_json(approval).decode("utf-8")
|