"""Orchestrator - 多 Agent 协作编排器 实现 Orchestrator-Worker 模式:中央编排器协调多 Agent 并行/串行执行。 """ from __future__ import annotations import asyncio import logging import uuid from dataclasses import dataclass, field from enum import Enum from typing import TYPE_CHECKING, Any from agentkit.core.protocol import TaskMessage, TaskResult, TaskStatus from agentkit.core.shared_workspace import SharedWorkspace if TYPE_CHECKING: from agentkit.core.goal_planner import GoalPlanner from agentkit.core.plan_executor import PlanExecutor from agentkit.core.plan_checker import PlanChecker logger = logging.getLogger(__name__) class AgentRole(str, Enum): """Agent 角色枚举""" ORCHESTRATOR = "orchestrator" WORKER = "worker" REVIEWER = "reviewer" class SubTaskStatus(str, Enum): """子任务状态""" PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" @dataclass class SubTask: """子任务定义""" task_id: str parent_task_id: str assigned_agent: str task_type: str input_data: dict[str, Any] status: SubTaskStatus = SubTaskStatus.PENDING result: dict[str, Any] | None = None error: str | None = None depends_on: list[str] = field(default_factory=list) @dataclass class OrchestrationPlan: """编排计划""" plan_id: str parent_task_id: str subtasks: list[SubTask] parallel_groups: list[list[str]] # 每组内的子任务可并行执行 @dataclass class OrchestrationResult: """编排结果""" plan_id: str parent_task_id: str subtask_results: dict[str, dict[str, Any]] aggregated_result: dict[str, Any] status: TaskStatus total_duration_ms: float metadata: dict[str, Any] = field(default_factory=dict) @dataclass class OrchestratorConfig: """Orchestrator 配置""" adaptive: bool = False max_iterations: int = 3 quality_threshold: float = 0.7 class Orchestrator: """多 Agent 协作编排器 Orchestrator-Worker 模式: 1. 接收复杂任务 2. LLM 驱动分解为子任务 3. 基于 Skill 能力匹配子任务到 Worker Agent 4. 并行/串行执行子任务 5. 汇总结果,生成最终输出 使用方式: orchestrator = Orchestrator(agent_pool=pool, workspace=workspace) result = await orchestrator.execute(task_message) """ def __init__( self, agent_pool: Any, workspace: SharedWorkspace | None = None, llm_gateway: Any = None, max_parallel: int = 5, subtask_timeout: float = 300.0, goal_planner: GoalPlanner | None = None, plan_executor: PlanExecutor | None = None, plan_checker: PlanChecker | None = None, config: OrchestratorConfig | None = None, message_bus: Any = None, ): """ Args: agent_pool: AgentPool 实例 workspace: 共享工作空间 llm_gateway: LLM Gateway,用于任务分解 max_parallel: 最大并行子任务数 subtask_timeout: 子任务超时时间(秒) goal_planner: GoalPlanner 实例,用于结构化目标分解(可选) plan_executor: PlanExecutor 实例,用于执行 ExecutionPlan(可选) plan_checker: PlanChecker 实例,用于检查和复盘(可选) config: Orchestrator 配置,包含自适应参数 message_bus: MessageBus 实例,用于 Agent 间通信 """ self._agent_pool = agent_pool self._workspace = workspace or SharedWorkspace() self._llm_gateway = llm_gateway self._max_parallel = max_parallel self._subtask_timeout = subtask_timeout self._goal_planner = goal_planner self._plan_executor = plan_executor self._plan_checker = plan_checker self._config = config or OrchestratorConfig() self._message_bus = message_bus async def execute(self, task: TaskMessage) -> OrchestrationResult: """执行编排任务 Args: task: 原始任务消息 Returns: OrchestrationResult: 编排结果 """ import time start_time = time.monotonic() # 1. Decompose task into subtasks plan = await self._decompose_task(task) if not plan.subtasks: return OrchestrationResult( plan_id=plan.plan_id, parent_task_id=task.task_id, subtask_results={}, aggregated_result={"error": "Failed to decompose task"}, status=TaskStatus.FAILED, total_duration_ms=0, ) # 2. Store plan in workspace await self._workspace.write( f"plan:{plan.plan_id}", {"task_id": task.task_id, "subtask_count": len(plan.subtasks)}, agent_id="orchestrator", ) # 3. Execute subtasks subtask_results = await self._execute_plan(plan, task) # 4. Aggregate results aggregated = await self._aggregate_results(plan, subtask_results, task) # 5. Determine overall status failed_count = sum( 1 for r in subtask_results.values() if r.get("status") == "failed" ) if failed_count == len(plan.subtasks): status = TaskStatus.FAILED elif failed_count > 0: status = TaskStatus.PARTIALLY_COMPLETED else: status = TaskStatus.COMPLETED duration_ms = (time.monotonic() - start_time) * 1000 return OrchestrationResult( plan_id=plan.plan_id, parent_task_id=task.task_id, subtask_results=subtask_results, aggregated_result=aggregated, status=status, total_duration_ms=duration_ms, ) async def _decompose_task(self, task: TaskMessage) -> OrchestrationPlan: """将复杂任务分解为子任务""" plan_id = str(uuid.uuid4())[:8] # If GoalPlanner available, use it for structured decomposition if self._goal_planner: try: execution_plan = await self._goal_planner.generate_plan( goal=str(task.input_data), context={"task_type": task.task_type, "agent_name": task.agent_name}, available_skills=self._get_available_skill_names(), ) subtasks = self._convert_execution_plan_to_subtasks( execution_plan, task.task_id, task.agent_name, task.task_type, task.input_data, ) if subtasks: parallel_groups = self._build_parallel_groups(subtasks) return OrchestrationPlan( plan_id=plan_id, parent_task_id=task.task_id, subtasks=subtasks, parallel_groups=parallel_groups, ) except Exception as e: logger.warning(f"GoalPlanner decomposition failed, falling back: {e}") # If LLM gateway available, use it for decomposition if self._llm_gateway: try: subtasks = await self._llm_decompose(task) if subtasks: parallel_groups = self._build_parallel_groups(subtasks) return OrchestrationPlan( plan_id=plan_id, parent_task_id=task.task_id, subtasks=subtasks, parallel_groups=parallel_groups, ) except Exception as e: logger.warning(f"LLM decomposition failed, falling back to simple: {e}") # Fallback: single subtask = original task subtask = SubTask( task_id=f"{plan_id}-0", parent_task_id=task.task_id, assigned_agent=task.agent_name, task_type=task.task_type, input_data=task.input_data, ) return OrchestrationPlan( plan_id=plan_id, parent_task_id=task.task_id, subtasks=[subtask], parallel_groups=[[subtask.task_id]], ) async def _llm_decompose(self, task: TaskMessage) -> list[SubTask]: """使用 LLM 分解任务""" # Get available agents and their capabilities agents_info = self._agent_pool.list_agents() agent_descriptions = "\n".join( f"- {a['name']} ({a['agent_type']}): {a.get('description', 'No description')}" for a in agents_info ) prompt = ( f"Decompose the following task into subtasks that can be assigned to available agents.\n\n" f"Task: {task.input_data}\n" f"Task Type: {task.task_type}\n\n" f"Available Agents:\n{agent_descriptions}\n\n" 'Respond ONLY with a JSON array: [{"agent_name": "...", "task_type": "...", ' '"input_data": {...}, "depends_on": []}]\n' "The depends_on field lists task indices (0-based) that must complete first.\n" "Do not include any other text." ) import json response = await self._llm_gateway.chat( messages=[{"role": "user", "content": prompt}], model="default", ) try: subtask_defs = json.loads(response.content) if not isinstance(subtask_defs, list): return [] subtasks = [] for i, defn in enumerate(subtask_defs): depends_on = [ f"task-{i}" for i in defn.get("depends_on", []) if isinstance(i, int) and 0 <= i < len(subtask_defs) ] subtasks.append(SubTask( task_id=f"task-{i}", parent_task_id=task.task_id, assigned_agent=defn.get("agent_name", task.agent_name), task_type=defn.get("task_type", task.task_type), input_data=defn.get("input_data", {}), depends_on=depends_on, )) return subtasks except (json.JSONDecodeError, KeyError) as e: logger.warning(f"Failed to parse LLM decomposition: {e}") return [] def _build_parallel_groups(self, subtasks: list[SubTask]) -> list[list[str]]: """构建并行执行组 基于依赖关系拓扑排序,无依赖的子任务分到同一组并行执行。 """ # Build dependency graph task_map = {st.task_id: st for st in subtasks} completed: set[str] = set() groups: list[list[str]] = [] remaining = set(st.task_id for st in subtasks) while remaining: # Find tasks with all dependencies satisfied ready = [] for tid in remaining: task = task_map[tid] if all(dep in completed for dep in task.depends_on): ready.append(tid) if not ready: # Circular dependency — put remaining in one group groups.append(list(remaining)) break # Limit group size group = ready[:self._max_parallel] groups.append(group) for tid in group: completed.add(tid) remaining.discard(tid) return groups async def _execute_plan( self, plan: OrchestrationPlan, original_task: TaskMessage ) -> dict[str, dict[str, Any]]: """执行编排计划""" subtask_results: dict[str, dict[str, Any]] = {} task_map = {st.task_id: st for st in plan.subtasks} for group in plan.parallel_groups: # Execute group in parallel tasks = [] for task_id in group: subtask = task_map[task_id] # Inject results from dependencies enriched_input = self._inject_dependency_results( subtask, subtask_results ) tasks.append(self._execute_subtask(subtask, enriched_input, original_task)) results = await asyncio.gather(*tasks, return_exceptions=True) for task_id, result in zip(group, results): if isinstance(result, Exception): subtask_results[task_id] = { "status": "failed", "error": str(result), } else: subtask_results[task_id] = result return subtask_results async def _execute_subtask( self, subtask: SubTask, input_data: dict[str, Any], original_task: TaskMessage, ) -> dict[str, Any]: """执行单个子任务""" agent = self._agent_pool.get_agent(subtask.assigned_agent) if agent is None: return {"status": "failed", "error": f"Agent '{subtask.assigned_agent}' not found"} sub_task_msg = TaskMessage( task_id=subtask.task_id, agent_name=subtask.assigned_agent, task_type=subtask.task_type, priority=original_task.priority, input_data=input_data, callback_url=None, created_at=original_task.created_at, timeout_seconds=int(self._subtask_timeout), ) try: result = await asyncio.wait_for( agent.execute(sub_task_msg), timeout=self._subtask_timeout, ) output = { "status": "completed", "output": result.output_data if hasattr(result, "output_data") else result, } # Publish progress via MessageBus if available if self._message_bus is not None: try: from agentkit.bus.message import AgentMessage await self._message_bus.publish(AgentMessage( sender=subtask.assigned_agent, recipient="orchestrator", topic="task.progress", payload={ "task_id": subtask.task_id, "status": "completed", }, )) except Exception as e: logger.warning(f"Failed to publish progress via MessageBus: {e}") return output except asyncio.TimeoutError: error_result = {"status": "failed", "error": "Subtask timed out"} if self._message_bus is not None: try: from agentkit.bus.message import AgentMessage await self._message_bus.publish(AgentMessage( sender=subtask.assigned_agent, recipient="orchestrator", topic="task.progress", payload={ "task_id": subtask.task_id, "status": "failed", "error": "Subtask timed out", }, )) except Exception as e: logger.warning(f"Failed to publish progress via MessageBus: {e}") return error_result except Exception as e: error_result = {"status": "failed", "error": str(e)} if self._message_bus is not None: try: from agentkit.bus.message import AgentMessage await self._message_bus.publish(AgentMessage( sender=subtask.assigned_agent, recipient="orchestrator", topic="task.progress", payload={ "task_id": subtask.task_id, "status": "failed", "error": str(e), }, )) except Exception as e: logger.warning(f"Failed to publish progress via MessageBus: {e}") return error_result def _inject_dependency_results( self, subtask: SubTask, subtask_results: dict[str, dict[str, Any]], ) -> dict[str, Any]: """将依赖子任务的结果注入到当前子任务的输入中""" enriched = dict(subtask.input_data) if subtask.depends_on: dep_results = {} for dep_id in subtask.depends_on: if dep_id in subtask_results: dep_results[dep_id] = subtask_results[dep_id] if dep_results: enriched["dependency_results"] = dep_results return enriched async def _aggregate_results( self, plan: OrchestrationPlan, subtask_results: dict[str, dict[str, Any]], original_task: TaskMessage, ) -> dict[str, Any]: """汇总子任务结果""" # Simple aggregation: collect all outputs outputs = {} errors = [] for subtask in plan.subtasks: result = subtask_results.get(subtask.task_id, {}) if result.get("status") == "completed": outputs[subtask.task_id] = result.get("output", {}) else: errors.append({ "task_id": subtask.task_id, "error": result.get("error", "Unknown error"), }) aggregated = { "outputs": outputs, "task_id": original_task.task_id, } if errors: aggregated["errors"] = errors aggregated["partial_success"] = True return aggregated def _get_available_skill_names(self) -> list[str]: """获取可用 Skill 名称列表""" try: agents_info = self._agent_pool.list_agents() return [a["name"] for a in agents_info] except Exception: return [] def _convert_execution_plan_to_subtasks( self, execution_plan: Any, parent_task_id: str, default_agent: str, default_task_type: str, original_input: dict[str, Any], ) -> list[SubTask]: """将 ExecutionPlan 的 PlanStep 转换为 SubTask 列表""" subtasks: list[SubTask] = [] for step in execution_plan.steps: # 尝试根据 required_skills 匹配 agent assigned_agent = default_agent if step.required_skills: matched_agent = self._match_agent_for_skills(step.required_skills) if matched_agent: assigned_agent = matched_agent subtasks.append(SubTask( task_id=step.step_id, parent_task_id=parent_task_id, assigned_agent=assigned_agent, task_type=default_task_type, input_data={ **original_input, "step_name": step.name, "step_description": step.description, }, depends_on=list(step.dependencies), )) return subtasks def _match_agent_for_skills(self, required_skills: list[str]) -> str | None: """根据所需 Skill 匹配 Agent""" try: agents_info = self._agent_pool.list_agents() for skill in required_skills: for agent in agents_info: name = agent.get("name", "") agent_type = agent.get("agent_type", "") description = agent.get("description", "").lower() if skill.lower() in name.lower() or skill.lower() in agent_type.lower() or skill.lower() in description: return name except Exception: pass return None async def execute_adaptive( self, task: TaskMessage, ) -> OrchestrationResult: """自适应编排:执行→评估→再分解循环。 与 execute() 不同,此方法在第一轮执行后评估子任务结果质量, 如果评估不通过且未达 max_iterations,则基于评估反馈重新分解 未达标的子任务,保留已完成的子任务结果,然后执行新分解的子任务。 Args: task: 原始任务消息 Returns: OrchestrationResult: 编排结果,metadata 中包含迭代历史 """ import time as _time start_time = _time.monotonic() iteration_history: list[dict[str, Any]] = [] # First execution result = await self.execute(task) # If adaptive not enabled or already succeeded, return directly if not self._config.adaptive or result.status == TaskStatus.COMPLETED: # Check quality even on success if self._config.adaptive and self._llm_gateway: quality = await self._evaluate_quality(task, result) if quality["score"] >= self._config.quality_threshold: result.metadata["quality_score"] = quality["score"] return result return result # Adaptive loop current_result = result for iteration in range(1, self._config.max_iterations + 1): # Evaluate quality quality = await self._evaluate_quality(task, current_result) iteration_history.append({ "iteration": iteration, "quality_score": quality["score"], "feedback": quality.get("feedback", ""), }) if quality["score"] >= self._config.quality_threshold: logger.info( f"Adaptive iteration {iteration}: quality " f"{quality['score']:.2f} >= {self._config.quality_threshold}" ) current_result.metadata["quality_score"] = quality["score"] current_result.metadata["iterations"] = iteration_history return current_result logger.info( f"Adaptive iteration {iteration}: quality " f"{quality['score']:.2f} < {self._config.quality_threshold}, " f"re-decomposing failed subtasks" ) # Re-decompose failed subtasks new_result = await self._reexecute_failed( task, current_result, quality, ) current_result = new_result # Exhausted iterations current_result.metadata["iterations"] = iteration_history return current_result async def _evaluate_quality( self, task: TaskMessage, result: OrchestrationResult, ) -> dict[str, Any]: """评估子任务结果质量。 Returns: Dict with "score" (0-1) and optional "feedback" string. """ # Rule-based evaluation when no LLM if self._llm_gateway is None: return self._rule_based_evaluate(result) try: return await self._llm_evaluate(task, result) except Exception as e: logger.warning(f"LLM evaluation failed, falling back to rule-based: {e}") return self._rule_based_evaluate(result) def _rule_based_evaluate( self, result: OrchestrationResult, ) -> dict[str, Any]: """基于规则的质量评估:根据完成率打分。""" total = len(result.subtask_results) if total == 0: return {"score": 0.0, "feedback": "No subtasks executed"} completed = sum( 1 for r in result.subtask_results.values() if r.get("status") == "completed" ) score = completed / total feedback = "" if score < 1.0: failed = [ tid for tid, r in result.subtask_results.items() if r.get("status") != "completed" ] feedback = f"Failed subtasks: {failed}" return {"score": score, "feedback": feedback} async def _llm_evaluate( self, task: TaskMessage, result: OrchestrationResult, ) -> dict[str, Any]: """使用 LLM 评估子任务结果质量。""" import json subtask_summary = [] for tid, r in result.subtask_results.items(): subtask_summary.append({ "task_id": tid, "status": r.get("status", "unknown"), "output_preview": str(r.get("output", ""))[:200], }) prompt = ( f"Evaluate the quality of the following orchestration result.\n\n" f"Original task: {task.input_data}\n" f"Subtask results:\n{json.dumps(subtask_summary, ensure_ascii=False)}\n\n" f'Respond ONLY with JSON: {{"score": 0.0-1.0, "feedback": "..."}}\n' f"Score 1.0 = perfect, 0.0 = completely failed." ) response = await self._llm_gateway.chat( messages=[{"role": "user", "content": prompt}], model="default", ) try: text = response.content.strip() if text.startswith("```"): lines = text.split("\n") text = "\n".join(lines[1:-1]) data = json.loads(text) return { "score": float(data.get("score", 0.0)), "feedback": data.get("feedback", ""), } except (json.JSONDecodeError, ValueError) as e: logger.warning(f"Failed to parse LLM evaluation: {e}") return self._rule_based_evaluate(result) async def _reexecute_failed( self, task: TaskMessage, previous_result: OrchestrationResult, quality: dict[str, Any], ) -> OrchestrationResult: """重新执行失败的子任务,保留已完成的结果。""" import time as _time start_time = _time.monotonic() # Identify failed subtasks failed_task_ids = [ tid for tid, r in previous_result.subtask_results.items() if r.get("status") != "completed" ] if not failed_task_ids: return previous_result # Create new subtasks for failed ones, incorporating feedback new_subtasks = [] for tid in failed_task_ids: old_result = previous_result.subtask_results[tid] new_subtasks.append(SubTask( task_id=f"retry-{tid}", parent_task_id=task.task_id, assigned_agent=task.agent_name, task_type=task.task_type, input_data={ **task.input_data, "previous_error": old_result.get("error", ""), "improvement_feedback": quality.get("feedback", ""), }, )) # Build a mini-plan for the retry subtasks plan = OrchestrationPlan( plan_id=f"retry-{previous_result.plan_id}", parent_task_id=task.task_id, subtasks=new_subtasks, parallel_groups=[[st.task_id for st in new_subtasks]], ) # Execute retry subtasks retry_results = await self._execute_plan(plan, task) # Merge: keep completed results, replace failed with retry results merged_results = {} for tid, r in previous_result.subtask_results.items(): if r.get("status") == "completed": merged_results[tid] = r for tid, r in retry_results.items(): # Map retry task IDs back to original original_tid = tid.replace("retry-", "", 1) merged_results[original_tid] = r # Re-aggregate all_subtasks = [] for tid, r in merged_results.items(): all_subtasks.append(SubTask( task_id=tid, parent_task_id=task.task_id, assigned_agent=task.agent_name, task_type=task.task_type, input_data=task.input_data, status=SubTaskStatus.COMPLETED if r.get("status") == "completed" else SubTaskStatus.FAILED, result=r.get("output"), )) retry_plan = OrchestrationPlan( plan_id=plan.plan_id, parent_task_id=task.task_id, subtasks=all_subtasks, parallel_groups=[], ) aggregated = await self._aggregate_results(retry_plan, merged_results, task) failed_count = sum( 1 for r in merged_results.values() if r.get("status") == "failed" ) if failed_count == len(merged_results): status = TaskStatus.FAILED elif failed_count > 0: status = TaskStatus.PARTIALLY_COMPLETED else: status = TaskStatus.COMPLETED duration_ms = (_time.monotonic() - start_time) * 1000 return OrchestrationResult( plan_id=plan.plan_id, parent_task_id=task.task_id, subtask_results=merged_results, aggregated_result=aggregated, status=status, total_duration_ms=duration_ms, )