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

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 Enum 

17from typing import Any, Dict, List, Optional, Set, Tuple 

18 

19from .base import BaseTool, ToolResult 

20from .validation import ToolErrorClassifier, ErrorCategory 

21 

22 

23class ExecutionStatus(str, Enum): 

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

25 PENDING = "pending" 

26 RUNNING = "running" 

27 SUCCESS = "success" 

28 TIMEOUT = "timeout" 

29 FAILED = "failed" 

30 CANCELLED = "cancelled" 

31 

32 

33class CircuitBreakerState(str, Enum): 

34 """熔断器状态。""" 

35 CLOSED = "closed" # 正常状态，允许执行 

36 OPEN = "open" # 熔断状态，拒绝执行 

37 HALF_OPEN = "half_open" # 半开状态，尝试恢复 

38 

39 

40@dataclass 

41class ExecutionMetrics: 

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

43 tool_name: str 

44 execution_count: int = 0 

45 success_count: int = 0 

46 failure_count: int = 0 

47 timeout_count: int = 0 

48 total_execution_time: float = 0.0 

49 last_execution_time: float = 0.0 

50 last_error: Optional[str] = None 

51 

52 @property 

53 def success_rate(self) -> float: 

54 if self.execution_count == 0: 

55 return 0.0 

56 return self.success_count / self.execution_count 

57 

58 @property 

59 def average_execution_time(self) -> float: 

60 if self.execution_count == 0: 

61 return 0.0 

62 return self.total_execution_time / self.execution_count 

63 

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

65 self.execution_count += 1 

66 self.success_count += 1 

67 self.total_execution_time += execution_time 

68 self.last_execution_time = execution_time 

69 self.last_error = None 

70 

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

72 self.execution_count += 1 

73 self.failure_count += 1 

74 self.total_execution_time += execution_time 

75 self.last_execution_time = execution_time 

76 self.last_error = error 

77 

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

79 self.execution_count += 1 

80 self.timeout_count += 1 

81 self.total_execution_time += execution_time 

82 self.last_execution_time = execution_time 

83 self.last_error = "timeout" 

84 

85 

86@dataclass 

87class CircuitBreaker: 

88 """熔断器：防止工具持续失败。""" 

89 

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

91 reset_timeout: float = 30.0 # 熔断恢复时间（秒） 

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

93 

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

95 _failure_count: int = 0 

96 _last_failure_time: float = 0.0 

97 _half_open_attempts: int = 0 

98 

99 def can_execute(self) -> bool: 

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

101 current_time = time.time() 

102 

103 if self._state == CircuitBreakerState.OPEN: 

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

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

106 self._state = CircuitBreakerState.HALF_OPEN 

107 self._half_open_attempts = 0 

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

109 return True 

110 return False 

111 

112 elif self._state == CircuitBreakerState.HALF_OPEN: 

113 if self._half_open_attempts >= self.half_open_max_attempts: 

114 return False 

115 return True 

116 

117 return True # CLOSED 状态 

118 

119 def record_success(self) -> None: 

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

121 if self._state == CircuitBreakerState.HALF_OPEN: 

122 # 半开状态成功，恢复正常 

123 self._state = CircuitBreakerState.CLOSED 

124 self._failure_count = 0 

125 self._half_open_attempts = 0 

126 else: 

127 self._failure_count = 0 

128 

129 def record_failure(self) -> None: 

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

131 self._failure_count += 1 

132 self._last_failure_time = time.time() 

133 

134 if self._state == CircuitBreakerState.HALF_OPEN: 

135 self._half_open_attempts += 1 

136 # 半开状态失败，重新熔断 

137 if self._half_open_attempts >= self.half_open_max_attempts: 

138 self._state = CircuitBreakerState.OPEN 

139 elif self._failure_count >= self.failure_threshold: 

140 self._state = CircuitBreakerState.OPEN 

141 

142 @property 

143 def state(self) -> CircuitBreakerState: 

144 return self._state 

145 

146 @property 

147 def time_until_reset(self) -> float: 

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

149 if self._state != CircuitBreakerState.OPEN: 

150 return 0.0 

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

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

153 

154 

155class AsyncToolExecutor: 

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

157 

158 def __init__( 

159 self, 

160 max_concurrent: int = 10, 

161 default_timeout: float = 30.0, 

162 enable_circuit_breaker: bool = True 

163 ): 

