Coverage for agentos/tools/fusion.py: 32%

1""" 

2Fusion Toolkit for NexusAgent. 

3 

4Multi-tool coordination system. Allows agents to use 

5multiple tools in sequence or parallel, with automatic 

6result fusion and conflict resolution. 

7""" 

8 

9from __future__ import annotations 

10 

11import asyncio 

12import time 

13import uuid 

14from collections.abc import Callable 

15from dataclasses import dataclass, field 

16from enum import StrEnum 

17from typing import Any 

18 

19 

20class FusionMode(StrEnum): 

21 """Tool fusion modes.""" 

22 

23 SEQUENTIAL = "sequential" # Run tools one by one 

24 PARALLEL = "parallel" # Run tools in parallel 

25 CHAIN = "chain" # Output of one feeds into next 

26 

27 

28@dataclass 

29class ToolSpec: 

30 """ 

31 Tool specification. 

32 

33 Attributes: 

34 name: Tool name 

35 description: Tool description 

36 func: Tool function 

37 parameters: Parameter schema 

38 timeout: Execution timeout 

39 retry_count: Number of retries 

40 """ 

41 

42 name: str 

43 description: str = "" 

44 func: Callable[..., Any] = None 

45 parameters: dict[str, Any] = field(default_factory=dict) 

46 timeout: float = 30.0 

47 retry_count: int = 0 

48 

49 def to_dict(self) -> dict[str, Any]: 

50 """Convert to dict.""" 

51 return { 

52 "name": self.name, 

53 "description": self.description, 

54 "parameters": self.parameters, 

55 "timeout": self.timeout, 

56 "retry_count": self.retry_count, 

57 } 

58 

59 

60@dataclass 

61class ToolResult: 

62 """ 

63 Result of a single tool execution. 

64 

65 Attributes: 

66 tool_name: Name of the tool 

67 success: Whether execution succeeded 

68 output: Tool output 

69 error: Error message (if failed) 

70 duration: Execution duration 

71 """ 

72 

73 tool_name: str 

74 success: bool 

75 output: Any = None 

76 error: str | None = None 

77 duration: float = 0.0 

78 

79 

80@dataclass 

81class FusionResult: 

82 """ 

83 Result of tool fusion. 

84 

85 Attributes: 

86 id: Unique identifier 

87 mode: Fusion mode used 

88 results: List of individual tool results 

89 fused_output: Fused final output 

90 total_duration: Total execution duration 

91 success: Whether fusion succeeded 

92 """ 

93 

94 id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) 

95 mode: FusionMode = FusionMode.SEQUENTIAL 

96 results: list[ToolResult] = field(default_factory=list) 

97 fused_output: Any = None 

98 total_duration: float = 0.0 

99 success: bool = True 

100 

101 def to_dict(self) -> dict[str, Any]: 

102 """Convert to dict.""" 

103 return { 

104 "id": self.id, 

105 "mode": self.mode.value, 

106 "results": [ 

107 { 

108 "tool_name": r.tool_name, 

109 "success": r.success, 

110 "output": r.output, 

111 "error": r.error, 

112 "duration": r.duration, 

113 } 

114 for r in self.results 

115 ], 

116 "fused_output": self.fused_output, 

117 "total_duration": self.total_duration, 

118 "success": self.success, 

119 } 

120 

121 

122class FusionToolkit: 

123 """ 

124 Multi-tool coordination system. 

125 

126 Allows agents to use multiple tools in different modes: 

127 - Sequential: Run tools one by one 

128 - Parallel: Run tools in parallel 

129 - Chain: Output of one feeds into next 

130 

131 Usage: 

132 toolkit = FusionToolkit() 

133 toolkit.register(ToolSpec(name="search", func=search_func)) 

134 toolkit.register(ToolSpec(name="summarize", func=summarize_func)) 

135 

136 # Sequential execution 

137 result = await toolkit.execute(["search", "summarize"], {"query": "AI"}) 

138 

139 # Parallel execution 

140 result = await toolkit.execute_parallel(["search", "summarize"], {"query": "AI"}) 

141 """ 

142 

143 def __init__(self, default_timeout: float = 30.0): 

144 """ 

145 Initialize fusion toolkit. 

146 

147 Args: 

148 default_timeout: Default tool timeout 

149 """ 

