# -*- 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, )