164 """ 

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

166  

167 Args: 

168 max_concurrent: 最大并发执行数 

169 default_timeout: 默认执行超时时间（秒） 

170 enable_circuit_breaker: 是否启用熔断器 

171 """ 

172 self.max_concurrent = max_concurrent 

173 self.default_timeout = default_timeout 

174 self.enable_circuit_breaker = enable_circuit_breaker 

175 

176 # 并发控制 

177 self._semaphore = asyncio.Semaphore(max_concurrent) 

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

179 

180 # 性能监控 

181 self._metrics: Dict[str, ExecutionMetrics] = {} 

182 self._circuit_breakers: Dict[str, CircuitBreaker] = {} 

183 

184 # 工具超时配置 

185 self._tool_timeouts: Dict[str, float] = {} 

186 

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

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

189 self._tool_timeouts[tool_name] = timeout 

190 

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

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

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

194 

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

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

197 if tool_name not in self._metrics: 

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

199 return self._metrics[tool_name] 

200 

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

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

203 if tool_name not in self._circuit_breakers: 

204 self._circuit_breakers[tool_name] = CircuitBreaker() 

205 return self._circuit_breakers[tool_name] 

206 

207 async def execute( 

208 self, 

209 tool: BaseTool, 

210 arguments: Dict[str, Any], 

211 call_id: Optional[str] = None, 

212 timeout: Optional[float] = None 

213 ) -> ToolResult: 

214 """ 

215 异步执行工具。 

216  

217 Args: 

218 tool: 要执行的工具 

219 arguments: 工具参数 

220 call_id: 调用ID（可选） 

221 timeout: 超时时间（可选，覆盖默认值） 

222  

223 Returns: 

224 ToolResult: 工具执行结果 

225 """ 

226 if call_id is None: 

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

228 

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

230 

231 # 检查熔断器 

232 if self.enable_circuit_breaker: 

233 circuit_breaker = self._get_or_create_circuit_breaker(tool_name) 

234 if not circuit_breaker.can_execute(): 

235 return ToolResult.fail( 

236 call_id=call_id, 

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

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

239 ) 

240 

241 # 获取超时时间 

242 exec_timeout = timeout or self.get_tool_timeout(tool_name) 

243 

244 # 获取性能指标 

245 metrics = self._get_or_create_metrics(tool_name) 

246 

247 # 创建任务 

248 task = asyncio.create_task( 

249 self._execute_with_semaphore( 

250 tool=tool, 

251 arguments=arguments, 

252 call_id=call_id, 

253 timeout=exec_timeout, 

254 tool_name=tool_name, 

255 metrics=metrics 

256 ) 

257 ) 

258 

259 self._active_tasks.add(task) 

260 task.add_done_callback(self._active_tasks.discard) 

261 

262 try: 

263 return await task 

264 except asyncio.CancelledError: 

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

266 

267 async def _execute_with_semaphore( 

268 self, 

269 tool: BaseTool, 

270 arguments: Dict[str, Any], 

271 call_id: str, 

272 timeout: float, 

273 tool_name: str, 

274 metrics: ExecutionMetrics 

275 ) -> ToolResult: 

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

277 start_time = time.time() 

278 

279 async with self._semaphore: 

280 try: 

281 # 执行工具（带超时） 

282 result = await asyncio.wait_for( 

283 tool.execute(arguments), 

284 timeout=timeout 

285 ) 

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 asyncio.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: Optional[int] = None, 

333 timeout_per_tool: Optional[float] = 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) 

353 for tool, args in batch 

354 ]) 

355 results.extend(batch_results) 

356 return results 

357 else: 

358 # 并发执行所有工具 

359 tasks = [ 

360 self.execute(tool, args, timeout=timeout_per_tool) 

361 for tool, args in tool_calls 

362 ] 

363 return await asyncio.gather(*tasks) 

364 

365 def get_metrics(self, tool_name: Optional[str] = None) -> Dict[str, ExecutionMetrics]: 

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

367 if tool_name: 

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

369 return self._metrics.copy() 

370 

371 def get_circuit_breaker_state(self, tool_name: str) -> Optional[CircuitBreakerState]: 

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

373 if tool_name in self._circuit_breakers: 

374 return self._circuit_breakers[tool_name].state 

375 return None 

376 

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

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

379 if tool_name in self._circuit_breakers: 

380 self._circuit_breakers[tool_name] = CircuitBreaker() 

381 return True 

382 return False 

383 

384 def reset_all_circuit_breakers(self) -> None: 

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

386 self._circuit_breakers.clear() 