150 self._tools: dict[str, ToolSpec] = {} 

151 self._default_timeout = default_timeout 

152 

153 def register(self, tool: ToolSpec) -> None: 

154 """ 

155 Register a tool. 

156 

157 Args: 

158 tool: Tool specification 

159 """ 

160 self._tools[tool.name] = tool 

161 

162 def unregister(self, tool_name: str) -> bool: 

163 """ 

164 Unregister a tool. 

165 

166 Args: 

167 tool_name: Tool name 

168 

169 Returns: 

170 True if unregistered, False if not found 

171 """ 

172 if tool_name in self._tools: 

173 del self._tools[tool_name] 

174 return True 

175 return False 

176 

177 def get_tool(self, tool_name: str) -> ToolSpec | None: 

178 """ 

179 Get a tool by name. 

180 

181 Args: 

182 tool_name: Tool name 

183 

184 Returns: 

185 ToolSpec if found, None otherwise 

186 """ 

187 return self._tools.get(tool_name) 

188 

189 def list_tools(self) -> list[ToolSpec]: 

190 """ 

191 List all registered tools. 

192 

193 Returns: 

194 List of ToolSpec 

195 """ 

196 return list(self._tools.values()) 

197 

198 async def execute( 

199 self, 

200 tool_names: list[str], 

201 inputs: dict[str, Any], 

202 mode: FusionMode = FusionMode.SEQUENTIAL, 

203 ) -> FusionResult: 

204 """ 

205 Execute multiple tools. 

206 

207 Args: 

208 tool_names: List of tool names 

209 inputs: Input parameters 

210 mode: Fusion mode 

211 

212 Returns: 

213 FusionResult 

214 """ 

215 start_time = time.time() 

216 

217 if mode == FusionMode.SEQUENTIAL: 

218 result = await self._execute_sequential(tool_names, inputs) 

219 elif mode == FusionMode.PARALLEL: 

220 result = await self._execute_parallel(tool_names, inputs) 

221 elif mode == FusionMode.CHAIN: 

222 result = await self._execute_chain(tool_names, inputs) 

223 else: 

224 raise ValueError(f"Unknown fusion mode: {mode}") 

225 

226 result.total_duration = time.time() - start_time 

227 

228 return result 

229 

230 async def _execute_sequential( 

231 self, 

232 tool_names: list[str], 

233 inputs: dict[str, Any], 

234 ) -> FusionResult: 

235 """Execute tools sequentially.""" 

236 result = FusionResult(mode=FusionMode.SEQUENTIAL) 

237 

238 for tool_name in tool_names: 

239 tool = self._tools.get(tool_name) 

240 if not tool: 

241 result.results.append( 

242 ToolResult( 

243 tool_name=tool_name, 

244 success=False, 

245 error=f"Tool not found: {tool_name}", 

246 ) 

247 ) 

248 result.success = False 

249 continue 

250 

251 try: 

252 tool_start = time.time() 

253 output = await self._execute_tool(tool, inputs) 

254 duration = time.time() - tool_start 

255 

256 result.results.append( 

257 ToolResult( 

258 tool_name=tool_name, 

259 success=True, 

260 output=output, 

261 duration=duration, 

262 ) 

263 ) 

264 except Exception as e: 

265 result.results.append( 

266 ToolResult( 

267 tool_name=tool_name, 

268 success=False, 

269 error=str(e), 

270 ) 

271 ) 

272 result.success = False 

273 

274 # Fuse outputs 

275 result.fused_output = self._fuse_outputs(result.results) 

276 

277 return result 

278 

279 async def _execute_parallel( 

280 self, 

281 tool_names: list[str], 

282 inputs: dict[str, Any], 

283 ) -> FusionResult: 

284 """Execute tools in parallel.""" 

285 result = FusionResult(mode=FusionMode.PARALLEL) 

286 

287 tasks = [] 

288 for tool_name in tool_names: 

289 tool = self._tools.get(tool_name) 

290 if tool: 

291 tasks.append(self._execute_tool_with_result(tool, inputs)) 

292 else: 

293 result.results.append( 

294 ToolResult( 

295 tool_name=tool_name, 

296 success=False, 

297 error=f"Tool not found: {tool_name}", 

298 ) 

299 ) 

300 

301 # Execute in parallel 

302 if tasks: 

