Coverage for agentos/orchestration/graph.py: 31%

1""" 

2Graph Orchestrator for NexusAgent. 

3 

4DAG-based workflow orchestration. Allows defining 

5complex workflows as graphs with nodes and edges. 

6""" 

7 

8from __future__ import annotations 

9 

10import asyncio 

11import time 

12import uuid 

13from collections.abc import Callable 

14from dataclasses import dataclass, field 

15from enum import StrEnum 

16from typing import Any 

17 

18 

19class NodeStatus(StrEnum): 

20 """Node execution status.""" 

21 

22 PENDING = "pending" 

23 RUNNING = "running" 

24 COMPLETED = "completed" 

25 FAILED = "failed" 

26 SKIPPED = "skipped" 

27 

28 

29@dataclass 

30class GraphNode: 

31 """ 

32 Node in execution graph. 

33 

34 Attributes: 

35 id: Unique identifier 

36 name: Node name 

37 func: Node function 

38 inputs: Input parameters 

39 outputs: Output values 

40 status: Execution status 

41 duration: Execution duration 

42 error: Error message (if failed) 

43 metadata: Additional metadata 

44 """ 

45 

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

47 name: str = "" 

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

49 inputs: dict[str, Any] = field(default_factory=dict) 

50 outputs: dict[str, Any] = field(default_factory=dict) 

51 status: NodeStatus = NodeStatus.PENDING 

52 duration: float = 0.0 

53 error: str | None = None 

54 metadata: dict[str, Any] = field(default_factory=dict) 

55 

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

57 """Convert to dict.""" 

58 return { 

59 "id": self.id, 

60 "name": self.name, 

61 "inputs": self.inputs, 

62 "outputs": self.outputs, 

63 "status": self.status.value, 

64 "duration": self.duration, 

65 "error": self.error, 

66 "metadata": self.metadata, 

67 } 

68 

69 

70@dataclass 

71class GraphEdge: 

72 """ 

73 Edge in execution graph. 

74 

75 Attributes: 

76 source: Source node ID 

77 target: Target node ID 

78 condition: Optional condition function 

79 metadata: Additional metadata 

80 """ 

81 

82 source: str 

83 target: str 

84 condition: Callable[[dict[str, Any]], bool] | None = None 

85 metadata: dict[str, Any] = field(default_factory=dict) 

86 

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

88 """Convert to dict.""" 

89 return { 

90 "source": self.source, 

91 "target": self.target, 

92 "metadata": self.metadata, 

93 } 

94 

95 

96@dataclass 

97class GraphResult: 

98 """ 

99 Result of graph execution. 

100 

101 Attributes: 

102 id: Unique identifier 

103 node_results: Node execution results 

104 total_duration: Total execution duration 

105 success: Whether execution succeeded 

106 error: Error message (if failed) 

107 """ 

108 

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

110 node_results: dict[str, dict[str, Any]] = field(default_factory=dict) 

111 total_duration: float = 0.0 

112 success: bool = True 

113 error: str | None = None 

114 

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

116 """Convert to dict.""" 

117 return { 

118 "id": self.id, 

119 "node_results": self.node_results, 

120 "total_duration": self.total_duration, 

121 "success": self.success, 

122 "error": self.error, 

123 } 

124 

125 

126class GraphOrchestrator: 

127 """ 

128 DAG-based workflow orchestrator. 

129 

130 Allows defining complex workflows as graphs: 

131 - Nodes represent tasks 

132 - Edges represent dependencies 

133 - Conditions for branching 

134 

135 Usage: 

136 orchestrator = GraphOrchestrator() 

137 

138 # Add nodes 

139 orchestrator.add_node("step1", step1_func) 

140 orchestrator.add_node("step2", step2_func) 

141 

142 # Add edges 

143 orchestrator.add_edge("step1", "step2") 

144 

145 # Execute 

146 result = await orchestrator.execute({"input": "data"}) 

147 """ 

148 

149 def __init__(self): 

150 """Initialize graph orchestrator.""" 

151 self._nodes: dict[str, GraphNode] = {} 

152 self._edges: list[GraphEdge] = [] 

153 self._start_nodes: list[str] = [] 

154 self._end_nodes: list[str] = [] 

155 

156 def add_node(self, name: str, func: Callable[..., Any], **metadata) -> GraphNode: 

