Python/产品需求文档AI生成/application/state/chat.py

# -*- coding: utf-8 -*-
"""
聊天页面状态
"""
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, Part, PartType, PREFIX_MAPING
from application.state.create_chat import CreateChatState
from application.state.database import DatabaseState

# 因 rx.state 需在进程与 WebSocket 间交互，非序列化对象会传输或储存失败，故需将智能体实例单独存储

# 绑定的智能体列表（键为聊天唯一标识，值为智能体实例）
agents: Dict[str, Agent] = {}


def get_agent(chat_id: str) -> Agent:
    """
    获取绑定的智能体
    :return: 绑定的智能体
    """
    # 若智能体唯一标识不存在则新建智能体实例
    if chat_id not in agents:
        agents[chat_id] = Agent(
            chat_id=chat_id,
        )  # 使用默认指令
    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_chats(self) -> Dict[str, Chat]:
        """
        获取聊天列表，用于前端渲染聊天列表
        :return: 聊天列表
        """
        return self.chats

    @rx.var
    def get_current_chat_description(self) -> str:
        """
        获取当前聊天描述，用于前端渲染导航栏中当前聊天描述
        :return: 当前聊天描述
        """
        # 当前聊天
        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 表示非正在流式输出）
        """
        # 当前聊天
        current_chat = self.chats.get(self.current_chat_id)
        return current_chat.is_streaming if current_chat else False

    @rx.var
    def get_dialogs(self) -> Dict[str, Dialog]:
        """
        获取对话列表，用于前端渲染对话列表
        :return: 对话列表
        """
        # 当前聊天
        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:
        """
        切换聊天
        :param chat_id: 聊天唯一标识
        :return: None
        """
        if chat_id not in self.chats:
            return
        self.current_chat_id = chat_id

    @rx.event
    async def create_chat(self, form_data: Dict[str, Any]) -> None:
        """
        新建聊天
        :param form_data: 表单数据
        :return: None
        """
        # 生成聊天唯一标识
        self.current_chat_id = uuid4().hex
        self.chats[self.current_chat_id] = Chat(
            description=form_data["chat_description"].strip() or "新聊天"
        )

        # 跨状态读取新建聊天模态窗状态（异步事件）
        create_chat_state = await self.get_state(CreateChatState)
        create_chat_state.is_open = False

    @rx.event
    def delete_chat(self, chat_id: str) -> None:
        """
        删除聊天
        :param chat_id: 聊天唯一标识
        :return: None
        """
        if chat_id not in self.chats:
            return
        del self.chats[chat_id]
        # 删除聊天后若聊天列表为空，则新建聊天
        if not self.chats:
            self.chats[uuid4().hex] = Chat()
        # 删除聊天后若当前聊天唯一标识不存在，则切换到第一个聊天
        if self.current_chat_id not in self.chats:
            self.current_chat_id = next(iter(self.chats))

    @rx.event
    async def run(self, form_data: dict[str, Any]) -> AsyncGenerator:
        """
        运行
        :param form_data: 表单数据
        :return: AsyncGenerator
        """
        # 获取用户提示词
        user_prompt = form_data["user_prompt"].strip()
        if not user_prompt:
            return

        # 当前聊天
        current_chat = self.chats[self.current_chat_id]
        if not current_chat:
            return

        # 生成新增对话唯一标识
        current_dialog_id = uuid4().hex
        # 新增对话
        current_chat.dialogs[current_dialog_id] = Dialog(user_prompt=user_prompt)
        # 当前对话
        current_dialog = current_chat.dialogs[current_dialog_id]

        # 设置当前聊天流式输出状态为正在流式输出
        current_chat.is_streaming = True
        yield  # 通知前端渲染输入消息

        # 获取当前聊天绑定的智能体
        agent = get_agent(chat_id=self.current_chat_id)

        # 数据库状态
        database_state = await self.get_state(DatabaseState)
        # 获取当前聊天消息历史
        message_history = await database_state.get_message_history(
            chat_id=self.current_chat_id
        )

        # 流式输出模型响应事件
        async for event in agent.run(
            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]

            # 获取当前片段
            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

                # 新增片段
                current_dialog.output[uuid4().hex] = (
                    new_part := Part(part_type=part_type, is_streaming=True)
                )
                current_part = new_part

            # 追加片段内容
            current_part.content += event.removeprefix(prefix)
            yield  # 通知前端渲染片段

        # 重置所有片段流式输出状态
        for part in current_dialog.output.values():
            part.is_streaming = False

        # 保存本轮对话消息
        await database_state.save_new_message(
            chat_id=self.current_chat_id, new_message=agent.new_messages
        )
        # 设置当前聊天流式输出状态非正在流式输出
        current_chat.is_streaming = False

    @rx.event
    def toggle_part_collapse(self, dialog_id: str, part_id: str) -> None:
        """
        打开/关闭片段折叠面板
        :param dialog_id: 对话唯一标识
        :param part_id: 片段唯一标识
        :return: None
        """
        # 当前聊天
        current_chat = self.chats.get(self.current_chat_id)
        if not current_chat:
            return

        # 当前对话
        current_dialog = current_chat.dialogs.get(dialog_id)
        if not current_dialog:
            return

        # 当前片段
        current_part = current_dialog.output.get(part_id)
        if not current_part:
            return

        # 切换折叠面板打开状态
        current_part.is_open = not current_part.is_open