Coverage for agentos/orchestration/task_decomposer.py: 28%

1""" 

2AgentOS v1.14.7 — Intelligent Task Decomposer 2.0. 

3 

4Replaces the simplistic 234-line single-prompt decomposer with: 

5- DAG cycle detection (Kahn's + DFS fallback) 

6- Dynamic re-planning on partial failure 

7- Task dependency validation 

8- Parallelism detection (independent sub-tasks) 

9- Confidence scoring per sub-task 

10- Observable decomposition trace 

11 

12Architecture: 

13 TaskInput → Decomposer.decompose() → TaskDAG 

14 Partial failure → Decomposer.replan(failed_node) → new_path 

15""" 

16 

17from __future__ import annotations 

18 

19import json 

20import logging 

21import uuid 

22from collections import deque 

23from collections.abc import Callable 

24from dataclasses import dataclass, field 

25from enum import StrEnum 

26from typing import Any 

27 

28logger = logging.getLogger(__name__) 

29 

30 

31# ── Types ──────────────────────────────────── 

32 

33 

34class TaskNodeStatus(StrEnum): 

35 PENDING = "pending" 

36 RUNNING = "running" 

37 COMPLETED = "completed" 

38 FAILED = "failed" 

39 SKIPPED = "skipped" 

40 

41 

42class DecompositionStrategy(StrEnum): 

43 """分解策略。""" 

44 

45 TOP_DOWN = "top_down" # 从目标逐层拆解 

46 BOTTOM_UP = "bottom_up" # 从子任务聚合 

47 RECURSIVE = "recursive" # 递归分解直到原子任务 

48 HEURISTIC = "heuristic" # 基于规则/模式匹配 

49 

50 

51@dataclass 

52class TaskEdge: 

53 """DAG 边：from_node 完成后才能执行 to_node。""" 

54 

55 from_node: str 

56 to_node: str 

57 dependency_type: str = "hard" # hard / soft 

58 

59 

60@dataclass 

61class TaskNode: 

62 """DAG 节点：单个可执行单元。""" 

63 

64 id: str = field(default_factory=lambda: f"tn-{uuid.uuid4().hex[:8]}") 

65 description: str = "" 

66 input_schema: dict = field(default_factory=dict) 

67 output_schema: dict = field(default_factory=dict) 

68 agent_type: str = "default" # 推荐执行 Agent 类型 

69 estimated_duration_s: float = 0.0 

70 confidence: float = 1.0 # 0~1，分解置信度 

71 retry_policy: str = "once" # once / retry_n / fallback 

72 max_retries: int = 1 

73 status: TaskNodeStatus = TaskNodeStatus.PENDING 

74 result: Any = None 

75 error: str = "" 

76 

77 

78@dataclass 

79class TaskDAG: 

80 """完整的任务 DAG。""" 

81 

82 dag_id: str = field(default_factory=lambda: f"dag-{uuid.uuid4().hex[:8]}") 

83 root_task: str = "" # 原始任务描述 

84 nodes: dict[str, TaskNode] = field(default_factory=dict) 

85 edges: list[TaskEdge] = field(default_factory=list) 

86 strategy: DecompositionStrategy = DecompositionStrategy.TOP_DOWN 

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

88 created_at: float = 0.0 

89 

90 def in_degree_map(self) -> dict[str, int]: 

91 """计算每个节点的入度。""" 

92 indeg: dict[str, int] = {nid: 0 for nid in self.nodes} 

93 for e in self.edges: 

94 indeg[e.to_node] = indeg.get(e.to_node, 0) + 1 

95 return indeg 

96 

97 def adjacency_map(self) -> dict[str, list[str]]: 

98 """邻接表。""" 

99 adj: dict[str, list[str]] = {nid: [] for nid in self.nodes} 

100 for e in self.edges: 

101 adj[e.from_node].append(e.to_node) 

102 return adj 

103 

104 def topological_order(self) -> list[str]: 

105 """Kahn 算法拓扑排序，遇循环抛 ValueError。""" 

106 indeg = self.in_degree_map() 

107 adj = self.adjacency_map() 

108 queue = deque([nid for nid, d in indeg.items() if d == 0]) 

109 order: list[str] = [] 

110 

111 while queue: 

112 node = queue.popleft() 

113 order.append(node) 

