临时存个代码

from datetime import datetime
from pathlib import Path
from typing import Literal, Optional

from sqlalchemy import (
    CheckConstraint,
    DateTime,
    ForeignKey,
    String,
    Text,
    create_engine,
    event,
    func,
    select,
)
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    Session,
    mapped_column,
    relationship,
)

# ============================================================
#  数据库配置
# ============================================================

BASE_DIR = Path(__file__).resolve().parent
DB_PATH = BASE_DIR / "chat.db"
DATABASE_URL = f"sqlite:///{DB_PATH}"

engine = create_engine(
    DATABASE_URL,
    echo=False,
    connect_args={"check_same_thread": False},
)

# 数据库事件监听器
# 每次 SQLAlchemy 新建一个数据库连接时，都会自动执行下该函数
# 这个为什么用cursor游标，因为是操作底层的数据库连接，而不是 ORM 层的会话
# 平时curd操作，都是在会话层进行的，而不是直接操作数据库连接
@event.listens_for(engine, "connect")
def enable_sqlite_foreign_keys(dbapi_connection, connection_record):
    """
    SQLite 默认不启用外键约束。

    如果不打开 PRAGMA foreign_keys=ON，
    ForeignKey(..., ondelete="CASCADE") 在 SQLite 中可能不会真正生效。
    """
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()


class Base(DeclarativeBase):
    """SQLAlchemy ORM 基类。"""

    pass


# ============================================================
#  ORM 模型
# ============================================================

class ConversationORM(Base):
    """
    会话表 ORM 模型。

    conversation_id 直接作为主键。
    每条记录代表一个独立的聊天会话。
    """

    __tablename__ = "conversations"

    # 会话唯一标识，直接作为主键
    conversation_id: Mapped[str] = mapped_column(
        String(100),
        primary_key=True,
        nullable=False,
    )

    # 会话标题
    title: Mapped[Optional[str]] = mapped_column(
        String(200),
        nullable=True,
    )

    # 创建时间
    created_at: Mapped[datetime] = mapped_column(
        DateTime,
        server_default=func.now(),
        nullable=False,
    )

    # 更新时间
    updated_at: Mapped[datetime] = mapped_column(
        DateTime,
        server_default=func.now(),
        onupdate=func.now(),
        nullable=False,
    )

    # 一对多：一个会话拥有多条消息
    messages: Mapped[list["ChatMessageORM"]] = relationship(
        back_populates="conversation",
        # cascade="all"	开启所有级联操作：添加、删除、合并、刷新
        # delete-orphan	删除所有子项（消息）时，也删除当前会话
        cascade="all, delete-orphan",
        # 让数据库自己处理级联删除，SQLAlchemy 不管
        passive_deletes=True,
        order_by="ChatMessageORM.id",
    )


