91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
数据库状态
|
|
"""
|
|
|
|
from typing import List, Dict
|
|
|
|
import reflex as rx
|
|
from sqlalchemy import desc, and_
|
|
from sqlmodel import select
|
|
|
|
from application.models import Conversations, RunResults
|
|
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter
|
|
|
|
|
|
class DatabaseState(rx.State):
|
|
"""
|
|
数据库状态
|
|
"""
|
|
|
|
async def get_conversations(self, user_id: str) -> Dict[str, Conversation]:
|
|
"""
|
|
获取指定用户的会话字典
|
|
:param user_id: 用户唯一标识
|
|
:return: 指定用户的会话字典
|
|
"""
|
|
conversations: Dict[str, Conversation] = {}
|
|
async with rx.asession() as session:
|
|
result = await session.exec(
|
|
select(Conversations, Dialogs, Parts)
|
|
.outerjoin(Dialogs, Conversations.id == Dialogs.conversation_id) # type: ignore
|
|
.outerjoin(Parts, Dialogs.id == Parts.dialog_id) # type: ignore
|
|
.where(
|
|
models.Conversation.user_id == user_id,
|
|
models.Conversation.is_deleted == False,
|
|
)
|
|
.order_by(
|
|
models.Conversation.id, models.Dialog.id, models.ReasoningPart.id
|
|
)
|
|
)
|
|
records = result.all()
|
|
for db_conversation, _, _ in records:
|
|
conversations[conversation.id] = models.Conversation(
|
|
user_id=user_id, description=conversation.description
|
|
)
|
|
|
|
return conversations
|
|
|
|
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(RunResults)
|
|
.where(RunResults.conversation_id == conversation_id)
|
|
.order_by(desc(RunResults.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:
|
|
record = RunResults.adapt(
|
|
conversation_id=conversation_id,
|
|
dialog_id=dialog_id,
|
|
new_messages=new_messages,
|
|
)
|
|
session.add(record)
|
|
await session.commit()
|