""" 速率限制策略基类 定义速率限制策略的接口 """ from abc import ABC, abstractmethod from dataclasses import dataclass from datetime import datetime from typing import Any, Dict, List, Optional from ..common import BasePlugin, HealthStatus, PluginMetadata @dataclass class RateLimitResult: """ 速率限制检查结果 """ allowed: bool remaining: int reset_at: Optional[datetime] = None retry_after: Optional[int] = None message: Optional[str] = None headers: Optional[Dict[str, str]] = None def __post_init__(self): if self.headers is None: self.headers = {} if self.remaining is not None: self.headers["X-RateLimit-Remaining"] = str(self.remaining) if self.reset_at: self.headers["X-RateLimit-Reset"] = str(int(self.reset_at.timestamp())) if self.retry_after: self.headers["Retry-After"] = str(self.retry_after) class RateLimitStrategy(BasePlugin): """ 速率限制策略基类 所有速率限制策略必须继承此类 """ def __init__( self, name: str, priority: int = 0, version: str = "1.0.0", author: str = "Unknown", description: str = "", api_version: str = "1.0", dependencies: List[str] = None, provides: List[str] = None, config: Dict[str, Any] = None, ): """ 初始化速率限制策略 Args: name: 策略名称 priority: 优先级(数字越大优先级越高) version: 插件版本 author: 插件作者 description: 插件描述 api_version: API版本 dependencies: 依赖的其他插件 provides: 提供的服务 config: 配置字典 """ super().__init__( name=name, priority=priority, version=version, author=author, description=description, api_version=api_version, dependencies=dependencies, provides=provides, config=config, ) @abstractmethod async def check_limit(self, key: str, **kwargs) -> RateLimitResult: """ 检查速率限制 Args: key: 限制键(如用户ID、API Key ID等) **kwargs: 额外参数 Returns: 速率限制检查结果 """ pass @abstractmethod async def consume(self, key: str, amount: int = 1, **kwargs) -> bool: """ 消费配额 Args: key: 限制键 amount: 消费数量 **kwargs: 额外参数 Returns: 是否成功消费 """ pass @abstractmethod async def reset(self, key: str): """ 重置限制 Args: key: 限制键 """ pass @abstractmethod async def get_stats(self, key: str) -> Dict[str, Any]: """ 获取统计信息 Args: key: 限制键 Returns: 统计信息字典 """ pass