1"""

2 AgentOS v0.60 Rate Limiter — 流量控制。

3 Token Bucket + Sliding Window + Concurrency Limiter + 多级配额。

4"""

5

6 from __future__ import annotations

7

8 import asyncio

9 import time

10 import math

11 from dataclasses import dataclass, field

12 from enum import Enum

13 from typing import Optional

14

15

16 class RateLimitStrategy(str, Enum):

17

18 """限流策略枚举。"""

19

20 TOKEN_BUCKET = "token_bucket"

21 SLIDING_WINDOW = "sliding_window"

22 FIXED_WINDOW = "fixed_window"

23

24

25@dataclass

26 class RateLimitConfig:

27 """限流配置。"""

28 strategy: RateLimitStrategy = RateLimitStrategy.TOKEN_BUCKET

29 max_requests: int = 60 # 每单位时间的最大请求数

30 per_seconds: float = 60.0 # 时间窗口（秒）

31 burst_size: int = 10 # 突发容量（token bucket 专用）

32 max_concurrent: int = 5 # 最大并发数

33 queue_timeout: float = 30.0 # 排队超时

34 retry_after_header: bool = True # 是否在拒绝时返回 Retry-After

35

36

37@dataclass

38 class RateLimitResult:

39 """限流检查结果。"""

40 allowed: bool

41 remaining: int = 0

42 reset_at: float = 0.0

43 retry_after: float = 0.0

44 limit: int = 0

45 reason: str = ""

46

47

48 class TokenBucket:

49 """令牌桶算法实现。"""

50

51 def __init__(self, rate: float, capacity: int):

52 self.rate = rate # 令牌填充速率（个/秒）

53 self.capacity = capacity # 桶容量（最大突发）

54 self.tokens = float(capacity)

55 self.last_refill = time.monotonic()

56 self._lock = asyncio.Lock()

57

58 async def consume(self, tokens: int = 1) -> bool:

59 async with self._lock:

60 self._refill()

61 if self.tokens >= tokens:

62 self.tokens -= tokens

63 return True

64 return False

65

66 def _refill(self):

67 now = time.monotonic()

68 elapsed = now - self.last_refill

69 self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)

70 self.last_refill = now

71

72 @property

73 def available(self) -> float:

74 self._refill()

75 return self.tokens

76

77

78 class SlidingWindow:

79 """滑动窗口计数器。"""

80

81 def __init__(self, max_requests: int, window_seconds: float):

82 self.max_requests = max_requests

83 self.window = window_seconds

84 self._timestamps: list[float] = []

85 self._lock = asyncio.Lock()

86

87 async def allow(self) -> bool:

88 async with self._lock:

89 now = time.monotonic()

90 cutoff = now - self.window

91 self._timestamps = [t for t in self._timestamps if t > cutoff]

92 if len(self._timestamps) < self.max_requests:

93 self._timestamps.append(now)

94 return True

95 return False

96

97 @property

98 def current_count(self) -> int:

99 cutoff = time.monotonic() - self.window

100 return sum(1 for t in self._timestamps if t > cutoff)

101

102

103 class ConcurrencyLimiter:

104 """并发请求限制器。"""

105

106 def __init__(self, max_concurrent: int):

107 self._semaphore = asyncio.Semaphore(max_concurrent)

108 self.max_concurrent = max_concurrent

109

110 async def acquire(self) -> bool:

111 return await self._semaphore.acquire()

112

113 def release(self):

114 self._semaphore.release()

115

116 @property

117 def available(self) -> int:

118 return self._semaphore._value

119

120

121 class RateLimiter:

122 """组合限流器：Token Bucket + Concurrency Limiter + 多级配额。"""

123

124 def __init__(self, config: RateLimitConfig | None = None):

125 cfg = config or RateLimitConfig()

126 self.config = cfg

127 self._bucket = TokenBucket(

128 rate=cfg.max_requests / cfg.per_seconds,

129 capacity=cfg.burst_size or cfg.max_requests

130 )

131 self._window = SlidingWindow(cfg.max_requests, cfg.per_seconds)

132 self._concurrency = ConcurrencyLimiter(cfg.max_concurrent)

133

134 async def acquire(self, weight: int = 1) -> RateLimitResult:

135 """尝试获取请求配额。先检查并发，再检查速率。"""

136 # 1. 并发检查

137 if not self._concurrency._semaphore.locked():

138 pass # 还有并发槽位

139

140 # 2. 速率检查

141 if self.config.strategy == RateLimitStrategy.TOKEN_BUCKET:

142 if await self._bucket.consume(weight):

143 return RateLimitResult(

144 allowed=True,

145 remaining=max(0, int(self._bucket.available)),

146 limit=self.config.max_requests,

147 )

148 wait = (weight - self._bucket.available) / self._bucket.rate

149 return RateLimitResult(

150 allowed=False,

151 remaining=0,

152 retry_after=wait,

153 limit=self.config.max_requests,

154 reason="rate_limit_exceeded",

155 )

156

157 elif self.config.strategy == RateLimitStrategy.SLIDING_WINDOW:

158 if await self._window.allow():

159 return RateLimitResult(

160 allowed=True,

161 remaining=self.config.max_requests - self._window.current_count,

162 limit=self.config.max_requests,

163 )

164 return RateLimitResult(

165 allowed=False, remaining=0,

166 retry_after=self.config.per_seconds,

167 limit=self.config.max_requests,

168 reason="window_exceeded",

169 )

170

171 # fixed window fallback

172 return RateLimitResult(allowed=True, limit=self.config.max_requests)

173

174 async def release(self):

175 self._concurrency.release()

176

177 def model_quota(self, model: str) -> RateLimitConfig:

178 """返回特定模型的配额配置。"""

179 quotas = {

180 "gpt-4o": RateLimitConfig(max_requests=50, per_seconds=60, burst_size=5),

181 "gpt-4o-mini": RateLimitConfig(max_requests=200, per_seconds=60, burst_size=20),

182 "claude-sonnet-4": RateLimitConfig(max_requests=40, per_seconds=60, burst_size=5),

183 "deepseek-v3.1": RateLimitConfig(max_requests=100, per_seconds=60, burst_size=15),

184 }

185 return quotas.get(model, self.config)

186

187

188 class QuotaManager:

189 """多租户配额管理。"""

190

191 def __init__(self):

192 self._limiters: dict[str, RateLimiter] = {}

193

194 def get(self, key: str, config: RateLimitConfig | None = None) -> RateLimiter:

195 if key not in self._limiters:

196 self._limiters[key] = RateLimiter(config)

197 return self._limiters[key]

198

199 def add_quota(self, key: str, config: RateLimitConfig):

200 self._limiters[key] = RateLimiter(config)

201

202 def clear_expired(self, ttl: float = 3600):

203 """清除超过TTL未使用的限流器（预留接口）。"""

204 pass