Coverage for agentos/queue/rate_limiter.py: 46%

1""" 

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

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

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import time 

10from dataclasses import dataclass 

11from enum import Enum 

12 

13 

14class RateLimitStrategy(str, Enum): 

15 

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

17 

18 TOKEN_BUCKET = "token_bucket" 

19 SLIDING_WINDOW = "sliding_window" 

20 FIXED_WINDOW = "fixed_window" 

21 

22 

23@dataclass 

24class RateLimitConfig: 

25 """限流配置。""" 

26 strategy: RateLimitStrategy = RateLimitStrategy.TOKEN_BUCKET 

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

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

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

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

31 queue_timeout: float = 30.0 # 排队超时 

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

33 

34 

35@dataclass 

36class RateLimitResult: 

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

38 allowed: bool 

39 remaining: int = 0 

40 reset_at: float = 0.0 

41 retry_after: float = 0.0 

42 limit: int = 0 

43 reason: str = "" 

44 

45 

46class TokenBucket: 

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

48 

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

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

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

52 self.tokens = float(capacity) 

53 self.last_refill = time.monotonic() 

54 self._lock = asyncio.Lock() 

55 

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

57 async with self._lock: 

58 self._refill() 

59 if self.tokens >= tokens: 

60 self.tokens -= tokens 

61 return True 

62 return False 

63 

64 def _refill(self): 

65 now = time.monotonic() 

66 elapsed = now - self.last_refill 

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

68 self.last_refill = now 

69 

70 @property 

71 def available(self) -> float: 

72 self._refill() 

73 return self.tokens 

74 

75 

76class SlidingWindow: 

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

78 

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

80 self.max_requests = max_requests 

81 self.window = window_seconds 

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

83 self._lock = asyncio.Lock() 

84 

85 async def allow(self) -> bool: 

86 async with self._lock: 

87 now = time.monotonic() 

88 cutoff = now - self.window 

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

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

91 self._timestamps.append(now) 

92 return True 

93 return False 

94 

95 @property 

96 def current_count(self) -> int: 

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

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

99 

100 

101class ConcurrencyLimiter: 

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

103 

104 def __init__(self, max_concurrent: int): 

105 self._semaphore = asyncio.Semaphore(max_concurrent) 

106 self.max_concurrent = max_concurrent 

107 

108 async def acquire(self) -> bool: 

109 return await self._semaphore.acquire() 

110 

111 def release(self): 

112 self._semaphore.release() 

113 

114 @property 

115 def available(self) -> int: 

116 return self._semaphore._value 

117 

118 

119class RateLimiter: 

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

121 

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

123 cfg = config or RateLimitConfig() 

124 self.config = cfg 

125 self._bucket = TokenBucket( 

126 rate=cfg.max_requests / cfg.per_seconds, 

127 capacity=cfg.burst_size or cfg.max_requests 

128 ) 

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

130 self._concurrency = ConcurrencyLimiter(cfg.max_concurrent) 

131 

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

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

134 # 1. 并发检查 

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

136 pass # 还有并发槽位 

137 

138 # 2. 速率检查 

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

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

141 return RateLimitResult( 

142 allowed=True, 

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

144 limit=self.config.max_requests, 

145 ) 

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

147 return RateLimitResult( 

148 allowed=False, 

149 remaining=0, 

150 retry_after=wait, 

151 limit=self.config.max_requests, 

152 reason="rate_limit_exceeded", 

153 ) 

154 

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

156 if await self._window.allow(): 

157 return RateLimitResult( 

158 allowed=True, 

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

160 limit=self.config.max_requests, 

161 ) 

162 return RateLimitResult( 

163 allowed=False, remaining=0, 

164 retry_after=self.config.per_seconds, 

165 limit=self.config.max_requests, 

166 reason="window_exceeded", 

167 ) 

168 

169 # fixed window fallback 

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

171 

172 async def release(self): 

173 self._concurrency.release() 

174 

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

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

177 quotas = { 

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

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

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

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

182 } 

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

184 

185 

186class QuotaManager: 

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

188 

189 def __init__(self): 

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

191 

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

193 if key not in self._limiters: 

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

195 return self._limiters[key] 

196 

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

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

199 

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

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

202 pass