# -*- coding: utf-8 -*- """ 任务状态 """ from datetime import date as Date, datetime from typing import AsyncGenerator 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 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_shown: bool = False # 当前数据库状态（私有变量） # 私有变量：reflex 约定以 _ 开头的变量为私有变量，后端不序列化，前端不可使用 _db_state: DatabaseState | None = None async def get_db_state(self) -> DatabaseState: """ 获取当前数据库状态 :return: 当前数据库状态 """ # 若数据库状态为 None 则获取数据库状态，否则直接返回当前数据库状态 if not self._db_state: self._db_state = await self.get_state(DatabaseState) return self._db_state async def resume_task_state(self, user_id: str) -> None: """ 恢复任务状态 :param user_id: 用户唯一标识 :return: None """ self.user_id = user_id if not self.user_id: return # 获取当前数据库状态 db_state = await self.get_db_state() # 查询当前用户的任务字典 self.tasks = await db_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 show_task_delete_confirm(self) -> None: """ 显示当前任务删除确认 :return: None """ self.is_task_delete_confirm_shown = True @rx.event async def hide_task_delete_confirm(self) -> None: """ 隐藏当前任务删除确认 :return: None """ self.is_task_delete_confirm_shown = False @rx.event async def confirm_task_delete(self) -> None: """ 确认当前任务删除 :return: None """ from application.states import SidebarState # 当前任务 task = self.tasks.get(self.task_id) if not task: return # 获取当前数据库状态 db_state = await self.get_db_state() # 更新任务记录 await db_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 = "" @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[None]: """ 运行 :param is_resume_run: 恢复运行 :param message_id: 消息唯一标识 :return: AsyncGenerator """ from application.states import SidebarState # 获取当前数据库状态 db_state = await self.get_db_state() # 若当前任务唯一标识为空则创建任务 if not self.task_id: # 获取当前侧边栏状态 sidebar_state = await self.get_state(SidebarState) # 将空设置为导航按钮 sidebar_state.nav_button = "" # 创建任务记录再添加任务 self.tasks.update( await db_state.create_task_record( self.user_id, self.temporary_user_prompt ) ) # 将最后一个任务的唯一标识作为当前任务唯一标识 self.task_id = next(reversed(self.tasks.keys())) # 当前任务 task = self.tasks[self.task_id] user_prompt = None # 用户提示词 deferred_tool_results = None # 延迟工具结果 # 若为恢复运行 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 yield # 通知前端更新渲染 # 构建延迟工具结果 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 db_state.create_message_record( task_id=task.id, message=Message( type=MessageType.USER_PROMPT, content=user_prompt, ), ) ) yield # 通知前端更新渲染 # 初始化片段索引映射为消息唯一标识字典 index_map_to_message_id: dict[int, str] = {} # 查询消息历史列表 message_history = await db_state.read_message_history(task_id=self.task_id) # 使用量 usage = task.usage # 初始化工作流 match task.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, Deps, ) 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 = {} # 上一个供应商详情 # 获取运行流式输出事件 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 db_state.update_task_record( id=task.id, last_run_at=task.last_run_at, usage=task.usage, workflow=task.workflow, ) # 创建运行记录 await db_state.create_run_record( task_id=task.id, new_messages=result.new_messages(), ) # 批量创建消息记录 await db_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 # 获取数据库状态 db_state = await self.get_db_state() # 更新任务记录 await db_state.update_task_record( id=task.id, workflow=task.workflow, )