Python/agent/application/states/database.py

# -*- coding: utf-8 -*-
"""
数据库状态
"""
from dataclasses import asdict
from datetime import datetime, timedelta
from random import choices

from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter, RunUsage
from pydantic_ai._uuid import uuid7
import reflex as rx
from sqlmodel import Field, SQLModel, select, update

from application.states.models import (
    Approval,
    Conversation,
    Message,
    MessageType,
    Workflow,
    approval_dump_json,
    approval_validate_json,
    usage_dump_json,
    usage_validate_json,
    workflow_dump_json,
    workflow_validate_json,
)


class CaptchaRecord(SQLModel, table=True):
    """
    验证码记录
    """

    email: str = Field(..., primary_key=True, description="邮箱")
    code: str = Field(
        default_factory=lambda: "".join(choices("0123456789", k=6)),
        primary_key=True,
        description="验证码",
    )
    is_valid: bool = Field(
        default=True,
        index=True,
        description="验证码有效：True 表示有效，False 表示无效",
    )
    is_verified: bool = Field(
        default=False,
        index=True,
        description="验证码已核验：True 表示已核验，False 表示未核验",
    )
    expired_at: datetime = Field(
        default_factory=lambda: datetime.now() + timedelta(minutes=30),
        primary_key=True,
        description="过期时间",
    )


class UserRecord(SQLModel, table=True):
    """
    用户记录
    """

    id: str = Field(
        default_factory=lambda: str(uuid7()),
        primary_key=True,
        description="用户唯一标识",
    )
    email: str = Field(..., index=True, description="邮箱")


class ConversationRecord(SQLModel, table=True):
    """
    会话记录
    """

    id: str = Field(
        default_factory=lambda: str(uuid7()),
        primary_key=True,
        description="会话唯一标识",
    )
    user_id: str = Field(..., index=True, description="用户唯一标识")
    description: str = Field(default="新会话", description="会话描述")
    created_at: datetime = Field(
        default_factory=datetime.now, description="会话创建时间"
    )
    is_deleted: bool = Field(
        default=False,
        index=True,
        description="会话已删除：True 表示已删除，False 表示未删除",
    )
    usage: str = Field(default="", description="会话使用量")
    workflow: str = Field(
        default="",
        description="会话工作流",
    )  # 原则上嵌套列表或字典若不涉及查询则以字符串储存


class RunRecord(SQLModel, table=True):
    """
    运行记录
    """

    id: str = Field(
        default_factory=lambda: str(uuid7()),
        primary_key=True,
        description="运行唯一标识",
    )
    conversation_id: str = Field(..., index=True, description="会话唯一标识")
    new_messages: str = Field(..., description="新增消息")


class MessageRecord(SQLModel, table=True):
    """
    消息记录
    """

    id: str = Field(
        default_factory=lambda: str(uuid7()),
        primary_key=True,
        description="消息唯一标识",
    )
    conversation_id: str = Field(..., index=True, description="会话唯一标识")
    type: MessageType = Field(..., description="消息类型")
    title: str = Field(default="", description="消息标题")
    content: str = Field(default="", description="消息内容")
    approval: str = Field(
        default="", description="审批"
    )  # 原则上嵌套列表或字典若不涉及查询则以字符串储存


