"""Quota system with registry pattern. Three subscription tiers with different limits: - free: 2GB storage, 5 videos/month, 3 concurrent, 3 templates, 50 titles, 10 voiceovers, no AI voice - basic: 20GB storage, 30 videos/month, 10 concurrent, 15 templates, 500 titles, 100 voiceovers, AI voice - premium: 100GB storage, 100 videos/month, 20 concurrent, unlimited templates, 500 titles, 100 voiceovers, AI voice Quota dimensions are registered by modules via the ModuleRegistry, and checked against the user's subscription plan. """ from __future__ import annotations from dataclasses import dataclass, field from enum import Enum from typing import Dict, List, Optional class QuotaDimension(str, Enum): """配额维度 - 所有可量化的资源限制""" STORAGE_GB = "storage_gb" # 存储空间 (GB) VIDEOS_PER_MONTH = "videos_per_month" # 每月生成视频数 MAX_CONCURRENT = "max_concurrent" # 最大并发任务数 MAX_TEMPLATES = "max_templates" # 最大模板数 MAX_TITLES = "max_titles" # 最大标题库条目数 MAX_VOICEOVERS = "max_voiceovers" # 最大配音库条目数 AI_VOICE_ENABLED = "ai_voice_enabled" # AI 配音是否可用 (0/1) # 以下维度由扩展模块注册，初始配额为 0（由模块注册时填充） AI_VOICE_CREDITS = "ai_voice_credits" # AI 配音积分（每月） BATCH_EXPORT_ENABLED = "batch_export_enabled" # 批量导出 MULTI_PLATFORM_ENABLED = "multi_platform_enabled" # 多平台发布 DEDUP_REPORT_ENABLED = "dedup_report_enabled" # 去重检测报告 @dataclass class QuotaTier: """一个套餐等级的配额定义""" name: str limits: Dict[str, float] = field(default_factory=dict) def get_limit(self, dimension: str) -> float: """获取指定维度的配额限制，未定义返回 0""" return self.limits.get(dimension, 0) def is_unlimited(self, dimension: str) -> bool: """指定维度是否不限量""" return self.limits.get(dimension, float("inf")) == float("inf") # 三个套餐等级的配额定义 QUOTA_TIERS: Dict[str, QuotaTier] = { "free": QuotaTier( name="free", limits={ QuotaDimension.STORAGE_GB: 2, QuotaDimension.VIDEOS_PER_MONTH: 5, QuotaDimension.MAX_CONCURRENT: 3, QuotaDimension.MAX_TEMPLATES: 3, QuotaDimension.MAX_TITLES: 50, QuotaDimension.MAX_VOICEOVERS: 10, QuotaDimension.AI_VOICE_ENABLED: 0, QuotaDimension.AI_VOICE_CREDITS: 0, QuotaDimension.BATCH_EXPORT_ENABLED: 0, QuotaDimension.MULTI_PLATFORM_ENABLED: 0, QuotaDimension.DEDUP_REPORT_ENABLED: 0, }, ), "basic": QuotaTier( name="basic", limits={ QuotaDimension.STORAGE_GB: 20, QuotaDimension.VIDEOS_PER_MONTH: 30, QuotaDimension.MAX_CONCURRENT: 10, QuotaDimension.MAX_TEMPLATES: 15, QuotaDimension.MAX_TITLES: 500, QuotaDimension.MAX_VOICEOVERS: 100, QuotaDimension.AI_VOICE_ENABLED: 1, QuotaDimension.AI_VOICE_CREDITS: 100, QuotaDimension.BATCH_EXPORT_ENABLED: 1, QuotaDimension.MULTI_PLATFORM_ENABLED: 0, QuotaDimension.DEDUP_REPORT_ENABLED: 0, }, ), "premium": QuotaTier( name="premium", limits={ QuotaDimension.STORAGE_GB: 100, QuotaDimension.VIDEOS_PER_MONTH: 100, QuotaDimension.MAX_CONCURRENT: 20, QuotaDimension.MAX_TEMPLATES: float("inf"), # 不限量 QuotaDimension.MAX_TITLES: 500, QuotaDimension.MAX_VOICEOVERS: 100, QuotaDimension.AI_VOICE_ENABLED: 1, QuotaDimension.AI_VOICE_CREDITS: 500, QuotaDimension.BATCH_EXPORT_ENABLED: 1, QuotaDimension.MULTI_PLATFORM_ENABLED: 1, QuotaDimension.DEDUP_REPORT_ENABLED: 1, }, ), } class QuotaWarningLevel: """配额告警级别""" NORMAL = "normal" # 使用量 < 80% WARNING = "warning" # 80% <= 使用量 < 100% CRITICAL = "critical" # 95% <= 使用量 < 100% EXCEEDED = "exceeded" # 使用量 >= 100% @dataclass class QuotaCheckResult: """配额检查结果""" allowed: bool dimension: str limit: float used: float remaining: float warning_level: str @property def usage_percent(self) -> float: if self.limit <= 0: return 100.0 if self.used > 0 else 0.0 if self.limit == float("inf"): return 0.0 return min(100.0, (self.used / self.limit) * 100) class QuotaRegistry: """配额注册表 管理所有配额维度的定义和套餐限制。 扩展模块可通过 register_dimension() 注册新的配额维度。 """ def __init__(self) -> None: self._dimensions: Dict[str, str] = {} # dimension_name -> description self._tiers: Dict[str, QuotaTier] = dict(QUOTA_TIERS) # 初始化内置维度 for dim in QuotaDimension: self._dimensions[dim.value] = dim.name def register_dimension( self, dimension: str, description: str, default_limits: Optional[Dict[str, float]] = None, ) -> None: """注册新的配额维度 Args: dimension: 维度名称 description: 人类可读描述 default_limits: 各套餐的默认限制 {plan_name: limit} """ if dimension in self._dimensions: return # 幂等 self._dimensions[dimension] = description # 为各套餐设置默认限制 if default_limits: for plan_name, limit in default_limits.items(): if plan_name in self._tiers: self._tiers[plan_name].limits[dimension] = limit else: # 默认所有套餐该维度为 0 for tier in self._tiers.values(): tier.limits.setdefault(dimension, 0) def get_tier(self, plan_name: str) -> Optional[QuotaTier]: """获取套餐配额定义""" return self._tiers.get(plan_name) def get_limit(self, plan_name: str, dimension: str) -> float: """获取指定套餐在指定维度的限制""" tier = self._tiers.get(plan_name) if tier is None: return 0 return tier.get_limit(dimension) def list_dimensions(self) -> Dict[str, str]: """列出所有已注册的配额维度""" return dict(self._dimensions) def list_tiers(self) -> List[str]: """列出所有套餐等级""" return list(self._tiers.keys()) class QuotaChecker: """配额检查器 检查用户在指定维度的使用量是否超出限制。 需要调用方提供当前使用量（used），本类不直接访问数据库。 """ def __init__(self, registry: Optional[QuotaRegistry] = None) -> None: self._registry = registry or QuotaRegistry() def check( self, plan_name: str, dimension: str, used: float, ) -> QuotaCheckResult: """检查指定维度的配额使用情况 Args: plan_name: 用户套餐等级 (free/basic/premium) dimension: 配额维度 used: 当前已使用量 Returns: QuotaCheckResult """ limit = self._registry.get_limit(plan_name, dimension) remaining = max(0, limit - used) if limit != float("inf") else float("inf") allowed = used < limit if limit != float("inf") else True warning_level = self._compute_warning_level(used, limit) return QuotaCheckResult( allowed=allowed, dimension=dimension, limit=limit, used=used, remaining=remaining, warning_level=warning_level, ) def check_multiple( self, plan_name: str, usage: Dict[str, float], ) -> List[QuotaCheckResult]: """批量检查多个维度的配额""" return [self.check(plan_name, dim, used) for dim, used in usage.items()] @staticmethod def _compute_warning_level(used: float, limit: float) -> str: """计算告警级别""" if limit <= 0: return QuotaWarningLevel.EXCEEDED if used > 0 else QuotaWarningLevel.NORMAL if limit == float("inf"): return QuotaWarningLevel.NORMAL ratio = used / limit if ratio >= 1.0: return QuotaWarningLevel.EXCEEDED if ratio >= 0.95: return QuotaWarningLevel.CRITICAL if ratio >= 0.80: return QuotaWarningLevel.WARNING return QuotaWarningLevel.NORMAL def get_warning_level(used: float, limit: float) -> str: """便捷函数：计算配额告警级别""" return QuotaChecker._compute_warning_level(used, limit) # 全局单例 quota_registry = QuotaRegistry() quota_checker = QuotaChecker(quota_registry)