Coverage for qreward / client / manager.py: 99.29%

1import asyncio (empty)

2import threading (empty)

3from typing import Any, AsyncIterator, Optional (empty)

4 

5from qreward.client.load_balancer import ( (empty)

6 LoadBalanceStrategy, 

7 RoundRobinSelector, 

8 WeightedRoundRobinSelector, 

9) 

10from qreward.client.model_router import ModelRouter (empty)

11from qreward.client.openai import OpenAIChatProxy (empty)

12 

13 

14class OpenAIChatProxyManager: (empty)

15 """管理多个 key -> OpenAIChatProxy 的映射。 

16 

17 提供代理的增删查改、负载均衡和故障转移功能，支持链式调用。 

18 

19 Example: 

20 manager = OpenAIChatProxyManager( 

21 strategy=LoadBalanceStrategy.ROUND_ROBIN 

22 ) 

23 manager.add_proxy_with_default( 

24 "key1", "http://api1.example.com", "sk-xxx" 

25 ) 

26 manager.add_proxy_with_default( 

27 "key2", "http://api2.example.com", "sk-yyy" 

28 ) 

29 proxy = manager.select_proxy() # round-robin selection 

30 """ 

31 

32 def __init__( (empty)

33 self, 

34 strategy: LoadBalanceStrategy = LoadBalanceStrategy.ROUND_ROBIN, 

35 ) -> None: 

36 self._proxies: dict[str, OpenAIChatProxy] = {} (empty)

37 self._insertion_order: list[str] = [] (empty)

38 self._healthy_keys: set[str] = set() (empty)

39 self._weights: dict[str, int] = {} (empty)

40 self._strategy = strategy (empty)

41 self._lock = asyncio.Lock() (empty)

42 self._health_lock = threading.Lock() (empty)

43 

44 self._rr_selector = RoundRobinSelector() (empty)

45 self._wrr_selector = WeightedRoundRobinSelector() (empty)

46 self._model_router: Optional[ModelRouter] = None (empty)

47 

48 def add_proxy_with_default( (empty)

49 self, key: str, base_url: str, api_key: str 

50 ) -> "OpenAIChatProxyManager": 

51 """添加一个使用默认配置的代理。 

52 

53 Args: 

54 key: 代理标识 

55 base_url: OpenAI API 基础 URL 

56 api_key: OpenAI API 密钥 

57 

58 Returns: 

59 self，支持链式调用 

60 """ 

61 self.add_proxy( (empty)

62 key=key, 

63 proxy=OpenAIChatProxy( 

64 base_url=base_url, 

65 api_key=api_key, 

66 ), 

67 ) 

68 return self (empty)

69 

70 def add_proxies_with_default( (empty)

71 self, proxies: dict[str, tuple[str, str]] 

72 ) -> "OpenAIChatProxyManager": 

73 """批量添加使用默认配置的代理。 

74 

75 Args: 

76 proxies: 代理映射，key 为代理标识，value 为 (base_url, api_key) 元组 

77 

78 Returns: 

79 self，支持链式调用 

80 """ 

81 for key, (base_url, api_key) in proxies.items(): (empty)

82 self.add_proxy_with_default( (empty)

83 key=key, 

84 base_url=base_url, 

85 api_key=api_key, 

86 ) 

87 return self (empty)

88 

89 def add_proxy( (empty)

90 self, 

91 key: str, 

92 proxy: OpenAIChatProxy, 

93 weight: int = 1, 

94 ) -> "OpenAIChatProxyManager": 

95 """添加一个自定义配置的代理。 

96 

97 Args: 

98 key: 代理标识 

99 proxy: OpenAIChatProxy 实例 

100 weight: 权重（用于加权轮询策略），默认为 1 

101 

102 Returns: 

103 self，支持链式调用 

104 

105 Raises: 

106 ValueError: 当 key 已存在时 

107 """ 

108 if key in self._proxies: (empty)

109 raise ValueError(f"Proxy with key={key!r} already exists.") (empty)

110 self._proxies[key] = proxy (empty)

111 self._insertion_order.append(key) (empty)

112 self._weights[key] = max(1, weight) (empty)

113 with self._health_lock: (empty)

114 self._healthy_keys.add(key) (empty)

115 self._wrr_selector.update_weights(self._weights) (empty)

116 return self (empty)

117 

118 def proxy(self, key: str) -> OpenAIChatProxy: (empty)

119 """获取指定 key 的代理。 

120 

121 Args: 

122 key: 代理标识 

123 

124 Returns: 

125 对应的 OpenAIChatProxy 实例 

126 

127 Raises: 

128 KeyError: 当 key 不存在时 

129 """ 

130 try: (empty)

131 return self._proxies[key] (empty)

132 except KeyError: (empty)

133 raise KeyError(f"Proxy with key={key!r} does not exist.") (empty)

134 

135 def exist_proxy(self, key: str) -> bool: (empty)

136 """检查指定 key 的代理是否存在。 

137 

138 Args: 

139 key: 代理标识 

140 

141 Returns: 

142 True 如果存在，否则 False 

143 """ 

144 return key in self._proxies (empty)

145 

146 def proxies(self) -> dict[str, OpenAIChatProxy]: (empty)

147 """获取所有代理的字典副本。 

148 

149 返回副本以防止外部修改内部状态。 

150 

151 Returns: 

152 key -> OpenAIChatProxy 的映射副本 

153 """ 

154 return dict(self._proxies) (empty)

155 

156 def register_model_route( (empty)

157 self, 

158 pattern: str, 

159 proxy_map: dict[str, OpenAIChatProxy], 

160 weights: Optional[dict[str, int]] = None, 

161 strategy: LoadBalanceStrategy = LoadBalanceStrategy.ROUND_ROBIN, 

162 ) -> "OpenAIChatProxyManager": 

