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 StrEnum 

12 

13 

14class RateLimitStrategy(StrEnum): 

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

16 

17 TOKEN_BUCKET = "token_bucket" 

18 SLIDING_WINDOW = "sliding_window" 

19 FIXED_WINDOW = "fixed_window" 

20 

21 

22@dataclass 

23class RateLimitConfig: 

24 """限流配置。""" 

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 

39 allowed: bool 

40 remaining: int = 0 

41 reset_at: float = 0.0 

42 retry_after: float = 0.0 

43 limit: int = 0 

44 reason: str = "" 

45 

46 

47class TokenBucket: 

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

49 

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

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

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

53 self.tokens = float(capacity) 

54 self.last_refill = time.monotonic() 

55 self._lock = asyncio.Lock() 

56 

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

58 async with self._lock: 

59 self._refill() 

60 if self.tokens >= tokens: 

61 self.tokens -= tokens 

62 return True 

63 return False 

64 

65 def _refill(self): 

66 now = time.monotonic() 

67 elapsed = now - self.last_refill 

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

69 self.last_refill = now 

70 

71 @property 

72 def available(self) -> float: 

73 self._refill() 

74 return self.tokens 

75 

76 

77class SlidingWindow: 

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

79 

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

81 self.max_requests = max_requests 

82 self.window = window_seconds 

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

84 self._lock = asyncio.Lock() 

85 

86 async def allow(self) -> bool: 

87 async with self._lock: 

88 now = time.monotonic() 

89 cutoff = now - self.window 

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

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

92 self._timestamps.append(now) 

93 return True 

94 return False 

95 

96 @property 

97 def current_count(self) -> int: 

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

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

100 

101 

102class ConcurrencyLimiter: 

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

104 

105 def __init__(self, max_concurrent: int): 

106 self._semaphore = asyncio.Semaphore(max_concurrent) 

107 self.max_concurrent = max_concurrent 

108 

109 async def acquire(self) -> bool: 

110 return await self._semaphore.acquire() 

111 

112 def release(self): 

113 self._semaphore.release() 

114 

115 @property 

116 def available(self) -> int: 

117 return self._semaphore._value 

118 

119 

120class RateLimiter: 

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

122 

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

124 cfg = config or RateLimitConfig() 

125 self.config = cfg 

126 self._bucket = TokenBucket( 

127 rate=cfg.max_requests / cfg.per_seconds, 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, 

164 remaining=0, 

165 retry_after=self.config.per_seconds, 

166 limit=self.config.max_requests, 

167 reason="window_exceeded", 

168 ) 

169 

170 # fixed window fallback 

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

172 

173 async def release(self): 

174 self._concurrency.release() 

175 

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

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

178 quotas = { 

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

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

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

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

183 } 

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

185 

186 

187class QuotaManager: 

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

189 

190 def __init__(self): 

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

192 

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

194 if key not in self._limiters: 

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

196 return self._limiters[key] 

197 

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

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

200 

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

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