From 2f51a9552c5410f432398c5de5672acacc3823a2 Mon Sep 17 00:00:00 2001 From: liubiren Date: Tue, 23 Jun 2026 20:09:54 +0800 Subject: [PATCH] 1 --- .../versions/{97c7b4855ab2_.py => e05e08cbaa82_.py} | 22 ++- 产品需求文档AI生成/application/application.py | 4 +- .../application/components/chat.py | 156 ++++++++++++------ 产品需求文档AI生成/application/models.py | 69 +++++--- 产品需求文档AI生成/application/state/chat.py | 38 ++++- .../application/state/global_.py | 14 +- 产品需求文档AI生成/application/utils/agent.py | 10 +- 7 files changed, 213 insertions(+), 100 deletions(-) rename 产品需求文档AI生成/alembic/versions/{97c7b4855ab2_.py => e05e08cbaa82_.py} (51%) diff --git a/产品需求文档AI生成/alembic/versions/97c7b4855ab2_.py b/产品需求文档AI生成/alembic/versions/e05e08cbaa82_.py similarity index 51% rename from 产品需求文档AI生成/alembic/versions/97c7b4855ab2_.py rename to 产品需求文档AI生成/alembic/versions/e05e08cbaa82_.py index 45524c6..a2eae52 100644 --- a/产品需求文档AI生成/alembic/versions/97c7b4855ab2_.py +++ b/产品需求文档AI生成/alembic/versions/e05e08cbaa82_.py @@ -1,8 +1,8 @@ """empty message -Revision ID: 97c7b4855ab2 +Revision ID: e05e08cbaa82 Revises: -Create Date: 2026-06-22 19:56:54.150570 +Create Date: 2026-06-23 10:49:22.696176 """ from typing import Sequence, Union @@ -12,7 +12,7 @@ import sqlalchemy as sa import sqlmodel # revision identifiers, used by Alembic. -revision: str = '97c7b4855ab2' +revision: str = 'e05e08cbaa82' down_revision: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -21,18 +21,26 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: """Upgrade schema.""" # ### commands auto generated by Alembic - please adjust! ### - op.create_table('historymessage', - sa.Column('id', sa.Integer(), nullable=False), + op.create_table('messagehistory', + sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False), sa.Column('chat_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False), sa.Column('new_message', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('timestamp', sa.Integer(), nullable=False), + sa.Column('create_at', sa.Integer(), nullable=False), sa.PrimaryKeyConstraint('id') ) + with op.batch_alter_table('messagehistory', schema=None) as batch_op: + batch_op.create_index(batch_op.f('ix_messagehistory_chat_id'), ['chat_id'], unique=False) + batch_op.create_index(batch_op.f('ix_messagehistory_create_at'), ['create_at'], unique=False) + # ### end Alembic commands ### def downgrade() -> None: """Downgrade schema.""" # ### commands auto generated by Alembic - please adjust! ### - op.drop_table('historymessage') + with op.batch_alter_table('messagehistory', schema=None) as batch_op: + batch_op.drop_index(batch_op.f('ix_messagehistory_create_at')) + batch_op.drop_index(batch_op.f('ix_messagehistory_chat_id')) + + op.drop_table('messagehistory') # ### end Alembic commands ### diff --git a/产品需求文档AI生成/application/application.py b/产品需求文档AI生成/application/application.py index 4975fa9..c5e91a8 100644 --- a/产品需求文档AI生成/application/application.py +++ b/产品需求文档AI生成/application/application.py @@ -5,6 +5,6 @@ import reflex as rx from application.pages.index import index_page -app = rx.App() +app = rx.App() # 此处变量名需使用 app ，具体原因目前尚不清楚 # 注册首页路由 -app.add_page(index_page, route="/") +app.add_page(component=index_page, route="/") diff --git a/产品需求文档AI生成/application/components/chat.py b/产品需求文档AI生成/application/components/chat.py index 5b5a03a..bb19070 100644 --- a/产品需求文档AI生成/application/components/chat.py +++ b/产品需求文档AI生成/application/components/chat.py @@ -1,90 +1,144 @@ # -*- coding: utf-8 -*- """ -会话相关组件 +聊天相关组件 """ import reflex as rx -from reflex.constants.colors import ColorType from application.models import Type_, Message, Dialog from application.state import ChatState -def input_message_component(message: str, color: ColorType) -> rx.Component: +def render_input_bubble(input_: str) -> rx.Component: """ - 输入消息组件 - :param message: 消息 - :param color: 颜色 + 渲染输入气泡 + :param content: 输入内容 :return: Component """ return rx.markdown( - message, - color=rx.color(color=color, shade=12), - background_color=rx.color(color=color, shade=4), - display="inline-block", - padding_inline="1em", - border_radius="8px", + input_, + color=rx.color("gray", 12), # 文字颜色 + background_color=rx.color("gray", 2), # 背景颜色 + display="inline-block", # 布局模式：自适应文本宽度 + max_width="85%", # 最大宽度 + padding_x="1.25em", # 水平内边距 + padding_y="0.5em", # 垂直内边距 + margin_left="auto", # 左侧外边距自动调整 + margin_bottom="8px", # 底部外边距 + border_radius="12px", # 圆角 ) -def output_message_component(message: Message) -> rx.Component: +def render_collapse_panel(dialog_id: str, message: Message) -> rx.Component: """ - 输出消息组件 - :param message: 消息 - :return: 气泡组件 + 渲染折叠面板 + :param dialog_id: 对话唯一标识 + :param message: 消息实例 + :return: Component """ - color = rx.cond( - message.type_ == Type_.TEXT, - "accent", + + return rx.box( rx.cond( - message.type_ == Type_.THINKING, - "iris", - rx.cond( - message.type_ == Type_.CALL, - "orange", - rx.cond( - message.type_ == Type_.TOOL_RETURN, - "teal", - rx.cond( - message.type_ == Type_.ERROR, - "red", - "mauve", # 兜底 - ), + 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( + rx.hstack( + rx.text( + loading_prefix, + 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", + ), # 折叠面板关闭 ), - ) - return rx.markdown( - message.content, - color=rx.color(color=color, shade=12), - background_color=rx.color(color=color, shade=4), - display="inline-block", - padding_inline="1em", - padding_block="0.5em", - border_radius="8px", - margin_bottom="4px", + rx.cond(message.is_open, detail_box, rx.fragment()), + width="100%", + max_width="85%", + margin_bottom="10px", key=message.id, + cursor="pointer", ) -def dialog_item_component(dialog: Dialog) -> rx.Component: +def render_output_container(dialog_id: str, message: Message) -> rx.Component: """ - 对话项组件 - :param dialog: 对话 + 渲染输出容器 + :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( - input_message_component(message=dialog.input_, color="mauve"), + render_input_bubble(input_=dialog.input_), text_align="right", + width="100%", margin_bottom="8px", ), rx.box( - rx.foreach(dialog.output, output_message_component), + rx.foreach( + dialog.output, + lambda message: render_output_container( + dialog_id=dialog.id, message=message + ), + ), text_align="left", + width="100%", margin_bottom="8px", ), - max_width="50em", - margin_inline="auto", + width="min(100%, 50em)", # 最大宽度：父级元素最大宽度和50em中较小值 + margin_x="auto", # 水平外边距：自动调整 key=dialog.id, ) @@ -125,7 +179,7 @@ def input_component() -> rx.Component: margin="0 auto", align_items="center", ), - on_submit=ChatState.process, # 处理输入，返回流式输出 + on_submit=ChatState.process_input, # 处理输入，返回流式输出 reset_on_submit=True, # 提交后清空输入框 ), rx.text( diff --git a/产品需求文档AI生成/application/models.py b/产品需求文档AI生成/application/models.py index f08a937..34e4ca0 100644 --- a/产品需求文档AI生成/application/models.py +++ b/产品需求文档AI生成/application/models.py @@ -5,38 +5,49 @@ from enum import StrEnum from time import time_ns from typing import List -from typing import Annotated from uuid import uuid4 from pydantic import BaseModel, Field from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter -from sqlmodel import SQLModel, Field as SQLField +from sqlmodel import Field as SqlField, SQLModel -# 数据库表模型 -class HistoryMessage(SQLModel, table=True): - id: int = SQLField(default_factory=int, primary_key=True) - chat_id: str - new_message: str - timestamp: int +# 消息历史数据表模型 +# 需使用 reflex db init 初始化数据库表，若重新初始化需手动删除 alembic 相关配置和文件夹 +class MessageHistory(SQLModel, table=True): + id: str = SqlField( + default_factory=lambda: uuid4().hex, + primary_key=True, + description="消息唯一标识", + ) + chat_id: str = SqlField(index=True, description="聊天唯一标识") + new_message: str = SqlField(description="新消息") + create_at: int = SqlField( + index=True, + description="创建时间戳（毫秒级）", + ) @staticmethod - def adapt(chat_id: str, new_message: List[ModelMessage]) -> "HistoryMessage": - return HistoryMessage( + def adapt(chat_id: str, new_message: List[ModelMessage]) -> "MessageHistory": + return MessageHistory( chat_id=chat_id, - new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode("utf-8"), - timestamp=time_ns() // 1000, # 微秒级时间戳 + new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode( + "utf-8" + ), # 序列化为 JSON 字符串 + create_at=time_ns() // 1_000_000, # 毫秒级时间戳 ) """ 聊天、对话和消息关系： -一次聊天包含若干轮对话，每轮对话包含一条输入消息（input_message）和若干条输出消息（output_messages）。其中，输入消息和输出消息合称消息（message）。 +一次聊天包含若干轮对话，每轮对话包含用户提示词（user_prompt）和输出（output） +其中， +输出包含若干片段（part） """ -class Type_(StrEnum): - """消息类型类""" +class PartType(StrEnum): + """片段类""" THINKING = "thinking" TEXT = "text" @@ -47,24 +58,34 @@ class Type_(StrEnum): ERROR = "error" -# 前缀映射表（动态生成） -PREFIX_MAPING = {f"{i:02d}:": t for i, t in enumerate(Type_)} +# 动态生成前缀和消息类型映射表 +PREFIX_MAPING = {f"{i:02d}:": t for i, t in enumerate(PartType)} -class Message(BaseModel): - """消息类""" +class Part(BaseModel): + """片段类（仅就输出消息）""" - id: str = Field(default_factory=lambda: uuid4().hex, description="消息唯一标识") - type_: Type_ = Field(..., description="消息类型") - content: str = Field(default="", description="消息内容") + id: str = Field(default_factory=lambda: uuid4().hex, description="片段唯一标识") + part_type: PartType = Field(..., alias="type", description="片段类型") + content: str = Field(default="", description="片段内容") + + is_streaming: bool = Field( + default=False, + description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出", + ) + + is_open: bool = Field( + default=False, + description="折叠面板打开状态，True表示打开，False表示关闭", + ) class Dialog(BaseModel): """对话类""" id: str = Field(default_factory=lambda: uuid4().hex, description="对话唯一标识") - input_: str = Field(..., description="输入消息") - output: List[Message] = Field(default_factory=list, description="输出消息") + user_prompt: str = Field(..., description="用户提示词") + output: List[Part] = Field(default_factory=list, description="输出") class Chat(BaseModel): diff --git a/产品需求文档AI生成/application/state/chat.py b/产品需求文档AI生成/application/state/chat.py index 1ca06a3..8d0d733 100644 --- a/产品需求文档AI生成/application/state/chat.py +++ b/产品需求文档AI生成/application/state/chat.py @@ -13,18 +13,19 @@ from application.state.global_ import GlobalState # 绑定的智能体列表（键为聊天唯一标识，值为智能体实例） +# 因 rx.state 需在 Python 进程与 WebSocket 之间传输和存储等操作，非序列化对象会失败，故需将智能体实例单独存储 agents: Dict[str, Agent] = {} def get_agent(chat_id: str) -> Agent: """ 获取绑定的智能体 - :return: 当前聊天绑定的智能体 + :return: 绑定的智能体 """ if chat_id not in agents: agents[chat_id] = Agent( chat_id=chat_id, - instructions="You are a friendly chatbot", + instructions="You are a friendly chatbot", # 智能体指令 ) return agents[chat_id] @@ -119,7 +120,7 @@ class ChatState(rx.State): self.current_chat_id = next(iter(self.chats)) @rx.event - async def process(self, form_data: dict[str, Any]) -> AsyncGenerator: + async def process_input(self, form_data: dict[str, Any]) -> AsyncGenerator: """ 处理输入，返回流式输出 :param form_data: 表单数据 @@ -131,11 +132,15 @@ class ChatState(rx.State): # 当前聊天 current_chat = self.chats[self.current_chat_id] + if not current_chat: + return + # 新增对话 current_chat.dialogs.append(Dialog(input_=input_)) # 当前对话 current_dialog = current_chat.dialogs[-1] + # 设置当前对话流式输出状态为正在流式输出 current_chat.is_streaming = True yield # 通知前端渲染输入消息 @@ -162,6 +167,7 @@ class ChatState(rx.State): if not prefix: continue + # 根据前缀匹配消息类型 type_ = PREFIX_MAPING[prefix] # 若当前对话输出消息为空或当前消息类型与新消息类型不相同则新增消息 if not current_dialog.output or current_dialog.output[-1].type_ != type_: @@ -175,4 +181,30 @@ class ChatState(rx.State): await global_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: + """ + 打开/关闭折叠面板 + :return: None + """ + # 当前聊天 + current_chat = self.chats[self.current_chat_id] + if not current_chat: + return + # 目标对话 + target_dialog = next( + (d for d in current_chat.dialogs if d.id == dialog_id), None + ) + if not target_dialog: + return + # 目标消息 + target_message = next( + (m for m in target_dialog.output if m.id == message_id), None + ) + if not target_message: + return + # 切换折叠面板打开状态 + target_message.is_open = not target_message.is_open diff --git a/产品需求文档AI生成/application/state/global_.py b/产品需求文档AI生成/application/state/global_.py index 8f6ba8b..b2455a1 100644 --- a/产品需求文档AI生成/application/state/global_.py +++ b/产品需求文档AI生成/application/state/global_.py @@ -3,20 +3,18 @@ 全局共享状态 """ -from typing import Dict, List +from typing import List from sqlmodel import select as sql_select import reflex as rx import sqlalchemy as sa from application.models import ( - HistoryMessage, + MessageHistory, ModelMessage, ModelMessagesTypeAdapter, ) - - class GlobalState(rx.State): """全局状态""" @@ -29,9 +27,9 @@ class GlobalState(rx.State): message_history = [] async with rx.asession() as session: records = await session.exec( - sql_select(HistoryMessage) - .where(HistoryMessage.chat_id == chat_id) - .order_by(sa.desc("timestamp")) + sql_select(MessageHistory) + .where(MessageHistory.chat_id == chat_id) + .order_by(sa.desc("create_at")) ) for record in records.all(): message_history.extend( @@ -49,6 +47,6 @@ class GlobalState(rx.State): :return: None """ async with rx.asession() as session: - record = HistoryMessage.adapt(chat_id=chat_id, new_message=new_message) + record = MessageHistory.adapt(chat_id=chat_id, new_message=new_message) session.add(record) await session.commit() diff --git a/产品需求文档AI生成/application/utils/agent.py b/产品需求文档AI生成/application/utils/agent.py index 4e5438a..7fea8c5 100644 --- a/产品需求文档AI生成/application/utils/agent.py +++ b/产品需求文档AI生成/application/utils/agent.py @@ -37,15 +37,15 @@ from pydantic_ai.providers.openai import OpenAIProvider from pydantic_ai.run import AgentRunResultEvent -class Buffer(BaseModel): +class PartBuffer(BaseModel): """ - 模型消息事件缓冲类 + 片段缓冲类 """ part_type: Literal["thinking", "text"] = Field(..., description="片段类型") content_delta: str = Field(default="", description="片段内容增量") task: Optional[Task] = Field(default=None, description="延迟刷新异步协程任务") - model_config = {"arbitrary_types_allowed": True} + model_config = {"arbitrary_types_allowed": True} # 允许任意类型 class Debouncer: @@ -61,8 +61,8 @@ class Debouncer: """ self.delay = delay - # 模型消息事件缓冲字典（数据类型为字典，键为片段索引，值为缓冲模型） - self.buffers: Dict[int, Buffer] = {} + # 片段缓冲字典（键为片段索引，值为片段缓冲实例） + self.part_buffers: Dict[int, PartBuffer] = {} # 待刷新异步队列 self.pending_flush_queue = Queue()