303 tool_results = await asyncio.gather(*tasks, return_exceptions=True) 

304 for tr in tool_results: 

305 if isinstance(tr, Exception): 

306 result.results.append( 

307 ToolResult( 

308 tool_name="unknown", 

309 success=False, 

310 error=str(tr), 

311 ) 

312 ) 

313 result.success = False 

314 else: 

315 result.results.append(tr) 

316 if not tr.success: 

317 result.success = False 

318 

319 # Fuse outputs 

320 result.fused_output = self._fuse_outputs(result.results) 

321 

322 return result 

323 

324 async def _execute_chain( 

325 self, 

326 tool_names: list[str], 

327 inputs: dict[str, Any], 

328 ) -> FusionResult: 

329 """Execute tools in chain (output feeds into next).""" 

330 result = FusionResult(mode=FusionMode.CHAIN) 

331 current_input = inputs.copy() 

332 

333 for tool_name in tool_names: 

334 tool = self._tools.get(tool_name) 

335 if not tool: 

336 result.results.append( 

337 ToolResult( 

338 tool_name=tool_name, 

339 success=False, 

340 error=f"Tool not found: {tool_name}", 

341 ) 

342 ) 

343 result.success = False 

344 break 

345 

346 try: 

347 tool_start = time.time() 

348 output = await self._execute_tool(tool, current_input) 

349 duration = time.time() - tool_start 

350 

351 result.results.append( 

352 ToolResult( 

353 tool_name=tool_name, 

354 success=True, 

355 output=output, 

356 duration=duration, 

357 ) 

358 ) 

359 

360 # Feed output into next tool 

361 current_input = {"input": output, **inputs} 

362 except Exception as e: 

363 result.results.append( 

364 ToolResult( 

365 tool_name=tool_name, 

366 success=False, 

367 error=str(e), 

368 ) 

369 ) 

370 result.success = False 

371 break 

372 

373 # Final output is last tool's output 

374 if result.results: 

375 last_result = result.results[-1] 

376 if last_result.success: 

377 result.fused_output = last_result.output 

378 

379 return result 

380 

381 async def _execute_tool( 

382 self, 

383 tool: ToolSpec, 

384 inputs: dict[str, Any], 

385 ) -> Any: 

386 """Execute a single tool.""" 

387 if not tool.func: 

388 raise ValueError(f"Tool {tool.name} has no function") 

389 

390 # Apply timeout 

391 try: 

392 if asyncio.iscoroutinefunction(tool.func): 

393 return await asyncio.wait_for( 

394 tool.func(**inputs), 

395 timeout=tool.timeout or self._default_timeout, 

396 ) 

397 else: 

398 return await asyncio.wait_for( 

399 asyncio.get_event_loop().run_in_executor(None, lambda: tool.func(**inputs)), 

400 timeout=tool.timeout or self._default_timeout, 

401 ) 

402 except TimeoutError: 

403 raise TimeoutError(f"Tool {tool.name} timed out") 

404 

405 async def _execute_tool_with_result( 

406 self, 

407 tool: ToolSpec, 

408 inputs: dict[str, Any], 

409 ) -> ToolResult: 

410 """Execute tool and return ToolResult.""" 

411 try: 

412 tool_start = time.time() 

413 output = await self._execute_tool(tool, inputs) 

414 duration = time.time() - tool_start 

415 

416 return ToolResult( 

417 tool_name=tool.name, 

418 success=True, 

419 output=output, 

420 duration=duration, 

421 ) 

422 except Exception as e: 

423 return ToolResult( 

424 tool_name=tool.name, 

425 success=False, 

426 error=str(e), 

427 ) 

428 

429 def _fuse_outputs(self, results: list[ToolResult]) -> Any: 

430 """Fuse multiple tool outputs.""" 

431 outputs = [r.output for r in results if r.success and r.output is not None] 

432 

433 if not outputs: 

434 return None 

435 

436 if len(outputs) == 1: 

437 return outputs[0] 

438 

439 # Default fusion: merge dicts, concatenate lists 

440 if all(isinstance(o, dict) for o in outputs): 

441 fused = {} 

442 for o in outputs: 

443 fused.update(o) 

444 return fused 

445 

446 if all(isinstance(o, list) for o in outputs): 

447 return [item for o in outputs for item in o] 

448 

449 # Default: return list of outputs 

450 return outputs