diff --git a/agent/application/workshop/__init__.py b/agent/application/agent/__init__.py similarity index 100% rename from agent/application/workshop/__init__.py rename to agent/application/agent/__init__.py diff --git a/agent/application/workshop/bank_customer_service.py b/agent/application/agent/bank_customer_service.py similarity index 98% rename from agent/application/workshop/bank_customer_service.py rename to agent/application/agent/bank_customer_service.py index be407a8..81f906c 100644 --- a/agent/application/workshop/bank_customer_service.py +++ b/agent/application/agent/bank_customer_service.py @@ -17,8 +17,8 @@ from pydantic_ai import ( from pydantic_ai.messages import AgentStreamEvent from pydantic_ai.run import AgentRunResultEvent -from application.workshop.database import connection_pool -from application.workshop.models import ( +from application.agent.database import connection_pool +from application.agent.models import ( DEEPSEEK_V4_FLASH_CHAT_MODEL, DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS, ) diff --git a/agent/application/workshop/book_flight.py b/agent/application/agent/book_flight.py similarity index 99% rename from agent/application/workshop/book_flight.py rename to agent/application/agent/book_flight.py index 5aa7eda..d03e4d4 100644 --- a/agent/application/workshop/book_flight.py +++ b/agent/application/agent/book_flight.py @@ -26,8 +26,8 @@ from pydantic_ai.messages import ( ) from pydantic_ai.run import AgentRunResultEvent, AgentRunResult -from application.workshop.events import AgentRuntimeEvent -from application.workshop.models import ( +from application.agent.events import AgentRuntimeEvent +from application.agent.models import ( DEEPSEEK_V4_FLASH_CHAT_MODEL, DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS, ) @@ -302,7 +302,6 @@ async def run_stream_events( deferred_tool_results=deferred_tool_results, ) as events: async for event in events: - print(event) match event: case PartStartEvent( part=part, diff --git a/agent/application/workshop/database.py b/agent/application/agent/database.py similarity index 100% rename from agent/application/workshop/database.py rename to agent/application/agent/database.py diff --git a/agent/application/workshop/events.py b/agent/application/agent/events.py similarity index 100% rename from agent/application/workshop/events.py rename to agent/application/agent/events.py diff --git a/agent/application/workshop/knowledge_base_retrieve.py b/agent/application/agent/knowledge_base_retrieve.py similarity index 98% rename from agent/application/workshop/knowledge_base_retrieve.py rename to agent/application/agent/knowledge_base_retrieve.py index 4e40f62..9348f9b 100644 --- a/agent/application/workshop/knowledge_base_retrieve.py +++ b/agent/application/agent/knowledge_base_retrieve.py @@ -13,8 +13,8 @@ from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.usage import RunUsage -from application.workshop.database import connection_pool -from application.workshop.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL +from application.agent.database import connection_pool +from application.agent.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL class Deps(BaseModel): diff --git a/agent/application/workshop/models.py b/agent/application/agent/models.py similarity index 100% rename from agent/application/workshop/models.py rename to agent/application/agent/models.py diff --git a/agent/application/agent/runtime.py b/agent/application/agent/runtime.py new file mode 100644 index 0000000..5bf6fe7 --- /dev/null +++ b/agent/application/agent/runtime.py @@ -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) diff --git a/agent/application/workshop/talk.py b/agent/application/agent/talk.py similarity index 95% rename from agent/application/workshop/talk.py rename to agent/application/agent/talk.py index 28ff9da..48152c6 100644 --- a/agent/application/workshop/talk.py +++ b/agent/application/agent/talk.py @@ -8,7 +8,7 @@ from pydantic_ai import Agent, ModelMessage, RunUsage from pydantic_ai.messages import AgentStreamEvent 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 = """ diff --git a/agent/application/workshop/web_retrieve.py b/agent/application/agent/web_retrieve.py similarity index 98% rename from agent/application/workshop/web_retrieve.py rename to agent/application/agent/web_retrieve.py index b38dc88..49a7f8a 100644 --- a/agent/application/workshop/web_retrieve.py +++ b/agent/application/agent/web_retrieve.py @@ -14,8 +14,8 @@ from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.usage import RunUsage -from application.workshop.database import connection_pool -from application.workshop.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL +from application.agent.database import connection_pool +from application.agent.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL """ 联网检索缓存表建表SQL diff --git a/agent/application/application.py b/agent/application/application.py index ed75d9a..8b7221f 100644 --- a/agent/application/application.py +++ b/agent/application/application.py @@ -5,7 +5,7 @@ import reflex as rx 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( diff --git a/agent/application/pages/login.py b/agent/application/pages/login.py index 2b889fd..5391a51 100644 --- a/agent/application/pages/login.py +++ b/agent/application/pages/login.py @@ -4,7 +4,7 @@ """ import reflex as rx -from application.states import AuthState +from application.states.auth_states import AuthState def login_page() -> rx.Component: diff --git a/agent/application/pages/task.py b/agent/application/pages/task.py index 8d55d00..760ac66 100644 --- a/agent/application/pages/task.py +++ b/agent/application/pages/task.py @@ -4,20 +4,24 @@ """ import reflex as rx -from application.states import AuthState, TaskState -from application.states.models import MessageHistoryItem, MessageType, WorkflowType +from application.states.task_states import TaskState +from application.states.models import ( + MessageHistoryItem, + MessageType, + CapabilityComponentType, +) -def shortcut(workflow_type: WorkflowType): +def shortcut(capability_component_type: CapabilityComponentType): """ 快捷方式 :param workflow_type: 工作类型 :return: rx.Component """ return rx.box( - workflow_type.value, - # 点击事件：设置工作流 - on_click=lambda: TaskState.set_workflow(workflow_type), + capability_component_type.value, + # 点击事件：设置功能组件 + on_click=lambda: TaskState.set_capability_component(capability_component_type), padding="10px 16px", background_color="var(--prismui-background-color-9)", border_radius="var(--prismui-border-radius-9)", @@ -75,14 +79,14 @@ def welcome() -> rx.Component: font_weight="var(--prismui-font-weight-3)", ), rx.hstack( - # 工作流按钮：预定航班 - shortcut(WorkflowType.BOOK_FLIGHT), - # 工作流按钮：银行客户服务 - shortcut(WorkflowType.BANK_CUSTOMER_SERVICE), - # 工作流按钮：联网检索 - shortcut(WorkflowType.WEB_RETRIEVE), - # 工作流按钮：知识库检索 - shortcut(WorkflowType.KNOWLEDGE_BASE_RETRIEVE), + # 快捷方式：预定航班 + shortcut(CapabilityComponentType.BOOK_FLIGHT), + # 快捷方式：银行客户服务 + shortcut(CapabilityComponentType.BANK_CUSTOMER_SERVICE), + # 快捷方式：联网检索 + shortcut(CapabilityComponentType.WEB_RETRIEVE), + # 快捷方式：知识库检索 + shortcut(CapabilityComponentType.KNOWLEDGE_BASE_RETRIEVE), flex_wrap="wrap", align_items="center", gap="12px", @@ -671,10 +675,10 @@ def message_history() -> rx.Component: rx.foreach(TaskState.message_history, message_history_item), # 若等待流式输出则显示加载动效 rx.cond( - TaskState.awaiting_stream, + TaskState.awaiting_stream_output, rx.box( rx.box( - class_name="awaiting_stream", + class_name="awaiting_stream_output", flex_shrink="0", ), display="flex", diff --git a/agent/application/pages/workbench.py b/agent/application/pages/workbench.py index 90fd0fc..9d85674 100644 --- a/agent/application/pages/workbench.py +++ b/agent/application/pages/workbench.py @@ -4,9 +4,11 @@ """ import reflex as rx -from application.pages import task_page, knowledge_base_page, login_page -from application.states import AuthState, TaskState, SidebarState +from application.pages import knowledge_base_page, login_page, task_page +from application.states.auth_states import AuthState from application.states.models import TaskHistoryItem +from application.states.sidebar_states import SidebarState +from application.states.task_states import TaskState # 导航按钮字典 diff --git a/agent/application/states/__init__.py b/agent/application/states/__init__.py index 785d232..40a96af 100644 --- a/agent/application/states/__init__.py +++ b/agent/application/states/__init__.py @@ -1,13 +1 @@ # -*- 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", -] diff --git a/agent/application/states/auth.py b/agent/application/states/auth_states.py similarity index 96% rename from agent/application/states/auth.py rename to agent/application/states/auth_states.py index e4212a1..d416603 100644 --- a/agent/application/states/auth.py +++ b/agent/application/states/auth_states.py @@ -12,8 +12,8 @@ from aiosmtplib import SMTP, SMTPException import reflex as rx from reflex.event import EventCallback -from application.states.database import DatabaseState -from application.states.task import TaskState +from application.states.task_states import TaskState +from application.states.database_states import DatabaseState # 验证码HTML模板 CAPTCHA_HTML_TEMPLATE = """ @@ -165,7 +165,7 @@ class AuthState(rx.State): # 标记验证码已发送 self.is_captcha_code_sent = True - # 并行处理：后台运行倒计时与后台运行发送验证码 + # 并行处理：将运行倒计时此事件提交至后台任务列队后，再提交运行发送验证码此事件至后台任务列队后 yield type(self).background_countdown() yield type(self).background_send_captcha_code() diff --git a/agent/application/states/database.py b/agent/application/states/database_states.py similarity index 93% rename from agent/application/states/database.py rename to agent/application/states/database_states.py index 951f1dd..b8170e1 100644 --- a/agent/application/states/database.py +++ b/agent/application/states/database_states.py @@ -20,13 +20,13 @@ from application.states.models import ( Message, MessageType, Task, - Workflow, + CapabilityComponent, approval_dump_json, approval_validate_json, usage_dump_json, usage_validate_json, - workflow_dump_json, - workflow_validate_json, + capability_component_dump_json, + capability_component_validate_json, ) @@ -92,9 +92,9 @@ class TaskRecord(SQLModel, table=True): description="任务已删除：True 表示已删除，False 表示未删除", ) usage: str = Field(default="", description="任务使用量") - workflow: str = Field( + capability_component: str = Field( default="", - description="任务工作流", + description="任务功能组件", ) # 原则上嵌套列表或字典若不涉及查询则以字符串储存 @@ -264,9 +264,9 @@ class DatabaseState(rx.State): usage=usage_validate_json( task_record.usage ), # 将 task_record.usage 由 JSON 字符串反序列化为 RunUsage - workflow=workflow_validate_json( - task_record.workflow - ), # 将 task_record.workflow 由 JSON 字符串反序列化为 Workflow, + capability_component=capability_component_validate_json( + task_record.capability_component + ), # 将 task_record.capability_component 由 JSON 字符串反序列化为 CapabilityComponent, ), ) if not message_record: @@ -290,7 +290,7 @@ class DatabaseState(rx.State): last_run_at: datetime | None = None, is_deleted: bool | None = None, usage: RunUsage | None = None, - workflow: Workflow | None = None, + capability_component: CapabilityComponent | None = None, ) -> None: """ 更新任务记录 @@ -299,7 +299,7 @@ class DatabaseState(rx.State): :param last_run_at: 上一次运行日期时间 :param is_deleted: 任务已删除 :param usage: 使用量 - :param workflow: 工作流 + :param capability_component: 功能组件 :return: None """ async with rx.asession() as session: @@ -316,10 +316,10 @@ class DatabaseState(rx.State): record.usage = usage_dump_json( RunUsage(**asdict(usage)) ) # 将 RunUsage 序列化为 JSON 字符串 - if isinstance(workflow, Workflow): - record.workflow = workflow_dump_json( - Workflow(**workflow.model_dump(mode="json")) - ) # 将 Workflow 序列化为 JSON 字符串（因 workflow 基于 pydantic 建模，故使用 model_dump 方法转为字典） + if isinstance(capability_component, CapabilityComponent): + record.capability_component = capability_component_dump_json( + CapabilityComponent(**capability_component.model_dump(mode="json")) + ) # 将 CapabilityComponent 序列化为 JSON 字符串（因 capability_component 基于 pydantic 建模，故使用 model_dump 方法转为字典） await session.commit() async def create_message_record( diff --git a/agent/application/states/models.py b/agent/application/states/models.py index 83ab0e9..1de82e1 100644 --- a/agent/application/states/models.py +++ b/agent/application/states/models.py @@ -53,9 +53,9 @@ class Message(BaseModel): ) -class WorkflowType(StrEnum): +class CapabilityComponentType(StrEnum): """ - 工作类型枚举 + 功能组件类型枚举 """ TALK = "对话" @@ -65,13 +65,15 @@ class WorkflowType(StrEnum): KNOWLEDGE_BASE_RETRIEVE = "知识库检索" -class Workflow(BaseModel): +class CapabilityComponent(BaseModel): """ - 工作流类 + 功能组件类 """ - type: WorkflowType = Field(default=WorkflowType.TALK, description="工作类型") - deps: Any = Field(default=None, description="工作依赖项") + 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="工具调用唯一标识映射为消息实例唯一标识字典" ) @@ -88,13 +90,15 @@ class Task(BaseModel): user_prompt: str = Field(default="", description="用户提示词") usage: RunUsage = Field(..., description="任务使用量") messages: dict[str, Message] = Field(default_factory=dict, description="消息字典") - workflow: Workflow = Field(default_factory=Workflow, description="工作，默认为对话") - is_running: bool = Field( - default=False, description="正在运行，True 表示正在运行， False 表示运行结束" + capability_component: CapabilityComponent = Field( + default_factory=CapabilityComponent, description="功能组件" ) - awaiting_stream: bool = Field( + is_running: bool = Field( + default=False, description="运行状态，True 表示正在运行， False 表示运行结束" + ) + awaiting_stream_output: bool = Field( default=False, - description="等待流式输出，True 表示等待流式输出，False 表示已开始流式输出或已完成", + description="等待流式输出状态，True 表示等待流式输出，False 表示已开始流式输出或已完成", ) @@ -120,26 +124,28 @@ class MessageHistoryItem(Message): 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: - return Workflow() - return WorkflowAdapter.validate_json(workflow) + if not capability_component: + return CapabilityComponent() + return CapabilityComponentAdapter.validate_json(capability_component) -def workflow_dump_json(workflow: Workflow) -> str: +def capability_component_dump_json(capability_component: CapabilityComponent) -> str: """ - 将工作流序列化 - :param workflow: 工作流 - :return: 字符串 + 将功能组件序列化 + :param capability_component: 功能组件 + :return: 序列化后功能组件 """ - return WorkflowAdapter.dump_json(workflow).decode("utf-8") + return CapabilityComponentAdapter.dump_json(capability_component).decode("utf-8") # 使用量（RunUsage）适配器 diff --git a/agent/application/states/sidebar.py b/agent/application/states/sidebar_states.py similarity index 93% rename from agent/application/states/sidebar.py rename to agent/application/states/sidebar_states.py index c2df10e..7a974af 100644 --- a/agent/application/states/sidebar.py +++ b/agent/application/states/sidebar_states.py @@ -4,7 +4,7 @@ """ import reflex as rx -from application.states.task import TaskState +from application.states.task_states import TaskState 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 diff --git a/agent/application/states/task.py b/agent/application/states/task.py deleted file mode 100644 index 497c8a8..0000000 --- a/agent/application/states/task.py +++ /dev/null @@ -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, - ) diff --git a/agent/application/states/task_states.py b/agent/application/states/task_states.py new file mode 100644 index 0000000..534a910 --- /dev/null +++ b/agent/application/states/task_states.py @@ -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, + ) diff --git a/agent/assets/styles.css b/agent/assets/styles.css index fd0a28e..c842466 100644 --- a/agent/assets/styles.css +++ b/agent/assets/styles.css @@ -67,13 +67,13 @@ pre, pre code { font-family: Consolas, "Microsoft YaHei", sans-serif !important; } -.awaiting_stream { +.awaiting_stream_output { width: 6px; aspect-ratio: 1; 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% { box-shadow: 12px 0 var(--prismui-color-5), -12px 0 var(--prismui-color-9); background: var(--prismui-color-5); diff --git a/agent/database.db b/agent/database.db index 03db19d..a4403dc 100644 Binary files a/agent/database.db and b/agent/database.db differ diff --git a/短视频自动合成_新/caches.db b/短视频自动合成_新/caches.db index 44eacdc..a903f55 100644 Binary files a/短视频自动合成_新/caches.db and b/短视频自动合成_新/caches.db differ diff --git a/短视频自动合成_新/main.py b/短视频自动合成_新/main.py index 00c47a0..50b37c6 100644 --- a/短视频自动合成_新/main.py +++ b/短视频自动合成_新/main.py @@ -14,5 +14,5 @@ if __name__ == "__main__": # 批量创建草稿 drafts_manager.batch_create( - draft_counts=150, + draft_counts=600, )