This commit is contained in:
liubiren 2026-10-23 20:09:17 +08:00
parent cd239926e4
commit d37517abb1
25 changed files with 1046 additions and 889 deletions

View File

@ -17,8 +17,8 @@ from pydantic_ai import (
from pydantic_ai.messages import AgentStreamEvent from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.run import AgentRunResultEvent
from application.workshop.database import connection_pool from application.agent.database import connection_pool
from application.workshop.models import ( from application.agent.models import (
DEEPSEEK_V4_FLASH_CHAT_MODEL, DEEPSEEK_V4_FLASH_CHAT_MODEL,
DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS, DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
) )

View File

@ -26,8 +26,8 @@ from pydantic_ai.messages import (
) )
from pydantic_ai.run import AgentRunResultEvent, AgentRunResult from pydantic_ai.run import AgentRunResultEvent, AgentRunResult
from application.workshop.events import AgentRuntimeEvent from application.agent.events import AgentRuntimeEvent
from application.workshop.models import ( from application.agent.models import (
DEEPSEEK_V4_FLASH_CHAT_MODEL, DEEPSEEK_V4_FLASH_CHAT_MODEL,
DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS, DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
) )
@ -302,7 +302,6 @@ async def run_stream_events(
deferred_tool_results=deferred_tool_results, deferred_tool_results=deferred_tool_results,
) as events: ) as events:
async for event in events: async for event in events:
print(event)
match event: match event:
case PartStartEvent( case PartStartEvent(
part=part, part=part,

View File

@ -13,8 +13,8 @@ from pydantic_ai.run import AgentRunResultEvent
from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai.usage import RunUsage from pydantic_ai.usage import RunUsage
from application.workshop.database import connection_pool from application.agent.database import connection_pool
from application.workshop.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL from application.agent.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL
class Deps(BaseModel): class Deps(BaseModel):

View File

@ -0,0 +1,155 @@
# -*- coding: utf-8 -*-
"""
智能体运行时模块
"""
from asyncio import Queue, Task, Event, Lock
from typing import AsyncGenerator, Optional
from application.agent.events import AgentRuntimeEvent
from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent
from pydantic import BaseModel, Field
from application.states.models import CapabilityComponentType
from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolApproved
class TaskRuntime(BaseModel):
queue: Queue[AgentStreamEvent | AgentRunResultEvent | None] = Field(
default_factory=lambda: Queue(maxsize=256), description="任务运行时队列"
)
worker: Task | None = Field(default=None, description="任务运行时工作器")
done: Event = Field(default_factory=Event, description="任务运行时完成状态")
class TaskRuntimeRegistrar:
"""
任务运行时注册器
"""
def __init__(self):
# 异步互斥锁
self.lock = Lock()
# 任务运行时字典,键为任务唯一标识,值为任务运行时实例
self.task_runtimes: dict[str, TaskRuntime] = {}
async def get(self, task_id: str) -> TaskRuntime:
"""
获取任务运行时实例
:param task_id: 任务唯一标识
:return: 任务运行时实例
"""
async with self.lock:
# 若任务唯一标识不存在则创建任务运行时实例
if task_id not in self.task_runtimes:
self.task_runtimes[task_id] = TaskRuntime()
return self.task_runtimes[task_id]
async def cancel(self, task_id: str):
"""
取消任务运行时实例
:param task_id: 任务唯一标识
"""
async with self.lock:
task_runtime = self.task_runtimes.pop(task_id, None)
if task_runtime and task_runtime.worker and not task_runtime.worker.done():
task_runtime.worker.cancel()
async def worker(
task_id: str,
user_prompt: str | None,
message_history,
usage,
capability_component_type: CapabilityComponentType,
queue: Queue,
deps=None,
deferred_tool_results: DeferredToolResults | None = None,
):
"""
工作器
"""
try:
# 根据能力组件的类型匹配流式输出方法
match capability_component_type:
case CapabilityComponentType.TALK:
from application.agent.talk import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 预定航班
case CapabilityComponentType.BOOK_FLIGHT:
from application.agent.book_flight import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
deferred_tool_results=deferred_tool_results,
)
# 银行客户服务
case CapabilityComponentType.BANK_CUSTOMER_SERVICE:
from application.agent.bank_customer_service import (
run_stream_events,
)
from application.agent.database import connection_pool
# 若连接池为空则直接返回
if not connection_pool:
return
# 运行并流式输出事件
stream_events = run_stream_events(
deps=deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 联网检索
case CapabilityComponentType.WEB_RETRIEVE:
from application.agent.web_retrieve import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 知识库检索
case CapabilityComponentType.KNOWLEDGE_BASE_RETRIEVE:
from application.agent.knowledge_base_retrieve import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
async for event in stream_events:
await queue.put(event)
# 哨兵:正常结束
await queue.put(None)
except Exception as e:
# 异常丢入队列给消费端处理
await queue.put(e)

View File

@ -8,7 +8,7 @@ from pydantic_ai import Agent, ModelMessage, RunUsage
from pydantic_ai.messages import AgentStreamEvent from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.run import AgentRunResultEvent
from application.workshop.models import DEEPSEEK_V4_FLASH_CHAT_MODEL from application.agent.models import DEEPSEEK_V4_FLASH_CHAT_MODEL
instruction = """ instruction = """

View File

@ -14,8 +14,8 @@ from pydantic_ai.run import AgentRunResultEvent
from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai.usage import RunUsage from pydantic_ai.usage import RunUsage
from application.workshop.database import connection_pool from application.agent.database import connection_pool
from application.workshop.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL from application.agent.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL
""" """
联网检索缓存表建表SQL 联网检索缓存表建表SQL

View File

@ -5,7 +5,7 @@
import reflex as rx import reflex as rx
from application.pages.workbench import workbench from application.pages.workbench import workbench
from application.workshop.database import connection_pool_init from application.agent.database import connection_pool_init
app = rx.App( app = rx.App(

View File

@ -4,7 +4,7 @@
""" """
import reflex as rx import reflex as rx
from application.states import AuthState from application.states.auth_states import AuthState
def login_page() -> rx.Component: def login_page() -> rx.Component:

View File

@ -4,20 +4,24 @@
""" """
import reflex as rx import reflex as rx
from application.states import AuthState, TaskState from application.states.task_states import TaskState
from application.states.models import MessageHistoryItem, MessageType, WorkflowType from application.states.models import (
MessageHistoryItem,
MessageType,
CapabilityComponentType,
)
def shortcut(workflow_type: WorkflowType): def shortcut(capability_component_type: CapabilityComponentType):
""" """
快捷方式 快捷方式
:param workflow_type: 工作类型 :param workflow_type: 工作类型
:return: rx.Component :return: rx.Component
""" """
return rx.box( return rx.box(
workflow_type.value, capability_component_type.value,
# 点击事件:设置工作流 # 点击事件:设置功能组件
on_click=lambda: TaskState.set_workflow(workflow_type), on_click=lambda: TaskState.set_capability_component(capability_component_type),
padding="10px 16px", padding="10px 16px",
background_color="var(--prismui-background-color-9)", background_color="var(--prismui-background-color-9)",
border_radius="var(--prismui-border-radius-9)", border_radius="var(--prismui-border-radius-9)",
@ -75,14 +79,14 @@ def welcome() -> rx.Component:
font_weight="var(--prismui-font-weight-3)", font_weight="var(--prismui-font-weight-3)",
), ),
rx.hstack( rx.hstack(
# 工作流按钮:预定航班 # 快捷方式:预定航班
shortcut(WorkflowType.BOOK_FLIGHT), shortcut(CapabilityComponentType.BOOK_FLIGHT),
# 工作流按钮:银行客户服务 # 快捷方式:银行客户服务
shortcut(WorkflowType.BANK_CUSTOMER_SERVICE), shortcut(CapabilityComponentType.BANK_CUSTOMER_SERVICE),
# 工作流按钮:联网检索 # 快捷方式:联网检索
shortcut(WorkflowType.WEB_RETRIEVE), shortcut(CapabilityComponentType.WEB_RETRIEVE),
# 工作流按钮:知识库检索 # 快捷方式:知识库检索
shortcut(WorkflowType.KNOWLEDGE_BASE_RETRIEVE), shortcut(CapabilityComponentType.KNOWLEDGE_BASE_RETRIEVE),
flex_wrap="wrap", flex_wrap="wrap",
align_items="center", align_items="center",
gap="12px", gap="12px",
@ -671,10 +675,10 @@ def message_history() -> rx.Component:
rx.foreach(TaskState.message_history, message_history_item), rx.foreach(TaskState.message_history, message_history_item),
# 若等待流式输出则显示加载动效 # 若等待流式输出则显示加载动效
rx.cond( rx.cond(
TaskState.awaiting_stream, TaskState.awaiting_stream_output,
rx.box( rx.box(
rx.box( rx.box(
class_name="awaiting_stream", class_name="awaiting_stream_output",
flex_shrink="0", flex_shrink="0",
), ),
display="flex", display="flex",

View File

@ -4,9 +4,11 @@
""" """
import reflex as rx import reflex as rx
from application.pages import task_page, knowledge_base_page, login_page from application.pages import knowledge_base_page, login_page, task_page
from application.states import AuthState, TaskState, SidebarState from application.states.auth_states import AuthState
from application.states.models import TaskHistoryItem from application.states.models import TaskHistoryItem
from application.states.sidebar_states import SidebarState
from application.states.task_states import TaskState
# 导航按钮字典 # 导航按钮字典

View File

@ -1,13 +1 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
from application.states.auth import AuthState
from application.states.database import DatabaseState
from application.states.sidebar import SidebarState
from application.states.task import TaskState
__all__ = [
"AuthState",
"DatabaseState",
"SidebarState",
"TaskState",
]

View File

@ -12,8 +12,8 @@ from aiosmtplib import SMTP, SMTPException
import reflex as rx import reflex as rx
from reflex.event import EventCallback from reflex.event import EventCallback
from application.states.database import DatabaseState from application.states.task_states import TaskState
from application.states.task import TaskState from application.states.database_states import DatabaseState
# 验证码HTML模板 # 验证码HTML模板
CAPTCHA_HTML_TEMPLATE = """ CAPTCHA_HTML_TEMPLATE = """
@ -165,7 +165,7 @@ class AuthState(rx.State):
# 标记验证码已发送 # 标记验证码已发送
self.is_captcha_code_sent = True self.is_captcha_code_sent = True
# 并行处理:后台运行倒计时与后台运行发送验证码 # 并行处理:将运行倒计时此事件提交至后台任务列队后,再提交运行发送验证码此事件至后台任务列队后
yield type(self).background_countdown() yield type(self).background_countdown()
yield type(self).background_send_captcha_code() yield type(self).background_send_captcha_code()

View File

@ -20,13 +20,13 @@ from application.states.models import (
Message, Message,
MessageType, MessageType,
Task, Task,
Workflow, CapabilityComponent,
approval_dump_json, approval_dump_json,
approval_validate_json, approval_validate_json,
usage_dump_json, usage_dump_json,
usage_validate_json, usage_validate_json,
workflow_dump_json, capability_component_dump_json,
workflow_validate_json, capability_component_validate_json,
) )
@ -92,9 +92,9 @@ class TaskRecord(SQLModel, table=True):
description="任务已删除:True 表示已删除,False 表示未删除", description="任务已删除:True 表示已删除,False 表示未删除",
) )
usage: str = Field(default="", description="任务使用量") usage: str = Field(default="", description="任务使用量")
workflow: str = Field( capability_component: str = Field(
default="", default="",
description="任务工作流", description="任务功能组件",
) # 原则上嵌套列表或字典若不涉及查询则以字符串储存 ) # 原则上嵌套列表或字典若不涉及查询则以字符串储存
@ -264,9 +264,9 @@ class DatabaseState(rx.State):
usage=usage_validate_json( usage=usage_validate_json(
task_record.usage task_record.usage
), # 将 task_record.usage 由 JSON 字符串反序列化为 RunUsage ), # 将 task_record.usage 由 JSON 字符串反序列化为 RunUsage
workflow=workflow_validate_json( capability_component=capability_component_validate_json(
task_record.workflow task_record.capability_component
), # 将 task_record.workflow 由 JSON 字符串反序列化为 Workflow, ), # 将 task_record.capability_component 由 JSON 字符串反序列化为 CapabilityComponent,
), ),
) )
if not message_record: if not message_record:
@ -290,7 +290,7 @@ class DatabaseState(rx.State):
last_run_at: datetime | None = None, last_run_at: datetime | None = None,
is_deleted: bool | None = None, is_deleted: bool | None = None,
usage: RunUsage | None = None, usage: RunUsage | None = None,
workflow: Workflow | None = None, capability_component: CapabilityComponent | None = None,
) -> None: ) -> None:
""" """
更新任务记录 更新任务记录
@ -299,7 +299,7 @@ class DatabaseState(rx.State):
:param last_run_at: 上一次运行日期时间 :param last_run_at: 上一次运行日期时间
:param is_deleted: 任务已删除 :param is_deleted: 任务已删除
:param usage: 使用量 :param usage: 使用量
:param workflow: 工作流 :param capability_component: 功能组件
:return: None :return: None
""" """
async with rx.asession() as session: async with rx.asession() as session:
@ -316,10 +316,10 @@ class DatabaseState(rx.State):
record.usage = usage_dump_json( record.usage = usage_dump_json(
RunUsage(**asdict(usage)) RunUsage(**asdict(usage))
) # 将 RunUsage 序列化为 JSON 字符串 ) # 将 RunUsage 序列化为 JSON 字符串
if isinstance(workflow, Workflow): if isinstance(capability_component, CapabilityComponent):
record.workflow = workflow_dump_json( record.capability_component = capability_component_dump_json(
Workflow(**workflow.model_dump(mode="json")) CapabilityComponent(**capability_component.model_dump(mode="json"))
) # 将 Workflow 序列化为 JSON 字符串(因 workflow 基于 pydantic 建模,故使用 model_dump 方法转为字典) ) # 将 CapabilityComponent 序列化为 JSON 字符串(因 capability_component 基于 pydantic 建模,故使用 model_dump 方法转为字典)
await session.commit() await session.commit()
async def create_message_record( async def create_message_record(

View File

@ -53,9 +53,9 @@ class Message(BaseModel):
) )
class WorkflowType(StrEnum): class CapabilityComponentType(StrEnum):
""" """
工作类型枚举 功能组件类型枚举
""" """
TALK = "对话" TALK = "对话"
@ -65,13 +65,15 @@ class WorkflowType(StrEnum):
KNOWLEDGE_BASE_RETRIEVE = "知识库检索" KNOWLEDGE_BASE_RETRIEVE = "知识库检索"
class Workflow(BaseModel): class CapabilityComponent(BaseModel):
""" """
工作流类 功能组件类
""" """
type: WorkflowType = Field(default=WorkflowType.TALK, description="工作类型") type: CapabilityComponentType = Field(
deps: Any = Field(default=None, description="工作依赖项") default=CapabilityComponentType.TALK, description="功能组件类型"
)
deps: Any = Field(default=None, description="功能组件依赖项")
tool_call_id_map_to_message_id: dict[str, str] = Field( tool_call_id_map_to_message_id: dict[str, str] = Field(
default_factory=dict, description="工具调用唯一标识映射为消息实例唯一标识字典" default_factory=dict, description="工具调用唯一标识映射为消息实例唯一标识字典"
) )
@ -88,13 +90,15 @@ class Task(BaseModel):
user_prompt: str = Field(default="", description="用户提示词") user_prompt: str = Field(default="", description="用户提示词")
usage: RunUsage = Field(..., description="任务使用量") usage: RunUsage = Field(..., description="任务使用量")
messages: dict[str, Message] = Field(default_factory=dict, description="消息字典") messages: dict[str, Message] = Field(default_factory=dict, description="消息字典")
workflow: Workflow = Field(default_factory=Workflow, description="工作,默认为对话") capability_component: CapabilityComponent = Field(
is_running: bool = Field( default_factory=CapabilityComponent, description="功能组件"
default=False, description="正在运行,True 表示正在运行, False 表示运行结束"
) )
awaiting_stream: bool = Field( is_running: bool = Field(
default=False, description="运行状态,True 表示正在运行, False 表示运行结束"
)
awaiting_stream_output: bool = Field(
default=False, default=False,
description="等待流式输出,True 表示等待流式输出,False 表示已开始流式输出或已完成", description="等待流式输出状态,True 表示等待流式输出,False 表示已开始流式输出或已完成",
) )
@ -120,26 +124,28 @@ class MessageHistoryItem(Message):
approval_content: str = Field(default="", description="审批内容") approval_content: str = Field(default="", description="审批内容")
# 工作流适配器 # 功能组件适配器
WorkflowAdapter = TypeAdapter(Workflow) CapabilityComponentAdapter = TypeAdapter(CapabilityComponent)
def workflow_validate_json(workflow: str) -> Any: def capability_component_validate_json(capability_component: str) -> Any:
""" """
将工作流反序列化 将功能组件反序列化
:param capability_component: 功能组件
:return: 反序列化后功能组件
""" """
if not workflow: if not capability_component:
return Workflow() return CapabilityComponent()
return WorkflowAdapter.validate_json(workflow) return CapabilityComponentAdapter.validate_json(capability_component)
def workflow_dump_json(workflow: Workflow) -> str: def capability_component_dump_json(capability_component: CapabilityComponent) -> str:
""" """
将工作流序列化 将功能组件序列化
:param workflow: 工作流 :param capability_component: 功能组件
:return: 字符串 :return: 序列化后功能组件
""" """
return WorkflowAdapter.dump_json(workflow).decode("utf-8") return CapabilityComponentAdapter.dump_json(capability_component).decode("utf-8")
# 使用量(RunUsage)适配器 # 使用量(RunUsage)适配器

View File

@ -4,7 +4,7 @@
""" """
import reflex as rx import reflex as rx
from application.states.task import TaskState from application.states.task_states import TaskState
class SidebarState(rx.State): class SidebarState(rx.State):
@ -13,7 +13,7 @@ class SidebarState(rx.State):
""" """
# 选中的导航按钮,默认为创建任务 # 选中的导航按钮,默认为创建任务
nav_button: str = "create_task" nav_button: str | None = "create_task"
# 退出登录确认打开状态 # 退出登录确认打开状态
is_logout_confirm_open: bool = False is_logout_confirm_open: bool = False

View File

@ -1,801 +0,0 @@
# -*- coding: utf-8 -*-
"""
任务状态
"""
from datetime import date as Date, datetime
from asyncio import create_task
from copy import copy
from typing import AsyncGenerator
from reflex_base.event import EventSpec
from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolApproved
from pydantic_ai.messages import (
FunctionToolResultEvent,
PartDeltaEvent,
PartEndEvent,
PartStartEvent,
TextPart,
TextPartDelta,
ThinkingPart,
ThinkingPartDelta,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.run import AgentRunResultEvent
import reflex as rx
from reflex.event import EventCallback
from application.states.database import DatabaseState
from application.states.models import (
Approval,
Task,
TaskHistoryItem,
Message,
MessageHistoryItem,
MessageType,
Workflow,
WorkflowType,
)
from application.workshop.events import AgentRuntimeEvent
class TaskState(rx.State):
"""
任务状态
"""
# 当前用户唯一标识
user_id: str = ""
# 当前用户的任务字典
tasks: dict[str, Task] = {} # 键为任务唯一标识,值为任务实例
# 当前任务唯一标识
task_id: str = ""
# 临时用户提示词
temporary_user_prompt: str = ""
# 任务删除确认打开状态
is_task_delete_confirm_open: bool = False
# 任务删除确认打开日期时间
task_delete_confirm_open_at: float = 0
async def resume(self, user_id: str) -> None:
"""
恢复任务状态
:param user_id: 用户唯一标识
:return: None
"""
self.user_id = user_id
if not self.user_id:
return
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 查询当前用户的任务字典
self.tasks = await database_state.read_tasks(self.user_id)
for task in self.tasks.values():
# 最后一个消息
message = next(iter(reversed(task.messages.values())), None)
# 若最后一个消息包含审批且审批结果为 None 则任务正在运行设置为是、消息正在运行设置为是
if (
message
and isinstance(approval := message.approval, Approval)
and approval.result is None
):
task.is_running = True
message.is_running = True
@rx.var
def task_history(self) -> list[TaskHistoryItem]:
"""
获取任务历史项列表(按照任务创建日期时间降序排序)
:return: 任务历史项列表
"""
def formatter(dt: datetime) -> str:
"""
格式器
:param dt: 日期时间
:return: 格式化后日期时间字符串
"""
# 时间间隔
delta = (datetime.now() - dt).total_seconds()
if delta < 60:
return "刚刚"
elif delta < 60 * 60:
return f"{round(delta / 60)}分钟前"
elif delta < 60 * 60 * 24:
return f"{round(delta / 60 / 60)}小时前"
else:
return f"{round(delta / 60 / 60 / 24)}天前"
return [
TaskHistoryItem(
id=task.id,
description=task.description,
last_run_at=formatter(task.last_run_at),
is_running=task.is_running,
)
for task in reversed(self.tasks.values())
]
@rx.var
def task_description(self) -> str:
"""
获取当前任务的描述
:return: 当前任务的描述
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return ""
return task.description
@rx.var
def task_last_run_at(self) -> str:
"""
获取当前任务的上一次运行日期时间
:return: 当前任务的上一次运行日期时间
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return ""
return f"更新于 {task.last_run_at.strftime("%Y-%m-%d %H:%M:%S")}"
@rx.event
async def toggle_task_delete_confirm_open(self, is_open: bool) -> None:
"""
切换任务删除确认打开状态
:param is_open: 任务删除确认打开状态
:return: None
"""
# 当前日期时间
current = datetime.now().timestamp()
# 根据当前时间与任务删除确认打开日期时间间隔防抖
if current - self.task_delete_confirm_open_at < 0.2:
return
self.is_task_delete_confirm_open = is_open
self.task_delete_confirm_open_at = current
@rx.event
async def delete_task(self) -> EventSpec | None:
"""
删除当前任务
:return: RX事件
"""
from application.states import SidebarState
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 更新任务记录
await database_state.update_task_record(task.id, is_deleted=True)
del self.tasks[task.id]
# 获取侧边栏状态
sidebar_state = await self.get_state(SidebarState)
# 将创建任务设置为导航按钮
sidebar_state.nav_button = ""
# 将空设置为当前任务唯一标识
self.task_id = ""
# 将任务删除确认显示设置为否
self.is_task_delete_confirm_open = False
return rx.toast.success("删除成功")
@rx.event
def set_task_id(self, task_id: str) -> None:
"""
将指定任务唯一标识设置为当前任务唯一标识
:param task_id: 任务唯一标识
:return: None
"""
self.task_id = task_id
@rx.var
def message_history(self) -> list[MessageHistoryItem]:
"""
获取当前任务的消息历史
:return: 当前任务的消息历史
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return []
message_history: list[MessageHistoryItem] = []
for message in task.messages.values():
message_history.append(
MessageHistoryItem(
**message.model_dump(),
has_approval=isinstance(approval := message.approval, Approval)
and approval.result is None,
approval_content=getattr(approval, "content", ""),
)
)
return message_history
@rx.var
def user_prompt(self) -> str:
"""
用户提示词
:return: 用户提示词
"""
# 当前任务
task = self.tasks.get(self.task_id)
# 若当前任务存在则返回当前任务的用户提示词,否则返回临时用户提示词
if task:
return task.user_prompt
else:
return self.temporary_user_prompt
@rx.event
def set_user_prompt(self, user_prompt: str) -> None:
"""
设置用户提示词
:param user_prompt: 用户提示词
:return: None
"""
# 当前任务
task = self.tasks.get(self.task_id)
# 若当前任务存在则将用户提示词设置为当前任务的用户提示词,否则将用户提示词设置为临时用户提示词
if task:
task.user_prompt = user_prompt
else:
self.temporary_user_prompt = user_prompt
@rx.var
def is_user_prompt_sending_button_disabled(self) -> bool:
"""
用户提示词发送按钮不可点击状态
:return: bool
"""
# 当前任务
task = self.tasks.get(self.task_id)
# 若当前任务存在则判断该任务是否正在运行或用户提示词非空
if task:
return task.is_running or not task.user_prompt
else:
return not self.temporary_user_prompt
@rx.event
async def run(
self,
is_resume_run: bool = False,
message_id: str | None = None,
) -> AsyncGenerator[EventCallback, None]:
"""
运行
:param is_resume_run: 恢复运行
:param message_id: 消息唯一标识
:return: None
"""
from application.states import SidebarState
user_prompt = None # 用户提示词
deferred_tool_results = None # 延迟工具结果
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
async with self:
# 若当前任务的唯一标识为空则创建任务并将此作为当前任务
if not self.task_id:
# 获取侧边栏状态
sidebar_state = await self.get_state(SidebarState)
# 将空设置为导航按钮
sidebar_state.nav_button = ""
# 先创建任务记录再将此添加至当前用户的任务字典
self.tasks.update(
await database_state.create_task_record(
self.user_id, self.temporary_user_prompt
)
)
# 将最后一个任务的唯一标识作为当前任务唯一标识
self.task_id = next(reversed(self.tasks.keys()))
# 当前任务的唯一标识
task_id = self.task_id
# 当前任务
task = self.tasks[task_id]
# 若为恢复运行
if is_resume_run:
if not message_id:
return
# 指定消息唯一标识的消息
message = task.messages.get(message_id)
if (
not message
or not isinstance(approval := message.approval, Approval)
or not approval.content
):
return
# 将等待流式输出设置为是
task.awaiting_stream = True
# 更新消息内容
message.content = f"您输入的是:\n\n{approval.content}\n\n"
# 更新审批结果
approval.result = True
# 构建延迟工具结果
deferred_tool_results = DeferredToolResults(
approvals={approval.tool_call_id: ToolApproved()}, # 审批通过
metadata={approval.tool_call_id: {"content": approval.content}},
)
else:
# 若用户提示词为空则直接返回
if not (
user_prompt := task.user_prompt.strip()
): # 去除用户提示词首尾空格
return
# 清空临时用户提示词和当前任务的用户提示词
self.temporary_user_prompt = task.user_prompt = ""
# 将当前任务设置为正在运行
task.is_running = True
# 将当前任务设置为等待流式输出
task.awaiting_stream = True
# 先创建消息记录再添加至当前任务的消息字典
task.messages.update(
await database_state.create_message_record(
task_id=task.id,
message=Message(
type=MessageType.USER_PROMPT,
content=user_prompt,
),
)
)
yield type(self).background_run(task_id, user_prompt, deferred_tool_results)
@rx.event(background=True)
async def background_run(
self,
task_id: str,
user_prompt: str | None,
deferred_tool_results: DeferredToolResults | None = None,
):
"""
后台运行运行
:param task_id: 任务唯一标识
:return: None
"""
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 查询消息历史列表
message_history = await database_state.read_message_history(task_id=task_id)
async with self:
# 当前任务
task = self.tasks[task_id]
# 工作流
workflow = task.workflow
# 使用量
usage = copy(task.usage)
if usage.details:
usage.details = dict(usage.details)
# 根据工作流的类型匹配流式输出
match workflow.type:
case WorkflowType.TALK:
from application.workshop.talk import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 预定航班
case WorkflowType.BOOK_FLIGHT:
from application.workshop.book_flight import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=task.workflow.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
deferred_tool_results=deferred_tool_results,
)
# 银行客户服务
case WorkflowType.BANK_CUSTOMER_SERVICE:
from application.workshop.bank_customer_service import (
run_stream_events,
)
from application.workshop.database import connection_pool
# 若连接池为空则直接返回
if not connection_pool:
return
# 运行并流式输出事件
stream_events = run_stream_events(
deps=task.workflow.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 联网检索
case WorkflowType.WEB_RETRIEVE:
from application.workshop.web_retrieve import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=task.workflow.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 知识库检索
case WorkflowType.KNOWLEDGE_BASE_RETRIEVE:
from application.workshop.knowledge_base_retrieve import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=task.workflow.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
messages: list[Message] = [] # 消息列表
previous_provider_details = {} # 上一个供应商详情
# 初始化片段索引映射为消息唯一标识字典
index_map_to_message_id: dict[int, str] = {}
# 获取运行流式输出事件
async for event in stream_events:
# 若等待流式输出则将等待流式输出设置为否
if task.awaiting_stream:
task.awaiting_stream = False
match event:
# ========== 开始事件 ==========
case PartStartEvent(
index=index,
part=part,
):
match part:
# 思考分片
case ThinkingPart(content=content):
# 构建消息
message = Message(
type=MessageType.THINKING,
title="正在思考",
content=content,
is_running=True,
)
# 添加至消息字典
task.messages[message.id] = message
# 将消息唯一标识与片段索引映射
index_map_to_message_id[index] = message.id
yield # 通知前端更新渲染
# 工具调用分片
case ToolCallPart(
tool_call_id=tool_call_id, tool_name=tool_name
):
if (
tool_call_id
not in task.workflow.tool_call_id_map_to_message_id
):
# 构建消息
message = Message(
type=MessageType.TOOL_CALL,
title=f"正在{tool_name}",
is_running=True,
)
# 添加至消息字典
task.messages[message.id] = message
# 将消息唯一标识与工具调用唯一标识映射
task.workflow.tool_call_id_map_to_message_id[
tool_call_id
] = message.id
yield # 通知前端更新渲染
# 文本分片
case TextPart(content=content):
# 构建消息
message = Message(type=MessageType.TEXT, content=content)
# 添加至消息字典
task.messages[message.id] = message
# 将消息唯一标识与片段索引映射
index_map_to_message_id[index] = message.id
yield # 通知前端更新渲染
# ========== 增量事件 ==========
case PartDeltaEvent(index=index, delta=delta):
# 获取消息
match delta:
# 思考分片增量
case ThinkingPartDelta(
content_delta=content_delta,
provider_details=provider_details,
):
message = task.messages[index_map_to_message_id[index]]
# 若内容增量非空则将内容追加至消息内容,否则若供应商详情可调用更新消息内容
if content_delta:
message.content += content_delta
yield # 通知前端更新渲染
else:
if callable(provider_details):
# 获取原始内容
raw_content = (
previous_provider_details := provider_details(
previous_provider_details
)
).get("raw_content")
if isinstance(raw_content, str):
...
elif isinstance(raw_content, list):
raw_content = "".join(raw_content)
else:
continue
message.content = raw_content
yield # 通知前端更新渲染
# 文本分片增量
case TextPartDelta(
content_delta=content_delta,
):
message = task.messages[index_map_to_message_id[index]]
# 增量更新
message.content += content_delta
yield # 通知前端更新渲染
# ========== 结束事件 ==========
case PartEndEvent(
index=index,
part=part,
):
match part:
# 思考分片
case ThinkingPart():
# 获取消息
message = task.messages[index_map_to_message_id[index]]
message.is_running = False
message.title = "思考已完成"
yield # 通知前端更新渲染
messages.append(message)
# 文本分片
case TextPart():
# 获取消息
message = task.messages[index_map_to_message_id[index]]
messages.append(message)
# ========== 工具调用结果事件 ==========
case FunctionToolResultEvent(part=part):
match part:
# 工具返回分片
case ToolReturnPart(
tool_call_id=tool_call_id,
tool_name=tool_name,
content=content,
):
# 获取消息
message = task.messages[
task.workflow.tool_call_id_map_to_message_id[
tool_call_id
]
]
message.is_running = False
message.title = f"{tool_name}已完成"
if isinstance(content, str) and content:
message.content = content
yield # 通知前端更新渲染
messages.append(message)
# ========== 智能体运行时事件 ==========
case AgentRuntimeEvent(tool_call_id=tool_call_id, content=content):
# 若工具调用唯一标识为空则创建消息并添加至消息字典,否则若消息内容非空则根据工具调用唯一标识获取消息并更新消息内容
if not tool_call_id:
message = Message(type=MessageType.TEXT, content=content)
task.messages[message.id] = message
yield # 通知前端更新渲染
messages.append(message)
else:
if content:
message = task.messages[
task.workflow.tool_call_id_map_to_message_id[
tool_call_id
]
]
message.content = content
yield # 通知前端更新渲染
# ========== 智能体运行结果事件 ==========
case AgentRunResultEvent(result=result):
if isinstance(output := result.output, DeferredToolRequests):
for appr in output.approvals:
# 获取消息
message = task.messages[
task.workflow.tool_call_id_map_to_message_id[
tool_call_id := appr.tool_call_id
]
]
message.content = output.metadata[tool_call_id]["content"]
message.approval = Approval(
tool_call_id=tool_call_id,
)
message.is_shown = True
yield # 通知前端更新渲染
messages.append(message)
else:
# 若工作流类型为预定航班则设置为对话
if task.workflow.type == WorkflowType.BOOK_FLIGHT:
task.workflow = Workflow()
# 更新上一次运行日期时间
task.last_run_at = datetime.now()
# 更新正在运行状态
task.is_running = False
yield # 通知前端更新渲染
# 更新使用量
task.usage = result.usage
# 更新任务记录
await database_state.update_task_record(
id=task.id,
last_run_at=task.last_run_at,
usage=task.usage,
workflow=task.workflow,
)
# 创建运行记录
await database_state.create_run_record(
task_id=task.id,
new_messages=result.new_messages(),
)
# 批量创建消息记录
await database_state.create_message_records(
task_id=task.id,
messages=messages,
)
@rx.var
def is_running(self) -> bool:
"""
当前会话正在运行
:return: 当前会话正在运行(True 表示正在运行,False 表示运行完成)
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return False
return task.is_running
@rx.var
def awaiting_stream(self) -> bool:
"""
当前会话正在等待流式输出
:return: 当前会话正在等待流式输出(True 表示等待流式输出,False 表示已开始流式输出或已完成)
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return False
return task.awaiting_stream
@rx.event
def toggle_message_history_item_shown(self, item_id: str) -> None:
"""
展示 / 隐藏消息历史项
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return
task.messages[item_id].is_shown ^= True
@rx.event
def set_approval_content(self, item_id: str, content: str) -> None:
"""
设置审批内容
:param item_id: 消息历史项唯一标识
:param content: 内容
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return
# 获取消息
message = task.messages.get(item_id)
if not message or not isinstance(approval := message.approval, Approval):
return
if not content.strip():
return
approval.content = content
@rx.event
async def set_workflow(self, type: WorkflowType) -> None:
"""
设置工作流
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return
match type:
# 预定航班
case WorkflowType.BOOK_FLIGHT:
from application.workshop.book_flight import Deps
# 初始化预定航班工作流
workflow = Workflow(
type=type,
deps=Deps(
date=Date(2025, 1, 10),
origin_airport_code="SFO",
destination_airport_code="ANC",
),
)
# 设置用户提示词
self.set_user_prompt(
f"帮我查询并预定在 {workflow.deps.date.strftime('%Y-%m-%d')} 从 {workflow.deps.origin_airport_code} 到 {workflow.deps.destination_airport_code} 的航班"
)
# 银行客户服务
case WorkflowType.BANK_CUSTOMER_SERVICE:
from application.workshop.bank_customer_service import Deps
# 初始化银行客户服务工作流
workflow = Workflow(type=type, deps=Deps(user_id=self.user_id))
# 设置用户提示词
self.set_user_prompt(f"帮我查询账户余额")
# 联网检索
case WorkflowType.WEB_RETRIEVE:
from application.workshop.web_retrieve import Deps
# 初始化联网检索工作流
workflow = Workflow(type=type, deps=Deps())
# 知识库检索
case WorkflowType.KNOWLEDGE_BASE_RETRIEVE:
from application.workshop.knowledge_base_retrieve import Deps
# 初始化知识库检索工作流
workflow = Workflow(type=type, deps=Deps())
# 其他
case _:
# 初始化对话工作流
workflow = Workflow()
# 设置工作流
task.workflow = workflow
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 更新任务记录
await database_state.update_task_record(
id=task.id,
workflow=task.workflow,
)

View File

@ -0,0 +1,804 @@
# -*- coding: utf-8 -*-
"""
任务状态
"""
from datetime import date as Date, datetime
from asyncio import create_task
from copy import copy
from typing import AsyncGenerator
from reflex_base.event import EventSpec
from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolApproved
from pydantic_ai.messages import (
FunctionToolResultEvent,
PartDeltaEvent,
PartEndEvent,
PartStartEvent,
TextPart,
TextPartDelta,
ThinkingPart,
ThinkingPartDelta,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.run import AgentRunResultEvent
import reflex as rx
from reflex.event import EventCallback
from application.states.database_states import DatabaseState
from application.states.models import (
Approval,
Task,
TaskHistoryItem,
Message,
MessageHistoryItem,
MessageType,
CapabilityComponent,
CapabilityComponentType,
)
from application.agent.events import AgentRuntimeEvent
class TaskState(rx.State):
"""
任务状态
"""
# 当前用户唯一标识
user_id: str = ""
# 当前用户的任务字典
tasks: dict[str, Task] = {} # 键为任务唯一标识,值为任务实例
# 当前任务唯一标识
task_id: str = ""
# 临时用户提示词
temporary_user_prompt: str = ""
# 任务删除确认打开状态
is_task_delete_confirm_open: bool = False
# 任务删除确认打开日期时间
task_delete_confirm_open_at: float = 0
async def resume(self, user_id: str) -> None:
"""
恢复任务状态
:param user_id: 用户唯一标识
:return: None
"""
self.user_id = user_id
if not self.user_id:
return
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 查询当前用户的任务字典
self.tasks = await database_state.read_tasks(self.user_id)
for task in self.tasks.values():
# 最后一个消息
message = next(iter(reversed(task.messages.values())), None)
# 若最后一个消息包含审批且审批结果为 None 则任务正在运行设置为是、消息正在运行设置为是
if (
message
and isinstance(approval := message.approval, Approval)
and approval.result is None
):
task.is_running = True
message.is_running = True
@rx.var
def task_history(self) -> list[TaskHistoryItem]:
"""
获取任务历史项列表(按照任务创建日期时间降序排序)
:return: 任务历史项列表
"""
def formatter(dt: datetime) -> str:
"""
格式器
:param dt: 日期时间
:return: 格式化后日期时间字符串
"""
# 时间间隔
delta = (datetime.now() - dt).total_seconds()
if delta < 60:
return "刚刚"
elif delta < 60 * 60:
return f"{round(delta / 60)}分钟前"
elif delta < 60 * 60 * 24:
return f"{round(delta / 60 / 60)}小时前"
else:
return f"{round(delta / 60 / 60 / 24)}天前"
return [
TaskHistoryItem(
id=task.id,
description=task.description,
last_run_at=formatter(task.last_run_at),
is_running=task.is_running,
)
for task in reversed(self.tasks.values())
]
@rx.var
def task_description(self) -> str:
"""
获取当前任务的描述
:return: 当前任务的描述
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return ""
return task.description
@rx.var
def task_last_run_at(self) -> str:
"""
获取当前任务的上一次运行日期时间
:return: 当前任务的上一次运行日期时间
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return ""
return f"更新于 {task.last_run_at.strftime("%Y-%m-%d %H:%M:%S")}"
@rx.event
async def toggle_task_delete_confirm_open(self, is_open: bool) -> None:
"""
切换任务删除确认打开状态
:param is_open: 任务删除确认打开状态
:return: None
"""
# 当前日期时间
current = datetime.now().timestamp()
# 根据当前时间与任务删除确认打开日期时间间隔防抖
if current - self.task_delete_confirm_open_at < 0.2:
return
self.is_task_delete_confirm_open = is_open
self.task_delete_confirm_open_at = current
@rx.event
async def delete_task(self) -> EventSpec | None:
"""
删除当前任务
:return: RX事件
"""
from application.states.sidebar_states import SidebarState
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 更新任务记录
await database_state.update_task_record(task.id, is_deleted=True)
del self.tasks[task.id]
# 获取侧边栏状态
sidebar_state = await self.get_state(SidebarState)
# 将创建任务设置为导航按钮
sidebar_state.nav_button = ""
# 将空设置为当前任务唯一标识
self.task_id = ""
# 将任务删除确认显示设置为否
self.is_task_delete_confirm_open = False
return rx.toast.success("删除成功")
@rx.event
def set_task_id(self, task_id: str) -> None:
"""
将指定任务唯一标识设置为当前任务唯一标识
:param task_id: 任务唯一标识
:return: None
"""
self.task_id = task_id
@rx.var
def message_history(self) -> list[MessageHistoryItem]:
"""
获取当前任务的消息历史
:return: 当前任务的消息历史
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return []
message_history: list[MessageHistoryItem] = []
for message in task.messages.values():
message_history.append(
MessageHistoryItem(
**message.model_dump(),
has_approval=isinstance(approval := message.approval, Approval)
and approval.result is None,
approval_content=getattr(approval, "content", ""),
)
)
return message_history
@rx.var
def user_prompt(self) -> str:
"""
用户提示词
:return: 用户提示词
"""
# 当前任务
task = self.tasks.get(self.task_id)
# 若当前任务存在则返回当前任务的用户提示词,否则返回临时用户提示词
if task:
return task.user_prompt
else:
return self.temporary_user_prompt
@rx.event
def set_user_prompt(self, user_prompt: str) -> None:
"""
设置用户提示词
:param user_prompt: 用户提示词
:return: None
"""
# 当前任务
task = self.tasks.get(self.task_id)
# 若当前任务存在则将用户提示词设置为当前任务的用户提示词,否则将用户提示词设置为临时用户提示词
if task:
task.user_prompt = user_prompt
else:
self.temporary_user_prompt = user_prompt
@rx.var
def is_user_prompt_sending_button_disabled(self) -> bool:
"""
用户提示词发送按钮不可点击状态
:return: bool
"""
# 当前任务
task = self.tasks.get(self.task_id)
# 若当前任务存在则判断该任务是否正在运行或用户提示词非空
if task:
return task.is_running or not task.user_prompt
else:
return not self.temporary_user_prompt
@rx.event
async def run(
self,
is_resume_run: bool = False,
message_id: str | None = None,
) -> AsyncGenerator[EventCallback, None]:
"""
运行
:param is_resume_run: 恢复运行
:param message_id: 消息唯一标识
:return: None
"""
from application.states.sidebar_states import SidebarState
user_prompt = None # 用户提示词
deferred_tool_results = None # 延迟工具结果
async with self:
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 若当前任务的唯一标识为空则创建任务并将此作为当前任务
if not self.task_id:
# 获取侧边栏状态
sidebar_state = await self.get_state(SidebarState)
# 将选中的导航按钮设置为空
sidebar_state.nav_button = None
# 先创建任务记录再将此添加至当前用户的任务字典
self.tasks.update(
await database_state.create_task_record(
self.user_id, self.temporary_user_prompt
)
)
# 将最后一个任务的唯一标识作为当前任务的唯一标识
self.task_id = next(reversed(self.tasks.keys()))
# 当前任务
task = self.tasks[task_id := self.task_id]
# 若为恢复运行
if is_resume_run:
if not message_id:
return
# 指定消息唯一标识的消息
message = task.messages.get(message_id)
if (
not message
or not isinstance(approval := message.approval, Approval)
or not approval.content
):
return
# 将当前任务的等待流式输出状态设置为等待流式输出
task.awaiting_stream_output = True
# 更新消息内容
message.content = f"您输入的是:\n\n{approval.content}\n\n"
# 更新审批结果
approval.result = True
# 构建延迟工具结果
deferred_tool_results = DeferredToolResults(
approvals={approval.tool_call_id: ToolApproved()}, # 审批通过
metadata={approval.tool_call_id: {"content": approval.content}},
)
else:
# 若用户提示词为空则直接返回
if not (user_prompt := task.user_prompt.strip()):
return
# 清空临时用户提示词和当前任务的用户提示词
self.temporary_user_prompt = task.user_prompt = ""
# 将当前任务的运行状态设置为正在运行
task.is_running = True
# 将当前任务的等待流式输出状态设置为等待流式输出
task.awaiting_stream_output = True
# 先创建消息记录再添加至当前任务的消息字典
task.messages.update(
await database_state.create_message_record(
task_id=task.id,
message=Message(
type=MessageType.USER_PROMPT,
content=user_prompt,
),
)
)
yield type(self).background_run(task_id, user_prompt, deferred_tool_results)
@rx.event(background=True)
async def background_run(
self,
task_id: str,
user_prompt: str | None,
deferred_tool_results: DeferredToolResults | None = None,
):
"""
后台运行运行
:param task_id: 任务唯一标识
:return: None
"""
async with self:
# 当前任务
task = self.tasks[task_id]
# 功能组件
capability_component = task.capability_component
# 使用量
usage = copy(task.usage)
if usage.details:
usage.details = dict(usage.details)
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 查询消息历史列表
message_history = await database_state.read_message_history(task_id)
# 根据功能组件的类型匹配流式输出
match capability_component.type:
case CapabilityComponentType.TALK:
from application.agent.talk import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 预定航班
case CapabilityComponentType.BOOK_FLIGHT:
from application.agent.book_flight import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=capability_component.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
deferred_tool_results=deferred_tool_results,
)
# 银行客户服务
case CapabilityComponentType.BANK_CUSTOMER_SERVICE:
from application.agent.bank_customer_service import (
run_stream_events,
)
from application.agent.database import connection_pool
# 若连接池为空则直接返回
if not connection_pool:
return
# 运行并流式输出事件
stream_events = run_stream_events(
deps=capability_component.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 联网检索
case CapabilityComponentType.WEB_RETRIEVE:
from application.agent.web_retrieve import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=capability_component.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
# 知识库检索
case CapabilityComponentType.KNOWLEDGE_BASE_RETRIEVE:
from application.agent.knowledge_base_retrieve import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
deps=capability_component.deps,
user_prompt=user_prompt,
message_history=message_history,
usage=usage,
)
messages: list[Message] = [] # 消息列表
previous_provider_details = {} # 上一个供应商详情
# 初始化片段索引映射为消息唯一标识字典
index_map_to_message_id: dict[int, str] = {}
# 获取运行流式输出事件
async for event in stream_events:
# 若为等待流式输出则将当前任务的等待流式输出状态设置为已开始流式输出
if task.awaiting_stream_output:
task.awaiting_stream_output = False
match event:
# ========== 开始事件 ==========
case PartStartEvent(
index=index,
part=part,
):
match part:
# 思考分片
case ThinkingPart(content=content):
# 构建消息
message = Message(
type=MessageType.THINKING,
title="正在思考",
content=content,
is_running=True,
)
# 添加至消息字典
task.messages[message.id] = message
# 将消息唯一标识与片段索引映射
index_map_to_message_id[index] = message.id
yield # 通知前端更新渲染
# 工具调用分片
case ToolCallPart(
tool_call_id=tool_call_id, tool_name=tool_name
):
if (
tool_call_id
not in capability_component.tool_call_id_map_to_message_id
):
# 构建消息
message = Message(
type=MessageType.TOOL_CALL,
title=f"正在{tool_name}",
is_running=True,
)
# 添加至消息字典
task.messages[message.id] = message
# 将消息唯一标识与工具调用唯一标识映射
capability_component.tool_call_id_map_to_message_id[
tool_call_id
] = message.id
yield # 通知前端更新渲染
# 文本分片
case TextPart(content=content):
# 构建消息
message = Message(
type=MessageType.TEXT, content=content
)
# 添加至消息字典
task.messages[message.id] = message
# 将消息唯一标识与片段索引映射
index_map_to_message_id[index] = message.id
yield # 通知前端更新渲染
# ========== 增量事件 ==========
case PartDeltaEvent(index=index, delta=delta):
# 获取消息
match delta:
# 思考分片增量
case ThinkingPartDelta(
content_delta=content_delta,
provider_details=provider_details,
):
message = task.messages[index_map_to_message_id[index]]
# 若内容增量非空则将内容追加至消息内容,否则若供应商详情可调用更新消息内容
if content_delta:
message.content += content_delta
yield # 通知前端更新渲染
else:
if callable(provider_details):
# 获取原始内容
raw_content = (
previous_provider_details := provider_details(
previous_provider_details
)
).get("raw_content")
if isinstance(raw_content, str):
...
elif isinstance(raw_content, list):
raw_content = "".join(raw_content)
else:
continue
message.content = raw_content
yield # 通知前端更新渲染
# 文本分片增量
case TextPartDelta(
content_delta=content_delta,
):
message = task.messages[index_map_to_message_id[index]]
# 增量更新
message.content += content_delta
yield # 通知前端更新渲染
# ========== 结束事件 ==========
case PartEndEvent(
index=index,
part=part,
):
match part:
# 思考分片
case ThinkingPart():
# 获取消息
message = task.messages[index_map_to_message_id[index]]
message.is_running = False
message.title = "思考已完成"
yield # 通知前端更新渲染
messages.append(message)
# 文本分片
case TextPart():
# 获取消息
message = task.messages[index_map_to_message_id[index]]
messages.append(message)
# ========== 工具调用结果事件 ==========
case FunctionToolResultEvent(part=part):
match part:
# 工具返回分片
case ToolReturnPart(
tool_call_id=tool_call_id,
tool_name=tool_name,
content=content,
):
# 获取消息
message = task.messages[
capability_component.tool_call_id_map_to_message_id[
tool_call_id
]
]
message.is_running = False
message.title = f"{tool_name}已完成"
if isinstance(content, str) and content:
message.content = content
yield # 通知前端更新渲染
messages.append(message)
# ========== 智能体运行时事件 ==========
case AgentRuntimeEvent(tool_call_id=tool_call_id, content=content):
# 若工具调用唯一标识为空则创建消息并添加至消息字典,否则若消息内容非空则根据工具调用唯一标识获取消息并更新消息内容
if not tool_call_id:
message = Message(type=MessageType.TEXT, content=content)
task.messages[message.id] = message
yield # 通知前端更新渲染
messages.append(message)
else:
if content:
message = task.messages[
capability_component.tool_call_id_map_to_message_id[
tool_call_id
]
]
message.content = content
yield # 通知前端更新渲染
# ========== 智能体运行结果事件 ==========
case AgentRunResultEvent(result=result):
if isinstance(output := result.output, DeferredToolRequests):
for appr in output.approvals:
# 获取消息
message = task.messages[
capability_component.tool_call_id_map_to_message_id[
tool_call_id := appr.tool_call_id
]
]
message.content = output.metadata[tool_call_id][
"content"
]
message.approval = Approval(
tool_call_id=tool_call_id,
)
message.is_shown = True
yield # 通知前端更新渲染
messages.append(message)
else:
# 若能力组件类型为预定航班则设置为对话
if (
capability_component.type
== CapabilityComponentType.BOOK_FLIGHT
):
capability_component = CapabilityComponent()
# 更新上一次运行日期时间
task.last_run_at = datetime.now()
# 更新正在运行状态
task.is_running = False
yield # 通知前端更新渲染
# 更新使用量
task.usage = result.usage
# 更新任务记录
await database_state.update_task_record(
id=task.id,
last_run_at=task.last_run_at,
usage=task.usage,
capability_component=capability_component,
)
# 创建运行记录
await database_state.create_run_record(
task_id=task.id,
new_messages=result.new_messages(),
)
# 批量创建消息记录
await database_state.create_message_records(
task_id=task.id,
messages=messages,
)
@rx.var
def is_running(self) -> bool:
"""
当前会话正在运行
:return: 当前会话正在运行(True 表示正在运行,False 表示运行完成)
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return False
return task.is_running
@rx.var
def awaiting_stream_output(self) -> bool:
"""
当前会话正在等待流式输出状态
:return: 当前会话正在等待流式输出状态(True 表示等待流式输出,False 表示已开始流式输出或已完成)
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return False
return task.awaiting_stream_output
@rx.event
def toggle_message_history_item_shown(self, item_id: str) -> None:
"""
展示 / 隐藏消息历史项
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return
task.messages[item_id].is_shown ^= True
@rx.event
def set_approval_content(self, item_id: str, content: str) -> None:
"""
设置审批内容
:param item_id: 消息历史项唯一标识
:param content: 内容
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return
# 获取消息
message = task.messages.get(item_id)
if not message or not isinstance(approval := message.approval, Approval):
return
if not content.strip():
return
approval.content = content
@rx.event
async def set_capability_component(self, type: CapabilityComponentType) -> None:
"""
设置能力组件
"""
# 当前会话
task = self.tasks.get(self.task_id)
if not task:
return
match type:
# 预定航班
case CapabilityComponentType.BOOK_FLIGHT:
from application.agent.book_flight import Deps
# 初始化预定航班能力组件
capability_component = CapabilityComponent(
type=type,
deps=Deps(
date=Date(2025, 1, 10),
origin_airport_code="SFO",
destination_airport_code="ANC",
),
)
# 设置用户提示词
self.set_user_prompt(
f"帮我查询并预定在 {capability_component.deps.date.strftime('%Y-%m-%d')} 从 {capability_component.deps.origin_airport_code} 到 {capability_component.deps.destination_airport_code} 的航班"
)
# 银行客户服务
case CapabilityComponentType.BANK_CUSTOMER_SERVICE:
from application.agent.bank_customer_service import Deps
# 初始化银行客户服务能力组件
capability_component = CapabilityComponent(
type=type, deps=Deps(user_id=self.user_id)
)
# 设置用户提示词
self.set_user_prompt(f"帮我查询账户余额")
# 联网检索
case CapabilityComponentType.WEB_RETRIEVE:
from application.agent.web_retrieve import Deps
# 初始化联网检索能力组件
capability_component = CapabilityComponent(type=type, deps=Deps())
# 知识库检索
case CapabilityComponentType.KNOWLEDGE_BASE_RETRIEVE:
from application.agent.knowledge_base_retrieve import Deps
# 初始化知识库检索能力组件
capability_component = CapabilityComponent(type=type, deps=Deps())
# 其他
case _:
# 初始化对话能力组件
capability_component = CapabilityComponent()
# 设置能力组件
task.capability_component = capability_component
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 更新任务记录
await database_state.update_task_record(
id=task.id,
capability_component=task.capability_component,
)

View File

@ -67,13 +67,13 @@ pre, pre code {
font-family: Consolas, "Microsoft YaHei", sans-serif !important; font-family: Consolas, "Microsoft YaHei", sans-serif !important;
} }
.awaiting_stream { .awaiting_stream_output {
width: 6px; width: 6px;
aspect-ratio: 1; aspect-ratio: 1;
border-radius: 50%; border-radius: 50%;
animation: awaiting_stream 0.9s infinite ease-in-out; animation: awaiting_stream_output 0.9s infinite ease-in-out;
} }
@keyframes awaiting_stream { @keyframes awaiting_stream_output {
0% { 0% {
box-shadow: 12px 0 var(--prismui-color-5), -12px 0 var(--prismui-color-9); box-shadow: 12px 0 var(--prismui-color-5), -12px 0 var(--prismui-color-9);
background: var(--prismui-color-5); background: var(--prismui-color-5);

Binary file not shown.

Binary file not shown.

View File

@ -14,5 +14,5 @@ if __name__ == "__main__":
# 批量创建草稿 # 批量创建草稿
drafts_manager.batch_create( drafts_manager.batch_create(
draft_counts=150, draft_counts=600,
) )