114 for neighbor in adj.get(node, []): 

115 indeg[neighbor] -= 1 

116 if indeg[neighbor] == 0: 

117 queue.append(neighbor) 

118 

119 if len(order) != len(self.nodes): 

120 remaining = set(self.nodes) - set(order) 

121 raise ValueError( 

122 f"Cycle detected in DAG. {len(remaining)} nodes in cycle: " 

123 f"{list(remaining)[:5]}..." 

124 ) 

125 

126 return order 

127 

128 def detect_cycles(self) -> set[str]: 

129 """检测并返回所有参与循环的节点 ID 集合。""" 

130 indeg = self.in_degree_map() 

131 adj = self.adjacency_map() 

132 queue = deque([nid for nid, d in indeg.items() if d == 0]) 

133 acyclic: set[str] = set() 

134 

135 while queue: 

136 node = queue.popleft() 

137 acyclic.add(node) 

138 for neighbor in adj.get(node, []): 

139 indeg[neighbor] -= 1 

140 if indeg[neighbor] == 0: 

141 queue.append(neighbor) 

142 

143 return set(self.nodes) - acyclic 

144 

145 def parallel_groups(self) -> list[list[str]]: 

146 """按拓扑层级分组，同一组内可并行执行。""" 

147 order = self.topological_order() 

148 indeg = self.in_degree_map() 

149 adj = self.adjacency_map() 

150 

151 groups: list[list[str]] = [] 

152 remaining = set(order) 

153 

154 while remaining: 

155 # 所有入度为 0 的当前批次 

156 batch = sorted([n for n in remaining if indeg.get(n, 0) == 0]) 

157 if not batch: 

158 break 

159 groups.append(batch) 

160 for n in batch: 

161 remaining.discard(n) 

162 for neighbor in adj.get(n, []): 

163 indeg[neighbor] -= 1 

164 

165 return groups 

166 

167 

168@dataclass 

169class DecompositionTrace: 

170 """分解过程可观测性记录。""" 

171 

172 iteration: int 

173 action: str # split / merge / refine / replan 

174 node_before: TaskNode | None = None 

175 nodes_after: list[TaskNode] = field(default_factory=list) 

176 reason: str = "" 

177 

178 

179# ── Decomposer ─────────────────────────────── 

180 

181 

182class TaskDecomposer: 

183 """智能任务分解器。 

184 

185 核心能力： 

186 1. 将复杂任务拆解为可执行的 DAG 

187 2. 检测并拒绝循环依赖 

188 3. 在部分失败时动态重规划 

189 4. 输出可观测的分解轨迹 

190 

191 Usage: 

192 decomposer = TaskDecomposer() 

193 dag = decomposer.decompose("从 10GB 日志中提取异常并生成日报") 

194 order = dag.topological_order() 

195 for nid in order: 

196 execute(dag.nodes[nid]) 

197 # 部分节点失败后 

198 new_dag = decomposer.replan(dag, failed_nodes=["tn-xxx"]) 

199 """ 

200 

201 MAX_DEPTH = 8 # 最大递归深度 

202 MIN_NODE_DURATION = 1.0 # 最小节点估算时长（秒），低于此不再分解 

203 MAX_NODES = 50 # 最多节点数 

204 

205 def __init__( 

206 self, 

207 strategy: DecompositionStrategy = DecompositionStrategy.RECURSIVE, 

208 llm_call: Callable | None = None, 

209 ): 

210 self._strategy = strategy 

211 self._llm_call = llm_call 

212 self._trace: list[DecompositionTrace] = [] 

213 self._iteration = 0 

214 

215 # ── Public API ───────────────────────── 

216 

217 def decompose( 

218 self, 

219 task: str, 

220 context: dict[str, Any] | None = None, 

221 ) -> TaskDAG: 

222 """将任务分解为 DAG。 

223 

224 Args: 

225 task: 任务描述 

226 context: 补充上下文（已训练的 Agent、可用工具等） 

227 

228 Returns: 

229 完整的 TaskDAG 

230 """ 

231 self._trace = [] 

232 self._iteration = 0 

233 

234 dag = TaskDAG( 

235 root_task=task, 

236 strategy=self._strategy, 

237 metadata=context or {}, 

238 created_at=__import__("time").time(), 

239 ) 

240 

241 root = self._create_node(task, confidence=1.0) 

