165 lines
5.8 KiB
Python
165 lines
5.8 KiB
Python
"""
|
||
配额检查业务逻辑。
|
||
|
||
「对话轮数」的定义：每条用户发出的消息（Segment 表的记录数）计为 1 轮。
|
||
"""
|
||
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.features.conversation.models import Conversation, Segment
|
||
from app.features.memoir.models import Chapter
|
||
from app.features.quota.schemas import QuotaCheckResponse
|
||
|
||
PLAN_QUOTAS = {
|
||
"free": {
|
||
"max_conversations": 50,
|
||
"max_chapters": 1,
|
||
"max_words": None,
|
||
},
|
||
"pro": {
|
||
"max_conversations": 2000,
|
||
"max_chapters": None,
|
||
"max_words": None,
|
||
},
|
||
"pro_plus": {
|
||
"max_conversations": 10000,
|
||
"max_chapters": None,
|
||
"max_words": None,
|
||
},
|
||
"premium": {
|
||
"max_conversations": None,
|
||
"max_chapters": None,
|
||
"max_words": None,
|
||
},
|
||
"test": {
|
||
"max_conversations": None,
|
||
"max_chapters": None,
|
||
"max_words": None,
|
||
},
|
||
}
|
||
|
||
|
||
async def get_segment_count(user_id: str, db: AsyncSession) -> int:
|
||
"""获取用户已消耗的对话轮数（= 该用户所有 Segment 记录数）。"""
|
||
stmt = (
|
||
select(func.count(Segment.id))
|
||
.join(Conversation, Segment.conversation_id == Conversation.id)
|
||
.where(Conversation.user_id == user_id)
|
||
)
|
||
result = await db.execute(stmt)
|
||
return result.scalar() or 0
|
||
|
||
|
||
async def get_chapter_count(user_id: str, db: AsyncSession) -> int:
|
||
"""获取用户当前章节数量。"""
|
||
stmt = select(func.count(Chapter.id)).where(Chapter.user_id == user_id)
|
||
result = await db.execute(stmt)
|
||
return result.scalar() or 0
|
||
|
||
|
||
async def get_conversation_count(user_id: str, db: AsyncSession) -> int:
|
||
"""别名：实际按 Segment 计数。"""
|
||
return await get_segment_count(user_id, db)
|
||
|
||
|
||
def check_can_send_message(
|
||
subscription_type: str,
|
||
segment_count: int,
|
||
) -> tuple[bool, str]:
|
||
"""检查用户是否还能发送消息（对话轮数）。返回 (是否允许, 提示信息)。"""
|
||
quotas = PLAN_QUOTAS.get(subscription_type, PLAN_QUOTAS["free"])
|
||
max_conv = quotas.get("max_conversations")
|
||
if max_conv is None:
|
||
return True, ""
|
||
if segment_count >= max_conv:
|
||
return (
|
||
False,
|
||
f"对话轮数已用完（{segment_count}/{max_conv}），请升级 Pro 或 Pro+ 继续使用",
|
||
)
|
||
return True, ""
|
||
|
||
|
||
def check_can_create_conversation(
|
||
subscription_type: str,
|
||
conversation_count: int,
|
||
) -> tuple[bool, str]:
|
||
"""兼容旧调用。"""
|
||
return check_can_send_message(subscription_type, conversation_count)
|
||
|
||
|
||
def check_can_submit_organize(
|
||
subscription_type: str,
|
||
chapter_count: int,
|
||
) -> tuple[bool, str]:
|
||
"""检查是否可以提交整理任务（生成新章节）。免费版仅允许 1 个章节。"""
|
||
quotas = PLAN_QUOTAS.get(subscription_type, PLAN_QUOTAS["free"])
|
||
max_ch = quotas.get("max_chapters")
|
||
if max_ch is None:
|
||
return True, ""
|
||
if chapter_count >= max_ch:
|
||
return False, "章节数量已达上限（免费版仅支持 1 个章节整理），请升级后继续"
|
||
return True, ""
|
||
|
||
|
||
class QuotaService:
|
||
def __init__(self, db: AsyncSession):
|
||
self._db = db
|
||
|
||
async def get_usage(self, user_id: str) -> tuple[int, int]:
|
||
"""Return (segment_count, chapter_count)."""
|
||
seg = await get_segment_count(user_id, self._db)
|
||
ch = await get_chapter_count(user_id, self._db)
|
||
return seg, ch
|
||
|
||
async def check_can_send_message(
|
||
self, user_id: str, subscription_type: str
|
||
) -> tuple[bool, str]:
|
||
"""检查用户是否还能发送消息（对话轮数）。返回 (是否允许, 提示信息)。"""
|
||
count = await get_segment_count(user_id, self._db)
|
||
return check_can_send_message(subscription_type, count)
|
||
|
||
async def check_can_submit_organize(
|
||
self, user_id: str, subscription_type: str
|
||
) -> tuple[bool, str]:
|
||
"""检查是否可以提交整理任务（生成新章节）。返回 (是否允许, 提示信息)。"""
|
||
chapter_count = await get_chapter_count(user_id, self._db)
|
||
return check_can_submit_organize(subscription_type, chapter_count)
|
||
|
||
async def check(self, user_id: str, subscription_type: str) -> QuotaCheckResponse:
|
||
"""检查用户配额使用情况。"""
|
||
quotas = PLAN_QUOTAS.get(subscription_type, PLAN_QUOTAS["free"])
|
||
segment_count, chapter_count = await self.get_usage(user_id)
|
||
max_conversations = quotas.get("max_conversations")
|
||
max_chapters = quotas.get("max_chapters")
|
||
max_words = quotas.get("max_words")
|
||
|
||
remaining_conversations = None
|
||
remaining_chapters = None
|
||
remaining_words = None
|
||
if max_conversations is not None:
|
||
remaining_conversations = max(0, max_conversations - segment_count)
|
||
if max_chapters is not None:
|
||
remaining_chapters = max(0, max_chapters - chapter_count)
|
||
|
||
has_quota = True
|
||
message = "配额充足"
|
||
if max_conversations is not None and segment_count >= max_conversations:
|
||
has_quota = False
|
||
message = f"对话轮数已用完（{segment_count}/{max_conversations}），请升级 Pro 或 Pro+ 继续使用"
|
||
elif max_chapters is not None and chapter_count >= max_chapters:
|
||
has_quota = False
|
||
message = "章节数量已达上限（免费版仅支持 1 个章节整理），请升级后继续"
|
||
|
||
return QuotaCheckResponse(
|
||
has_quota=has_quota,
|
||
remaining_conversations=remaining_conversations,
|
||
remaining_chapters=remaining_chapters,
|
||
remaining_words=remaining_words,
|
||
used_conversations=segment_count,
|
||
used_chapters=chapter_count,
|
||
max_conversations=max_conversations,
|
||
max_chapters=max_chapters,
|
||
message=message,
|
||
)
|