229 lines
7.3 KiB
Python
229 lines
7.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据库状态
|
||
"""
|
||
|
||
from datetime import datetime
|
||
from typing import Dict, List, cast, TypeAlias
|
||
|
||
from pydantic import BaseModel, Field, TypeAdapter
|
||
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter
|
||
import reflex as rx
|
||
from sqlalchemy import ColumnElement, desc
|
||
from sqlmodel import select
|
||
|
||
from application.models import Conversation, Dialog, RunResult
|
||
|
||
|
||
class ConversationMemory(BaseModel):
|
||
"""
|
||
会话内存模型
|
||
"""
|
||
|
||
description: str = Field(default="新会话", description="会话描述")
|
||
is_running: bool = Field(
|
||
default=False, description="运行状态，True 表示运行中，False 表示运行完成"
|
||
)
|
||
dialogs: Dict[str, DialogMemory] = Field(
|
||
default_factory=dict, description="对话字典"
|
||
)
|
||
created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
|
||
|
||
|
||
class DialogMemory(BaseModel):
|
||
"""
|
||
对话内存模型
|
||
"""
|
||
|
||
question: str = Field(..., description="问题")
|
||
is_thinking: bool = Field(
|
||
default=False, description="思考状态，True 表示思考中，False 表示思考完成"
|
||
)
|
||
is_collapse_expanded: bool = Field(
|
||
default=False, description="思考折叠面板展开状态，True 表示展开，False 表示折叠"
|
||
)
|
||
thought_nodes: list[ThoughtNodeMemory] = Field(
|
||
default_factory=list, description="思考节点列表"
|
||
)
|
||
answer: str = Field(default="", description="回答")
|
||
|
||
|
||
class ThoughtNodeMemory(BaseModel):
|
||
"""
|
||
思考节点内存模型
|
||
"""
|
||
|
||
kind: str
|
||
content: str
|
||
|
||
|
||
# 会话字典类型别名
|
||
Conversations: TypeAlias = Dict[str, ConversationMemory]
|
||
|
||
# 思考节点内存模型类型适配器
|
||
ThoughtNodeMemoryTypeAdapter = TypeAdapter(list[ThoughtNodeMemory])
|
||
|
||
|
||
class DatabaseState(rx.State):
|
||
"""
|
||
数据库状态
|
||
"""
|
||
|
||
async def get_conversations(self, user_id: str) -> Conversations:
|
||
"""
|
||
获取会话字典
|
||
:param user_id: 用户唯一标识
|
||
:return: 会话字典
|
||
"""
|
||
conversations: Conversations = {}
|
||
async with rx.asession() as session:
|
||
result = await session.exec(
|
||
select(Conversation, Dialog)
|
||
.outerjoin(
|
||
Dialog,
|
||
cast(
|
||
ColumnElement[bool],
|
||
Conversation.id == Dialog.conversation_id,
|
||
),
|
||
)
|
||
.where(
|
||
Conversation.user_id == user_id,
|
||
Conversation.is_deleted == False,
|
||
)
|
||
.order_by(Conversation.id, Dialog.id)
|
||
)
|
||
for conversation, dialog in result.all():
|
||
conversation_memory = conversations.setdefault(
|
||
conversation.id,
|
||
ConversationMemory(
|
||
description=conversation.description,
|
||
created_at=conversation.created_at,
|
||
),
|
||
)
|
||
|
||
if not dialog:
|
||
continue
|
||
conversation_memory.dialogs[dialog.id] = DialogMemory(
|
||
question=dialog.question,
|
||
thought_nodes=ThoughtNodeMemoryTypeAdapter.validate_python(
|
||
dialog.thought_nodes
|
||
), # 反序列化为思考节点内存实例列表
|
||
answer=dialog.answer,
|
||
)
|
||
return conversations
|
||
|
||
async def create_conversation(
|
||
self, user_id: str, description: str = "新会话"
|
||
) -> tuple[str, datetime]:
|
||
"""
|
||
新建会话
|
||
:param user_id: 用户唯一标识
|
||
:param description: 会话描述
|
||
:return: 新建会话唯一标识、创建时间
|
||
"""
|
||
async with rx.asession() as session:
|
||
conversation = Conversation(user_id=user_id, description=description)
|
||
session.add(conversation)
|
||
await session.commit()
|
||
await session.refresh(conversation)
|
||
return conversation.id, conversation.created_at
|
||
|
||
async def delete_conversation(self, conversation_id: str) -> None:
|
||
"""
|
||
删除指定会话（逻辑删除）
|
||
:param conversation_id: 指定会话唯一标识
|
||
:return: None
|
||
"""
|
||
async with rx.asession() as session:
|
||
conversation = await session.get(Conversation, conversation_id)
|
||
if not conversation:
|
||
return
|
||
conversation.is_deleted = True
|
||
await session.commit()
|
||
|
||
async def create_dialog(self, conversation_id: str, question: str) -> str:
|
||
"""
|
||
新建对话
|
||
:param conversation_id: 会话唯一标识
|
||
:param question: 提问
|
||
:return: 新建对话唯一标识
|
||
"""
|
||
async with rx.asession() as session:
|
||
dialog = Dialog(
|
||
conversation_id=conversation_id,
|
||
question=question,
|
||
)
|
||
session.add(dialog)
|
||
await session.commit()
|
||
await session.refresh(dialog)
|
||
return dialog.id
|
||
|
||
async def complete_dialog(
|
||
self,
|
||
dialog_id: str,
|
||
thought_nodes: list[ThoughtNodeMemory],
|
||
answer: str,
|
||
) -> None:
|
||
"""
|
||
完成对话
|
||
:param dialog_id: 对话唯一标识
|
||
:param thought_nodes: 思考节点列表
|
||
:param answer: 回答
|
||
:return: None
|
||
"""
|
||
async with rx.asession() as session:
|
||
dialog = await session.get(Dialog, dialog_id)
|
||
if not dialog:
|
||
return
|
||
dialog.thought_nodes = ThoughtNodeMemoryTypeAdapter.dump_python(
|
||
thought_nodes
|
||
)
|
||
dialog.answer = answer
|
||
await session.commit()
|
||
|
||
async def get_message_history(self, conversation_id: str) -> List[ModelMessage]:
|
||
"""
|
||
获取消息历史
|
||
:param conversation_id: 会话唯一标识
|
||
:return: 消息历史
|
||
"""
|
||
message_history: List[ModelMessage] = []
|
||
async with rx.asession() as session:
|
||
result = await session.exec(
|
||
select(RunResult)
|
||
.where(RunResult.conversation_id == conversation_id)
|
||
.order_by(desc(RunResult.dialog_id))
|
||
)
|
||
for record in result.all():
|
||
message_history.extend(
|
||
ModelMessagesTypeAdapter.validate_json(record.new_messages)
|
||
)
|
||
return message_history
|
||
|
||
async def save_new_messages(
|
||
self,
|
||
conversation_id: str,
|
||
dialog_id: str,
|
||
new_messages: List[ModelMessage],
|
||
) -> None:
|
||
"""
|
||
保存新增消息
|
||
:param conversation_id: 会话唯一标识
|
||
:param dialog_id: 对话唯一标识
|
||
:param new_messages: 新增消息
|
||
:return: None
|
||
"""
|
||
async with rx.asession() as session:
|
||
session.add(
|
||
RunResult(
|
||
conversation_id=conversation_id,
|
||
dialog_id=dialog_id,
|
||
new_messages=ModelMessagesTypeAdapter.dump_json(
|
||
new_messages
|
||
).decode(
|
||
"utf-8"
|
||
), # 序列化为 JSON 字符串
|
||
)
|
||
)
|
||
await session.commit()
|