242 dag.nodes[root.id] = root 

243 

244 # 递归分解 

245 self._decompose_recursive(dag, root.id, depth=0) 

246 

247 # 验证无循环 

248 cycles = dag.detect_cycles() 

249 if cycles: 

250 logger.warning(f"Decomposition produced cycle: {cycles}. Re-resolving.") 

251 dag = self._break_cycles(dag, cycles) 

252 

253 return dag 

254 

255 def replan( 

256 self, 

257 dag: TaskDAG, 

258 failed_nodes: list[str], 

259 ) -> TaskDAG: 

260 """在部分节点执行失败后动态重规划。 

261 

262 Args: 

263 dag: 当前 DAG（含已执行的节点状态） 

264 failed_nodes: 失败节点 ID 列表 

265 

266 Returns: 

267 重规划后的新 DAG（仅影响失败节点及其下游） 

268 """ 

269 self._iteration += 1 

270 

271 affected = self._collect_downstream(dag, failed_nodes) 

272 

273 new_dag = TaskDAG( 

274 dag_id=f"{dag.dag_id}-replan-{self._iteration}", 

275 root_task=dag.root_task, 

276 metadata=dag.metadata, 

277 created_at=__import__("time").time(), 

278 ) 

279 

280 # 保留不受影响的节点和边 

281 for nid, node in dag.nodes.items(): 

282 if nid not in affected: 

283 new_dag.nodes[nid] = node 

284 

285 for edge in dag.edges: 

286 if edge.from_node not in affected and edge.to_node not in affected: 

287 new_dag.edges.append(edge) 

288 

289 # 为每个失败节点构建替代路径 

290 for nid in failed_nodes: 

291 node = dag.nodes[nid] 

292 original_desc = node.description 

293 

294 alt_node = self._create_node( 

295 f"[Retry Plan] {original_desc}", 

296 confidence=node.confidence * 0.8, # 降信心 

297 ) 

298 alt_node.retry_policy = "retry_n" 

299 alt_node.max_retries = node.max_retries + 1 

300 new_dag.nodes[alt_node.id] = alt_node 

301 

302 # 重连边：失败节点上游 → 新节点，新节点 → 失败节点下游 

303 incoming = [e for e in dag.edges if e.to_node == nid] 

304 outgoing = [e for e in dag.edges if e.from_node == nid] 

305 

306 for e in incoming: 

307 if e.from_node not in affected: 

308 new_dag.edges.append(TaskEdge(from_node=e.from_node, to_node=alt_node.id)) 

309 for e in outgoing: 

310 if e.to_node not in affected: 

311 new_dag.edges.append(TaskEdge(from_node=alt_node.id, to_node=e.to_node)) 

312 

313 self._trace.append( 

314 DecompositionTrace( 

315 iteration=self._iteration, 

316 action="replan", 

317 node_before=node, 

318 nodes_after=[alt_node], 

319 reason=f"Node {nid} failed: {node.error or 'unknown'}", 

320 ) 

321 ) 

322 

323 return new_dag 

324 

325 def get_trace(self) -> list[DecompositionTrace]: 

326 """获取完整的分解轨迹（用于可观测性）。""" 

327 return list(self._trace) 

328 

329 def validate_dag(self, dag: TaskDAG) -> tuple[bool, str]: 

330 """验证 DAG 的结构完整性。 

331 

332 Returns: 

333 (is_valid, error_message) 

334 """ 

335 # 空 DAG 

336 if not dag.nodes: 

337 return False, "DAG has no nodes" 

338 

339 # 循环检测 

340 cycles = dag.detect_cycles() 

341 if cycles: 

342 return False, f"DAG contains cycles: {cycles}" 

343 

344 # 孤立节点检测 

345 connected: set[str] = set() 

346 for e in dag.edges: 

347 connected.add(e.from_node) 

348 connected.add(e.to_node) 

349 isolated = set(dag.nodes) - connected 

350 if isolated and len(dag.nodes) > 1: 

351 logger.warning(f"Isolated nodes: {isolated}") 

352 

353 # 拓扑可达性（至少存在一条从源到汇的路径） 

354 try: 

355 dag.topological_order() 

356 except ValueError as e: 

357 return False, str(e) 

358 

359 return True, "valid" 

360 

361 # ── Internal ──────────────────────────── 