157 """ 

158 Add a node to the graph. 

159 

160 Args: 

161 name: Node name 

162 func: Node function 

163 **metadata: Additional metadata 

164 

165 Returns: 

166 Created GraphNode 

167 """ 

168 node = GraphNode( 

169 name=name, 

170 func=func, 

171 metadata=metadata, 

172 ) 

173 self._nodes[name] = node 

174 

175 # If first node, mark as start 

176 if len(self._nodes) == 1: 

177 self._start_nodes.append(name) 

178 

179 return node 

180 

181 def remove_node(self, name: str) -> bool: 

182 """ 

183 Remove a node from the graph. 

184 

185 Args: 

186 name: Node name 

187 

188 Returns: 

189 True if removed, False if not found 

190 """ 

191 if name not in self._nodes: 

192 return False 

193 

194 del self._nodes[name] 

195 

196 # Remove edges 

197 self._edges = [e for e in self._edges if e.source != name and e.target != name] 

198 

199 # Update start/end nodes 

200 if name in self._start_nodes: 

201 self._start_nodes.remove(name) 

202 if name in self._end_nodes: 

203 self._end_nodes.remove(name) 

204 

205 return True 

206 

207 def add_edge( 

208 self, 

209 source: str, 

210 target: str, 

211 condition: Callable[[dict[str, Any]], bool] | None = None, 

212 **metadata, 

213 ) -> GraphEdge: 

214 """ 

215 Add an edge to the graph. 

216 

217 Args: 

218 source: Source node name 

219 target: Target node name 

220 condition: Optional condition function 

221 **metadata: Additional metadata 

222 

223 Returns: 

224 Created GraphEdge 

225 """ 

226 if source not in self._nodes: 

227 raise ValueError(f"Source node not found: {source}") 

228 if target not in self._nodes: 

229 raise ValueError(f"Target node not found: {target}") 

230 

231 edge = GraphEdge( 

232 source=source, 

233 target=target, 

234 condition=condition, 

235 metadata=metadata, 

236 ) 

237 self._edges.append(edge) 

238 

239 # Update start/end nodes 

240 if target in self._start_nodes: 

241 self._start_nodes.remove(target) 

242 if source in self._end_nodes: 

243 self._end_nodes.remove(source) 

244 

245 if source not in [e.target for e in self._edges]: 

246 if source not in self._start_nodes: 

247 self._start_nodes.append(source) 

248 

249 if target not in [e.source for e in self._edges]: 

250 if target not in self._end_nodes: 

251 self._end_nodes.append(target) 

252 

253 return edge 

254 

255 def remove_edge(self, source: str, target: str) -> bool: 

256 """ 

257 Remove an edge from the graph. 

258 

259 Args: 

260 source: Source node name 

261 target: Target node name 

262 

263 Returns: 

264 True if removed, False if not found 

265 """ 

266 for edge in self._edges: 

267 if edge.source == source and edge.target == target: 

268 self._edges.remove(edge) 

269 return True 

270 return False 

271 

272 def get_node(self, name: str) -> GraphNode | None: 

273 """ 

274 Get a node by name. 

275 

276 Args: 

277 name: Node name 

278 

279 Returns: 

280 GraphNode if found, None otherwise 

281 """ 

282 return self._nodes.get(name) 

283 

284 def list_nodes(self) -> list[str]: 

285 """ 

286 List all nodes. 

287 

288 Returns: 

289 List of node names 

290 """ 

291 return list(self._nodes.keys()) 

292 

293 def list_edges(self) -> list[tuple[str, str]]: 

294 """ 

295 List all edges. 

296 

297 Returns: 

298 List of (source, target) tuples 

299 """ 

300 return [(e.source, e.target) for e in self._edges] 

301 

302 async def execute(self, inputs: dict[str, Any], **metadata) -> GraphResult: 

303 """ 

304 Execute the graph. 

305 

306 Args: 

307 inputs: Input parameters 

308 **metadata: Additional metadata 

309 

310 Returns: 

311 GraphResult 

312 """ 

313 start_time = time.time() 

314 result = GraphResult() 

315 

316 # Reset node status 

317 for node in self._nodes.values(): 

318 node.status = NodeStatus.PENDING 

319 node.outputs = {} 

320 node.error = None 

