470 lines
17 KiB
Python
470 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
会话状态
|
||
"""
|
||
from typing import AsyncGenerator
|
||
from datetime import datetime
|
||
|
||
from pydantic_ai.messages import (
|
||
FunctionToolCallEvent,
|
||
FunctionToolResultEvent,
|
||
LoadCapabilityCallPart,
|
||
PartDeltaEvent,
|
||
PartEndEvent,
|
||
PartStartEvent,
|
||
TextPart,
|
||
TextPartDelta,
|
||
ThinkingPart,
|
||
ThinkingPartDelta,
|
||
ToolCallPart,
|
||
ToolSearchCallPart,
|
||
)
|
||
from pydantic_ai.run import AgentRunResultEvent
|
||
import reflex as rx
|
||
from pydantic_ai._uuid import uuid7
|
||
from application.states.database import DatabaseState
|
||
from application.states.models import (
|
||
Conversation,
|
||
Message,
|
||
MessageType,
|
||
WorkFlow,
|
||
WorkFlowType,
|
||
ConversationHistoryItem,
|
||
usage_validate_python,
|
||
)
|
||
|
||
|
||
def format_conversation_created_at(created_at: datetime) -> str:
|
||
"""
|
||
格式化会话创建日期时间
|
||
:param created_at: 创建日期时间
|
||
:return: 格式化后的日期时间字符串
|
||
"""
|
||
match (datetime.now().date() - created_at.date()).days:
|
||
case 0:
|
||
formatted_created_at = f"{created_at.strftime('%H:%M')}"
|
||
case 1:
|
||
formatted_created_at = f"昨天 {created_at.strftime('%H:%M')}"
|
||
case _:
|
||
formatted_created_at = created_at.strftime("%Y-%m-%d %H:%M")
|
||
return formatted_created_at
|
||
|
||
|
||
class ConversationState(rx.State):
|
||
"""
|
||
会话状态
|
||
"""
|
||
|
||
# 当前用户唯一标识
|
||
user_id: str = ""
|
||
# 当前用户的会话字典（私有变量）
|
||
|
||
# 当前用户的会话字典
|
||
conversations: dict[str, Conversation] = {} # 键为会话唯一标识，值为会话实例
|
||
# 激活会话唯一标识
|
||
actived_conversation_id: str = ""
|
||
|
||
# 显示会话历史，True 表示显示，False 表示隐藏
|
||
is_conversation_history_shown: bool = False
|
||
# 会话历史项显示悬停气泡
|
||
is_conversation_history_item_popover_shown: 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_conversation_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.conversations = await db_state.get_conversations(self.user_id)
|
||
# 若当前用户的会话字典为空则先创建会话记录再添加会话实例
|
||
if not self.conversations:
|
||
self.conversations.update(
|
||
await db_state.create_conversation_record(self.user_id)
|
||
)
|
||
# 将最后一个会话的唯一标识作为激活会话唯一标识
|
||
self.actived_conversation_id = next(reversed(self.conversations.keys()))
|
||
|
||
@rx.event
|
||
def toggle_conversation_history_shown(self) -> None:
|
||
"""
|
||
显示 / 隐藏会话历史
|
||
:return: None
|
||
"""
|
||
self.is_conversation_history_shown = not self.is_conversation_history_shown
|
||
|
||
@rx.var
|
||
def conversation_history(self) -> list[ConversationHistoryItem]:
|
||
"""
|
||
获取会话历史项列表（按照会话创建日期时间降序排序）
|
||
:return: 会话历史项列表
|
||
"""
|
||
return [
|
||
ConversationHistoryItem(
|
||
id=conversation.id,
|
||
description=conversation.description,
|
||
created_at=format_conversation_created_at(conversation.created_at),
|
||
)
|
||
for conversation in reversed(self.conversations.values())
|
||
]
|
||
|
||
@rx.event
|
||
def set_conversation_history_item_popover_shown(
|
||
self, conversation_id: str, is_shown: bool
|
||
) -> None:
|
||
"""
|
||
设置会话历史项悬停气泡显示 / 隐藏
|
||
:return: None
|
||
"""
|
||
self.is_conversation_history_item_popover_shown = (
|
||
conversation_id if is_shown else ""
|
||
)
|
||
|
||
@rx.event
|
||
async def delete_conversation(self, conversation_id: str) -> None:
|
||
"""
|
||
删除会话
|
||
:param conversation_id: 需删除的会话唯一标识
|
||
:return: None
|
||
"""
|
||
# 获取当前数据库状态
|
||
db_state = await self.get_db_state()
|
||
# 更新会话记录
|
||
await db_state.update_conversation_record(conversation_id, is_deleted=True)
|
||
del self.conversations[conversation_id]
|
||
|
||
# 删除后，若当前用户的会话字典为空则创建会话
|
||
if not self.conversations:
|
||
self.conversations.update(
|
||
await db_state.create_conversation_record(self.user_id)
|
||
)
|
||
|
||
# 删除后，若激活会话唯一标识不存在则将最后一个会话的唯一标识作为激活会话唯一标识
|
||
if self.actived_conversation_id not in self.conversations:
|
||
self.actived_conversation_id = next(reversed(self.conversations.keys()))
|
||
|
||
@rx.event
|
||
def set_actived_conversation(self, conversation_id: str) -> None:
|
||
"""
|
||
设置激活会话
|
||
:param conversation_id: 会话唯一标识
|
||
:return: None
|
||
"""
|
||
self.actived_conversation_id = conversation_id
|
||
|
||
@rx.var
|
||
def message_history(self) -> list[Message]:
|
||
"""
|
||
获取当前会话的消息历史
|
||
:return: 当前会话的消息历史
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return []
|
||
return list(conversation.messages.values())
|
||
|
||
@rx.event
|
||
async def create_conversation(self) -> None:
|
||
"""
|
||
创建会话
|
||
:return: None
|
||
"""
|
||
# 获取当前数据库状态
|
||
db_state = await self.get_db_state()
|
||
# 创建会话记录再添加会话实例
|
||
self.conversations.update(
|
||
await db_state.create_conversation_record(self.user_id)
|
||
)
|
||
# 将最后一个会话的唯一标识作为激活会话唯一标识
|
||
self.actived_conversation_id = next(reversed(self.conversations.keys()))
|
||
|
||
@rx.var
|
||
def user_prompt(self) -> str:
|
||
"""
|
||
用户提示词
|
||
:return: 用户提示词
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return ""
|
||
return conversation.user_prompt
|
||
|
||
@rx.event
|
||
def set_user_prompt(self, user_prompt: str) -> None:
|
||
"""
|
||
设置用户提示词
|
||
:param user_prompt: 用户提示词
|
||
:return: None
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return
|
||
conversation.user_prompt = user_prompt.strip()
|
||
|
||
@rx.var
|
||
def is_user_prompt_sending_button_disabled(self) -> bool:
|
||
"""
|
||
用户提示词发送按钮不可点击状态
|
||
:return: bool
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return True
|
||
return conversation.is_running or not conversation.user_prompt
|
||
|
||
@rx.var
|
||
def is_running(self) -> bool:
|
||
"""
|
||
当前会话正在运行
|
||
:return: 当前会话正在运行（True 表示正在运行，False 表示运行完成）
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return False
|
||
return conversation.is_running
|
||
|
||
@rx.var
|
||
def awaiting_stream(self) -> bool:
|
||
"""
|
||
当前会话正在等待流式输出
|
||
:return: 当前会话正在等待流式输出（True 表示等待流式输出，False 表示已开始流式输出或已完成）
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return False
|
||
return conversation.awaiting_stream
|
||
|
||
@rx.event
|
||
async def run(self) -> AsyncGenerator[None]:
|
||
"""
|
||
运行
|
||
:return: AsyncGenerator[None]
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return
|
||
|
||
# 若用户提示词为空则直接返回
|
||
if not conversation.user_prompt:
|
||
return
|
||
|
||
# 将正在运行设置为是
|
||
conversation.is_running = True
|
||
# 将等待流式输出设置为是
|
||
conversation.awaiting_stream = True
|
||
yield # 通知前端更新渲染
|
||
|
||
# 获取数据库状态
|
||
db_state = await self.get_db_state()
|
||
# 先创建消息记录再添加消息实例
|
||
conversation.messages.update(
|
||
await db_state.create_message_record(
|
||
conversation_id=conversation.id,
|
||
message=Message(
|
||
type=MessageType.USER_PROMPT,
|
||
content=(user_prompt := conversation.user_prompt),
|
||
),
|
||
)
|
||
)
|
||
# 清空用户提示词
|
||
conversation.user_prompt = ""
|
||
yield # 通知前端更新渲染
|
||
|
||
# 初始化片段索引映射为消息实例唯一标识字典
|
||
index_map_to_message_id: dict[int, str] = {}
|
||
# 初始化工具调用唯一标识集合
|
||
tool_call_ids: set[str] = set()
|
||
|
||
# 匹配工作流类型
|
||
match conversation.work_flow:
|
||
# 预定航班
|
||
case WorkFlowType.BOOK_FLIGHT:
|
||
from application.workshop.book_flight import run_stream_events
|
||
|
||
# 非结构化对话
|
||
case _:
|
||
from application.workshop.unstructured_dialogue import (
|
||
run_stream_events,
|
||
)
|
||
|
||
# 消息列表
|
||
messages: list[Message] = []
|
||
# 运行并流式输出事件
|
||
stream_events = run_stream_events(
|
||
user_prompt=user_prompt,
|
||
message_history=await db_state.get_message_history(
|
||
conversation_id=self.actived_conversation_id
|
||
),
|
||
usage=usage_validate_python(dict(conversation.usage)),
|
||
)
|
||
try:
|
||
# 获取运行流式输出事件
|
||
async for event in stream_events:
|
||
# 将等待流式输出设置为否
|
||
conversation.awaiting_stream = False
|
||
message: Message | None = None
|
||
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,
|
||
)
|
||
# 添加至消息字典
|
||
conversation.messages[message.id] = message
|
||
# 将消息实例唯一标识与片段索引映射
|
||
index_map_to_message_id[index] = message.id
|
||
yield # 通知前端更新渲染
|
||
|
||
# 文本分片开始事件
|
||
case TextPart(content=content):
|
||
# 构建消息实例
|
||
message = Message(type=MessageType.RESULT_OUTPUT)
|
||
message.content = content
|
||
# 添加至消息字典
|
||
conversation.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,
|
||
):
|
||
# 若内容增量不为空则增量更新
|
||
if content_delta:
|
||
conversation.messages[
|
||
index_map_to_message_id[index]
|
||
].content += content_delta
|
||
yield # 通知前端更新渲染
|
||
|
||
# 文本分片增量事件
|
||
case TextPartDelta(
|
||
content_delta=content_delta,
|
||
):
|
||
# 增量更新
|
||
conversation.messages[
|
||
index_map_to_message_id[index]
|
||
].content += content_delta
|
||
yield # 通知前端更新渲染
|
||
|
||
# ========== 结束事件 ==========
|
||
case PartEndEvent(
|
||
index=index,
|
||
part=part,
|
||
):
|
||
match part:
|
||
# 思考分片结束事件
|
||
case ThinkingPart(content=content):
|
||
# 获取消息实例
|
||
message = conversation.messages[
|
||
index_map_to_message_id[index]
|
||
]
|
||
message.is_running = False
|
||
message.title = "思考完成"
|
||
yield # 通知前端更新渲染
|
||
messages.append(message)
|
||
|
||
case TextPart(content=content):
|
||
# 获取消息实例
|
||
message = conversation.messages[
|
||
index_map_to_message_id[index]
|
||
]
|
||
messages.append(message)
|
||
|
||
# ========== 智能体运行结果事件 ==========
|
||
case AgentRunResultEvent(result=result):
|
||
# 更新使用量
|
||
await db_state.update_conversation_record(
|
||
conversation.id, usage=result.usage
|
||
)
|
||
# 创建运行记录
|
||
await db_state.create_run_record(
|
||
conversation_id=conversation.id,
|
||
new_messages=result.new_messages(),
|
||
)
|
||
|
||
except Exception as e:
|
||
...
|
||
finally:
|
||
# 将正在运行设置为否
|
||
conversation.is_running = False
|
||
# 将等待流式输出设置为否
|
||
conversation.awaiting_stream = False
|
||
self.conversations[self.actived_conversation_id] = conversation
|
||
yield
|
||
|
||
# 批量创建消息记录
|
||
await db_state.create_message_records(
|
||
conversation_id=conversation.id,
|
||
messages=messages,
|
||
)
|
||
|
||
@rx.event
|
||
async def init_work_flow(self, work_flow_type: WorkFlowType) -> None:
|
||
"""
|
||
初始化工作流
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return
|
||
|
||
match work_flow_type:
|
||
# 预定航班工作流
|
||
case WorkFlowType.BOOK_FLIGHT:
|
||
from application.workshop.book_flight import init_work_flow
|
||
|
||
# 初始化预定航班任务
|
||
conversation.work_flow = (work_flow := init_work_flow())
|
||
|
||
@rx.event
|
||
def toggle_message_history_item_shown(self, message_id: str) -> None:
|
||
"""
|
||
展示 / 隐藏消息历史项
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.actived_conversation_id)
|
||
if not conversation:
|
||
return
|
||
conversation.messages[message_id].is_shown ^= True
|