class DatabaseState(rx.State):
    """
    数据库状态
    """

    async def create_captcha_record(self, email: str) -> str:
        """
        创建验证记录
        :param email: 邮箱
        :return: 验证码
        """
        async with rx.asession() as session:
            # 先将该邮箱有效、未核验的验证码记录设置为无效
            await session.exec(
                update(CaptchaRecord)
                .where(
                    CaptchaRecord.email == email,  # type: ignore
                    CaptchaRecord.is_valid == True,  # type: ignore
                    CaptchaRecord.is_verified == False,  # type: ignore
                )
                .values(is_valid=False)
            )
            await session.flush()
            # 创建验证记录
            record = CaptchaRecord(
                email=email,
            )
            session.add(record)
            await session.commit()
            await session.refresh(record)
            return record.code

    async def verify_captcha_code(self, email: str, captcha_code: str) -> bool:
        """
        核验验证码
        :param email: 邮箱
        :param code: 验证码
        :return: 核验是否成功，True 表示核验成功，False 表示核验失败
        """
        async with rx.asession() as session:
            result = await session.exec(
                select(CaptchaRecord).where(
                    CaptchaRecord.email == email,
                    CaptchaRecord.code == captcha_code,
                    CaptchaRecord.is_valid == True,
                    CaptchaRecord.is_verified == False,
                    CaptchaRecord.expired_at > datetime.now(),
                )
            )
            record = result.first()
            if not record:
                return False
            record.is_verified = True
            await session.commit()
            await session.refresh(record)
            return True

    async def create_user_record(self, email: str) -> str:
        """
        创建用户记录
        :param email: 邮箱
        :return: 用户唯一标识
        """
        async with rx.asession() as session:
            result = await session.exec(
                select(UserRecord).where(UserRecord.email == email)
            )
            record = result.first()
            # 若用户记录已存在则返回用户唯一标识，否则先创建用户记录再返回用户唯一标识
            if record:
                return record.id
            record = UserRecord(email=email)
            session.add(record)
            await session.commit()
            await session.refresh(record)
            return record.id

    async def create_conversation_record(self, user_id: str) -> dict[str, Conversation]:
        """
        创建会话记录
        :param user_id: 用户唯一标识
        :return: 会话实例
        """
        async with rx.asession() as session:
            record = ConversationRecord(user_id=user_id)
            session.add(record)
            await session.commit()
            await session.refresh(record)
        return {
            record.id: Conversation(
                id=record.id,
                description=record.description,
                usage=usage_validate_json(
                    record.usage
                ),  # 将 record.usage 由 JSON 字符串反序列化为 RunUsage
                created_at=record.created_at,
            )
        }

    async def read_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(ConversationRecord, MessageRecord)
                .outerjoin(MessageRecord, MessageRecord.conversation_id == ConversationRecord.id)  # type: ignore
                .where(
                    ConversationRecord.user_id == user_id,
                    ConversationRecord.is_deleted == False,
                )
                .order_by(ConversationRecord.id, MessageRecord.id)
            )
            for (
                conversation_record,
                message_record,
            ) in result.all():
                conversation = conversations.setdefault(
                    conversation_record.id,
                    Conversation(
                        id=conversation_record.id,
                        description=conversation_record.description,
                        usage=usage_validate_json(
                            conversation_record.usage
                        ),  # 将 conversation_record.usage 由 JSON 字符串反序列化为 RunUsage
                        created_at=conversation_record.created_at,
                        workflow=workflow_validate_json(
                            conversation_record.workflow
                        ),  # 将 conversation_record.workflow 由 JSON 字符串反序列化为 Workflow,
                    ),
                )
                if not message_record:
                    continue
                conversation.messages.setdefault(
                    message_record.id,
                    Message(
                        id=message_record.id,
                        type=message_record.type,
                        title=message_record.title,
                        content=message_record.content,
                        approval=approval_validate_json(message_record.approval),
                    ),
                )
        return conversations

    async def update_conversation_record(
        self,
        conversation_id: str,
        description: str | None = None,
        is_deleted: bool | None = None,
        usage: RunUsage | None = None,
        workflow: Workflow | None = None,
    ) -> None:
        """
        更新会话记录
        :param conversation_id: 指定会话唯一标识
        :param description: 会话描述
        :param is_deleted: 会话已删除
        :param usage: 使用量
        :param workflow: 工作流
        :return: None
        """
        async with rx.asession() as session:
            record = await session.get(ConversationRecord, conversation_id)
            if not record:
                return
            if isinstance(description, str):
                record.description = description
            if isinstance(is_deleted, bool):
                record.is_deleted = is_deleted
            if isinstance(usage, RunUsage):
                record.usage = usage_dump_json(
                    RunUsage(**asdict(usage))
                )  # 将 RunUsage 序列化为 JSON 字符串
            if isinstance(workflow, Workflow):
                record.workflow = workflow_dump_json(
                    Workflow(**workflow.model_dump(mode="json"))
                )  # 将 Workflow 序列化为 JSON 字符串（因 workflow 基于 pydantic 建模，故使用 model_dump 方法转为字典）
            await session.commit()

    async def create_message_record(
        self,
        conversation_id: str,
        message: Message,
    ) -> dict[str, Message]:
        """
        创建消息记录
        :param conversation_id: 会话唯一标识
        :param message: 消息实例
        :return: 消息实例
        """
        async with rx.asession() as session:
            record = MessageRecord(
                id=message.id,
                conversation_id=conversation_id,
                type=message.type,
                title=message.title,
                content=message.content,
                approval=approval_dump_json(message.approval),
            )
            session.add(record)
            await session.commit()
            await session.refresh(record)
        return {message.id: message}

    async def create_message_records(
        self,
        conversation_id: str,
        messages: list[Message],
    ) -> None:
        """
        创建消息记录
        :param conversation_id: 会话唯一标识
        :param messages: 消息实例列表
        :return: None
        """
        async with rx.asession() as session:
            for message in messages:
                result = await session.exec(
                    select(MessageRecord).where(MessageRecord.id == message.id)
                )
                record = result.first()
                if record:
                    record.title = message.title
                    record.content = message.content
                    record.approval = approval_dump_json(message.approval)
                else:
                    session.add(
                        MessageRecord(
                            id=message.id,
                            conversation_id=conversation_id,
                            type=message.type,
                            title=message.title,
                            content=message.content,
                            approval=approval_dump_json(message.approval),
                        )
                    )
            await session.commit()


    async def update_message_record(
        self,
        message_id: str,
        approval: Approval | None = None,
    ) -> None:
        """
        更新消息记录
        :param message_id: 消息唯一标识
        :param approval: 审批实例
        """
        async with rx.asession() as session:
            record = await session.get(MessageRecord, message_id)
            if not record:
                return
            if isinstance(approval, Approval):
                record.approval = approval_dump_json(approval)
            await session.commit()

    async def read_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(RunRecord)
                .where(RunRecord.conversation_id == conversation_id)
                .order_by(RunRecord.id)
            )
            for run in result.all():
                message_history.extend(
                    ModelMessagesTypeAdapter.validate_json(
                        run.new_messages
                    )  # 将 run.new_messages 由 JSON 字符串反序列化为 List[ModelMessage]
                )
        return message_history

    async def create_run_record(
        self,
        conversation_id: str,
        new_messages: list[ModelMessage],
    ) -> None:
        """
        创建运行记录
        :param conversation_id: 会话唯一标识
        :param new_messages: 新增消息
        :return: None
        """
        async with rx.asession() as session:
            session.add(
                RunRecord(
                    conversation_id=conversation_id,
                    new_messages=ModelMessagesTypeAdapter.dump_json(
                        new_messages
                    ).decode("utf-8"),
                    # 将 messages 由 List[ModelMessage] 序列化为 JSON 字符串
                )
            )
            await session.commit()