Coverage for agentos/tools/async_executor.py: 0%

248 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-09 09:19 +0800

1""" 

2v1.15.1 — 异步工具执行优化:并发控制 + 超时熔断 + 性能监控。 

3 

4核心功能: 

51. 并发执行控制:限制同时执行的工具数量 

62. 超时熔断:工具执行超时自动中断 

73. 性能监控:记录工具执行时间、成功率 

84. 智能重试:根据错误类型自动重试 

9""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14import time 

15from dataclasses import dataclass, field 

16from enum import StrEnum 

17from typing import Any 

18 

19from .base import BaseTool, ToolResult 

20from .validation import ErrorCategory, ToolErrorClassifier 

21 

22 

23class ExecutionStatus(StrEnum): 

24 """工具执行状态。""" 

25 

26 PENDING = "pending" 

27 RUNNING = "running" 

28 SUCCESS = "success" 

29 TIMEOUT = "timeout" 

30 FAILED = "failed" 

31 CANCELLED = "cancelled" 

32 

33 

34class CircuitBreakerState(StrEnum): 

35 """熔断器状态。""" 

36 

37 CLOSED = "closed" # 正常状态,允许执行 

38 OPEN = "open" # 熔断状态,拒绝执行 

39 HALF_OPEN = "half_open" # 半开状态,尝试恢复 

40 

41 

42@dataclass 

43class ExecutionMetrics: 

44 """工具执行性能指标。""" 

45 

46 tool_name: str 

47 execution_count: int = 0 

48 success_count: int = 0 

49 failure_count: int = 0 

50 timeout_count: int = 0 

51 total_execution_time: float = 0.0 

52 last_execution_time: float = 0.0 

53 last_error: str | None = None 

54 

55 @property 

56 def success_rate(self) -> float: 

57 if self.execution_count == 0: 

58 return 0.0 

59 return self.success_count / self.execution_count 

60 

61 @property 

62 def average_execution_time(self) -> float: 

63 if self.execution_count == 0: 

64 return 0.0 

65 return self.total_execution_time / self.execution_count 

66 

67 def record_success(self, execution_time: float) -> None: 

68 self.execution_count += 1 

69 self.success_count += 1 

70 self.total_execution_time += execution_time 

71 self.last_execution_time = execution_time 

72 self.last_error = None 

73 

74 def record_failure(self, execution_time: float, error: str) -> None: 

75 self.execution_count += 1 

76 self.failure_count += 1 

77 self.total_execution_time += execution_time 

78 self.last_execution_time = execution_time 

79 self.last_error = error 

80 

81 def record_timeout(self, execution_time: float) -> None: 

82 self.execution_count += 1 

83 self.timeout_count += 1 

84 self.total_execution_time += execution_time 

85 self.last_execution_time = execution_time 

86 self.last_error = "timeout" 

87 

88 

89@dataclass 

90class CircuitBreaker: 

91 """熔断器:防止工具持续失败。""" 

92 

93 failure_threshold: int = 5 # 连续失败次数阈值 

94 reset_timeout: float = 30.0 # 熔断恢复时间(秒) 

95 half_open_max_attempts: int = 3 # 半开状态最大尝试次数 

96 

97 _state: CircuitBreakerState = field(default=CircuitBreakerState.CLOSED) 

98 _failure_count: int = 0 

99 _last_failure_time: float = 0.0 

100 _half_open_attempts: int = 0 

101 

102 def can_execute(self) -> bool: 

103 """检查是否允许执行。""" 

104 current_time = time.time() 

105 

106 if self._state == CircuitBreakerState.OPEN: 

107 # 检查是否应该进入半开状态 

108 if current_time - self._last_failure_time >= self.reset_timeout: 

109 self._state = CircuitBreakerState.HALF_OPEN 

110 self._half_open_attempts = 0 

111 self._failure_count = 0 # 重置失败计数 

112 return True 

113 return False 

114 

115 elif self._state == CircuitBreakerState.HALF_OPEN: 

116 if self._half_open_attempts >= self.half_open_max_attempts: 

117 return False 

118 return True 

119 

120 return True # CLOSED 状态 

121 

122 def record_success(self) -> None: 

123 """记录成功执行。""" 

124 if self._state == CircuitBreakerState.HALF_OPEN: 

125 # 半开状态成功,恢复正常 

126 self._state = CircuitBreakerState.CLOSED 

127 self._failure_count = 0 

128 self._half_open_attempts = 0 

129 else: 

130 self._failure_count = 0 

131 

132 def record_failure(self) -> None: 

133 """记录失败执行。""" 

134 self._failure_count += 1 

135 self._last_failure_time = time.time() 

136 

137 if self._state == CircuitBreakerState.HALF_OPEN: 

138 self._half_open_attempts += 1 

139 # 半开状态失败,重新熔断 

140 if self._half_open_attempts >= self.half_open_max_attempts: 

141 self._state = CircuitBreakerState.OPEN 

142 elif self._failure_count >= self.failure_threshold: 

143 self._state = CircuitBreakerState.OPEN 

144 

145 @property 

146 def state(self) -> CircuitBreakerState: 

147 return self._state 

148 

149 @property 

150 def time_until_reset(self) -> float: 

151 """距离熔断恢复的剩余时间。""" 

152 if self._state != CircuitBreakerState.OPEN: 

153 return 0.0 

154 elapsed = time.time() - self._last_failure_time 

155 return max(0.0, self.reset_timeout - elapsed) 

156 

157 

158class AsyncToolExecutor: 

159 """异步工具执行器,支持并发控制和熔断。""" 

160 

161 def __init__( 

162 self, 

163 max_concurrent: int = 10, 

164 default_timeout: float = 30.0, 

165 enable_circuit_breaker: bool = True, 

166 ): 

167 """ 

