211 lines
6.7 KiB
Python
211 lines
6.7 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
聊天页面状态
|
||
"""
|
||
from typing import Any, AsyncGenerator, Dict, List
|
||
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
|
||
|
||
|
||
# 绑定的智能体列表（键为聊天唯一标识，值为智能体实例）
|
||
# 因 rx.state 需在 Python 进程与 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,
|
||
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]:
|
||
"""
|
||
获取聊天唯一标识列表
|
||
:return: 聊天唯一标识列表
|
||
"""
|
||
return list(self.chats)
|
||
|
||
@rx.var
|
||
def get_chat_descriptions(self) -> Dict[str, str]:
|
||
"""
|
||
获取聊天描述列表
|
||
:return: 聊天描述列表
|
||
"""
|
||
return {chat_id: chat.description for chat_id, chat in self.chats.items()}
|
||
|
||
@rx.var
|
||
def get_current_chat_status(self) -> bool:
|
||
"""
|
||
获取当前聊天流式输出状态
|
||
:return: 当前聊天流式输出状态（True 表示正在流式输出，False 表示非正在流式输出）
|
||
"""
|
||
chat = self.chats.get(self.current_chat_id)
|
||
return chat.is_streaming if chat else False
|
||
|
||
@rx.var
|
||
def get_current_chat_dialogs(self) -> List[Dialog]:
|
||
"""
|
||
获取当前聊天对话列表
|
||
:return: 当前聊天对话列表
|
||
"""
|
||
chat = self.chats.get(self.current_chat_id)
|
||
return chat.dialogs if 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_modal_state = await self.get_state(CreateChatModalState)
|
||
create_chat_modal_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 process_input(self, form_data: dict[str, Any]) -> AsyncGenerator:
|
||
"""
|
||
处理输入，返回流式输出
|
||
:param form_data: 表单数据
|
||
:return: AsyncGenerator
|
||
"""
|
||
input_ = form_data["input"].strip()
|
||
if not input_:
|
||
return
|
||
|
||
# 当前聊天
|
||
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 # 通知前端渲染输入消息
|
||
|
||
# 全局状态
|
||
global_state = await self.get_state(GlobalState)
|
||
# 获取当前聊天绑定的智能体
|
||
agent = get_agent(chat_id=self.current_chat_id)
|
||
# 获取当前聊天消息历史
|
||
message_history = await global_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
|
||
):
|
||
# 若消息事件为空则跳过
|
||
if not event:
|
||
continue
|
||
|
||
# 匹配前缀
|
||
prefix = next((p for p in PREFIX_MAPING if event.startswith(p)), None)
|
||
# 若未匹配到前缀则跳过
|
||
if not prefix:
|
||
continue
|
||
|
||
# 根据前缀匹配消息类型
|
||
type_ = PREFIX_MAPING[prefix]
|
||
# 若当前对话输出消息为空或当前消息类型与新消息类型不相同则新增消息
|
||
if not current_dialog.output or current_dialog.output[-1].type_ != type_:
|
||
current_dialog.output.append(Message(type_=type_))
|
||
|
||
# 追加消息内容
|
||
current_dialog.output[-1].content += event.removeprefix(prefix)
|
||
yield # 通知前端渲染输出消息
|
||
|
||
# 保存本轮对话消息
|
||
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
|