163 """Register a model route for model-based proxy selection. 

164 

165 Args: 

166 pattern: Model name or glob pattern (e.g. "gpt-4" or "gpt-*"). 

167 proxy_map: Mapping of key → OpenAIChatProxy for this model. 

168 weights: Optional per-key weights for weighted round-robin. 

169 strategy: Load balancing strategy for this model group. 

170 

171 Returns: 

172 self, for chaining. 

173 """ 

174 if self._model_router is None: (empty)

175 self._model_router = ModelRouter() (empty)

176 self._model_router.register(pattern, proxy_map, weights, strategy) (empty)

177 return self (empty)

178 

179 def list_routes(self) -> list[dict]: (empty)

180 """List all registered model routes. 

181 

182 Returns: 

183 List of route info dicts, or empty list if no router configured. 

184 """ 

185 if self._model_router is None: (empty)

186 return [] (empty)

187 return self._model_router.list_routes() (empty)

188 

189 def select_proxy(self, model: Optional[str] = None) -> OpenAIChatProxy: (empty)

190 """根据负载均衡策略选择一个可用代理。 

191 

192 Args: 

193 model: Optional model name for model-based routing. 

194 

195 Returns: 

196 选中的 OpenAIChatProxy 实例。 

197 

198 Raises: 

199 RuntimeError: 当没有可用的健康代理时。 

200 """ 

201 if model and self._model_router is not None: (empty)

202 group = self._model_router.resolve(model) (empty)

203 if group is not None: (empty)

204 return self._model_router.select_from_group(group) (empty)

205 

206 with self._health_lock: (empty)

207 healthy_snapshot = set(self._healthy_keys) (empty)

208 

209 selected_key: str | None = None (empty)

210 

211 if self._strategy == LoadBalanceStrategy.ROUND_ROBIN: (empty)

212 selected_key = self._rr_selector.select( (empty)

213 self._insertion_order, healthy_snapshot 

214 ) 

215 elif self._strategy == LoadBalanceStrategy.WEIGHTED_ROUND_ROBIN: (empty)

216 selected_key = self._wrr_selector.select( (empty)

217 self._insertion_order, healthy_snapshot, self._weights 

218 ) 

219 elif self._strategy == LoadBalanceStrategy.LEAST_CONNECTIONS: 219 ↛ 224line 219 didn't jump to line 224 because the condition on line 219 was always true(empty)

220 selected_key = self._rr_selector.select( (empty)

221 self._insertion_order, healthy_snapshot 

222 ) 

223 

224 if selected_key is None: (empty)

225 raise RuntimeError("No healthy proxy available.") (empty)

226 

227 return self._proxies[selected_key] (empty)

228 

229 def mark_unhealthy(self, key: str) -> None: (empty)

230 """标记代理为不健康状态。 

231 

232 Args: 

233 key: 代理标识。 

234 

235 Raises: 

236 KeyError: 当 key 不存在时。 

237 """ 

238 if key not in self._proxies: (empty)

239 raise KeyError(f"Proxy with key={key!r} does not exist.") (empty)

240 with self._health_lock: (empty)

241 self._healthy_keys.discard(key) (empty)

242 

243 def mark_healthy(self, key: str) -> None: (empty)

244 """标记代理为健康状态。 

245 

246 Args: 

247 key: 代理标识。 

248 

249 Raises: 

250 KeyError: 当 key 不存在时。 

251 """ 

252 if key not in self._proxies: (empty)

253 raise KeyError(f"Proxy with key={key!r} does not exist.") (empty)

254 with self._health_lock: (empty)

255 self._healthy_keys.add(key) (empty)

256 

257 def healthy_proxies(self) -> dict[str, OpenAIChatProxy]: (empty)

258 """获取所有健康代理的字典副本。 

259 

260 Returns: 

261 key -> OpenAIChatProxy 的映射副本（仅健康代理）。 

262 """ 

263 with self._health_lock: (empty)

264 healthy_snapshot = set(self._healthy_keys) (empty)

265 return { (empty)

266 k: v 

267 for k, v in self._proxies.items() 

268 if k in healthy_snapshot 

269 } 

270 

271 async def remove_proxy(self, key: str) -> None: (empty)

272 """移除并关闭指定 key 的代理。 

273 

274 Args: 

275 key: 代理标识 

276 """ 

277 async with self._lock: (empty)

278 proxy = self._proxies.pop(key, None) (empty)

279 if key in self._insertion_order: (empty)

280 self._insertion_order.remove(key) (empty)

281 self._weights.pop(key, None) (empty)

282 with self._health_lock: (empty)

283 self._healthy_keys.discard(key) (empty)

284 if proxy: (empty)

285 await proxy.client.close() (empty)

286 

287 async def batch_stream_chat_completion( (empty)

288 self, 

289 batch_messages: list[list[dict[str, Any]]], 

290 model: str, 

291 **kwargs: Any, 

292 ) -> "AsyncIterator[tuple[int, str]]": 

293 """通过负载均衡选择代理并执行批量流式聊天补全。""" 

294 proxy = self.select_proxy(model=model) (empty)

295 async for item in proxy.batch_stream_chat_completion( (empty)

296 batch_messages, model, **kwargs 

297 ): 

298 yield item (empty)

299 

300 async def close(self) -> None: (empty)

301 """关闭所有代理并清空映射。""" 

302 async with self._lock: (empty)

303 proxies_to_close = list(self._proxies.values()) (empty)

304 self._proxies.clear() (empty)

305 self._insertion_order.clear() (empty)

306 self._weights.clear() (empty)

307 with self._health_lock: (empty)

308 self._healthy_keys.clear() (empty)

309 for proxy in proxies_to_close: (empty)

310 await proxy.client.close() (empty)