325 lines
9.0 KiB
Python
325 lines
9.0 KiB
Python
"""
|
||
套餐额度预警服务
|
||
|
||
监控用户的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
|