168 初始化异步工具执行器。 

169 

170 Args: 

171 max_concurrent: 最大并发执行数 

172 default_timeout: 默认执行超时时间(秒) 

173 enable_circuit_breaker: 是否启用熔断器 

174 """ 

175 self.max_concurrent = max_concurrent 

176 self.default_timeout = default_timeout 

177 self.enable_circuit_breaker = enable_circuit_breaker 

178 

179 # 并发控制 

180 self._semaphore = asyncio.Semaphore(max_concurrent) 

181 self._active_tasks: set[asyncio.Task] = set() 

182 

183 # 性能监控 

184 self._metrics: dict[str, ExecutionMetrics] = {} 

185 self._circuit_breakers: dict[str, CircuitBreaker] = {} 

186 

187 # 工具超时配置 

188 self._tool_timeouts: dict[str, float] = {} 

189 

190 def set_tool_timeout(self, tool_name: str, timeout: float) -> None: 

191 """为特定工具设置超时时间。""" 

192 self._tool_timeouts[tool_name] = timeout 

193 

194 def get_tool_timeout(self, tool_name: str) -> float: 

195 """获取工具的超时时间。""" 

196 return self._tool_timeouts.get(tool_name, self.default_timeout) 

197 

198 def _get_or_create_metrics(self, tool_name: str) -> ExecutionMetrics: 

199 """获取或创建性能指标。""" 

200 if tool_name not in self._metrics: 

201 self._metrics[tool_name] = ExecutionMetrics(tool_name=tool_name) 

202 return self._metrics[tool_name] 

203 

204 def _get_or_create_circuit_breaker(self, tool_name: str) -> CircuitBreaker: 

205 """获取或创建熔断器。""" 

206 if tool_name not in self._circuit_breakers: 

207 self._circuit_breakers[tool_name] = CircuitBreaker() 

208 return self._circuit_breakers[tool_name] 

209 

210 async def execute( 

211 self, 

212 tool: BaseTool, 

213 arguments: dict[str, Any], 

214 call_id: str | None = None, 

215 timeout: float | None = None, 

216 ) -> ToolResult: 

217 """ 

218 异步执行工具。 

219 

220 Args: 

221 tool: 要执行的工具 

222 arguments: 工具参数 

223 call_id: 调用ID(可选) 

224 timeout: 超时时间(可选,覆盖默认值) 

225 

226 Returns: 

227 ToolResult: 工具执行结果 

