Merge branch 'refactor/backend-architecture' into development

api/app/features/quota/service.py

@@ -0,0 +1,160 @@
+"""
+配额检查业务逻辑。
+
+「对话轮数」的定义：每条用户发出的消息（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,
+        )