445 lines
15 KiB
Python
445 lines
15 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据库状态
|
||
|
||
更新数据库方法：
|
||
1. reflex db makemigrations
|
||
2. reflex db migrate
|
||
"""
|
||
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,
|
||
Message,
|
||
MessageType,
|
||
Task,
|
||
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 TaskRecord(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="任务描述")
|
||
last_run_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="运行唯一标识",
|
||
)
|
||
task_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="消息唯一标识",
|
||
)
|
||
task_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 captcha_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_task_record(
|
||
self, user_id: str, user_prompt: str = ""
|
||
) -> dict[str, Task]:
|
||
"""
|
||
创建任务记录
|
||
:param user_id: 用户唯一标识
|
||
:param user_prompt: 用户提示词，默认为空
|
||
:return: 任务实例
|
||
"""
|
||
async with rx.asession() as session:
|
||
record = TaskRecord(user_id=user_id)
|
||
session.add(record)
|
||
await session.commit()
|
||
await session.refresh(record)
|
||
return {
|
||
record.id: Task(
|
||
id=record.id,
|
||
description=record.description,
|
||
user_prompt=user_prompt,
|
||
usage=usage_validate_json(
|
||
record.usage
|
||
), # 将 record.usage 由 JSON 字符串反序列化为 RunUsage
|
||
last_run_at=record.last_run_at,
|
||
)
|
||
}
|
||
|
||
async def read_tasks(self, user_id: str) -> dict[str, Task]:
|
||
"""
|
||
查询当前用户的任务字典
|
||
:param user_id: 用户唯一标识
|
||
:return: 当前用户的任务字典
|
||
"""
|
||
tasks: dict[str, Task] = {}
|
||
async with rx.asession() as session:
|
||
result = await session.exec(
|
||
select(TaskRecord, MessageRecord)
|
||
.outerjoin(MessageRecord, MessageRecord.task_id == TaskRecord.id) # type: ignore
|
||
.where(
|
||
TaskRecord.user_id == user_id,
|
||
TaskRecord.is_deleted == False,
|
||
)
|
||
.order_by(TaskRecord.id, MessageRecord.id)
|
||
)
|
||
for (
|
||
task_record,
|
||
message_record,
|
||
) in result.all():
|
||
task = tasks.setdefault(
|
||
task_record.id,
|
||
Task(
|
||
id=task_record.id,
|
||
description=task_record.description,
|
||
last_run_at=task_record.last_run_at,
|
||
usage=usage_validate_json(
|
||
task_record.usage
|
||
), # 将 task_record.usage 由 JSON 字符串反序列化为 RunUsage
|
||
workflow=workflow_validate_json(
|
||
task_record.workflow
|
||
), # 将 task_record.workflow 由 JSON 字符串反序列化为 Workflow,
|
||
),
|
||
)
|
||
if not message_record:
|
||
continue
|
||
task.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 tasks
|
||
|
||
async def update_task_record(
|
||
self,
|
||
id: str,
|
||
description: str | None = None,
|
||
last_run_at: datetime | None = None,
|
||
is_deleted: bool | None = None,
|
||
usage: RunUsage | None = None,
|
||
workflow: Workflow | None = None,
|
||
) -> None:
|
||
"""
|
||
更新任务记录
|
||
:param id: 任务唯一标识
|
||
:param description: 任务描述
|
||
:param last_run_at: 上一次运行日期时间
|
||
:param is_deleted: 任务已删除
|
||
:param usage: 使用量
|
||
:param workflow: 工作流
|
||
:return: None
|
||
"""
|
||
async with rx.asession() as session:
|
||
record = await session.get(TaskRecord, id)
|
||
if not record:
|
||
return
|
||
if isinstance(description, str):
|
||
record.description = description
|
||
if isinstance(last_run_at, datetime):
|
||
record.last_run_at = last_run_at
|
||
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,
|
||
task_id: str,
|
||
message: Message,
|
||
) -> dict[str, Message]:
|
||
"""
|
||
创建消息记录
|
||
:param task_id: 任务唯一标识
|
||
:param message: 消息实例
|
||
:return: 消息实例
|
||
"""
|
||
async with rx.asession() as session:
|
||
record = MessageRecord(
|
||
id=message.id,
|
||
task_id=task_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,
|
||
task_id: str,
|
||
messages: list[Message],
|
||
) -> None:
|
||
"""
|
||
创建消息记录
|
||
:param task_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,
|
||
task_id=task_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, task_id: str) -> list[ModelMessage]:
|
||
"""
|
||
查询消息历史列表
|
||
:param task_id: 任务唯一标识
|
||
:return: 消息历史
|
||
"""
|
||
message_history: list[ModelMessage] = []
|
||
async with rx.asession() as session:
|
||
result = await session.exec(
|
||
select(RunRecord)
|
||
.where(RunRecord.task_id == task_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,
|
||
task_id: str,
|
||
new_messages: list[ModelMessage],
|
||
) -> None:
|
||
"""
|
||
创建运行记录
|
||
:param task_id: 任务唯一标识
|
||
:param new_messages: 新增消息
|
||
:return: None
|
||
"""
|
||
async with rx.asession() as session:
|
||
session.add(
|
||
RunRecord(
|
||
task_id=task_id,
|
||
new_messages=ModelMessagesTypeAdapter.dump_json(
|
||
new_messages
|
||
).decode("utf-8"),
|
||
# 将 messages 由 List[ModelMessage] 序列化为 JSON 字符串
|
||
)
|
||
)
|
||
await session.commit()
|