228 """ 

229 if call_id is None: 

230 call_id = f"call_{int(time.time() * 1000)}" 

231 

232 tool_name = tool.name or tool.__class__.__name__ 

233 

234 # 检查熔断器 

235 if self.enable_circuit_breaker: 

236 circuit_breaker = self._get_or_create_circuit_breaker(tool_name) 

237 if not circuit_breaker.can_execute(): 

238 return ToolResult.fail( 

239 call_id=call_id, 

240 error=f"Circuit breaker is OPEN for tool '{tool_name}'. " 

241 f"Try again in {circuit_breaker.time_until_reset:.1f}s.", 

242 ) 

243 

244 # 获取超时时间 

245 exec_timeout = timeout or self.get_tool_timeout(tool_name) 

246 

247 # 获取性能指标 

248 metrics = self._get_or_create_metrics(tool_name) 

249 

250 # 创建任务 

251 task = asyncio.create_task( 

252 self._execute_with_semaphore( 

253 tool=tool, 

254 arguments=arguments, 

255 call_id=call_id, 

256 timeout=exec_timeout, 

257 tool_name=tool_name, 

258 metrics=metrics, 

259 ) 

260 ) 

261 

262 self._active_tasks.add(task) 

263 task.add_done_callback(self._active_tasks.discard) 

264 

265 try: 

266 return await task 

267 except asyncio.CancelledError: 

268 return ToolResult.fail(call_id=call_id, error="Execution cancelled") 

269 

270 async def _execute_with_semaphore( 

271 self, 

272 tool: BaseTool, 

273 arguments: dict[str, Any], 

274 call_id: str, 

275 timeout: float, 

276 tool_name: str, 

277 metrics: ExecutionMetrics, 

278 ) -> ToolResult: 

279 """使用信号量控制并发执行。""" 

280 start_time = time.time() 

281 

282 async with self._semaphore: 

283 try: 

284 # 执行工具(带超时) 

285 result = await asyncio.wait_for(tool.execute(arguments), timeout=timeout) 

286 

287 execution_time = time.time() - start_time 

288 

289 # 检查结果是否失败 

290 if result.error is not None: 

291 # 工具执行失败 

292 metrics.record_failure(execution_time, result.error) 

293 if self.enable_circuit_breaker: 

294 circuit_breaker = self._get_or_create_circuit_breaker(tool_name) 

295 circuit_breaker.record_failure() 

296 else: 

297 # 工具执行成功 

298 metrics.record_success(execution_time) 

299 if self.enable_circuit_breaker: 

300 circuit_breaker = self._get_or_create_circuit_breaker(tool_name) 

301 circuit_breaker.record_success() 

302 

303 return result 

304 

305 except TimeoutError: 

306 execution_time = time.time() - start_time 

307 metrics.record_timeout(execution_time) 

308 

309 if self.enable_circuit_breaker: 

310 circuit_breaker = self._get_or_create_circuit_breaker(tool_name) 

311 circuit_breaker.record_failure() 

312 

313 return ToolResult.fail( 

314 call_id=call_id, 

315 error=f"Tool '{tool_name}' execution timed out after {timeout}s", 

316 ) 

317 

318 except Exception as e: 

319 execution_time = time.time() - start_time 

320 error_msg = str(e) 

321 metrics.record_failure(execution_time, error_msg) 

322 

323 if self.enable_circuit_breaker: 

324 circuit_breaker = self._get_or_create_circuit_breaker(tool_name) 

325 circuit_breaker.record_failure() 

326 

327 return ToolResult.fail(call_id=call_id, error=error_msg) 

328 

329 async def execute_batch( 

330 self, 

331 tool_calls: list[tuple[BaseTool, dict[str, Any]]], 

332 max_batch_size: int | None = None, 

333 timeout_per_tool: float | None = None, 

334 ) -> list[ToolResult]: 

335 """ 

336 批量执行工具。 

337 

338 Args: 

339 tool_calls: 工具调用列表 [(tool, arguments), ...] 

340 max_batch_size: 最大批量大小(None表示无限制) 

341 timeout_per_tool: 每个工具的超时时间 

342 

343 Returns: 

344 List[ToolResult]: 工具执行结果列表 

345 """ 

346 if max_batch_size is not None: 

347 # 分批执行 

348 results = [] 

349 for i in range(0, len(tool_calls), max_batch_size): 

350 batch = tool_calls[i : i + max_batch_size] 

351 batch_results = await asyncio.gather( 

352 *[self.execute(tool, args, timeout=timeout_per_tool) for tool, args in batch] 

353 ) 

354 results.extend(batch_results) 

355 return results 

356 else: 

357 # 并发执行所有工具 

358 tasks = [ 

359 self.execute(tool, args, timeout=timeout_per_tool) for tool, args in tool_calls 

360 ] 

361 return await asyncio.gather(*tasks) 

362 

363 def get_metrics(self, tool_name: str | None = None) -> dict[str, ExecutionMetrics]: 

364 """获取性能指标。""" 

365 if tool_name: 

366 return {tool_name: self._metrics.get(tool_name)} 

367 return self._metrics.copy() 

368 

369 def get_circuit_breaker_state(self, tool_name: str) -> CircuitBreakerState | None: 

370 """获取熔断器状态。""" 

371 if tool_name in self._circuit_breakers: 

372 return self._circuit_breakers[tool_name].state 

373 return None 

374 

375 def reset_circuit_breaker(self, tool_name: str) -> bool: 

376 """重置指定工具的熔断器。""" 

377 if tool_name in self._circuit_breakers: 

378 self._circuit_breakers[tool_name] = CircuitBreaker() 

379 return True 

380 return False 

381 

382 def reset_all_circuit_breakers(self) -> None: 

383 """重置所有熔断器。""" 

384 self._circuit_breakers.clear() 

385 

386 async def shutdown(self, timeout: float = 5.0) -> None: 

387 """优雅关闭执行器。""" 

388 # 取消所有正在执行的任务 

389 for task in self._active_tasks.copy(): 

390 task.cancel() 

391 

392 # 等待任务完成或超时 

393 if self._active_tasks: 

394 try: 

395 await asyncio.wait_for( 

396 asyncio.gather(*self._active_tasks, return_exceptions=True), timeout=timeout 

397 ) 

398 except TimeoutError: 

399 pass 

400 

401 @property 

402 def active_task_count(self) -> int: 

403 """当前活跃任务数量。""" 

404 return len(self._active_tasks) 

405 

406 @property 

407 def available_slots(self) -> int: 

408 """可用并发槽位数量。""" 

409 return self.max_concurrent - self.active_task_count 

410 

411 

412class SmartRetryExecutor: 

413 """智能重试执行器:根据错误类型自动重试。""" 

414 

415 def __init__( 

416 self, 

417 max_retries: int = 3, 

418 retry_delay: float = 1.0, 

419 backoff_factor: float = 2.0, 

420 retryable_categories: list[ErrorCategory] | None = None, 

421 ): 

422 """ 

