53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
数据库状态
|
|
"""
|
|
|
|
from typing import List
|
|
from sqlmodel import select as sql_select
|
|
import reflex as rx
|
|
import sqlalchemy as sa
|
|
|
|
from application.models import (
|
|
MessageHistory,
|
|
ModelMessage,
|
|
ModelMessagesTypeAdapter,
|
|
)
|
|
|
|
|
|
class DatabaseState(rx.State):
|
|
"""数据库状态"""
|
|
|
|
async def get_message_history(self, chat_id: str) -> List[ModelMessage]:
|
|
"""
|
|
获取聊天消息历史
|
|
:param chat_id: 聊天唯一标识
|
|
:return: 聊天消息历史
|
|
"""
|
|
message_history = []
|
|
async with rx.asession() as session:
|
|
records = await session.exec(
|
|
sql_select(MessageHistory)
|
|
.where(MessageHistory.chat_id == chat_id)
|
|
.order_by(sa.desc("create_at"))
|
|
)
|
|
for record in records.all():
|
|
message_history.extend(
|
|
ModelMessagesTypeAdapter.validate_json(record.new_message)
|
|
)
|
|
return message_history
|
|
|
|
async def save_new_message(
|
|
self, chat_id: str, new_message: List[ModelMessage]
|
|
) -> None:
|
|
"""
|
|
保存一轮对话消息
|
|
:param chat_id: 聊天唯一标识
|
|
:param new_message: 一轮对话消息
|
|
:return: None
|
|
"""
|
|
async with rx.asession() as session:
|
|
record = MessageHistory.adapt(chat_id=chat_id, new_message=new_message)
|
|
session.add(record)
|
|
await session.commit()
|