Python/agent/application/states/task.py

768 lines
28 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# -*- 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
def format_datetime(datetime: datetime) -> str:
"""
格式化日期时间
:param datetime: 日期时间
:return: 格式化后日期时间字符串
"""
match (datetime.now().date() - datetime.date()).days:
case 0:
formatted_datetime = f"{datetime.strftime('%H:%M')}"
case 1:
formatted_datetime = f"昨天 {datetime.strftime('%H:%M')}"
case _:
formatted_datetime = datetime.strftime("%Y-%m-%d %H:%M")
return formatted_datetime
class TaskState(rx.State):
"""
任务状态
"""
# 当前用户唯一标识
user_id: str = ""
# 当前用户的任务字典
tasks: dict[str, Task] = {} # 键为任务唯一标识,值为任务实例
# 当前任务唯一标识
task_id: str = ""
# 悬停气泡显示的任务历史项唯一标识(一个任务只允许存在一个悬停气泡任务历史项)
popover_shown_task_history_item_id: str = ""
# 当前数据库状态(私有变量)
# 私有变量: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.strip()
if not self.user_id:
return
# 获取当前数据库状态
db_state = await self.get_db_state()
# 查询当前用户的任务字典
self.tasks = await db_state.read_tasks(self.user_id)
# 若当前用户的任务字典为空则先创建任务记录再添加任务
if not self.tasks:
self.tasks.update(
await db_state.create_task_record(self.user_id)
)
# 将最后一个任务的唯一标识作为当前任务唯一标识
self.task_id = next(reversed(self.tasks.keys()))
# 当前任务
task = self.tasks[self.task_id]
# 最后一个消息
message = next(iter(reversed(task.messages.values())), None)
# 若最后一个消息包含审批且审批结果为 None 则任务正在运行设置为是、消息正在运行设置为是
if (
message is not None
and isinstance(approval := message.approval, Approval)
and approval.result is None
):
task.is_running = True
message.is_running = True
@rx.event
def toggle_task_history_shown(self) -> None:
"""
显示 / 隐藏任务历史
:return: None
"""
self.is_task_history_shown = not self.is_task_history_shown
@rx.var
def task_history(self) -> list[TaskHistoryItem]:
"""
获取任务历史项列表(按照任务创建日期时间降序排序)
:return: 任务历史项列表
"""
return [
TaskHistoryItem(
id=task.id,
description=task.description,
last_run_at=format_datetime(task.last_run_at),
is_running=task.is_running,
)
for task in reversed(self.tasks.values())
]
@rx.event
def set_task_history_item_popover_shown(
self, item_id: str, is_shown: bool
) -> None:
"""
设置会话历史项悬停气泡显示 / 隐藏
:return: None
"""
self.popover_shown_conversation_history_item_id = item_id if is_shown else ""
@rx.event
async def delete_task(self, task_id: str) -> None:
"""
删除任务
:param task_id: 需删除的任务唯一标识
:return: None
"""
# 获取当前数据库状态
db_state = await self.get_db_state()
# 更新任务记录
await db_state.update_task_record(task_id, is_deleted=True)
del self.tasks[task_id]
# 删除后,若当前用户的任务字典为空则创建任务
if not self.tasks:
self.tasks.update(
await db_state.create_task_record(self.user_id)
)
# 删除后,若当前任务唯一标识不存在则将最后一个任务的唯一标识作为当前任务唯一标识
if self.task_id not in self.tasks:
self.task_id = next(reversed(self.tasks.keys()))
@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.event
async def create_task(self) -> None:
"""
创建任务
:return: None
"""
# 获取当前数据库状态
db_state = await self.get_db_state()
# 创建任务记录再添加任务
self.tasks.update(
await db_state.create_task_record(self.user_id)
)
# 将最后一个任务的唯一标识作为当前任务唯一标识
self.task_id = next(reversed(self.tasks.keys()))
@rx.var
def user_prompt(self) -> str:
"""
用户提示词
:return: 用户提示词
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return ""
return task.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 not task:
return
if not user_prompt.strip():
return
task.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 not task:
return True
return task.is_running or not task.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
"""
# 当前任务
task = self.tasks.get(self.task_id)
if not task:
return
# 获取当前数据库状态
db_state = await self.get_db_state()
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):
return
# 清空用户提示词
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,
)