class ChatMessageORM(Base):
    """
    聊天消息表 ORM 模型。

    每条消息属于一个会话。
    """

    __tablename__ = "chat_messages"

    __table_args__ = (
        CheckConstraint(
            "role IN ('system', 'user', 'assistant')",
            name="check_chat_message_role",
        ),
    )

    # 消息自增主键
    # 注意：这个 id 是消息 id，不是会话 id
    id: Mapped[int] = mapped_column(
        primary_key=True,
        autoincrement=True,
    )

    # 外键：关联 conversations.conversation_id
    conversation_id: Mapped[str] = mapped_column(
        String(100),
        ForeignKey("conversations.conversation_id", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )

    # 消息角色
    role: Mapped[str] = mapped_column(
        String(20),
        nullable=False,
    )

    # 消息内容
    content: Mapped[str] = mapped_column(
        Text,
        nullable=False,
    )

    # 创建时间
    created_at: Mapped[datetime] = mapped_column(
        DateTime,
        server_default=func.now(),
        nullable=False,
    )

    # 多对一：多条消息属于一个会话
    conversation: Mapped["ConversationORM"] = relationship(
        back_populates="messages",
    )


# ============================================================
#  序列化工具
# ============================================================

def conversation_to_dict(conversation: ConversationORM) -> dict:
    """
    将 ConversationORM 转成普通 dict。

    避免 Session 关闭后返回 ORM 对象导致 DetachedInstanceError。
    """
    return {
        "conversation_id": conversation.conversation_id,
        "title": conversation.title,
        "created_at": conversation.created_at,
        "updated_at": conversation.updated_at,
    }


def message_to_dict(message: ChatMessageORM) -> dict:
    """
    将 ChatMessageORM 转成普通 dict。
    """
    return {
        "id": message.id,
        "conversation_id": message.conversation_id,
        "role": message.role,
        "content": message.content,
        "created_at": message.created_at,
    }


# ============================================================
#  数据库初始化 & Session 工具
# ============================================================

def init_db() -> None:
    """
    初始化数据库。

    如果 chat.db 不存在，会自动创建。
    如果 conversations / chat_messages 表不存在，也会自动创建。
    """
    Base.metadata.create_all(bind=engine)


def create_session() -> Session:
    """
    创建一个数据库 Session。

    用法：
        with create_session() as session:
            ...
    """
    return Session(engine)


# ============================================================
#  会话 Conversation 相关操作
# ============================================================

def create_conversation(
    conversation_id: str,
    title: Optional[str] = None,
) -> dict:
    """
    创建一个新会话。

    如果 conversation_id 已存在，会抛出 IntegrityError。
    """
    with create_session() as session:
        conversation = ConversationORM(
            conversation_id=conversation_id,
            title=title,
        )
        session.add(conversation)
        session.commit()
        session.refresh(conversation)

        return conversation_to_dict(conversation)


def get_conversation(conversation_id: str) -> Optional[dict]:
    """
    根据 conversation_id 获取会话。

    如果不存在，返回 None。
    """
    with create_session() as session:
        stmt = select(ConversationORM).where(
            ConversationORM.conversation_id == conversation_id
        )
        conversation = session.scalar(stmt)

        if conversation is None:
            return None

        return conversation_to_dict(conversation)


def get_or_create_conversation(
    conversation_id: str,
    title: Optional[str] = None,
) -> dict:
    """
    获取已有会话，不存在则自动创建。

    适合首次发消息时自动建会话的场景。
    """
    with create_session() as session:
        stmt = select(ConversationORM).where(
            ConversationORM.conversation_id == conversation_id
        )
        conversation = session.scalar(stmt)

        if conversation is None:
            conversation = ConversationORM(
                conversation_id=conversation_id,
                title=title,
            )
            session.add(conversation)
            session.commit()
            session.refresh(conversation)

        return conversation_to_dict(conversation)


def update_conversation_title(
    conversation_id: str,
    title: str,
) -> Optional[dict]:
    """
    更新会话标题。

    如果会话不存在，返回 None。
    """
    with create_session() as session:
        stmt = select(ConversationORM).where(
            ConversationORM.conversation_id == conversation_id
        )
        conversation = session.scalar(stmt)

        if conversation is None:
            return None

        conversation.title = title
        session.commit()
        session.refresh(conversation)

        return conversation_to_dict(conversation)


def list_conversations(limit: int = 50) -> list[dict]:
    """
    列出所有会话，按最近更新时间倒序。
    """
    with create_session() as session:
        stmt = (
            select(ConversationORM)
            .order_by(ConversationORM.updated_at.desc())
            .limit(limit)
        )
        conversations = session.scalars(stmt).all()

        return [
            conversation_to_dict(conversation)
            for conversation in conversations
        ]


def delete_conversation(conversation_id: str) -> bool:
    """
    删除一个会话及其所有消息。

    返回：
        True: 删除成功
        False: 会话不存在
    """
    with create_session() as session:
        stmt = select(ConversationORM).where(
            ConversationORM.conversation_id == conversation_id
        )
        conversation = session.scalar(stmt)

        if conversation is None:
            return False

        session.delete(conversation)
        session.commit()

        return True


# ============================================================
#  消息 ChatMessage 相关操作
# ============================================================

def save_message(
    conversation_id: str,
    role: Literal["system", "user", "assistant"],
    content: str,
    auto_create_conversation: bool = True,
) -> dict:
    """
    保存一条聊天消息。

    参数：
        conversation_id:
            会话 ID。

        role:
            消息角色，只允许 system / user / assistant。

        content:
            消息内容。

        auto_create_conversation:
            如果会话不存在，是否自动创建。
    """
    with create_session() as session:
        stmt = select(ConversationORM).where(
            ConversationORM.conversation_id == conversation_id
        )
        conversation = session.scalar(stmt)

        if conversation is None:
            if not auto_create_conversation:
                raise ValueError(f"会话 {conversation_id} 不存在")

            conversation = ConversationORM(
                conversation_id=conversation_id,
            )
            session.add(conversation)
            session.flush()

        message = ChatMessageORM(
            conversation_id=conversation_id,
            role=role,
            content=content,
        )
        session.add(message)

        # 让会话更新时间刷新
        conversation.updated_at = func.now()

        session.commit()
        session.refresh(message)

        return message_to_dict(message)


def get_history(
    conversation_id: str,
    limit: int = 20,
) -> list[dict]:
    """
    获取某个会话最近的聊天历史。

    返回格式兼容 OpenAI / LangChain messages：

    [
        {"role": "user", "content": "你好"},
        {"role": "assistant", "content": "你好，有什么可以帮你？"}
    ]
    """
    with create_session() as session:
        stmt = (
            select(ChatMessageORM)
            .where(ChatMessageORM.conversation_id == conversation_id)
            .order_by(ChatMessageORM.id.desc())
            .limit(limit)
        )
        messages = session.scalars(stmt).all()

        messages = list(reversed(messages))

        return [
            {
                "role": message.role,
                "content": message.content,
            }
            for message in messages
        ]


def get_messages(
    conversation_id: str,
    limit: int = 20,
) -> list[dict]:
    """
    获取某个会话最近的完整消息列表。

    和 get_history 不同，这个函数会返回 message_id、created_at 等信息。
    """
    with create_session() as session:
        stmt = (
            select(ChatMessageORM)
            .where(ChatMessageORM.conversation_id == conversation_id)
            .order_by(ChatMessageORM.id.desc())
            .limit(limit)
        )
        messages = session.scalars(stmt).all()

        messages = list(reversed(messages))

        return [
            message_to_dict(message)
            for message in messages
        ]


def clear_messages(conversation_id: str) -> int:
    """
    清空某个会话下的所有消息，但不删除会话本身。

    返回删除的消息数量。
    """
    with create_session() as session:
        stmt = select(ChatMessageORM).where(
            ChatMessageORM.conversation_id == conversation_id
        )
        messages = session.scalars(stmt).all()

        count = len(messages)

        for message in messages:
            session.delete(message)

        # 清空消息后也刷新会话 updated_at
        conversation_stmt = select(ConversationORM).where(
            ConversationORM.conversation_id == conversation_id
        )
        conversation = session.scalar(conversation_stmt)
        if conversation is not None:
            conversation.updated_at = func.now()

        session.commit()

        return count