362 

363 def _create_node(self, description: str, confidence: float = 1.0) -> TaskNode: 

364 return TaskNode( 

365 description=description, 

366 confidence=confidence, 

367 estimated_duration_s=max(1.0, len(description.split()) * 0.5), 

368 ) 

369 

370 def _decompose_recursive(self, dag: TaskDAG, node_id: str, depth: int): 

371 """递归分解节点直到达到原子粒度。""" 

372 if depth >= self.MAX_DEPTH: 

373 return 

374 

375 node = dag.nodes.get(node_id) 

376 if not node: 

377 return 

378 

379 # 判断是否继续分解 

380 if self._is_atomic(node, depth): 

381 return 

382 

383 self._iteration += 1 

384 

385 # 调用 LLM 或启发式规则生成子任务 

386 sub_tasks = self._generate_sub_tasks(node, dag.metadata) 

387 

388 if not sub_tasks or len(sub_tasks) <= 1: 

389 return # 无法继续分解 

390 

391 # 移除原节点，插入子节点和边 

392 dag.nodes.pop(node_id) 

393 for i, sub in enumerate(sub_tasks): 

394 dag.nodes[sub.id] = sub 

395 # 子任务按顺序或并行链接 

396 if i > 0: 

397 dag.edges.append( 

398 TaskEdge( 

399 from_node=sub_tasks[i - 1].id, 

400 to_node=sub.id, 

401 ) 

402 ) 

403 

404 # 重连原节点的入边和出边 

405 incoming = [e for e in dag.edges if e.to_node == node_id] 

406 outgoing = [e for e in dag.edges if e.from_node == node_id] 

407 

408 # 移除旧边 

409 dag.edges = [e for e in dag.edges if e.to_node != node_id and e.from_node != node_id] 

410 

411 if sub_tasks: 

412 first = sub_tasks[0] 

413 for e in incoming: 

414 dag.edges.append(TaskEdge(from_node=e.from_node, to_node=first.id)) 

415 last = sub_tasks[-1] 

416 for e in outgoing: 

417 dag.edges.append(TaskEdge(from_node=last.id, to_node=e.to_node)) 

418 

419 self._trace.append( 

420 DecompositionTrace( 

421 iteration=self._iteration, 

422 action="split", 

423 node_before=node, 

424 nodes_after=sub_tasks, 

425 reason=f"Decomposed at depth {depth}", 

426 ) 

427 ) 

428 

429 # 递归分解子任务 

430 if len(dag.nodes) < self.MAX_NODES: 

431 for sub in sub_tasks: 

432 if sub.id in dag.nodes: 

433 self._decompose_recursive(dag, sub.id, depth + 1) 

434 

435 def _is_atomic(self, node: TaskNode, depth: int) -> bool: 

436 """判断节点是否已达到原子粒度，无需进一步分解。""" 

437 # 规则 1：估算时长够短 

438 if node.estimated_duration_s < self.MIN_NODE_DURATION: 

439 return True 

440 # 规则 2：节点数已接近上限 

441 if depth > self.MAX_DEPTH - 1: 

442 return True 

443 # 规则 3：描述过于简单（单步骤） 

444 if len(node.description.split()) < 5: 

445 return True 

446 return False 

447 

448 def _generate_sub_tasks(self, node: TaskNode, context: dict[str, Any]) -> list[TaskNode]: 

449 """生成节点的子任务列表。 

450 

451 优先使用 LLM 调用，降级为启发式规则。 

452 """ 

453 if self._llm_call: 

454 return self._llm_generate(node, context) 

455 return self._heuristic_generate(node) 

456 

457 def _llm_generate(self, node: TaskNode, context: dict[str, Any]) -> list[TaskNode]: 

458 """通过 LLM 调用生成子任务。""" 

459 prompt = f"""Break down the following task into 2-5 subtasks. 

460 

461Task: {node.description} 

462Context: {json.dumps(context, default=str) if context else 'None'} 

463 

464Output JSON array of subtasks, each with: 

465- description: string 

466- agent_type: string (default/planner/executor/analyst) 

467- estimated_duration_s: float 

468 

469Only respond with the JSON array, no other text.""" 

470 try: 

471 result = self._llm_call(prompt) 

472 items = json.loads(result) if isinstance(result, str) else result 

