""" 配额检查 API 路由 """ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from typing import Optional from middleware.auth import get_current_user from database.models import User from sqlalchemy.ext.asyncio import AsyncSession from database import get_async_db from sqlalchemy import select, func router = APIRouter(prefix="/api/quota", tags=["quota"]) class QuotaCheckResponse(BaseModel): """配额检查响应""" has_quota: bool # 是否有配额 remaining_conversations: Optional[int] = None # 剩余对话次数 remaining_chapters: Optional[int] = None # 剩余章节数 remaining_words: Optional[int] = None # 剩余字数 message: str # 提示信息 # 计划配额限制 PLAN_QUOTAS = { "free": { "max_conversations": 3, "max_chapters": 10, "max_words": 50000 }, "premium": { "max_conversations": None, # 无限制 "max_chapters": None, "max_words": None } } @router.get("/check", response_model=QuotaCheckResponse) async def check_quota( current_user: User = Depends(get_current_user) ): """ 检查用户配额使用情况 根据用户的订阅计划检查是否还有配额可以使用 """ plan_type = current_user.subscription_type quotas = PLAN_QUOTAS.get(plan_type, PLAN_QUOTAS["free"]) # 如果是高级版,无限制 if plan_type == "premium": return QuotaCheckResponse( has_quota=True, remaining_conversations=None, remaining_chapters=None, remaining_words=None, message="高级版用户,无使用限制" ) # 统计使用情况 async for db in get_async_db(): # 统计对话数量 from database.models import Conversation stmt = select(func.count(Conversation.id)).where( Conversation.user_id == current_user.id ) result = await db.execute(stmt) conversation_count = result.scalar() or 0 # 统计章节数量 from database.models import Chapter stmt = select(func.count(Chapter.id)).where( Chapter.user_id == current_user.id ) result = await db.execute(stmt) chapter_count = result.scalar() or 0 # 统计总字数 stmt = select(func.sum(func.length(Chapter.content))).where( Chapter.user_id == current_user.id ) result = await db.execute(stmt) total_words = result.scalar() or 0 # 计算剩余配额 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 - conversation_count) if max_chapters is not None: remaining_chapters = max(0, max_chapters - chapter_count) if max_words is not None: remaining_words = max(0, max_words - total_words) # 检查是否有配额 has_quota = True message = "配额充足" if max_conversations is not None and conversation_count >= max_conversations: has_quota = False message = "对话次数已用完,请升级到高级版" elif max_chapters is not None and chapter_count >= max_chapters: has_quota = False message = "章节数量已达上限,请升级到高级版" elif max_words is not None and total_words >= max_words: has_quota = False message = "字数已达上限,请升级到高级版" return QuotaCheckResponse( has_quota=has_quota, remaining_conversations=remaining_conversations, remaining_chapters=remaining_chapters, remaining_words=remaining_words, message=message )