Python/agent/application/states/database.py

# -*- coding: utf-8 -*-
"""
数据库状态
"""
from datetime import datetime
from typing import Dict, List, cast

from pydantic import 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,
    Conversations,
    Dialog,
    Dialogs,
    RunResults,
    ThoughtNode,
)


# 思考节点内存模型类型适配器
ThoughtNodeTypeAdapter = TypeAdapter(dict[int, ThoughtNode])


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)
                .outerjoin(
                    Dialogs,
                    cast(
                        ColumnElement[bool],
                        Conversations.id == Dialogs.conversation_id,
                    ),
                )
                .where(
                    Conversations.user_id == user_id,
                    Conversations.is_deleted == False,
                )
                .order_by(Conversations.id, Dialogs.id)
            )
            for conversation, dialog in result.all():
                conversation_memory = conversations.setdefault(
                    conversation.id,
                    Conversation(
                        description=conversation.description,
                        created_at=conversation.created_at,
                    ),
                )

                if not dialog:
                    continue
                conversation_memory.dialogs[dialog.id] = Dialog(
                    question=dialog.question,
                    thought_nodes=ThoughtNodeTypeAdapter.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 = Conversations(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(Conversations, 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 = Dialogs(
                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: dict[int, ThoughtNode],
        answer: str,
    ) -> None:
        """
        完成对话
        :param dialog_id: 对话唯一标识
        :param thought_nodes: 思考节点列表
        :param answer: 回答
        :return: None
        """
        async with rx.asession() as session:
            dialog = await session.get(Dialogs, dialog_id)
            if not dialog:
                return
            dialog.thought_nodes = ThoughtNodeTypeAdapter.dump_python(
                {k: v for k, v in thought_nodes.items()}
            )
            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(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:
            session.add(
                RunResults(
                    conversation_id=conversation_id,
                    dialog_id=dialog_id,
                    new_messages=ModelMessagesTypeAdapter.dump_json(
                        new_messages
                    ).decode(
                        "utf-8"
                    ),  # 序列化为 JSON 字符串
                )
            )
            await session.commit()