""" 配额检查业务逻辑。 「对话轮数」的定义:每条用户发出的消息(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, Conversation.deleted_at.is_(None), ) ) 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_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, )