Coverage for qreward / client / openai.py: 99.49%

1import asyncio (empty)

2import logging (empty)

3import os (empty)

4import warnings (empty)

5from collections.abc import Callable (empty)

6from datetime import datetime (empty)

7from typing import ( (empty)

8 Any, 

9 AsyncIterator, 

10 Optional, 

11 Sequence, 

12 Union, 

13) 

14 

15from aiolimiter import AsyncLimiter (empty)

16from httpx import Limits, Timeout, TimeoutException (empty)

17from openai import ( (empty)

18 AsyncOpenAI, 

19 APIConnectionError, 

20 APIStatusError, 

21 APITimeoutError, 

22 DefaultAioHttpClient, 

23 RateLimitError, 

24) 

25from openai.types.chat import ChatCompletion (empty)

26from tenacity import ( (empty)

27 retry, 

28 retry_if_exception_type, 

29 stop_after_attempt, 

30 wait_exponential, 

31) 

32 

33from qreward.client.patch_openai import patch_openai_embeddings (empty)

34from qreward.types import RequestHook, ResponseHook (empty)

35from qreward.utils import patch_httpx (empty)

36 

37logger = logging.getLogger(__name__) (empty)

38 

39MAX_RETRIES = 5 (empty)

40_ERROR_GROUP = ( (empty)

41 TimeoutError, 

42 asyncio.TimeoutError, 

43 TimeoutException, 

44 RateLimitError, 

45 APIStatusError, 

46 APITimeoutError, 

47 APIConnectionError, 

48) 

49 

50 

51class OpenAIChatProxy: (empty)

52 def __init__( (empty)

53 self, 

54 base_url: str, 

55 *, 

56 api_key: str | None = None, 

57 debug: bool = False, 

58 max_concurrent: int = 64, 

59 chat_process_func: Callable | None = None, 

60 error_process_func: Callable | None = None, 

61 chat_process_fuc: Callable | None = None, 

62 error_process_fuc: Callable | None = None, 

63 timeout: float | Timeout = None, 

64 rate_limiter_bucket_size: float = 50.0, 

65 rate_limiter_bucket_period: float = 1.0, 

66 httpx_request_hook: RequestHook | None = None, 

67 httpx_response_hook: ResponseHook | None = None, 

68 is_hack_embedding_method: bool = False, 

69 verify_ssl: bool = False, 

70 ): 

71 """初始化 OpenAI Chat 代理。""" 

72 # Lazy patch: apply httpx JSON optimization on first proxy creation 

73 patch_httpx() (empty)

74 

75 # Handle deprecated parameters 

76 if chat_process_fuc is not None: (empty)

77 warnings.warn( (empty)

78 "chat_process_fuc is deprecated, use " 

79 "chat_process_func instead", 

80 DeprecationWarning, 

81 stacklevel=2, 

82 ) 

83 if chat_process_func is None: (empty)

84 chat_process_func = chat_process_fuc (empty)

85 

86 if error_process_fuc is not None: (empty)

87 warnings.warn( (empty)

88 "error_process_fuc is deprecated, use " 

89 "error_process_func instead", 

90 DeprecationWarning, 

91 stacklevel=2, 

92 ) 

93 if error_process_func is None: (empty)

94 error_process_func = error_process_fuc (empty)

95 

96 self.debug = debug (empty)

97 self._max_concurrent = max_concurrent (empty)

98 self._api_key = api_key or self.get_openai_key() (empty)

99 self.client = AsyncOpenAI( (empty)

100 api_key=self._api_key, 

101 base_url=base_url, 

102 timeout=timeout, 

103 max_retries=0, 

104 http_client=self._default_http_client( 

105 request_hook=httpx_request_hook, 

106 response_hook=httpx_response_hook, 

107 verify=verify_ssl, 

108 ), 

109 ) 

110 

111 self.semaphore = asyncio.Semaphore(value=self._max_concurrent) (empty)

112 

113 # default token bucket size is 50 QPS 

114 self.rate_limiter = AsyncLimiter( (empty)

115 max_rate=rate_limiter_bucket_size, 

116 time_period=rate_limiter_bucket_period, 

117 ) 

118 

119 self._default_temperature = 0.0 (empty)

120 self._default_timeout = 60 (empty)

121 self._default_chat_process_func = chat_process_func (empty)

122 self._default_error_process_func = error_process_func (empty)

123 

124 # for hacking embedding 

125 self._is_hack_embedding = is_hack_embedding_method (empty)

126 if self._is_hack_embedding: (empty)

127 patch_openai_embeddings() (empty)

128 

129 async def __aenter__(self): (empty)

130 return self (empty)

131 

132 async def __aexit__(self, exc_type, exc_val, exc_tb): (empty)

133 await self.client.close() (empty)

134 

135 def _default_http_client( (empty)

136 self, 

137 request_hook: RequestHook | None = None, 

138 response_hook: ResponseHook | None = None, 

139 verify: bool = False, 

140 ) -> DefaultAioHttpClient: 

