237 lines
8.5 KiB
Python
237 lines
8.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
会话页面状态
|
||
"""
|
||
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, cast
|
||
|
||
from pydantic_ai._uuid import uuid7
|
||
import reflex as rx
|
||
|
||
from application.models import Conversation, EventKind, PartKind, Run
|
||
from application.state.create_conversation import CreateConversationState
|
||
from application.state.database import DatabaseState
|
||
from application.utils.agent import AIAgent
|
||
|
||
|
||
class ConversationState(rx.State):
|
||
"""
|
||
会话状态
|
||
"""
|
||
|
||
# 当前会话唯一标识
|
||
current_conversation_id: str = str(uuid7())
|
||
# 会话字典
|
||
conversations: Dict[str, Conversation] = {current_conversation_id: Conversation()}
|
||
|
||
# 当前运行唯一标识
|
||
current_run_id: Optional[str] = None
|
||
|
||
# 初始化智能体
|
||
agent: AIAgent = AIAgent()
|
||
|
||
@rx.var
|
||
def get_conversations(self) -> Dict[str, Conversation]:
|
||
"""
|
||
获取会话字典，用于前端渲染会话列表
|
||
:return: 会话字典
|
||
"""
|
||
return self.conversations
|
||
|
||
@rx.var
|
||
def get_current_conversation_description(self) -> str:
|
||
"""
|
||
获取当前会话描述，用于前端渲染导航栏中当前会话描述
|
||
:return: 当前会话描述
|
||
"""
|
||
# 若当前会话唯一标识不存在则创建会话
|
||
if self.current_conversation_id not in self.conversations:
|
||
self.conversations[self.current_conversation_id] = Conversation()
|
||
|
||
# 当前会话
|
||
current_conversation = self.conversations[self.current_conversation_id]
|
||
return current_conversation.description
|
||
|
||
@rx.event
|
||
def set_current_conversation_id(self, conversation_id: str) -> None:
|
||
"""
|
||
将指定会话唯一标识设置为当前会话唯一标识
|
||
:param conversation_id: 指定会话唯一标识
|
||
:return: None
|
||
"""
|
||
# 若指定会话唯一标识不存在则创建会话
|
||
if conversation_id not in self.conversations:
|
||
self.conversations[conversation_id] = Conversation()
|
||
|
||
# 将指定会话唯一标识设置为当前会话唯一标识
|
||
self.current_conversation_id = conversation_id
|
||
|
||
@rx.event
|
||
async def create_conversation(self, form_data: Dict[str, Any]) -> None:
|
||
"""
|
||
创建会话
|
||
:param form_data: 表单数据
|
||
:return: None
|
||
"""
|
||
# 获取描述
|
||
description = form_data["conversation_description"].strip()
|
||
if not description:
|
||
description = "新会话"
|
||
|
||
# 创建会话唯一标识
|
||
self.current_conversation_id = str(uuid7())
|
||
self.conversations[self.current_conversation_id] = Conversation(
|
||
description=description
|
||
)
|
||
|
||
# 跨状态读取创建会话模态窗状态
|
||
create_conversation_state = await self.get_state(CreateConversationState)
|
||
create_conversation_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.conversations[str(uuid7())] = Conversation()
|
||
# 删除后若当前会话唯一标识不存在则将第一个会话唯一标识设置为当前会话唯一标识
|
||
if self.current_conversation_id not in self.conversations:
|
||
self.current_conversation_id = next(iter(self.conversations))
|
||
|
||
@rx.var
|
||
def get_runs(self) -> Dict[str, Run]:
|
||
"""
|
||
获取当前会话运行字典，用于前端渲染运行列表
|
||
:return: 当前会话运行字典
|
||
"""
|
||
# 若当前会话唯一标识不存在则创建会话
|
||
if self.current_conversation_id not in self.conversations:
|
||
self.conversations[self.current_conversation_id] = Conversation()
|
||
|
||
# 当前会话
|
||
current_conversation = self.conversations[self.current_conversation_id]
|
||
return current_conversation.runs
|
||
|
||
@rx.var
|
||
def get_current_run_streaming_status(self) -> bool:
|
||
"""
|
||
获取当前运行流式输出状态
|
||
:return: 当前运行流式输出状态（True 表示正在流式输出，False 表示非正在流式输出）
|
||
"""
|
||
# 若当前会话唯一标识不存在则创建会话
|
||
if self.current_conversation_id not in self.conversations:
|
||
self.conversations[self.current_conversation_id] = Conversation()
|
||
|
||
# 当前会话
|
||
current_conversation = self.conversations[self.current_conversation_id]
|
||
|
||
# 若当前运行唯一标识为空或不存在则返回 False
|
||
if (
|
||
not self.current_run_id
|
||
or self.current_run_id not in current_conversation.runs
|
||
):
|
||
return False
|
||
|
||
# 当前运行
|
||
current_run = current_conversation.runs[self.current_run_id]
|
||
return current_run.is_streaming
|
||
|
||
@rx.event
|
||
async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
|
||
"""
|
||
运行
|
||
:param form_data: 表单数据
|
||
:return: AsyncGenerator
|
||
"""
|
||
# 若当前会话唯一标识不存在则创建会话
|
||
if self.current_conversation_id not in self.conversations:
|
||
self.conversations[self.current_conversation_id] = Conversation()
|
||
|
||
# 当前会话
|
||
current_conversation = self.conversations[self.current_conversation_id]
|
||
|
||
# 获取用户提示词
|
||
user_prompt = form_data["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.current_conversation_id
|
||
)
|
||
|
||
async for event in self.agent.run(
|
||
conversation_id=self.current_conversation_id,
|
||
user_prompt=user_prompt,
|
||
message_history=message_history,
|
||
):
|
||
|
||
# 若事件为空或当前运行唯一标识为空或当前运行唯一标识与事件运行唯一标识不相同则跳过
|
||
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
|
||
# 创建运行
|
||
current_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
|