473 return [ 

474 TaskNode( 

475 description=item["description"], 

476 agent_type=item.get("agent_type", "default"), 

477 estimated_duration_s=item.get("estimated_duration_s", 5.0), 

478 confidence=0.7, 

479 ) 

480 for item in items 

481 ] 

482 except Exception as e: 

483 logger.warning(f"LLM decomposition failed: {e}, falling back to heuristic") 

484 return self._heuristic_generate(node) 

485 

486 def _heuristic_generate(self, node: TaskNode) -> list[TaskNode]: 

487 """启发式任务分解 — 基于关键词和模式匹配。""" 

488 desc = node.description.lower() 

489 subtasks: list[TaskNode] = [] 

490 

491 # 模式 1：提取/收集 → 分析 → 生成 

492 if any(kw in desc for kw in ("extract", "collect", "fetch", "retrieve", "提取", "收集")): 

493 subtasks.append( 

494 self._create_node(f"Phase 1: Collect data for: {node.description[:60]}") 

495 ) 

496 subtasks.append(self._create_node("Phase 2: Analyze/process collected data")) 

497 subtasks.append( 

498 self._create_node(f"Phase 3: Generate output/report for: {node.description[:60]}") 

499 ) 

500 

501 # 模式 2：对比/比较 

502 elif any(kw in desc for kw in ("compare", "vs", "对比", "比较", "versus")): 

503 parts = desc.replace("compare ", "").replace("对比 ", "").split(" vs ") 

504 if len(parts) < 2: 

505 parts = desc.replace("compare ", "").split(" and ") 

506 if len(parts) >= 2: 

507 subtasks.append(self._create_node(f"Analyze: {parts[0].strip()}")) 

508 subtasks.append(self._create_node(f"Analyze: {parts[1].strip()}")) 

509 subtasks.append(self._create_node("Synthesize comparison results")) 

510 

511 # 模式 3：transform/convert/migrate 

512 elif any(kw in desc for kw in ("transform", "convert", "migrate", "转换", "迁移")): 

513 subtasks.append(self._create_node("Validate source data integrity")) 

514 subtasks.append(self._create_node(f"Execute transformation: {node.description[:60]}")) 

515 subtasks.append(self._create_node("Verify output correctness")) 

516 

517 # 模式 4：default — 按步骤拆 

518 else: 

519 subtasks.append(self._create_node(f"Plan: outline steps for '{node.description[:60]}'")) 

520 subtasks.append(self._create_node(f"Execute: carry out '{node.description[:60]}'")) 

521 subtasks.append( 

522 self._create_node(f"Validate: check results of '{node.description[:60]}'") 

523 ) 

524 

525 for sub in subtasks: 

526 sub.confidence = 0.6 # 启发式分解信心较低 

527 return subtasks 

528 

529 def _collect_downstream(self, dag: TaskDAG, failed_nodes: list[str]) -> set[str]: 

530 """收集失败节点及所有下游节点。""" 

531 adj = dag.adjacency_map() 

532 affected: set[str] = set() 

533 

534 queue = deque(failed_nodes) 

535 while queue: 

536 nid = queue.popleft() 

537 if nid in affected: 

538 continue 

539 affected.add(nid) 

540 for neighbor in adj.get(nid, []): 

541 if neighbor not in affected: 

542 queue.append(neighbor) 

543 

544 return affected 

545 

546 def _break_cycles(self, dag: TaskDAG, cycles: set[str]) -> TaskDAG: 

547 """打破循环 — 移除循环中置信度最低的边。""" 

548 cycle_edges = [e for e in dag.edges if e.from_node in cycles and e.to_node in cycles] 

549 if cycle_edges: 

550 # 移除第一条循环边（可改进为最小置信度边） 

551 dag.edges.remove(cycle_edges[0]) 

552 logger.info( 

553 f"Removed edge {cycle_edges[0].from_node}→{cycle_edges[0].to_node} to break cycle" 

554 ) 

555 return dag 

556 

557 

558# ── Quick Start ────────────────────────────── 

559 

560 

561def create_decomposer( 

562 strategy: DecompositionStrategy = DecompositionStrategy.RECURSIVE, 

563 llm_call: Callable | None = None, 

564) -> TaskDecomposer: 

565 return TaskDecomposer(strategy=strategy, llm_call=llm_call)