From d0eb2932dcec7d51801c2a7432c4580bd09355e8 Mon Sep 17 00:00:00 2001 From: liubiren Date: Wed, 24 Jun 2026 20:07:17 +0800 Subject: [PATCH] =?UTF-8?q?260624=E6=9B=B4=E6=96=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../application/components/chat.py | 177 ++++++++------- .../application/components/navbar.py | 90 ++++---- 产品需求文档AI生成/application/models.py | 22 +- 产品需求文档AI生成/application/pages/index.py | 6 +- 产品需求文档AI生成/application/state/chat.py | 133 ++++++----- .../state/{create_chat_modal.py => create_chat.py} | 4 +- .../application/state/{global_.py => database.py} | 6 +- 产品需求文档AI生成/application/utils/agent.py | 211 ++++++++++-------- 8 files changed, 351 insertions(+), 298 deletions(-) rename 产品需求文档AI生成/application/state/{create_chat_modal.py => create_chat.py} (83%) rename 产品需求文档AI生成/application/state/{global_.py => database.py} (95%) diff --git a/产品需求文档AI生成/application/components/chat.py b/产品需求文档AI生成/application/components/chat.py index bb19070..76e9223 100644 --- a/产品需求文档AI生成/application/components/chat.py +++ b/产品需求文档AI生成/application/components/chat.py @@ -1,21 +1,21 @@ # -*- coding: utf-8 -*- """ -聊天相关组件 +聊天页面相关组件 """ import reflex as rx -from application.models import Type_, Message, Dialog +from application.models import PartType, Part, Dialog from application.state import ChatState -def render_input_bubble(input_: str) -> rx.Component: +def render_user_prompt(user_prompt: str) -> rx.Component: """ - 渲染输入气泡 - :param content: 输入内容 + 渲染用户提示词 + :param user_prompt: 用户提示词 :return: Component """ return rx.markdown( - input_, + user_prompt, color=rx.color("gray", 12), # 文字颜色 background_color=rx.color("gray", 2), # 背景颜色 display="inline-block", # 布局模式：自适应文本宽度 @@ -28,100 +28,94 @@ def render_input_bubble(input_: str) -> rx.Component: ) -def render_collapse_panel(dialog_id: str, message: Message) -> rx.Component: +def render_part(dialog_id: str, part_id: str, part: Part): """ - 渲染折叠面板 + 渲染片段 :param dialog_id: 对话唯一标识 - :param message: 消息实例 + :param part_id: 片段唯一标识 + :param part: 片段实例 :return: Component """ - - return rx.box( - rx.cond( - message.is_open, - rx.hstack( - rx.text( - message.content, - font_size="0.90rem", - bold=True, - color=rx.color("gray", 10), - ), - rx.spacer(), - rx.icon("chevron_up", size=16, color=rx.color("gray", 6)), - width="100%", - cursor="pointer", - on_click=ChatState.toggle_collapse_panel(dialog_id, message.id), - padding_x="1em", - padding_y="0.6em", - background_color=rx.color("gray", 2), - border_radius="8px", - ), # 折叠面板打开 - rx.hstack( + return rx.match( + part.part_type, + ( + PartType.TEXT, + rx.markdown( + part.content, + font_size="0.90rem", + color=rx.color("gray", 12), # 字体颜色 + background_color="transparent", # 背景颜色：设置为透明以继承父元素背景颜色 + display="block", # 布局模式：铺满 + max_width="85%", # 最大宽度 + padding="0", # 内边距 + margin_right="auto", # 右侧外边距自动调整 + margin_bottom="12px", # 底部外边距 + key=part_id, + ), + ), # 文本片段 + (PartType.FINISHED, rx.fragment(key=part_id)), # 结束片段 + rx.box( + rx.cond( + part.is_open, rx.hstack( rx.text( - loading_prefix, + part.content, font_size="0.90rem", bold=True, color=rx.color("gray", 10), ), - loading_dot_group(message.is_streaming), - spacing="8", # 文字和圆点拉开一点间距，更美观 - align_items="center", - ), - rx.spacer(), - rx.icon("chevron_down", size=16, color=rx.color("gray", 6)), - width="100%", - cursor="pointer", - on_click=ChatState.toggle_collapse_panel(dialog_id, message.id), - padding_x="1em", - padding_y="0.6em", - background_color=rx.color("gray", 2), - border_radius="8px", - ), # 折叠面板关闭 + rx.spacer(), + rx.icon("chevron_up", size=16, color=rx.color("gray", 6)), + width="100%", + cursor="pointer", + padding_x="1em", + padding_y="0.6em", + background_color=rx.color("gray", 2), + border_radius="8px", + on_click=ChatState.toggle_collapse_panel(dialog_id, part_id), + ), # 折叠面板打开时标题栏 + rx.hstack( + rx.hstack( + rx.text( + loading_prefix, + font_size="0.90rem", + bold=True, + color=rx.color("gray", 10), + ), + loading_dot_group(part.is_streaming), + spacing="8", # 间距 + align_items="center", + ), + rx.spacer(), + rx.icon("chevron_down", size=16, color=rx.color("gray", 6)), + width="100%", + cursor="pointer", + on_click=ChatState.toggle_collapse_panel(dialog_id, part_id), + padding_x="1em", + padding_y="0.6em", + background_color=rx.color("gray", 2), + border_radius="8px", + ), # 折叠面板关闭时标题栏 + ), + rx.cond(part.is_open, detail_box, rx.fragment()), + width="100%", + max_width="85%", + margin_bottom="10px", + cursor="pointer", + key=part_id, ), - rx.cond(message.is_open, detail_box, rx.fragment()), - width="100%", - max_width="85%", - margin_bottom="10px", - key=message.id, - cursor="pointer", ) -def render_output_container(dialog_id: str, message: Message) -> rx.Component: +def render_dialog_item(dialog_id: str, dialog: Dialog) -> rx.Component: """ - 渲染输出容器 - :param dialog_id: 对话唯一标识 - :param message: 消息实例 - :return: Component - """ - return rx.cond( - message.type_ == Type_.TEXT, - rx.markdown( - message.content, - font_size="0.90rem", - color=rx.color("gray", 12), # 字体颜色 - background_color="transparent", # 背景颜色：设置为透明以继承父元素背景颜色 - display="block", # 布局模式：铺满 - max_width="85%", # 最大宽度 - padding="0", # 内边距 - margin_right="auto", # 右侧外边距自动调整 - margin_bottom="12px", # 底部外边距 - key=message.id, - ), - render_collapse_panel(dialog_id=dialog_id, message=message), - ) - - -def render_dialog_container(dialog: Dialog) -> rx.Component: - """ - 渲染对话容器 + 渲染对话项 :param dialog: 对话实例 :return: Component """ return rx.box( rx.box( - render_input_bubble(input_=dialog.input_), + render_user_prompt(user_prompt=dialog.user_prompt), text_align="right", width="100%", margin_bottom="8px", @@ -129,8 +123,8 @@ def render_dialog_container(dialog: Dialog) -> rx.Component: rx.box( rx.foreach( dialog.output, - lambda message: render_output_container( - dialog_id=dialog.id, message=message + lambda part_id, part: render_part( + dialog_id=dialog_id, part_id=part_id, part=part ), ), text_align="left", @@ -139,33 +133,38 @@ def render_dialog_container(dialog: Dialog) -> rx.Component: ), width="min(100%, 50em)", # 最大宽度：父级元素最大宽度和50em中较小值 margin_x="auto", # 水平外边距：自动调整 - key=dialog.id, + key=dialog_id, ) -def dialog_list_component() -> rx.Component: +def render_dialog_list() -> rx.Component: """ - 对话列表组件 + 渲染对话列表 :return: Component """ return rx.auto_scroll( - rx.foreach(ChatState.get_current_chat_dialogs, dialog_item_component), + rx.foreach( + ChatState.get_dialogs, + lambda dialog_id, dialog: render_dialog_item( + dialog_id=dialog_id, dialog=dialog + ), + ), flex="1", padding="8px", overflow_y="auto", ) -def input_component() -> rx.Component: +def render_input_bar() -> rx.Component: """ - 输入组件 + 渲染输入栏 """ return rx.center( rx.vstack( rx.form( rx.hstack( rx.input( - name="input", + name="user_prompt", placeholder="请输入...", flex="auto", ), diff --git a/产品需求文档AI生成/application/components/navbar.py b/产品需求文档AI生成/application/components/navbar.py index 033cb6b..2e55236 100644 --- a/产品需求文档AI生成/application/components/navbar.py +++ b/产品需求文档AI生成/application/components/navbar.py @@ -1,53 +1,25 @@ # -*- coding: utf-8 -*- """ -导航栏组件 +导航栏相关组件 """ import reflex as rx from application.state.chat import ChatState -from application.state.create_chat_modal import CreateChatModalState +from application.state.create_chat import CreateChatState -def create_chat_modal_component(trigger) -> rx.Component: +def render_chat_item(chat_id: str, chat_description: str) -> rx.Component: """ - 新建聊天模态窗组件 - """ - return rx.dialog.root( - rx.dialog.trigger(trigger), - rx.dialog.content( - rx.form( - rx.hstack( - rx.input( - name="chat_description", - placeholder="聊天描述", - flex="auto", - min_width="20ch", - ), - rx.button("新建聊天"), - spacing="2", - wrap="wrap", - width="100%", - ), - on_submit=ChatState.create_chat, - ), - background_color=rx.color("mauve", 1), - ), - open=CreateChatModalState.is_open, - on_open_change=CreateChatModalState.toggle, - ) - - -def chat_item_component(chat_id: str) -> rx.Component: - """ - 聊天项组件 + 渲染聊天项 :param chat_id: 聊天唯一标识 + :param description: 聊天描述 :return: Component """ return rx.drawer.close( rx.hstack( rx.button( - ChatState.get_chat_descriptions[chat_id], + chat_description, on_click=lambda: ChatState.switch_chat(chat_id), # 点击按钮将切换会话 width="80%", variant="surface", @@ -68,9 +40,9 @@ def chat_item_component(chat_id: str) -> rx.Component: ) -def chat_list_component(trigger) -> rx.Component: +def render_chat_list(trigger) -> rx.Component: """ - 聊天列表组件 + 渲染聊天列表 """ return rx.drawer.root( rx.drawer.trigger(trigger), @@ -81,9 +53,10 @@ def chat_list_component(trigger) -> rx.Component: rx.heading("聊天列表", color=rx.color("mauve", 11)), rx.divider(), rx.foreach( - ChatState.get_chat_ids, # 获取所有聊天唯一标识 - lambda chat_id: chat_item_component( - chat_id=chat_id + ChatState.get_chats, # 获取聊天列表 + lambda chat_id, chat: render_chat_item( + chat_id=chat_id, + chat_description=chat.description, ), # 创建聊天组件 ), align_items="stretch", @@ -102,21 +75,50 @@ def chat_list_component(trigger) -> rx.Component: ) -def navbar_component() -> rx.Component: +def render_create_chat_modal(trigger) -> rx.Component: """ - 导航栏组件 + 渲染新建聊天模态窗 + """ + return rx.dialog.root( + rx.dialog.trigger(trigger), + rx.dialog.content( + rx.form( + rx.hstack( + rx.input( + name="chat_description", + placeholder="请输入聊天描述（可选）", + flex="auto", + min_width="20ch", + ), + rx.button("新建"), + spacing="2", + wrap="wrap", + width="100%", + ), + on_submit=ChatState.create_chat, + ), + background_color=rx.color("mauve", 1), + ), # 模态窗的内容容器 + open=CreateChatState.is_open, + on_open_change=CreateChatState.toggle, + ) + + +def render_navbar() -> rx.Component: + """ + 渲染导航栏 """ return rx.hstack( rx.badge( - ChatState.get_chat_descriptions[ChatState.current_chat_id], + ChatState.get_current_chat_description, size="3", variant="soft", margin_inline_end="auto", ), - create_chat_modal_component( + render_create_chat_modal( rx.box(rx.tooltip(rx.icon("message-square-plus"), content="新建聊天")) ), - chat_list_component( + render_chat_list( rx.box( rx.tooltip( rx.icon("messages-square"), diff --git a/产品需求文档AI生成/application/models.py b/产品需求文档AI生成/application/models.py index 34e4ca0..9f8b81a 100644 --- a/产品需求文档AI生成/application/models.py +++ b/产品需求文档AI生成/application/models.py @@ -4,7 +4,7 @@ """ from enum import StrEnum from time import time_ns -from typing import List +from typing import List, Dict from uuid import uuid4 from pydantic import BaseModel, Field @@ -47,26 +47,25 @@ class MessageHistory(SQLModel, table=True): class PartType(StrEnum): - """片段类""" + """片段类型""" THINKING = "thinking" TEXT = "text" CALL = "call" TOOL_ARGS = "tool_args" TOOL_RETURN = "tool_return" - RESULT = "result" + FINISHED = "finished" ERROR = "error" -# 动态生成前缀和消息类型映射表 +# 动态生成前缀和片段类型映射表 PREFIX_MAPING = {f"{i:02d}:": t for i, t in enumerate(PartType)} class Part(BaseModel): """片段类（仅就输出消息）""" - id: str = Field(default_factory=lambda: uuid4().hex, description="片段唯一标识") - part_type: PartType = Field(..., alias="type", description="片段类型") + part_type: PartType = Field(..., description="片段类型") content: str = Field(default="", description="片段内容") is_streaming: bool = Field( @@ -83,17 +82,20 @@ class Part(BaseModel): class Dialog(BaseModel): """对话类""" - id: str = Field(default_factory=lambda: uuid4().hex, description="对话唯一标识") user_prompt: str = Field(..., description="用户提示词") - output: List[Part] = Field(default_factory=list, description="输出") + output: Dict[str, Part] = Field( + default_factory=dict, description="输出，键为片段唯一标识，值为片段对象" + ) class Chat(BaseModel): """聊天类""" - description: str = Field(default="新聊天", description="描述") + description: str = Field(default="新聊天", description="聊天描述") is_streaming: bool = Field( default=False, description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出", ) - dialogs: List[Dialog] = Field(default_factory=list, description="对话列表") + dialogs: Dict[str, Dialog] = Field( + default_factory=dict, description="对话列表，键为对话唯一标识，值为对话对象" + ) diff --git a/产品需求文档AI生成/application/pages/index.py b/产品需求文档AI生成/application/pages/index.py index 0df9e16..87afbb5 100644 --- a/产品需求文档AI生成/application/pages/index.py +++ b/产品需求文档AI生成/application/pages/index.py @@ -5,14 +5,14 @@ import reflex as rx -from ..components.navbar import navbar_component +from ..components.navbar import render_navbar from ..components.chat import dialog_list_component, input_component def index_page() -> rx.Component: - """根页面""" + """根页面（聊天页面）""" return rx.vstack( - navbar_component(), # 导航栏组件 + render_navbar(), # 导航栏组件 dialog_list_component(), # 对话列表组件 input_component(), # 输入组件 color=rx.color(color="mauve", shade=12), diff --git a/产品需求文档AI生成/application/state/chat.py b/产品需求文档AI生成/application/state/chat.py index 8d0d733..c6744e7 100644 --- a/产品需求文档AI生成/application/state/chat.py +++ b/产品需求文档AI生成/application/state/chat.py @@ -2,18 +2,18 @@ """ 聊天页面状态 """ -from typing import Any, AsyncGenerator, Dict, List +from typing import Any, AsyncGenerator, Dict, List, Tuple from uuid import uuid4 import reflex as rx from application.utils.agent import Agent -from application.models import Chat, Dialog, Message, PREFIX_MAPING -from application.state.create_chat_modal import CreateChatModalState -from application.state.global_ import GlobalState +from application.models import Chat, Dialog, Part, PartType, PREFIX_MAPING +from application.state.create_chat import CreateChatState +from application.state.database import DatabaseState +# 因 rx.state 需在进程与 WebSocket 间交互，非序列化对象会传输或储存失败，故需将智能体实例单独存储 # 绑定的智能体列表（键为聊天唯一标识，值为智能体实例） -# 因 rx.state 需在 Python 进程与 WebSocket 之间传输和存储等操作，非序列化对象会失败，故需将智能体实例单独存储 agents: Dict[str, Agent] = {} @@ -22,57 +22,61 @@ def get_agent(chat_id: str) -> Agent: 获取绑定的智能体 :return: 绑定的智能体 """ + # 若智能体唯一标识不存在则新建智能体实例 if chat_id not in agents: agents[chat_id] = Agent( chat_id=chat_id, - instructions="You are a friendly chatbot", # 智能体指令 - ) + ) # 使用默认指令 return agents[chat_id] class ChatState(rx.State): """ - 聊天页面状态 + 聊天状态 """ - # 当前聊天唯一标识 + # 初始化当前聊天唯一标识 current_chat_id: str = uuid4().hex # 聊天列表 chats: Dict[str, Chat] = {current_chat_id: Chat()} @rx.var - def get_chat_ids(self) -> List[str]: + def get_chats(self) -> Dict[str, Chat]: """ - 获取聊天唯一标识列表 - :return: 聊天唯一标识列表 + 获取聊天列表，用于前端渲染聊天列表 + :return: 聊天列表 """ - return list(self.chats) + return self.chats @rx.var - def get_chat_descriptions(self) -> Dict[str, str]: + def get_current_chat_description(self) -> str: """ - 获取聊天描述列表 - :return: 聊天描述列表 + 获取当前聊天描述，用于前端渲染导航栏中当前聊天描述 + :return: 当前聊天描述 """ - return {chat_id: chat.description for chat_id, chat in self.chats.items()} + # 当前聊天 + current_chat = self.chats.get(self.current_chat_id) + return current_chat.description if current_chat else "新聊天" @rx.var def get_current_chat_status(self) -> bool: """ - 获取当前聊天流式输出状态 - :return: 当前聊天流式输出状态（True 表示正在流式输出，False 表示非正在流式输出） + 获取当前聊天状态 + :return: 当前聊天状态（True 表示正在流式输出，False 表示非正在流式输出） """ - chat = self.chats.get(self.current_chat_id) - return chat.is_streaming if chat else False + # 当前聊天 + current_chat = self.chats.get(self.current_chat_id) + return current_chat.is_streaming if current_chat else False @rx.var - def get_current_chat_dialogs(self) -> List[Dialog]: + def get_dialogs(self) -> Dict[str, Dialog]: """ - 获取当前聊天对话列表 - :return: 当前聊天对话列表 + 获取对话列表，用于前端渲染对话列表 + :return: 对话列表 """ - chat = self.chats.get(self.current_chat_id) - return chat.dialogs if chat else [] + # 当前聊天 + current_chat = self.chats.get(self.current_chat_id) + return current_chat.dialogs if current_chat else {} @rx.event def switch_chat(self, chat_id: str) -> None: @@ -99,8 +103,8 @@ class ChatState(rx.State): ) # 跨状态读取新建聊天模态窗状态（异步事件） - create_chat_modal_state = await self.get_state(CreateChatModalState) - create_chat_modal_state.is_open = False + create_chat_state = await self.get_state(CreateChatState) + create_chat_state.is_open = False @rx.event def delete_chat(self, chat_id: str) -> None: @@ -126,8 +130,9 @@ class ChatState(rx.State): :param form_data: 表单数据 :return: AsyncGenerator """ - input_ = form_data["input"].strip() - if not input_: + # 获取用户提示词 + user_prompt = form_data["user_prompt"].strip() + if not user_prompt: return # 当前聊天 @@ -135,64 +140,82 @@ class ChatState(rx.State): if not current_chat: return + # 生成新增对话唯一标识 + current_dialog_id = uuid4().hex # 新增对话 - current_chat.dialogs.append(Dialog(input_=input_)) + current_chat.dialogs[current_dialog_id] = Dialog(user_prompt=user_prompt) # 当前对话 - current_dialog = current_chat.dialogs[-1] + current_dialog = current_chat.dialogs[current_dialog_id] - # 设置当前对话流式输出状态为正在流式输出 + # 设置当前聊天流式输出状态为正在流式输出 current_chat.is_streaming = True yield # 通知前端渲染输入消息 - # 全局状态 - global_state = await self.get_state(GlobalState) # 获取当前聊天绑定的智能体 agent = get_agent(chat_id=self.current_chat_id) + + # 数据库状态 + database_state = await self.get_state(DatabaseState) # 获取当前聊天消息历史 - message_history = await global_state.get_message_history( + message_history = await database_state.get_message_history( chat_id=self.current_chat_id ) - # 流式输出消息事件 - async for event in agent.stream_messages_events( - user_prompt=input_, message_history=message_history + # 流式输出模型响应事件 + async for event in agent.stream_events( + user_prompt=user_prompt, message_history=message_history ): - # 若消息事件为空则跳过 + # 若模型响应事件为空则跳过 if not event: continue - # 匹配前缀 + # 根据前缀和片段类型映射表匹配前缀 prefix = next((p for p in PREFIX_MAPING if event.startswith(p)), None) # 若未匹配到前缀则跳过 if not prefix: continue + # 根据前缀匹配片段类型 + part_type = PREFIX_MAPING[prefix] - # 根据前缀匹配消息类型 - type_ = PREFIX_MAPING[prefix] - # 若当前对话输出消息为空或当前消息类型与新消息类型不相同则新增消息 - if not current_dialog.output or current_dialog.output[-1].type_ != type_: - current_dialog.output.append(Message(type_=type_)) + # 获取当前片段 + current_part = ( + next(iter(reversed(current_dialog.output.values()))) + if current_dialog.output + else None + ) + # 若当前片段为空或其片段类型与当前事件的片段类型不相同则新增片段 + if not current_part or current_part.part_type != part_type: + # 若当前片段非空则设置片段流式输出状态非正在流式输出 + if current_part: + current_part.is_streaming = False + new_part = Part(part_type=part_type, is_streaming=True) + current_dialog.output[uuid4().hex] = new_part + current_part = new_part - # 追加消息内容 - current_dialog.output[-1].content += event.removeprefix(prefix) - yield # 通知前端渲染输出消息 + # 追加片段内容 + current_part.content += event.removeprefix(prefix) + yield # 通知前端渲染片段 + + # 重置所有片段流式输出状态 + for part in current_dialog.output.values(): + part.is_streaming = False # 保存本轮对话消息 - await global_state.save_new_message( + await database_state.save_new_message( chat_id=self.current_chat_id, new_message=agent.new_messages ) - # 设置当前对话流式输出状态非正在流式输出 + # 设置当前聊天流式输出状态非正在流式输出 current_chat.is_streaming = False @rx.event - async def toggle_collapse_panel(self, dialog_id: str, message_id: str) -> None: + async def toggle_collapse_panel(self, dialog_id: str, part_id: str) -> None: """ 打开/关闭折叠面板 :return: None """ - # 当前聊天 - current_chat = self.chats[self.current_chat_id] - if not current_chat: + # 当前片段 + current_part = self.chats.get(self.current_chat_id, {}).dialogs.get(dialog_id, {}).get(part_id, None) + if not current_part: return # 目标对话 target_dialog = next( diff --git a/产品需求文档AI生成/application/state/create_chat_modal.py b/产品需求文档AI生成/application/state/create_chat.py similarity index 83% rename from 产品需求文档AI生成/application/state/create_chat_modal.py rename to 产品需求文档AI生成/application/state/create_chat.py index 0afbbcd..9eefc75 100644 --- a/产品需求文档AI生成/application/state/create_chat_modal.py +++ b/产品需求文档AI生成/application/state/create_chat.py @@ -5,9 +5,9 @@ import reflex as rx -class CreateChatModalState(rx.State): +class CreateChatState(rx.State): """ - 新建聊天模态窗状态 + 新建聊天状态 """ # 新建聊天模态窗打开状态，True表示打开，False表示关闭 diff --git a/产品需求文档AI生成/application/state/global_.py b/产品需求文档AI生成/application/state/database.py similarity index 95% rename from 产品需求文档AI生成/application/state/global_.py rename to 产品需求文档AI生成/application/state/database.py index b2455a1..fa164ee 100644 --- a/产品需求文档AI生成/application/state/global_.py +++ b/产品需求文档AI生成/application/state/database.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- """ -全局共享状态 +数据库状态 """ from typing import List @@ -15,8 +15,8 @@ from application.models import ( ) -class GlobalState(rx.State): - """全局状态""" +class DatabaseState(rx.State): + """数据库状态""" async def get_message_history(self, chat_id: str) -> List[ModelMessage]: """ diff --git a/产品需求文档AI生成/application/utils/agent.py b/产品需求文档AI生成/application/utils/agent.py index 7fea8c5..fcb1e95 100644 --- a/产品需求文档AI生成/application/utils/agent.py +++ b/产品需求文档AI生成/application/utils/agent.py @@ -39,52 +39,56 @@ from pydantic_ai.run import AgentRunResultEvent class PartBuffer(BaseModel): """ - 片段缓冲类 + 模型响应片段缓冲类 """ - part_type: Literal["thinking", "text"] = Field(..., description="片段类型") - content_delta: str = Field(default="", description="片段内容增量") - task: Optional[Task] = Field(default=None, description="延迟刷新异步协程任务") + part_type: Literal["thinking", "text"] = Field( + ..., description="片段类型，仅支持思考和文本片段" + ) + part_content_delta: str = Field(default="", description="片段内容增量") + flush_task: Optional[Task] = Field( + default=None, description="片段绑定的延迟刷新任务" + ) model_config = {"arbitrary_types_allowed": True} # 允许任意类型 class Debouncer: """ - 防抖器 - 用于处理 run_stream_events 返回的模型消息事件 + 模型响应事件防抖器 + 用于平滑 run_stream_events 返回的高频模型响应事件，降低前端渲染频率 """ - def __init__(self, delay: float = 0.25): + def __init__(self, flush_delay: float): """ 初始化 - :param delay: 延迟时长（单位为秒），停顿超该时长则刷新 + :param flush_delay: 刷新延迟时长（单位为秒） """ - self.delay = delay + self.flush_delay = flush_delay # 片段缓冲字典（键为片段索引，值为片段缓冲实例） self.part_buffers: Dict[int, PartBuffer] = {} - # 待刷新异步队列 - self.pending_flush_queue = Queue() + # 待推送刷新队列 + self.flush_queue = Queue() async def handle_event( self, event: Union[AgentStreamEvent, AgentRunResultEvent] ) -> AsyncGenerator[Union[AgentStreamEvent, AgentRunResultEvent]]: """ - 处理模型消息事件 - 在 Pydantic-AI 中模型流式输出包含若干类型片段，例如思考片段、文本片段等。每类型片段包含若干模型消息事件，例如片段开始事件、片段增量事件、片段结束事件等 - :param event: 模型消息事件 + 处理模型响应事件 + 在 Pydantic-AI 中模型流式输出包含若干类型片段，例如思考片段、文本片段等。每类型片段包含若干模型响应事件，例如片段开始事件、片段增量事件、片段结束事件等 + :param event: 模型响应事件 :yield: AsyncGenerator """ - # 返回待刷新异步队列中模型消息事件 + # 优先推送队列中积压的刷新任务 while True: try: - yield self.pending_flush_queue.get_nowait() + yield self.flush_queue.get_nowait() except QueueEmpty: break - # 若为思考或文本的片段开始事件则先缓存片段再返回，其它事件直接返回 + # 若为思考或文本的片段开始事件则初始化片段缓冲，其它事件透传 if isinstance(event, PartStartEvent): - index, part = event.index, event.part # 片段索引、片段 + part_index, part = event.index, event.part # 片段索引、片段 match part: case ThinkingPart(): part_type = "thinking" @@ -93,113 +97,133 @@ class Debouncer: case _: yield event return - # 新增片段缓存 - self.buffers[index] = Buffer( + # 初始化片段缓冲 + self.part_buffers[part_index] = PartBuffer( part_type=part_type, ) yield event return - """ - 原理： - 收到思考或文本片段增量事件时取消未完成的延迟刷新任务、追加增量并重设；若模型持续输出则先缓存，停顿超过阈值再返回，实现防抖 - """ + # 若为思考或文本片段增量事件则先取消延迟刷新任务，更新片段缓冲中的片段内容增量并重新绑定延迟刷新任务，其它事件透传 if isinstance(event, PartDeltaEvent): - index, delta = event.index, event.delta - buffer = self.buffers.get(index) - if buffer and isinstance(delta, (ThinkingPartDelta, TextPartDelta)): - # 追加增量 - buffer.content_delta += delta.content_delta or "" - - # 取消上一轮未完成的延迟刷新任务 - self._cancel_task(index=index) - # 创建延迟刷新任务 - buffer.task = create_task(coro=self._delay_flush(index=index)) + part_index, delta = event.index, event.delta + part_buffer = self.part_buffers.get(part_index) + if part_buffer and isinstance(delta, (ThinkingPartDelta, TextPartDelta)): + # 取消延迟刷新任务 + self._cancel_task(part_index=part_index) + # 更新片段缓冲中的片段内容增量 + part_buffer.part_content_delta += delta.content_delta or "" + # 重新绑定延迟刷新任务 + part_buffer.flush_task = create_task( + coro=self._flush_after_delay(part_index=part_index) + ) else: yield event return - # 若为其它事件则先批量刷新片段增量事件再返回该事件 - async for event_ in self._batch_flush(): - yield event_ - # 若为片段结束事件则删除该片段缓存 + # 若为其它事件则先刷新所有片段缓冲再透传当前事件 + async for flush_event in self._flush_all(): + yield flush_event + # 若为片段结束事件则删除相应片段缓冲 if isinstance(event, PartEndEvent): - self.buffers.pop(event.index, None) + self.part_buffers.pop(event.index, None) yield event - def _cancel_task(self, index: int) -> None: + def _cancel_task(self, part_index: int) -> None: """ - 取消延迟刷新异步协程任务 - :param index: 片段索引 + 取消延迟刷新任务 + :param part_index: 片段索引 :return: None """ - buffer = self.buffers.get(index) - if not buffer or not buffer.task: + part_buffer = self.part_buffers.get(part_index) + if not part_buffer or not part_buffer.flush_task: return - if not buffer.task.done(): - buffer.task.cancel() - buffer.task = None - async def _delay_flush(self, index: int) -> None: + # 若延迟刷新任务未推送则取消 + if not part_buffer.flush_task.done(): + part_buffer.flush_task.cancel() + part_buffer.flush_task = None + + async def _flush_after_delay(self, part_index: int) -> None: """ 延迟刷新 - :param index: 片段索引 + :param part_index: 片段索引 :return: None """ - await sleep(delay=self.delay) + await sleep(delay=self.flush_delay) - # 刷新片段增量事件 - event = await self._flush(index=index) - if event: - await self.pending_flush_queue.put(item=event) + # 刷新单个片段缓冲 + flush_event = await self._flush(part_index=part_index) + if flush_event: + await self.flush_queue.put(item=flush_event) - async def _flush(self, index: int) -> Optional[PartDeltaEvent]: + async def _flush(self, part_index: int) -> Optional[PartDeltaEvent]: """ - 刷新片段增量事件 - :param index: 片段索引 + 刷新单个片段缓冲 + :param part_index: 片段索引 :return: Optional[PartDeltaEvent] """ - buffer = self.buffers.get(index) - if not buffer or not buffer.content_delta: + part_buffer = self.part_buffers.get(part_index) + if not part_buffer or not part_buffer.part_content_delta: return - # 构建片段增量事件中增量部分 - match buffer.part_type: + # 构建片段增量事件中增量 + match part_buffer.part_type: case "thinking": - delta = ThinkingPartDelta(content_delta=buffer.content_delta) + delta = ThinkingPartDelta(content_delta=part_buffer.part_content_delta) case "text": - delta = TextPartDelta(content_delta=buffer.content_delta) + delta = TextPartDelta(content_delta=part_buffer.part_content_delta) - buffer.content_delta = "" - buffer.task = None - return PartDeltaEvent(index=index, delta=delta) + # 重置片段缓冲中片段内容增量 + part_buffer.part_content_delta = "" + # 重置片段缓冲中延迟刷新任务 + part_buffer.flush_task = None + return PartDeltaEvent(index=part_index, delta=delta) - async def _batch_flush(self) -> AsyncGenerator[PartDeltaEvent, None]: + async def _flush_all(self) -> AsyncGenerator[PartDeltaEvent]: """ - 批量刷新片段增量事件 + 刷新所有片段缓冲 :yield: AsyncGenerator """ - for index in list(self.buffers.keys()): - event = await self._flush(index=index) - if event: - yield event + for part_index in list(self.part_buffers.keys()): + flush_event = await self._flush(part_index=part_index) + if flush_event: + yield flush_event + + +DEFAULT_INSTRUCTIONS: str = """ +# 角色 +专业友好AI助手，结构化解答各类问题。 + +# 输出硬性规则 +1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容； +2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表； +3. 代码块用 ```语言名``` 包裹； +4. 重点内容标注 **粗体**/*斜体*； +5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签； +6. 内容分点拆分，排版整洁适配前端Markdown渲染。 + +# 行文要求 +语言通俗，逻辑完整简洁，无多余废话。 +""" class Agent: """ - Pydantic AI 智能体 + 基于 Pydantic AI 封装的智能体，支持： + 1、流式输出模型响应事件 """ def __init__( self, chat_id: str, - instructions: str, + instructions: Optional[str] = None, output_type: OutputSpec = str, capabilities: Optional[List[AgentCapability]] = None, retries: int = 1, ): """ - 初始化智能体 + 初始化 :param chat_id: 聊天唯一标识 :param instructions: 指令 :param capabilities: 技能列表，默认为不使用技能 @@ -210,11 +234,16 @@ class Agent: # 聊天唯一标识 self.chat_id = chat_id - # 本轮对话新增消息 - # 一次聊天（chat）包含若干论对话（dialog），每一轮对话由输入消息（message）和输出消息（message）组成 + # 一次聊天（chat）包含若干论对话（dialog），每一轮对话由用户提示词（user_prompt）和输出（output）组成，两者统称为消息（message） + + # 本轮对话新增消息列表 self.new_messages: List[ModelMessage] = [] - # 实例智能体 + # 若指令为空则使用默认指令 + if not instructions: + instructions = DEFAULT_INSTRUCTIONS + + # 初始化智能体 self.agent = PydanticAIAgent( model=OpenAIChatModel( model_name="deepseek-v4-flash", @@ -229,27 +258,27 @@ class Agent: retries=retries, ) - async def stream_messages_events( + async def stream_events( self, user_prompt: str | List[str], message_history: Optional[List[ModelMessage]] = None, - delay: float = 0.2, + flush_delay: float = 0.15, ) -> AsyncGenerator[str, None]: """ - 流式输出消息事件 - :param user_prompt: 用户提示词（用户输入消息） - :param delay: 延迟时长（单位为秒），停顿超该时长则刷新 + 流式输出模型响应事件 + :param user_prompt: 用户提示词（用户提示词） + :param flush_delay: 刷新延迟时长（单位为秒） :yield: AsyncGenerator """ - # 实例防抖器 - debouncer = Debouncer(delay=delay) + # 初始化防抖器 + debouncer = Debouncer(flush_delay=flush_delay) async with self.agent.run_stream_events( user_prompt=user_prompt, message_history=message_history, ) as events: async for event in events: - # 处理模型消息事件 + # 处理模型响应事件 async for event in debouncer.handle_event(event): match event: # 片段开始事件 @@ -266,11 +295,11 @@ class Agent: | ToolCallPart(tool_name=tool_name) | LoadCapabilityCallPart(tool_name=tool_name) ): - yield f"02:正在使用技能：{tool_name}" + yield f"02:{tool_name}" case NativeToolReturnPart( content=content ) | ToolReturnPart(content=content): - yield f"04:技能返回：{content}" + yield f"04:{content}" case _: yield f"06:未知片段类型{type(part).__name__}" # 片段增量事件 @@ -281,9 +310,7 @@ class Agent: case TextPartDelta(content_delta=content_delta): yield f"01:{content_delta}" case ToolCallPartDelta(args_delta=args_delta): - yield f"03:技能参数：{args_delta}" - case _: - yield f"06:未知片段类型{type(delta).__name__}" + yield f"03:{args_delta}" # 片段结束事件、最终结果事件无需处理 case PartEndEvent(): continue @@ -291,6 +318,6 @@ class Agent: continue case AgentRunResultEvent(result=result): self.new_messages = result.new_messages() - yield "05:FinalResultEvent" + yield "05:" case _: yield f"06:未知事件类型{type(event).__name__}"