387 

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

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

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

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

392 task.cancel() 

393 

394 # 等待任务完成或超时 

395 if self._active_tasks: 

396 try: 

397 await asyncio.wait_for( 

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

399 timeout=timeout 

400 ) 

401 except asyncio.TimeoutError: 

402 pass 

403 

404 @property 

405 def active_task_count(self) -> int: 

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

407 return len(self._active_tasks) 

408 

409 @property 

410 def available_slots(self) -> int: 

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

412 return self.max_concurrent - self.active_task_count 

413 

414 

415class SmartRetryExecutor: 

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

417 

418 def __init__( 

419 self, 

420 max_retries: int = 3, 

421 retry_delay: float = 1.0, 

422 backoff_factor: float = 2.0, 

423 retryable_categories: Optional[List[ErrorCategory]] = None 

424 ): 

425 """ 

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

427  

428 Args: 

429 max_retries: 最大重试次数 

430 retry_delay: 初始重试延迟（秒） 

431 backoff_factor: 退避因子 

432 retryable_categories: 可重试的错误类别 

433 """ 

434 self.max_retries = max_retries 

435 self.retry_delay = retry_delay 

436 self.backoff_factor = backoff_factor 

437 

438 if retryable_categories is None: 

439 self.retryable_categories = [ 

440 ErrorCategory.NETWORK, 

441 ErrorCategory.TIMEOUT, 

442 ErrorCategory.RATE_LIMIT, 

443 ErrorCategory.UNKNOWN 

444 ] 

445 else: 

446 self.retryable_categories = retryable_categories 

447 

448 async def execute_with_retry( 

449 self, 

450 tool: BaseTool, 

451 arguments: Dict[str, Any], 

452 call_id: Optional[str] = None, 

453 base_executor: Optional[AsyncToolExecutor] = None 

454 ) -> ToolResult: 

455 """ 

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

457  

458 Args: 

459 tool: 要执行的工具 

460 arguments: 工具参数 

461 call_id: 调用ID 

462 base_executor: 基础执行器（可选） 

463  

464 Returns: 

465 ToolResult: 最终执行结果 

466 """ 

467 if call_id is None: 

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

469 

470 if base_executor is None: 

471 base_executor = AsyncToolExecutor() 

472 

473 last_result = None 

474 delay = self.retry_delay 

475 

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

477 if attempt > 0: 

478 # 等待重试延迟 

479 await asyncio.sleep(delay) 

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

481 

482 # 执行工具 

483 result = await base_executor.execute( 

484 tool=tool, 

485 arguments=arguments, 

486 call_id=f"{call_id}_attempt{attempt}" 

487 ) 

488 

489 if result.error is None: 

490 # 执行成功 

491 return result 

492 

493 # 检查是否可重试 

494 last_result = result 

495 error_category = ToolErrorClassifier.classify(result) 

496 

497 if error_category not in self.retryable_categories: 

498 # 不可重试的错误 

499 break 

500 

501 if attempt == self.max_retries: 

502 # 达到最大重试次数 

503 break 

504 

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

506 return last_result or ToolResult.fail( 

507 call_id=call_id, 

508 error="Execution failed after retries" 

509 ) 

510 

511 

512# 便捷函数 

513async def execute_tool_with_retry( 

514 tool: BaseTool, 

515 arguments: Dict[str, Any], 

516 max_retries: int = 3, 

517 call_id: Optional[str] = None 

518) -> ToolResult: 

519 """ 

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

521  

522 Args: 

523 tool: 要执行的工具 

524 arguments: 工具参数 

525 max_retries: 最大重试次数 

526 call_id: 调用ID 

527  

528 Returns: 

529 ToolResult: 执行结果 

530 """ 

531 executor = SmartRetryExecutor(max_retries=max_retries) 

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

533 

534 

535async def execute_tools_concurrently( 

536 tool_calls: List[Tuple[BaseTool, Dict[str, Any]]], 

537 max_concurrent: int = 10, 

538 timeout_per_tool: Optional[float] = None 

539) -> List[ToolResult]: 

540 """ 

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

542  

543 Args: 

544 tool_calls: 工具调用列表 

545 max_concurrent: 最大并发数 

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

547  

548 Returns: 

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

550 """ 

551 executor = AsyncToolExecutor(max_concurrent=max_concurrent) 

552 return await executor.execute_batch( 

553 tool_calls=tool_calls, 

554 timeout_per_tool=timeout_per_tool 

555 )