141 _http_client = DefaultAioHttpClient( (empty)

142 verify=verify, 

143 # 默认值是 100 / 20，现在需要根据 self._max_concurrent 的值来调整配比 

144 limits=Limits( 

145 max_connections=self._max_concurrent, 

146 max_keepalive_connections=self._max_concurrent, 

147 ), 

148 ) 

149 

150 if request_hook and isinstance(request_hook, Callable): (empty)

151 _http_client.event_hooks.get("request").append(request_hook) (empty)

152 

153 if response_hook and isinstance(response_hook, Callable): (empty)

154 _http_client.event_hooks.get("response").append(response_hook) (empty)

155 

156 return _http_client (empty)

157 

158 @staticmethod (empty)

159 def get_openai_key() -> str | None: (empty)

160 return os.getenv("OPENAI_API_KEY") (empty)

161 

162 def with_max_concurrent(self, max_concurrent: int): (empty)

163 self._max_concurrent = max_concurrent (empty)

164 return self (empty)

165 

166 def with_temperature(self, temperature: float): (empty)

167 self._default_temperature = temperature (empty)

168 return self (empty)

169 

170 def with_timeout(self, timeout: int): (empty)

171 self.client.timeout = timeout (empty)

172 self._default_timeout = timeout (empty)

173 return self (empty)

174 

175 def with_error_process_func(self, error_process_func: Callable): (empty)

176 self._default_error_process_func = error_process_func (empty)

177 return self (empty)

178 

179 def with_error_process_fuc(self, error_process_fuc: Callable): (empty)

180 warnings.warn( (empty)

181 "with_error_process_fuc is deprecated, use " 

182 "with_error_process_func instead", 

183 DeprecationWarning, 

184 stacklevel=2, 

185 ) 

186 return self.with_error_process_func(error_process_fuc) (empty)

187 

188 @staticmethod (empty)

189 def _default_chat_completion_func(completion: ChatCompletion): (empty)

190 return completion.choices[0].message.content (empty)

191 

192 @retry( (empty)

193 retry=retry_if_exception_type(exception_types=_ERROR_GROUP), 

194 wait=wait_exponential(multiplier=1, min=2, max=4), 

195 stop=stop_after_attempt(max_attempt_number=MAX_RETRIES), 

196 reraise=True, 

197 ) 

198 async def chat_completion( (empty)

199 self, 

200 messages: list[dict[str, Any]], 

201 model: str, 

202 **kwargs: Any, 

203 ) -> str: 

204 """调用 OpenAI Chat 接口。""" 

205 if self.debug: (empty)

206 logger.debug( (empty)

207 "[time: %s] - [Begin] - Call model: %s", 

208 datetime.now(), 

209 model, 

210 ) 

211 

212 async with self.semaphore, self.rate_limiter: (empty)

213 completion: ChatCompletion = await asyncio.wait_for( (empty)

214 self.client.chat.completions.create( 

215 model=model, 

216 messages=messages, 

217 temperature=kwargs.get( 

218 "temperature", 

219 self._default_temperature, 

220 ), 

221 stream=kwargs.get("stream"), 

222 ), 

223 timeout=kwargs.get("timeout", self._default_timeout), 

224 ) 

225 if self.debug: (empty)

226 logger.debug( (empty)

227 "[time: %s] - [End] - Call model: %s success!", 

228 datetime.now(), 

229 model, 

230 ) 

231 

232 if self._default_chat_process_func: (empty)

233 completion = self._default_chat_process_func(completion) (empty)

234 

235 return self._default_chat_completion_func(completion=completion) (empty)

236 

237 async def batch_chat_completion( (empty)

238 self, 

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

240 model: str, 

241 **kwargs: Any, 

242 ) -> list[Any]: 

243 """批量调用 OpenAI Chat 接口。""" 

244 tasks = [] (empty)

245 for i in range(len(batch_messages)): (empty)

246 tasks.append( (empty)

247 asyncio.create_task( 

248 self.chat_completion( 

249 model=model, 

250 messages=batch_messages[i], 

251 **kwargs, 

252 ), 

253 ), 

254 ) 

255 

256 results = await asyncio.gather(*tasks, return_exceptions=True) (empty)

257 

258 processed_results = [] (empty)

259 for index, result in enumerate(results): (empty)

260 if isinstance(result, Exception): (empty)

261 logger.warning( (empty)

262 "batch_chat_completion task %d/%d failed: %s: %s", 

263 index + 1, 

264 len(results), 

265 type(result).__name__, 

266 result, 

267 ) 

268 if self._default_error_process_func: (empty)

269 result = self._default_error_process_func(result) (empty)

270 

271 processed_results.append(result) (empty)

272 

273 return processed_results (empty)

274 

275 async def stream_chat_completion( (empty)

276 self, 

277 messages: list[dict[str, Any]], 

278 model: str, 

279 **kwargs: Any, 

280 ) -> AsyncIterator[str]: 

281 """流式聊天补全，逐 token 返回内容。 

282 

283 Args: 

284 messages: 消息列表 

285 model: 使用的模型 

286 **kwargs: 其他参数，支持 temperature / timeout 等 

287 

288 Yields: 

289 每个 chunk 的文本内容 

290 """ 

291 if self.debug: (empty)

