443 lines
17 KiB
Python
443 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
会话状态
|
||
"""
|
||
from datetime import datetime
|
||
from typing import Any, AsyncGenerator, Dict, List
|
||
from pydantic_ai import Agent, ThinkingPartDelta
|
||
from pydantic_ai.messages import (
|
||
FunctionToolCallEvent,
|
||
FunctionToolResultEvent,
|
||
LoadCapabilityCallPart,
|
||
PartDeltaEvent,
|
||
PartEndEvent,
|
||
PartStartEvent,
|
||
TextPart,
|
||
TextPartDelta,
|
||
ThinkingPart,
|
||
ToolCallPart,
|
||
ToolSearchCallPart,
|
||
)
|
||
from pydantic_ai.models.openai import OpenAIChatModel
|
||
from pydantic_ai.providers.openai import OpenAIProvider
|
||
from pydantic_ai.run import AgentRunResultEvent
|
||
import reflex as rx
|
||
|
||
from application.states.database import DatabaseState
|
||
from application.domain_models import Conversation, Dialog, ConversationHistoryItem
|
||
|
||
instructions: str = """
|
||
# 角色
|
||
专业友好AI助手，结构化解答各类问题。
|
||
|
||
# 输出硬性规则
|
||
1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容；
|
||
2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表；
|
||
3. 代码块用 ```语言名``` 包裹；
|
||
4. 重点内容标注 **粗体**/*斜体*；
|
||
5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签；
|
||
6. 内容分点拆分，排版整洁适配前端Markdown渲染。
|
||
|
||
# 行文要求
|
||
语言通俗，逻辑完整简洁，无多余废话。
|
||
"""
|
||
|
||
# 实例化智能体（因无法序列化故剥离出状态管理）
|
||
agent: Agent = Agent(
|
||
model=OpenAIChatModel(
|
||
model_name="deepseek-v4-flash",
|
||
provider=OpenAIProvider(
|
||
base_url="https://tokenhub.tencentmaas.com/v1",
|
||
api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
|
||
),
|
||
),
|
||
instructions=instructions,
|
||
capabilities=None,
|
||
output_type=str,
|
||
retries=1,
|
||
)
|
||
|
||
|
||
class ConversationState(rx.State):
|
||
"""
|
||
会话状态
|
||
"""
|
||
|
||
# 当前用户唯一标识
|
||
user_id: str = ""
|
||
# 键为会话唯一标识，值为会话实例的会话字典
|
||
conversations: Dict[str, Conversation] = {} # 按照会话唯一标识顺序排序
|
||
# 当前会话唯一标识
|
||
conversation_id: str = ""
|
||
|
||
# 会话历史展示状态，True表示展示，False表示隐藏
|
||
is_conversation_history_shown: bool = False
|
||
|
||
# 会话创建状态，True表示正在创建，False表示未正在创建
|
||
is_conversation_creating: bool = False
|
||
|
||
async def set_user_id(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.conversations = await database_state.get_conversations(
|
||
user_id=self.user_id
|
||
)
|
||
# 若会话字典为空则先创建会话记录再在会话字典中添加会话实例
|
||
if not self.conversations:
|
||
self.conversations.update(
|
||
await database_state.create_conversations_record(user_id=self.user_id)
|
||
)
|
||
# 将最后一个会话作为当前会话并更新会话唯一标识
|
||
self.conversation_id = next(reversed(self.conversations))
|
||
|
||
@rx.event
|
||
def toggle_conversation_history_shown(self) -> None:
|
||
"""
|
||
切换会话历史展示状态，用于前端是否渲染会话历史
|
||
:return: None
|
||
"""
|
||
self.is_conversation_history_shown = not self.is_conversation_history_shown
|
||
|
||
@staticmethod
|
||
def format_created_at(created_at: datetime) -> str:
|
||
"""
|
||
格式化创建时间
|
||
:param created_at: 创建时间
|
||
:return: 格式化后的创建时间
|
||
"""
|
||
match (datetime.now().date() - created_at.date()).days:
|
||
case 0:
|
||
formated_created_at = f"今天 {created_at.strftime('%H:%M')}"
|
||
case 1:
|
||
formated_created_at = f"昨天 {created_at.strftime('%H:%M')}"
|
||
case _:
|
||
formated_created_at = created_at.strftime("%Y-%m-%d %H:%M")
|
||
return formated_created_at
|
||
|
||
@rx.var
|
||
def conversation_history_items(self) -> List[ConversationHistoryItem]:
|
||
"""
|
||
获取会话历史项列表，用于前端渲染会话历史
|
||
:return: 会话历史项列表
|
||
"""
|
||
items: List[ConversationHistoryItem] = []
|
||
for conversation in reversed(
|
||
self.conversations.values()
|
||
): # 按照会话唯一标识倒序排序
|
||
items.append(
|
||
ConversationHistoryItem(
|
||
id=conversation.id,
|
||
description=conversation.description,
|
||
created_at=self.format_created_at(
|
||
created_at=conversation.created_at
|
||
),
|
||
)
|
||
)
|
||
return items
|
||
|
||
@rx.event
|
||
async def create_conversation(self, form_data: Dict[str, Any]) -> None:
|
||
"""
|
||
新建会话
|
||
:param form_data: 表单数据
|
||
:return: None
|
||
"""
|
||
# 解析会话描述
|
||
description = form_data["description"].strip() or "新会话"
|
||
|
||
# 创建会话记录再添加会话实例
|
||
database_state = await self.get_state(DatabaseState)
|
||
self.conversations.update(
|
||
await database_state.create_conversations_record(
|
||
user_id=self.user_id, description=description
|
||
)
|
||
)
|
||
# 将最后一个会话作为当前会话并更新会话唯一标识
|
||
self.conversation_id = next(reversed(self.conversations))
|
||
|
||
# 会话创建状态设置为未正在创建
|
||
self.is_conversation_creating = False
|
||
|
||
@rx.event
|
||
async def delete_conversation(self, conversation_id: str) -> None:
|
||
"""
|
||
删除指定会话
|
||
:param conversation_id: 指定会话唯一标识
|
||
:return: None
|
||
"""
|
||
# 先逻辑删除会话记录再在会话列表删除会话实例
|
||
database_state = await self.get_state(DatabaseState)
|
||
await database_state.delete_conversations_record(
|
||
conversation_id=conversation_id
|
||
)
|
||
del self.conversations[idx]
|
||
|
||
# 删除后，若会话列表为空则先创建会话记录再添加会话实例
|
||
if not self.conversations:
|
||
self.conversations.append(
|
||
await database_state.create_conversations_record(user_id=self.user_id)
|
||
)
|
||
|
||
# 删除后，若当前会话唯一标识不存在则将最后一个会话的唯一标识设置为当前会话唯一标识
|
||
if self.conversation_id not in self.conversations:
|
||
self.conversation_id = next(reversed(self.conversations.keys()))
|
||
|
||
@rx.event
|
||
def switch_conversation(self, conversation_id: str) -> None:
|
||
"""
|
||
将指定会话的唯一标识设置为当前会话唯一标识
|
||
:param conversation_id: 指定会话唯一标识
|
||
:return: None
|
||
"""
|
||
if conversation_id not in self.conversations:
|
||
return
|
||
self.conversation_id = conversation_id
|
||
|
||
@rx.event
|
||
def toggle_conversation_creating(self) -> None:
|
||
"""
|
||
切换正在新建会话状态
|
||
"""
|
||
self.is_conversation_creating = not self.is_conversation_creating
|
||
|
||
@rx.var
|
||
def dialog_history(self) -> Dict[str, Dialog]:
|
||
"""
|
||
获取当前会话的对话历史，用于前端渲染对话历史
|
||
:return: 当前会话的对话历史
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
return conversation.dialogs if conversation else {}
|
||
|
||
@rx.var
|
||
def running_status(self) -> bool:
|
||
"""
|
||
获取当前会话的运行状态
|
||
:return: 当前会话的运行状态（True 表示正在运行，False 表示运行完成）
|
||
"""
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
if not conversation:
|
||
return False
|
||
return conversation.is_running
|
||
|
||
@rx.event
|
||
async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
|
||
"""
|
||
运行
|
||
:param form_data: 表单数据
|
||
:return: AsyncGenerator
|
||
"""
|
||
# 获取问题
|
||
question = form_data.get("question", "").strip()
|
||
if not question:
|
||
return
|
||
|
||
# 当前会话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
if not conversation:
|
||
return
|
||
|
||
# 将当前会话的运行状态设置为正在运行
|
||
conversation.is_running = True
|
||
|
||
# 获取数据库状态
|
||
database_state = await self.get_state(DatabaseState)
|
||
|
||
# 检索运行结果记录并作为消息历史
|
||
message_history = await database_state.retrieve_message_history(
|
||
conversation_id=self.conversation_id
|
||
)
|
||
|
||
# 先创建对话记录再在当前对话字典中新增对话
|
||
dialog_id = await database_state.create_dialogs_record(
|
||
conversation_id=self.conversation_id, question=question
|
||
)
|
||
dialog = conversation.dialogs.setdefault(dialog_id, Dialog(question=question))
|
||
|
||
# 初始化工具调用唯一标识和片段索引映射字典
|
||
tool_call_ids: Dict[str, int] = {}
|
||
|
||
yield # 通知前端渲染
|
||
|
||
async with agent.run_stream_events(
|
||
conversation_id=self.conversation_id,
|
||
user_prompt=question,
|
||
message_history=message_history,
|
||
) as events:
|
||
async for event in events:
|
||
match event:
|
||
# ========== 开始事件 ==========
|
||
case PartStartEvent(
|
||
index=index,
|
||
part=part,
|
||
previous_part_kind=previous_part_kind,
|
||
):
|
||
match part:
|
||
# 思考分片开始事件
|
||
case ThinkingPart(content=content):
|
||
# 若上一分片种类为空则将思考状态设置为正在思考、思考折叠面板展开状态设置为展开
|
||
if not previous_part_kind:
|
||
dialog.is_thinking = True
|
||
dialog.is_collapse_expanded = True
|
||
|
||
dialog.thought_nodes[index] = ThoughtNode(
|
||
kind="thinking", content=content
|
||
)
|
||
yield
|
||
|
||
# 工具检索分片开始事件
|
||
case ToolSearchCallPart(tool_call_id=tool_call_id):
|
||
# 创建工具调用唯一标识与片段索引映射
|
||
tool_call_ids[tool_call_id] = index
|
||
|
||
dialog.thought_nodes[index] = ThoughtNode(
|
||
kind="tool_search",
|
||
content="正在生成检索关键词",
|
||
)
|
||
yield
|
||
|
||
# 能力加载分片开始事件
|
||
case LoadCapabilityCallPart(tool_call_id=tool_call_id):
|
||
tool_call_ids[tool_call_id] = index
|
||
|
||
dialog.thought_nodes[index] = ThoughtNode(
|
||
kind="capability_load",
|
||
content="正在生成加载参数",
|
||
)
|
||
yield
|
||
|
||
# 工具调用分片开始事件
|
||
case ToolCallPart(tool_call_id=tool_call_id):
|
||
tool_call_ids[tool_call_id] = index
|
||
|
||
dialog.thought_nodes[index] = ThoughtNode(
|
||
kind="tool_call",
|
||
content="正在生成调用参数",
|
||
)
|
||
yield
|
||
|
||
# 文本分片开始事件
|
||
case TextPart(content=content):
|
||
dialog.answer = content
|
||
yield
|
||
|
||
# ========== 增量事件 ==========
|
||
case PartDeltaEvent(index=index, delta=delta):
|
||
match delta:
|
||
# 思考分片增量事件
|
||
case ThinkingPartDelta(
|
||
content_delta=content_delta,
|
||
):
|
||
dialog.thought_nodes[index].content += (
|
||
content_delta or ""
|
||
)
|
||
yield
|
||
|
||
# 文本分片增量事件
|
||
case TextPartDelta(
|
||
content_delta=content_delta,
|
||
):
|
||
dialog.answer += content_delta
|
||
yield
|
||
|
||
# ========== 结束事件 ==========
|
||
case PartEndEvent(
|
||
index=index,
|
||
part=part,
|
||
next_part_kind=next_part_kind,
|
||
):
|
||
match part:
|
||
# 思考分片结束事件
|
||
case ThinkingPart(part_kind=part_kind, content=content):
|
||
# 若下一分片种类为文本则将思考状态设置为思考完成、思考面板展开状态设置为折叠
|
||
if next_part_kind == "text":
|
||
dialog.is_thinking = False
|
||
dialog.is_collapse_expanded = False
|
||
yield
|
||
|
||
# ========== 函数工具调用事件 ==========
|
||
case FunctionToolCallEvent(tool_call_id=tool_call_id, part=part):
|
||
# 获取分片索引
|
||
index = tool_call_ids[tool_call_id]
|
||
match dialog.thought_nodes[index].kind:
|
||
# 工具检索
|
||
case "tool_search":
|
||
dialog.thought_nodes[index].content = "正在检索"
|
||
yield
|
||
|
||
# 能力加载
|
||
case "capability_load":
|
||
dialog.thought_nodes[index].content = (
|
||
f"正在加载能力 {part.tool_name}"
|
||
)
|
||
yield
|
||
|
||
# 工具调用
|
||
case "tool_call":
|
||
dialog.thought_nodes[index].content = (
|
||
f"正在调用工具 {part.tool_name}"
|
||
)
|
||
yield
|
||
|
||
# ========== 函数工具结果事件 ==========
|
||
case FunctionToolResultEvent(
|
||
tool_call_id=tool_call_id,
|
||
content=content,
|
||
):
|
||
index = tool_call_ids[tool_call_id]
|
||
match dialog.thought_nodes[index].kind:
|
||
# 工具检索
|
||
case "tool_search":
|
||
dialog.thought_nodes[index].content = (
|
||
content if isinstance(content, str) else ""
|
||
) # 暂仅考虑文本内容
|
||
yield
|
||
|
||
# 能力加载
|
||
case "capability_load":
|
||
dialog.thought_nodes[index].content = "已加载"
|
||
yield
|
||
|
||
# 工具调用
|
||
case "tool_call":
|
||
dialog.thought_nodes[index].content = f"已调用"
|
||
yield
|
||
|
||
# ========== 智能体运行结果事件 ==========
|
||
case AgentRunResultEvent(result=result):
|
||
# 更新对话记录
|
||
await database_state.update_dialog_record(
|
||
dialog_id=dialog_id,
|
||
thought_nodes=dialog.thought_nodes,
|
||
answer=dialog.answer,
|
||
)
|
||
# 创建运行结果记录
|
||
await database_state.create_run_result_record(
|
||
conversation_id=self.conversation_id,
|
||
dialog_id=dialog_id,
|
||
new_messages=result.new_messages(),
|
||
)
|
||
# 将当前会话的运行状态设置为运行完成
|
||
conversation.is_running = False
|
||
yield
|
||
|
||
@rx.event
|
||
def toggle_collapse(self, dialog_id: str) -> None:
|
||
"""
|
||
展开/折叠思考折叠面板
|
||
"""
|
||
# 指定运行
|
||
dialog = self.conversations[self.conversation_id].dialogs[dialog_id]
|
||
dialog.is_collapse_expanded = not dialog.is_collapse_expanded
|