Coverage for agentos/tools/orchestrator.py: 33%
300 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 23:17 +0800
1"""
2AgentOS v1.1.7 — 工具链编排引擎（Checkpoint/恢复）。
3基因来源: Airflow DAG + LangChain Tool Composition
5支持:
6- 顺序链 (chain): 工具A → 工具B → 工具C
7- 并行分支 (parallel): A + B 同时 → C
8- 条件执行 (conditional): if X then A else B
9- 重试策略 (retry): 指数退避 / 固定间隔
10- 超时控制 (timeout): 单工具 / 全链
11- Checkpoint/恢复: 长时间DAG断点保存与续跑
12"""
14from __future__ import annotations
16import asyncio
17import json
18import time
19from collections.abc import Callable
20from dataclasses import dataclass, field
21from enum import StrEnum
22from typing import Any
24# ── Core Types ──────────────────────────────────
27class NodeState(StrEnum):
28 """DAG 节点状态。"""
30 PENDING = "pending"
31 RUNNING = "running"
32 SUCCESS = "success"
33 FAILED = "failed"
34 SKIPPED = "skipped"
35 TIMEOUT = "timeout"
38@dataclass
39class NodeResult:
40 """DAG 节点执行结果。"""
42 node_id: str
43 state: NodeState
44 output: Any = None
45 error: str | None = None
46 duration_ms: float = 0.0
47 retries: int = 0
50@dataclass
51class DAGResult:
52 """DAG 执行结果。"""
54 nodes: dict[str, NodeResult]
55 final_output: Any = None
56 total_duration_ms: float = 0.0
57 success: bool = False
58 error: str | None = None
61@dataclass
62class RetryPolicy:
63 """重试策略类。"""
65 max_retries: int = 3
66 base_delay: float = 1.0 # seconds
67 max_delay: float = 30.0 # seconds
68 backoff: str = "exponential" # exponential | fixed | linear
69 retry_on: tuple = (Exception,)
72# ── DAG Node Types ─────────────────────────────
75@dataclass
76class ToolNode:
77 """工具执行节点。"""
79 tool_name: str
80 tool_args: dict[str, Any] = field(default_factory=dict)
81 depends_on: list[str] = field(default_factory=list) # upstream node IDs
82 timeout: float = 60.0
83 retry: RetryPolicy | None = None
85 # Optional transform: map upstream outputs → tool_args
86 input_transform: Callable[[dict[str, Any]], dict[str, Any]] | None = None
89@dataclass
90class ConditionNode:
91 """条件分支节点。"""
93 condition: Callable[[dict[str, Any]], str] # returns target node_id
94 depends_on: list[str] = field(default_factory=list)
97@dataclass
98class ParallelGroup:
99 """并行执行组 — 所有节点同时执行。"""
101 node_ids: list[str]
102 depends_on: list[str] = field(default_factory=list)
103 max_concurrency: int = 5
106@dataclass
107class DAGSpec:
108 """DAG编排规格。"""
110 name: str
111 nodes: dict[str, ToolNode] = field(default_factory=dict)
112 parallels: list[ParallelGroup] = field(default_factory=list)
113 conditions: dict[str, ConditionNode] = field(default_factory=dict)
114 entry: list[str] = field(default_factory=list)
115 global_timeout: float = 300.0
118# ── Checkpoint Data (v1.1.7) ──────────────────
121@dataclass
122class CheckpointData:
123 """DAG执行快照，支持断点续跑。"""
125 dag_name: str
126 completed_nodes: dict[str, dict] = field(default_factory=dict)
127 pending_nodes: list[str] = field(default_factory=list)
128 timestamp: float = 0.0
129 version: str = "1.0"
131 def to_dict(self) -> dict:
132 return {
133 "dag_name": self.dag_name,
134 "completed_nodes": self.completed_nodes,
135 "pending_nodes": self.pending_nodes,
136 "timestamp": self.timestamp,
137 "version": self.version,
138 }
140 @classmethod
141 def from_dict(cls, data: dict) -> CheckpointData:
142 return cls(
143 dag_name=data.get("dag_name", ""),
144 completed_nodes=data.get("completed_nodes", {}),
145 pending_nodes=data.get("pending_nodes", []),
146 timestamp=data.get("timestamp", 0.0),
147 version=data.get("version", "1.0"),
148 )
150 def to_json(self) -> str:
151 return json.dumps(self.to_dict(), ensure_ascii=False, indent=2)
153 @classmethod
154 def from_json(cls, json_str: str) -> CheckpointData:
155 return cls.from_dict(json.loads(json_str))
158# ── Orchestrator Engine ────────────────────────
161class ToolOrchestrator:
162 """
163 工具链编排引擎 — DAG执行、并行调度、条件分支、Checkpoint恢复。
164 """
166 def __init__(self, tool_registry: Any):
167 self.registry = tool_registry
168 self._results: dict[str, NodeResult] = {}
169 self._aborted: bool = False
171 async def execute(self, dag: DAGSpec) -> DAGResult:
172 """执行完整DAG。若从checkpoint恢复，会跳过已完成节点。"""
173 was_restored = bool(self._results)
174 if not was_restored:
175 self._results = {}
176 self._aborted = False
177 start = time.time()
179 try:
180 await asyncio.wait_for(
181 self._execute_entry(dag),
182 timeout=dag.global_timeout,
183 )
184 except TimeoutError:
185 return DAGResult(
186 nodes=self._results,
187 total_duration_ms=(time.time() - start) * 1000,
188 error=f"DAG timeout ({dag.global_timeout}s)",
189 )
191 duration_ms = (time.time() - start) * 1000
192 success = all(
193 r.state == NodeState.SUCCESS
194 for r in self._results.values()
195 if not self._is_conditional(dag, r.node_id)
196 )
198 # Final output = last entry node (or last successful)
199 last = None
200 for nid in reversed(dag.entry):
201 if nid in self._results and self._results[nid].state == NodeState.SUCCESS:
202 last = self._results[nid].output
203 break
205 return DAGResult(
206 nodes=self._results,
207 final_output=last,
208 total_duration_ms=duration_ms,
209 success=success,
210 )
212 async def _execute_entry(self, dag: DAGSpec):
213 """递归执行入口节点。"""
214 pending = set(dag.entry or dag.nodes.keys())
215 # v1.1.7: 跳过checkpoint恢复后已完成的节点，并发现其下游节点
216 completed = {
217 nid
218 for nid in pending
219 if nid in self._results and self._results[nid].state == NodeState.SUCCESS
220 }
221 pending -= completed
222 for nid in completed:
223 for other_nid, other_node in dag.nodes.items():
224 if (
225 nid in other_node.depends_on
226 and other_nid not in pending
227 and other_nid not in self._results
228 ):
229 pending.add(other_nid)
231 while pending:
232 # Find nodes ready to execute (all deps satisfied)
233 ready = []
234 for nid in list(pending):
235 node = dag.nodes.get(nid)
236 if not node:
237 continue
238 if self._deps_ready(node.depends_on):
239 ready.append(nid)
241 if not ready:
242 # Check for deadlocks
243 stuck = [nid for nid in pending if not self._can_proceed(dag, nid)]
244 if stuck:
245 for nid in stuck:
246 self._results[nid] = NodeResult(
247 node_id=nid,
248 state=NodeState.FAILED,
249 error="Deadlock: dependencies not met",
250 )
251 break
252 await asyncio.sleep(0.01)
253 continue
255 # Execute ready nodes (parallel groups first)
256 parallel_nodes = []
257 sequential_nodes = []
258 for nid in ready:
259 if any(nid in pg.node_ids for pg in dag.parallels):
260 parallel_nodes.append(nid)
261 else:
262 sequential_nodes.append(nid)
264 # Run sequential nodes concurrently
265 tasks = []
266 for nid in sequential_nodes:
267 tasks.append(self._run_node(dag, nid))
268 if tasks:
269 await asyncio.gather(*tasks)
271 # Run parallel groups
272 for pg in dag.parallels:
273 group_ready = [nid for nid in pg.node_ids if nid in ready]
274 if group_ready:
275 sem = asyncio.Semaphore(pg.max_concurrency)
277 async def bounded(nid):
278 async with sem:
279 await self._run_node(dag, nid)
281 await asyncio.gather(*[bounded(nid) for nid in group_ready])
283 pending -= set(ready)
285 # v1.1.7: 发现已完成节点的下游节点（支持checkpoint恢复 & 多节点链）
286 for nid in list(ready):
287 for other_nid, other_node in dag.nodes.items():
288 if (
289 nid in other_node.depends_on
290 and other_nid not in pending
291 and other_nid not in self._results
292 ):
293 pending.add(other_nid)
295 # Process conditions
296 for cid, cond in dag.conditions.items():
297 if self._deps_ready(cond.depends_on):
298 upstream = {nid: self._results.get(nid) for nid in cond.depends_on}
299 target = cond.condition(upstream)
300 if target and target in dag.nodes:
301 pending.add(target)
303 async def _run_node(self, dag: DAGSpec, nid: str) -> NodeResult:
304 """执行单个节点（带重试）。"""
305 node = dag.nodes.get(nid)
306 if not node:
307 return NodeResult(nid, NodeState.FAILED, error=f"Unknown node: {nid}")
309 # Gather upstream outputs
310 upstream = {}
311 for dep in node.depends_on:
312 dep_result = self._results.get(dep)
313 if dep_result and dep_result.state == NodeState.SUCCESS:
314 upstream[dep] = dep_result.output
316 # Transform inputs if needed
317 args = dict(node.tool_args)
318 if node.input_transform and upstream:
319 try:
320 transformed = node.input_transform(upstream)
321 args.update(transformed)
322 except Exception as e:
323 result = NodeResult(nid, NodeState.FAILED, error=f"Input transform error: {e}")
324 self._results[nid] = result
325 return result
327 # Execute with retry
328 retry_policy = node.retry or RetryPolicy(max_retries=0)
329 last_error = None
331 for attempt in range(retry_policy.max_retries + 1):
332 try:
333 step_start = time.time()
334 output = await asyncio.wait_for(
335 self._execute_tool(node.tool_name, args, upstream),
336 timeout=node.timeout,
337 )
338 duration_ms = (time.time() - step_start) * 1000
339 result = NodeResult(
340 node_id=nid,
341 state=NodeState.SUCCESS,
342 output=output,
343 duration_ms=duration_ms,
344 retries=attempt,
345 )
346 self._results[nid] = result
347 return result
348 except TimeoutError:
349 last_error = f"Tool timeout ({node.timeout}s)"
350 result = NodeResult(nid, NodeState.TIMEOUT, error=last_error, retries=attempt)
351 except Exception as e:
352 last_error = str(e)
353 if attempt < retry_policy.max_retries:
354 delay = self._calc_retry_delay(retry_policy, attempt)
355 await asyncio.sleep(delay)
357 result = NodeResult(
358 nid, NodeState.FAILED, error=last_error, retries=retry_policy.max_retries
359 )
360 self._results[nid] = result
361 return result
363 async def _execute_tool(
364 self,
365 tool_name: str,
366 args: dict,
367 upstream: dict[str, Any],
368 ) -> Any:
369 """执行具体工具。"""
370 tool = self.registry.get(tool_name)
371 if not tool:
372 raise ValueError(f"Tool not found: {tool_name}")
374 # Inject upstream results into args
375 full_args = {**args}
376 for dep_id, dep_output in upstream.items():
377 full_args[f"_{dep_id}"] = dep_output
378 full_args["_upstream"] = upstream
380 if asyncio.iscoroutinefunction(tool.execute):
381 return await tool.execute(**full_args)
382 else:
383 return tool.execute(**full_args)
385 def _deps_ready(self, deps: list[str]) -> bool:
386 """所有依赖是否成功完成。"""
387 for dep in deps:
388 r = self._results.get(dep)
389 if not r or r.state != NodeState.SUCCESS:
390 return False
391 return True
393 def _can_proceed(self, dag: DAGSpec, nid: str) -> bool:
394 """节点是否有可能继续执行（未永久失败）。"""
395 node = dag.nodes.get(nid)
396 if not node:
397 return False
398 for dep in node.depends_on:
399 r = self._results.get(dep)
400 if r and r.state in (NodeState.FAILED, NodeState.TIMEOUT):
401 return False
402 return True
404 def _is_conditional(self, dag: DAGSpec, nid: str) -> bool:
405 return nid in dag.conditions
407 def _calc_retry_delay(self, policy: RetryPolicy, attempt: int) -> float:
408 if policy.backoff == "fixed":
409 return policy.base_delay
410 elif policy.backoff == "linear":
411 return min(policy.base_delay * (attempt + 1), policy.max_delay)
412 else: # exponential
413 return min(policy.base_delay * (2**attempt), policy.max_delay)
415 @property
416 def results(self) -> dict[str, NodeResult]:
417 return dict(self._results)
419 # ── Checkpoint / Restore (v1.1.7) ──────────────
421 def checkpoint(self, dag: DAGSpec) -> CheckpointData:
422 """保存当前DAG执行进度为快照。"""
423 completed = {}
424 for nid, result in self._results.items():
425 completed[nid] = {
426 "node_id": result.node_id,
427 "state": result.state.value,
428 "output": result.output,
429 "error": result.error,
430 "duration_ms": result.duration_ms,
431 "retries": result.retries,
432 }
433 # 未完成的节点（在dag中但不在results里）
434 pending = [nid for nid in dag.nodes if nid not in self._results]
435 return CheckpointData(
436 dag_name=dag.name,
437 completed_nodes=completed,
438 pending_nodes=pending,
439 timestamp=time.time(),
440 )
442 def restore_from_checkpoint(self, dag: DAGSpec, cp: CheckpointData) -> dict[str, NodeResult]:
443 """从快照恢复已完成的节点状态，返回可继续执行的results基础。"""
444 restored = {}
445 for nid, data in cp.completed_nodes.items():
446 restored[nid] = NodeResult(
447 node_id=data["node_id"],
448 state=NodeState(data["state"]),
449 output=data.get("output"),
450 error=data.get("error"),
451 duration_ms=data.get("duration_ms", 0),
452 retries=data.get("retries", 0),
453 )
454 self._results = restored
455 return restored
457 async def execute_with_checkpoint(
458 self,
459 dag: DAGSpec,
460 checkpoint_callback: Callable[[CheckpointData], None] = None,
461 checkpoint_interval: float = 60.0,
462 ) -> dict[str, NodeResult]:
463 """
464 执行DAG并周期保存快照。超时或异常时保留已执行结果。
466 Args:
467 dag: DAG规格
468 checkpoint_callback: 快照回调，收到最新的CheckpointData
469 checkpoint_interval: 快照保存间隔（秒）
470 Returns:
471 最终执行结果
472 """
473 try:
474 await self.execute(dag)
475 except (TimeoutError, Exception):
476 # 异常时保存当前状态
477 cp = self.checkpoint(dag)
478 if checkpoint_callback:
479 checkpoint_callback(cp)
480 raise
481 else:
482 # 最终完成快照
483 cp = self.checkpoint(dag)
484 if checkpoint_callback:
485 checkpoint_callback(cp)
486 return self._results
489# ── DAG Builder (Fluent API) ────────────────────
492class DAGBuilder:
493 """流式构建DAG。"""
495 def __init__(self, name: str = "unnamed"):
496 self.name = name
497 self._nodes: dict[str, ToolNode] = {}
498 self._parallels: list[ParallelGroup] = []
499 self._conditions: dict[str, ConditionNode] = {}
500 self._entry: list[str] = []
502 def node(
503 self,
504 node_id: str,
505 tool_name: str,
506 tool_args: dict | None = None,
507 depends_on: list[str] | None = None,
508 timeout: float = 60.0,
509 retry: RetryPolicy | None = None,
510 input_transform: Callable | None = None,
511 ) -> DAGBuilder:
512 self._nodes[node_id] = ToolNode(
513 tool_name=tool_name,
514 tool_args=tool_args or {},
515 depends_on=depends_on or [],
516 timeout=timeout,
517 retry=retry,
518 input_transform=input_transform,
519 )
520 if not depends_on:
521 self._entry.append(node_id)
522 return self
524 def parallel(
525 self, node_ids: list[str], depends_on: list[str] | None = None, max_concurrency: int = 5
526 ) -> DAGBuilder:
527 self._parallels.append(
528 ParallelGroup(
529 node_ids=node_ids,
530 depends_on=depends_on or [],
531 max_concurrency=max_concurrency,
532 )
533 )
534 return self
536 def condition(self, cond_id: str, condition: Callable, depends_on: list[str]) -> DAGBuilder:
537 self._conditions[cond_id] = ConditionNode(
538 condition=condition,
539 depends_on=depends_on,
540 )
541 return self
543 def build(self, global_timeout: float = 300.0) -> DAGSpec:
544 return DAGSpec(
545 name=self.name,
546 nodes=self._nodes,
547 parallels=self._parallels,
548 conditions=self._conditions,
549 entry=self._entry,
550 global_timeout=global_timeout,
551 )
554# ── Pre-built Chains ────────────────────────────
557def chain_builder(name: str, tool_names: list[str]) -> DAGSpec:
558 """构建简单顺序链。"""
559 builder = DAGBuilder(name)
560 for i, tool_name in enumerate(tool_names):
561 nid = f"step_{i}"
562 deps = [f"step_{i - 1}"] if i > 0 else []
563 builder.node(nid, tool_name, depends_on=deps)
564 return builder.build()
567def parallel_then_merge(name: str, parallel_tools: list[str], merge_tool: str) -> DAGSpec:
568 """构建 并行→合并 模式。"""
569 builder = DAGBuilder(name)
570 pids = []
571 for i, tool_name in enumerate(parallel_tools):
572 nid = f"par_{i}"
573 builder.node(nid, tool_name)
574 pids.append(nid)
575 builder.parallel(pids)
576 builder.node("merge", merge_tool, depends_on=pids)
577 return builder.build()
580def if_then_else(name: str, check_tool: str, true_tool: str, false_tool: str) -> DAGSpec:
581 """构建 if-then-else 条件分支。"""
582 builder = DAGBuilder(name)
583 builder.node("check", check_tool)
584 builder.condition(
585 "cond",
586 lambda up: "true_branch" if up.get("check", {}).get("output") else "false_branch",
587 depends_on=["check"],
588 )
589 builder.node("true_branch", true_tool, depends_on=["check"])
590 builder.node("false_branch", false_tool, depends_on=["check"])
591 return builder.build()