292 logger.debug( (empty)

293 "[time: %s] - [Begin] - Stream model: %s", 

294 datetime.now(), 

295 model, 

296 ) 

297 

298 async with self.semaphore, self.rate_limiter: (empty)

299 stream = await self.client.chat.completions.create( (empty)

300 model=model, 

301 messages=messages, 

302 temperature=kwargs.get( 

303 "temperature", self._default_temperature 

304 ), 

305 stream=True, 

306 timeout=kwargs.get("timeout", self._default_timeout), 

307 ) 

308 try: (empty)

309 async for chunk in stream: (empty)

310 if ( (empty)

311 chunk.choices 

312 and chunk.choices[0].delta.content is not None 

313 ): 

314 yield chunk.choices[0].delta.content (empty)

315 finally: 

316 await stream.close() (empty)

317 

318 if self.debug: (empty)

319 logger.debug( (empty)

320 "[time: %s] - [End] - Stream model: %s done!", 

321 datetime.now(), 

322 model, 

323 ) 

324 

325 @retry( (empty)

326 retry=retry_if_exception_type(exception_types=_ERROR_GROUP), 

327 wait=wait_exponential(multiplier=1, min=2, max=4), 

328 stop=stop_after_attempt(max_attempt_number=MAX_RETRIES), 

329 reraise=True, 

330 ) 

331 async def embeddings( (empty)

332 self, 

333 *, 

334 sentences: Optional[Union[str, Sequence]] = None, 

335 model: str | None = None, 

336 extra_body: object | None = None, 

337 **kwargs: Any, 

338 ) -> list[Any]: 

339 """调用 OpenAI Embeddings 接口。""" 

340 if self.debug: (empty)

341 logger.debug( (empty)

342 "[time: %s] - [Begin] - Call embedding: %s", 

343 datetime.now(), 

344 model, 

345 ) 

346 

347 embedding_resp = await asyncio.wait_for( (empty)

348 self.client.embeddings.create( 

349 input=sentences, 

350 model=model, 

351 extra_body=extra_body, 

352 ), 

353 timeout=kwargs.get("timeout", self._default_timeout), 

354 ) 

355 

356 if self.debug: (empty)

357 logger.debug( (empty)

358 "[time: %s] - [End] - Call embedding: %s success!", 

359 datetime.now(), 

360 model, 

361 ) 

362 

363 return ( (empty)

364 embedding_resp.embeddings 

365 if self._is_hack_embedding 

366 else embedding_resp.data 

367 ) 

368 

369 async def batch_stream_chat_completion( (empty)

370 self, 

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

372 model: str, 

373 max_concurrent_streams: int = 0, 

374 on_stream_error: Callable[[int, Exception], None] | None = None, 

375 **kwargs: Any, 

376 ) -> AsyncIterator[tuple[int, str]]: 

377 """并发多个流式请求，yield (batch_index, token) 元组。""" 

378 if not batch_messages: (empty)

379 return (empty)

380 queue: asyncio.Queue[tuple[int, str] | None] = asyncio.Queue() (empty)

381 sem = ( (empty)

382 asyncio.Semaphore(max_concurrent_streams) 

383 if max_concurrent_streams > 0 

384 else None 

385 ) 

386 

387 async def _run(idx: int, msgs: list[dict[str, Any]]) -> None: (empty)

388 try: (empty)

389 if sem: (empty)

390 await sem.acquire() (empty)

391 async for tok in self.stream_chat_completion( (empty)

392 msgs, model, **kwargs 

393 ): 

394 await queue.put((idx, tok)) (empty)

395 except Exception as exc: (empty)

396 if on_stream_error: 396 ↛ 399line 396 didn't jump to line 399 because the condition on line 396 was always true(empty)

397 on_stream_error(idx, exc) (empty)

398 finally: 

399 if sem: (empty)

400 sem.release() (empty)

401 

402 tasks = [ (empty)

403 asyncio.create_task(_run(i, m)) 

404 for i, m in enumerate(batch_messages) 

405 ] 

406 

407 async def _drain() -> None: (empty)

408 await asyncio.gather(*tasks, return_exceptions=True) (empty)

409 await queue.put(None) (empty)

410 

411 sentinel = asyncio.create_task(_drain()) (empty)

412 while (item := await queue.get()) is not None: (empty)

413 yield item (empty)

414 await sentinel (empty)

415 

416 async def batch_embeddings( (empty)

417 self, 

418 *, 

419 batch_sentences: list[Union[str, Sequence]], 

420 model: str | None = None, 

421 extra_bodies: list[object] | None = None, 

422 **kwargs: Any, 

423 ) -> list[list[Any]]: 

424 """批量调用 OpenAI Embeddings 接口。""" 

425 tasks = [ (empty)

426 asyncio.create_task( 

427 self.embeddings( 

428 model=model, 

429 sentences=batch_sentences[i], 

430 extra_body=extra_bodies[i] if extra_bodies else None, 

431 **kwargs, 

432 ) 

433 ) 

434 for i in range(len(batch_sentences)) 

435 ] 

436 return await asyncio.gather(*tasks) (empty)