321 

322 # Execute start nodes 

323 try: 

324 await self._execute_nodes(self._start_nodes, inputs, result, metadata) 

325 

326 # Execute remaining nodes in topological order 

327 executed = set(self._start_nodes) 

328 while len(executed) < len(self._nodes): 

329 next_nodes = self._get_next_nodes(executed) 

330 if not next_nodes: 

331 break 

332 await self._execute_nodes(next_nodes, inputs, result, metadata) 

333 executed.update(next_nodes) 

334 

335 except Exception as e: 

336 result.success = False 

337 result.error = str(e) 

338 

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

340 

341 return result 

342 

343 async def _execute_nodes( 

344 self, 

345 node_names: list[str], 

346 inputs: dict[str, Any], 

347 result: GraphResult, 

348 metadata: dict[str, Any], 

349 ) -> None: 

350 """Execute multiple nodes.""" 

351 tasks = [] 

352 for name in node_names: 

353 node = self._nodes.get(name) 

354 if node: 

355 tasks.append(self._execute_node(node, inputs, result, metadata)) 

356 

357 if tasks: 

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

359 

360 async def _execute_node( 

361 self, 

362 node: GraphNode, 

363 inputs: dict[str, Any], 

364 result: GraphResult, 

365 metadata: dict[str, Any], 

366 ) -> None: 

367 """Execute a single node.""" 

368 # Check conditions 

369 for edge in self._edges: 

370 if edge.target == node.name and edge.condition: 

371 if not edge.condition(inputs): 

372 node.status = NodeStatus.SKIPPED 

373 result.node_results[node.name] = node.to_dict() 

374 return 

375 

376 # Execute node 

377 node.status = NodeStatus.RUNNING 

378 start_time = time.time() 

379 

380 try: 

381 if asyncio.iscoroutinefunction(node.func): 

382 output = await node.func(**inputs, **metadata) 

383 else: 

384 output = await asyncio.get_event_loop().run_in_executor( 

385 None, lambda: node.func(**inputs, **metadata) 

386 ) 

387 

388 node.outputs = output if isinstance(output, dict) else {"result": output} 

389 node.status = NodeStatus.COMPLETED 

390 node.duration = time.time() - start_time 

391 

392 except Exception as e: 

393 node.status = NodeStatus.FAILED 

394 node.error = str(e) 

395 node.duration = time.time() - start_time 

396 result.success = False 

397 

398 result.node_results[node.name] = node.to_dict() 

399 

400 # Update inputs for next nodes 

401 inputs.update(node.outputs) 

402 

403 def _get_next_nodes(self, executed: set[str]) -> list[str]: 

404 """Get next nodes to execute.""" 

405 next_nodes = [] 

406 

407 for edge in self._edges: 

408 if edge.source in executed and edge.target not in executed: 

409 # Check if all dependencies are executed 

410 deps = [e.source for e in self._edges if e.target == edge.target] 

411 if all(d in executed for d in deps): 

412 next_nodes.append(edge.target) 

413 

414 return next_nodes 

415 

416 def get_execution_order(self) -> list[str]: 

417 """ 

418 Get topological execution order. 

419 

420 Returns: 

421 List of node names in execution order 

422 """ 

423 order = [] 

424 visited = set() 

425 

426 def visit(node_name: str): 

427 if node_name in visited: 

428 return 

429 visited.add(node_name) 

430 

431 # Visit dependencies first 

432 for edge in self._edges: 

433 if edge.target == node_name: 

434 visit(edge.source) 

435 

436 order.append(node_name) 

437 

438 for node_name in self._nodes.keys(): 

439 visit(node_name) 

440 

441 return order 

442 

443 def validate(self) -> bool: 

444 """ 

445 Validate the graph. 

446 

447 Returns: 

448 True if valid, False otherwise 

449 """ 

450 # Check for cycles 

451 try: 

452 self.get_execution_order() 

453 except Exception: 

454 return False 

455 

456 # Check for disconnected nodes 

457 if not self._start_nodes or not self._end_nodes: 

458 return False 

459 

460 return True 

461 

462 def clear(self) -> None: 

463 """Clear the graph.""" 

464 self._nodes.clear() 

465 self._edges.clear() 

466 self._start_nodes.clear() 

467 self._end_nodes.clear()