423 初始化智能重试执行器。 

424 

425 Args: 

426 max_retries: 最大重试次数 

427 retry_delay: 初始重试延迟(秒) 

428 backoff_factor: 退避因子 

429 retryable_categories: 可重试的错误类别 

430 """ 

431 self.max_retries = max_retries 

432 self.retry_delay = retry_delay 

433 self.backoff_factor = backoff_factor 

434 

435 if retryable_categories is None: 

436 self.retryable_categories = [ 

437 ErrorCategory.NETWORK, 

438 ErrorCategory.TIMEOUT, 

439 ErrorCategory.RATE_LIMIT, 

440 ErrorCategory.UNKNOWN, 

441 ] 

442 else: 

443 self.retryable_categories = retryable_categories 

444 

445 async def execute_with_retry( 

446 self, 

447 tool: BaseTool, 

448 arguments: dict[str, Any], 

449 call_id: str | None = None, 

450 base_executor: AsyncToolExecutor | None = None, 

451 ) -> ToolResult: 

452 """ 

453 带智能重试的工具执行。 

454 

455 Args: 

456 tool: 要执行的工具 

457 arguments: 工具参数 

458 call_id: 调用ID 

459 base_executor: 基础执行器(可选) 

460 

461 Returns: 

462 ToolResult: 最终执行结果 

463 """ 

464 if call_id is None: 

465 call_id = f"retry_{int(time.time() * 1000)}" 

466 

467 if base_executor is None: 

468 base_executor = AsyncToolExecutor() 

469 

470 last_result = None 

471 delay = self.retry_delay 

472 

473 for attempt in range(self.max_retries + 1): 

474 if attempt > 0: 

475 # 等待重试延迟 

476 await asyncio.sleep(delay) 

477 delay *= self.backoff_factor # 指数退避 

478 

479 # 执行工具 

480 result = await base_executor.execute( 

481 tool=tool, arguments=arguments, call_id=f"{call_id}_attempt{attempt}" 

482 ) 

483 

484 if result.error is None: 

485 # 执行成功 

486 return result 

487 

488 # 检查是否可重试 

489 last_result = result 

490 error_category = ToolErrorClassifier.classify(result) 

491 

492 if error_category not in self.retryable_categories: 

493 # 不可重试的错误 

494 break 

495 

496 if attempt == self.max_retries: 

497 # 达到最大重试次数 

498 break 

499 

500 # 返回最后一次失败的结果 

501 return last_result or ToolResult.fail( 

502 call_id=call_id, error="Execution failed after retries" 

503 ) 

504 

505 

506# 便捷函数 

507async def execute_tool_with_retry( 

508 tool: BaseTool, arguments: dict[str, Any], max_retries: int = 3, call_id: str | None = None 

509) -> ToolResult: 

510 """ 

511 带重试的工具执行便捷函数。 

512 

513 Args: 

514 tool: 要执行的工具 

515 arguments: 工具参数 

516 max_retries: 最大重试次数 

517 call_id: 调用ID 

518 

519 Returns: 

520 ToolResult: 执行结果 

521 """ 

522 executor = SmartRetryExecutor(max_retries=max_retries) 

523 return await executor.execute_with_retry(tool, arguments, call_id) 

524 

525 

526async def execute_tools_concurrently( 

527 tool_calls: list[tuple[BaseTool, dict[str, Any]]], 

528 max_concurrent: int = 10, 

529 timeout_per_tool: float | None = None, 

530) -> list[ToolResult]: 

531 """ 

532 并发执行多个工具的便捷函数。 

533 

534 Args: 

535 tool_calls: 工具调用列表 

536 max_concurrent: 最大并发数 

537 timeout_per_tool: 每个工具的超时时间 

538 

539 Returns: 

540 List[ToolResult]: 执行结果列表 

541 """ 

542 executor = AsyncToolExecutor(max_concurrent=max_concurrent) 

543 return await executor.execute_batch(tool_calls=tool_calls, timeout_per_tool=timeout_per_tool)