523 lines
20 KiB
Python
523 lines
20 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
对话状态
|
||
"""
|
||
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, cast
|
||
|
||
from pydantic_ai import Agent, ThinkingPartDelta
|
||
from pydantic_ai._uuid import uuid7
|
||
from pydantic_ai.capabilities import AgentCapability
|
||
from pydantic_ai.messages import (
|
||
FunctionToolCallEvent,
|
||
FunctionToolResultEvent,
|
||
LoadCapabilityCallPart,
|
||
ModelMessage,
|
||
PartDeltaEvent,
|
||
PartEndEvent,
|
||
PartStartEvent,
|
||
TextPart,
|
||
TextPartDelta,
|
||
ThinkingPart,
|
||
ToolCallPart,
|
||
ToolSearchCallPart,
|
||
)
|
||
from pydantic_ai.models.openai import OpenAIChatModel
|
||
from pydantic_ai.providers.openai import OpenAIProvider
|
||
import reflex as rx
|
||
|
||
from application.models import Conversation, EventKind, PartKind, Run, Event, EventKind, PartKind
|
||
from application.state.create_conversation_modal import CreateConversationModalState
|
||
from application.state.database import DatabaseState
|
||
|
||
|
||
instructions: str = """
|
||
# 角色
|
||
专业友好AI助手，结构化解答各类问题。
|
||
|
||
# 输出硬性规则
|
||
1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容；
|
||
2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表；
|
||
3. 代码块用 ```语言名``` 包裹；
|
||
4. 重点内容标注 **粗体**/*斜体*；
|
||
5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签；
|
||
6. 内容分点拆分，排版整洁适配前端Markdown渲染。
|
||
|
||
# 行文要求
|
||
语言通俗，逻辑完整简洁，无多余废话。
|
||
"""
|
||
|
||
|
||
class ConversationState(rx.State):
|
||
"""
|
||
对话状态
|
||
"""
|
||
|
||
# 初始化当前对话唯一标识
|
||
conversation_id: str = str(uuid7())
|
||
# 初始化对话字典
|
||
conversations: Dict[str, Conversation] = {conversation_id: Conversation()}
|
||
|
||
# 初始化当前运行唯一标识
|
||
run_id: Optional[str] = None
|
||
# 初始化当前运行工具调用唯一标识与片段索引映射表
|
||
run_mapping: dict[str, int] = {}
|
||
|
||
# 初始化智能体
|
||
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,
|
||
)
|
||
|
||
@rx.var
|
||
def get_conversations(self) -> Dict[str, Conversation]:
|
||
"""
|
||
获取对话字典，用于前端按照创建倒序渲染对话历史
|
||
:return: 对话字典（倒序）
|
||
"""
|
||
return dict(reversed(self.conversations.items()))
|
||
|
||
@rx.event
|
||
async def create_conversation(self, form_data: Dict[str, Any]) -> None:
|
||
"""
|
||
创建对话
|
||
:param form_data: 表单数据
|
||
:return: None
|
||
"""
|
||
# 获取描述
|
||
description = form_data["description"].strip()
|
||
if not description:
|
||
description = "新对话"
|
||
|
||
# 生成并设置为当前对话唯一标识
|
||
self.conversation_id = str(uuid7())
|
||
# 创建对话
|
||
self.conversations[self.conversation_id] = Conversation(description=description)
|
||
# 初始化当前运行唯一标识
|
||
self.run_id = None
|
||
|
||
# 获取创建对话模态窗状态
|
||
create_conversation_modal_state = await self.get_state(
|
||
CreateConversationModalState
|
||
)
|
||
# 关闭创建对话模态窗
|
||
create_conversation_modal_state.is_open = False
|
||
|
||
@rx.event
|
||
def delete_conversation(self, conversation_id: str) -> None:
|
||
"""
|
||
删除对话
|
||
:param conversation_id: 对话唯一标识
|
||
:return: None
|
||
"""
|
||
if conversation_id not in self.conversations:
|
||
return
|
||
del self.conversations[conversation_id]
|
||
|
||
if not self.conversations:
|
||
# 生成并设置为当前对话唯一标识
|
||
self.conversation_id = str(uuid7())
|
||
# 创建对话
|
||
self.conversations[self.conversation_id] = Conversation()
|
||
# 初始化当前运行唯一标识
|
||
self.run_id = None
|
||
|
||
if self.conversation_id not in self.conversations:
|
||
# 将末位对话唯一标识设置为当前对话唯一标识
|
||
self.conversation_id = next(reversed(self.conversations.keys()))
|
||
# 初始化当前运行唯一标识
|
||
self.run_id = next(reversed(self.conversations[self.conversation_id].runs.keys()), None)
|
||
|
||
@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.var
|
||
def get_conversation_description(self) -> str:
|
||
"""
|
||
获取当前对话描述，用于前端渲染对话标题
|
||
:return: 当前对话描述
|
||
"""
|
||
# 当前对话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
return conversation.description if conversation else "新对话"
|
||
|
||
@rx.var
|
||
def get_runs(self) -> Dict[str, Run]:
|
||
"""
|
||
获取运行字典，用于前端渲染运行历史
|
||
:return: 运行字典
|
||
"""
|
||
# 当前对话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
return conversation.runs if conversation else {}
|
||
|
||
@rx.var
|
||
def get_run_streaming_status(self) -> bool:
|
||
"""
|
||
获取当前运行流式输出状态
|
||
:return: 当前运行流式输出状态（True 表示正在流式输出，False 表示非正在流式输出）
|
||
"""
|
||
# 当前对话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
if not conversation or not conversation.runs:
|
||
return False
|
||
|
||
# 末位运行
|
||
run = next(reversed(conversation.runs.values()))
|
||
return run.is_streaming
|
||
|
||
@rx.event
|
||
async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
|
||
"""
|
||
运行
|
||
:param form_data: 表单数据
|
||
:return: AsyncGenerator
|
||
"""
|
||
# 当前对话
|
||
conversation = self.conversations.get(self.conversation_id)
|
||
if not conversation:
|
||
return
|
||
|
||
# 更新当前运行唯一标识
|
||
self.run_id = str(uuid7())
|
||
|
||
# 获取用户提示词
|
||
user_prompt = form_data.get("user_prompt", "").strip()
|
||
if not user_prompt:
|
||
return
|
||
|
||
# 获取数据库状态
|
||
database_state = await self.get_state(DatabaseState)
|
||
# 获取消息历史
|
||
message_history = await database_state.get_message_history(
|
||
conversation_id=self.conversation_id
|
||
)
|
||
|
||
# 重置当前运行工具调用唯一标识与片段索引映射表
|
||
self.run_mapping.clear()
|
||
|
||
async for event in self.agent.run_stream_events(
|
||
conversation_id=self.conversation_id,
|
||
user_prompt=user_prompt,
|
||
message_history=message_history,
|
||
):
|
||
match event:
|
||
# ========== 分片开始事件 ==========
|
||
case PartStartEvent(
|
||
event_kind=event_kind,
|
||
index=index,
|
||
part=part,
|
||
previous_part_kind=previous_part_kind,
|
||
):
|
||
match part:
|
||
# 思考分片开始事件
|
||
case ThinkingPart(
|
||
part_kind=part_kind, content=content
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=content,
|
||
part_index=index,
|
||
previous_part_kind=PartKind(
|
||
previous_part_kind
|
||
),
|
||
part_kind=PartKind(part_kind),
|
||
)
|
||
# 检索工具分片开始事件
|
||
case ToolSearchCallPart(
|
||
tool_kind=tool_kind,
|
||
tool_call_id=tool_call_id,
|
||
tool_name=tool_name,
|
||
):
|
||
# 记录工具调用唯一标识与片段索引映射
|
||
self.tool_call_ids[tool_call_id] = index
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=index,
|
||
previous_part_kind=PartKind(
|
||
previous_part_kind
|
||
),
|
||
part_kind=PartKind(tool_kind),
|
||
tool_name=tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
# 加载能力分片开始事件
|
||
case LoadCapabilityCallPart(
|
||
tool_kind=tool_kind,
|
||
tool_call_id=tool_call_id,
|
||
tool_name=tool_name,
|
||
):
|
||
# 记录工具调用唯一标识与片段索引映射
|
||
self.tool_call_ids[tool_call_id] = index
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=index,
|
||
previous_part_kind=PartKind(
|
||
previous_part_kind
|
||
),
|
||
part_kind=PartKind(tool_kind),
|
||
tool_name=tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
# 调用工具分片开始事件
|
||
case ToolCallPart(
|
||
part_kind=part_kind,
|
||
tool_call_id=tool_call_id,
|
||
tool_name=tool_name,
|
||
):
|
||
# 记录工具调用唯一标识与片段索引映射
|
||
self.tool_call_ids[tool_call_id] = index
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=index,
|
||
previous_part_kind=PartKind(
|
||
previous_part_kind
|
||
),
|
||
part_kind=PartKind(part_kind),
|
||
tool_name=tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
# 文本分片开始事件
|
||
case TextPart(
|
||
part_kind=part_kind, content=content
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=content,
|
||
part_index=index,
|
||
previous_part_kind=PartKind(
|
||
previous_part_kind
|
||
),
|
||
part_kind=PartKind(part_kind),
|
||
run_id=current_run_id,
|
||
)
|
||
|
||
# ========== 分片增量事件（前端仅就文本片段实现打字机效果） ==========
|
||
case PartDeltaEvent(
|
||
event_kind=event_kind, index=index, delta=delta
|
||
):
|
||
match delta:
|
||
# 思考分片增量事件
|
||
case ThinkingPartDelta(
|
||
part_delta_kind=part_delta_kind,
|
||
content_delta=content_delta,
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=content_delta or "",
|
||
part_index=index,
|
||
part_kind=PartKind(part_delta_kind),
|
||
run_id=current_run_id,
|
||
)
|
||
# 文本分片增量事件
|
||
case TextPartDelta(
|
||
part_delta_kind=part_delta_kind,
|
||
content_delta=content_delta,
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=content_delta,
|
||
part_index=index,
|
||
part_kind=PartKind(part_delta_kind),
|
||
run_id=current_run_id,
|
||
)
|
||
|
||
# ========== 分片结束事件 ==========
|
||
case PartEndEvent(
|
||
event_kind=event_kind,
|
||
index=index,
|
||
part=part,
|
||
next_part_kind=next_part_kind,
|
||
):
|
||
match part:
|
||
# 思考分片结束事件
|
||
case ThinkingPart(
|
||
part_kind=part_kind, content=content
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=content,
|
||
part_index=index,
|
||
part_kind=PartKind(part_kind),
|
||
next_part_kind=PartKind(next_part_kind),
|
||
run_id=current_run_id,
|
||
)
|
||
# 检索工具分片结束事件
|
||
case ToolSearchCallPart(
|
||
tool_kind=tool_kind, tool_name=tool_name
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=index,
|
||
part_kind=PartKind(tool_kind),
|
||
next_part_kind=PartKind(next_part_kind),
|
||
tool_name=tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
# 加载能力分片结束事件
|
||
case LoadCapabilityCallPart(
|
||
tool_kind=tool_kind, tool_name=tool_name
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=index,
|
||
part_kind=PartKind(tool_kind),
|
||
next_part_kind=PartKind(next_part_kind),
|
||
tool_name=tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
# 调用工具分片结束事件
|
||
case ToolCallPart(
|
||
part_kind=part_kind, tool_name=tool_name
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=index,
|
||
part_kind=PartKind(part_kind),
|
||
next_part_kind=PartKind(next_part_kind),
|
||
tool_name=tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
# 文本分片结束事件
|
||
case TextPart(
|
||
part_kind=part_kind, content=content
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=content,
|
||
part_index=index,
|
||
part_kind=PartKind(part_kind),
|
||
next_part_kind=PartKind(next_part_kind),
|
||
run_id=current_run_id,
|
||
)
|
||
|
||
# ========== 函数工具调用事件 ==========
|
||
case FunctionToolCallEvent(
|
||
event_kind=event_kind,
|
||
part=part,
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
part_index=self.tool_call_ids.get(
|
||
part.tool_call_id
|
||
),
|
||
part_kind=PartKind(part.part_kind),
|
||
tool_name=part.tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
|
||
# ========== 函数工具回调事件 ==========
|
||
case FunctionToolResultEvent(
|
||
event_kind=event_kind,
|
||
content=content,
|
||
part=part,
|
||
):
|
||
yield Event(
|
||
event_kind=EventKind(event_kind),
|
||
event_content=(
|
||
content if isinstance(content, str) else ""
|
||
), # 暂仅处理文本类型，后续处理多模态和缓存
|
||
part_index=self.tool_call_ids.get(
|
||
part.tool_call_id
|
||
),
|
||
part_kind=PartKind(part.part_kind),
|
||
tool_name=part.tool_name,
|
||
run_id=current_run_id,
|
||
)
|
||
|
||
yield Event(
|
||
event_kind=EventKind.RUN_END,
|
||
run_id=current_run_id,
|
||
run_new_messages=agent_run.new_messages(),
|
||
)
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
# 运行，处理分片生命周期事件
|
||
async for self.agent.run(
|
||
conversation_id=self.conversation_id,
|
||
user_prompt=user_prompt,
|
||
message_history=message_history,
|
||
) as async_event:
|
||
if (
|
||
not event
|
||
or not self.current_run_id
|
||
or self.current_run_id != event.run_id
|
||
):
|
||
continue
|
||
|
||
# ========== 运行开始 ==========
|
||
if event.event_kind == EventKind.RUN_START:
|
||
# 将事件运行唯一标识设置为当前运行唯一标识
|
||
self.current_run_id = event.run_id
|
||
# 创建运行
|
||
conversation.runs[self.current_run_id] = Run(
|
||
user_prompt=user_prompt, is_streaming=True
|
||
)
|
||
yield # 通知前端渲染
|
||
continue
|
||
|
||
# 当前运行
|
||
current_run = current_conversation.runs[self.current_run_id]
|
||
|
||
# ========== 推理 ==========
|
||
if (
|
||
event.event_kind == EventKind.PART_START
|
||
and event.part_kind == PartKind.THINKING
|
||
):
|
||
current_run.thinking += event.event_content
|
||
yield # 通知前端渲染
|
||
continue
|
||
|
||
if part_index := event.part_index:
|
||
if part_index not in current_run.reasonings:
|
||
current_run.reasonings[part_index] += event.event_content
|
||
yield # 通知前端渲染
|
||
continue
|
||
|
||
# ========== 回答 ==========
|
||
if (
|
||
event.part_kind == PartKind.TEXT
|
||
and event.event_content != EventKind.PART_END
|
||
):
|
||
current_run.answer += event.event_content
|
||
yield # 通知前端渲染
|
||
continue
|
||
|
||
# ========== 运行结束 ==========
|
||
if event.event_kind == EventKind.RUN_END:
|
||
# 保存运行新增消息
|
||
await database_state.save_new_message(
|
||
conversation_id=self.current_conversation_id,
|
||
run_id=self.current_run_id,
|
||
run_new_message=event.run_new_messages,
|
||
)
|
||
# 设置当前运行流式输出状态非正在流式输出
|
||
current_run.is_streaming = False
|
||
yield # 通知前端渲染
|
||
continue
|