""" 套餐额度预警服务 监控用户的API调用额度、查询次数等使用情况,在接近限制时触发预警。 功能: - 额度使用查询: 获取用户当前额度使用情况 - 使用率计算: 计算额度使用百分比 - 预警阈值检查: 80%警告、90%严重、100%限制 - 预警消息生成: 生成带建议操作的预警消息 - 额度重置: 支持额度重置功能 额度类型: - api_calls: API调用次数 - queries: 查询次数 - content_generation: 内容生成次数 - storage: 存储空间(MB) 预警阈值: - warning: 80% - 警告 - critical: 90% - 严重 - exhausted: 100% - 耗尽 """ from __future__ import annotations from dataclasses import dataclass from enum import Enum class QuotaType(str, Enum): """额度类型枚举""" API_CALLS = "api_calls" QUERIES = "queries" CONTENT_GENERATION = "content_generation" STORAGE = "storage" QUOTA_LIMITS = { "free": { "api_calls": 1000, "queries": 50, "content_generation": 10, "storage": 100, }, "basic": { "api_calls": 10000, "queries": 500, "content_generation": 100, "storage": 1000, }, "pro": { "api_calls": 100000, "queries": 5000, "content_generation": 1000, "storage": 10000, }, "unlimited": { "api_calls": -1, "queries": -1, "content_generation": -1, "storage": -1, }, } WARNING_THRESHOLDS = { "warning": 0.80, "critical": 0.90, "exhausted": 1.00, } WARNING_MESSAGES = { "warning": "{quota_type}额度已使用{percentage:.0f}%", "critical": "{quota_type}额度已使用{percentage:.0f}%,即将耗尽", "exhausted": "{quota_type}额度已使用{percentage:.0f}%,已耗尽", } RECOMMENDED_ACTIONS = { "warning": "请关注使用情况,考虑升级套餐", "critical": "额度即将耗尽,建议立即升级套餐", "exhausted": "额度已耗尽,请升级套餐或等待重置", } @dataclass class QuotaUsage: """额度使用数据结构 Attributes: quota_type: 额度类型 used: 已使用量 limit: 额度限制(-1表示无限) usage_percentage: 使用百分比(0-100,无限时为0) status: 状态(ok/warning/critical/exhausted/unlimited) remaining: 剩余额度(-1表示无限) """ quota_type: str used: int limit: int usage_percentage: float status: str remaining: int @dataclass class QuotaWarning: """预警数据结构 Attributes: quota_type: 额度类型 status: 预警状态 usage_percentage: 使用百分比 message: 预警消息 recommended_action: 建议操作 """ quota_type: str status: str usage_percentage: float message: str recommended_action: str class QuotaService: """套餐额度预警服务 提供额度查询、使用率计算、预警检查等功能。 """ def calculate_usage_percentage(self, used: int, limit: int) -> float: """计算额度使用百分比 Args: used: 已使用量 limit: 额度限制(-1表示无限) Returns: 使用百分比(0-100),无限额度时返回0.0 """ if limit == -1 or limit == 0: return 0.0 return (used / limit) * 100.0 def get_quota_status(self, usage_percentage: float, limit: int = 0) -> str: """根据使用率获取额度状态 Args: usage_percentage: 使用百分比 limit: 额度限制(-1表示无限) Returns: 状态字符串: ok/warning/critical/exhausted/unlimited """ if limit == -1: return "unlimited" if usage_percentage >= 100.0: return "exhausted" if usage_percentage >= 90.0: return "critical" if usage_percentage >= 80.0: return "warning" return "ok" def get_quota_limit(self, plan: str, quota_type: QuotaType) -> int: """获取指定套餐的额度限制 Args: plan: 套餐名称(free/basic/pro/unlimited) quota_type: 额度类型 Returns: 额度限制值(-1表示无限) """ return QUOTA_LIMITS.get(plan, {}).get(quota_type.value, 0) def get_remaining(self, used: int, limit: int) -> int: """计算剩余额度 Args: used: 已使用量 limit: 额度限制(-1表示无限) Returns: 剩余额度(-1表示无限) """ if limit == -1: return -1 return limit - used def _format_warning_message(self, quota_type: str, status: str, percentage: float) -> str: """格式化预警消息 Args: quota_type: 额度类型 status: 预警状态 percentage: 使用百分比 Returns: 格式化后的预警消息 """ template = WARNING_MESSAGES.get(status, "{quota_type}额度使用情况") return template.format(quota_type=quota_type, percentage=percentage) def _get_recommended_action(self, status: str) -> str: """获取建议操作 Args: status: 预警状态 Returns: 建议操作描述 """ return RECOMMENDED_ACTIONS.get(status, "请关注使用情况") def generate_warning( self, quota_type: QuotaType, status: str, usage_percentage: float, ) -> QuotaWarning: """生成预警信息 Args: quota_type: 额度类型 status: 预警状态 usage_percentage: 使用百分比 Returns: QuotaWarning预警数据对象 """ return QuotaWarning( quota_type=quota_type.value, status=status, usage_percentage=usage_percentage, message=self._format_warning_message(quota_type.value, status, usage_percentage), recommended_action=self._get_recommended_action(status), ) def check_quota( self, plan: str, quota_type: QuotaType, used: int, ) -> QuotaUsage: """检查指定套餐的额度使用情况 Args: plan: 套餐名称 quota_type: 额度类型 used: 已使用量 Returns: QuotaUsage额度使用数据对象 """ limit = self.get_quota_limit(plan, quota_type) percentage = self.calculate_usage_percentage(used, limit) status = self.get_quota_status(percentage, limit) remaining = self.get_remaining(used, limit) return QuotaUsage( quota_type=quota_type.value, used=used, limit=limit, usage_percentage=percentage, status=status, remaining=remaining, ) def reset_quota(self, used: int, reset_to: int = 0) -> int: """重置额度使用量 Args: used: 当前使用量(未使用,仅为接口兼容) reset_to: 重置到的值,默认为0 Returns: 重置后的使用量 """ return reset_to def get_all_quota_usage( self, plan: str, api_calls_used: int = 0, queries_used: int = 0, content_generation_used: int = 0, storage_used: int = 0, ) -> list[QuotaUsage]: """获取所有额度类型的使用情况 Args: plan: 套餐名称 api_calls_used: API调用已使用量 queries_used: 查询已使用量 content_generation_used: 内容生成已使用量 storage_used: 存储已使用量(MB) Returns: 所有额度类型的使用情况列表 """ return [ self.check_quota(plan, QuotaType.API_CALLS, api_calls_used), self.check_quota(plan, QuotaType.QUERIES, queries_used), self.check_quota(plan, QuotaType.CONTENT_GENERATION, content_generation_used), self.check_quota(plan, QuotaType.STORAGE, storage_used), ] def get_warnings(self, usage_list: list[QuotaUsage]) -> list[QuotaWarning]: """从使用情况列表中生成预警 Args: usage_list: 额度使用情况列表 Returns: 预警信息列表(仅包含需要预警的项) """ warnings = [] for usage in usage_list: if usage.status in ["warning", "critical", "exhausted"]: warning = self.generate_warning( quota_type=QuotaType(usage.quota_type), status=usage.status, usage_percentage=usage.usage_percentage, ) warnings.append(warning) return warnings