第八章：任务编排与工作流引擎¶

from typing import Dict, List, Set, Optional
from dataclasses import dataclass, field
from enum import Enum
import json

class TaskType(Enum):
    """7种Task类型"""
    LOCAL_BASH = "local_bash"
    LOCAL_AGENT = "local_agent"
    REMOTE_AGENT = "remote_agent"
    IN_PROCESS_TEAMMATE = "in_process_teammate"
    WORKFLOW = "workflow"
    MONITOR_MCP = "monitor_mcp"
    DREAM = "dream"

class TaskState(Enum):
    """任务生命周期状态"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    KILLED = "killed"

@dataclass
class TaskDefinition:
    """任务定义"""
    task_id: str
    task_type: TaskType
    description: str
    dependencies: List[str] = field(default_factory=list)
    timeout_seconds: int = 300
    max_retries: int = 1
    execution_config: Dict = field(default_factory=dict)

    def add_dependency(self, task_id: str) -> None:
        """添加任务依赖"""
        if task_id not in self.dependencies:
            self.dependencies.append(task_id)

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "task_type": self.task_type.value,
            "description": self.description,
            "dependencies": self.dependencies,
            "timeout_seconds": self.timeout_seconds,
            "max_retries": self.max_retries,
            "execution_config": self.execution_config,
        }

class TaskDAG:
    """任务依赖图"""

    def __init__(self):
        self.tasks: Dict[str, TaskDefinition] = {}
        self.execution_order: List[str] = []

    def add_task(self, task_def: TaskDefinition) -> None:
        """添加任务定义"""
        self.tasks[task_def.task_id] = task_def

    def validate(self) -> bool:
        """验证DAG的合法性"""
        # 1. 检查所有依赖都存在
        for task_id, task in self.tasks.items():
            for dep_id in task.dependencies:
                if dep_id not in self.tasks:
                    raise ValueError(f"任务{task_id}依赖的任务{dep_id}不存在")

        # 2. 检查是否存在循环依赖
        if self._has_cycle():
            raise ValueError("检测到循环依赖")

        return True

    def _has_cycle(self) -> bool:
        """DFS检测循环依赖"""
        visited = set()
        rec_stack = set()

        def dfs(node: str) -> bool:
            visited.add(node)
            rec_stack.add(node)

            for neighbor in self.tasks[node].dependencies:
                if neighbor not in visited:
                    if dfs(neighbor):
                        return True
                elif neighbor in rec_stack:
                    return True

            rec_stack.remove(node)
            return False

        for task_id in self.tasks:
            if task_id not in visited:
                if dfs(task_id):
                    return True
        return False

    def topological_sort(self) -> List[str]:
        """拓扑排序,生成执行顺序"""
        self.validate()

        # 入度表
        in_degree = {task_id: 0 for task_id in self.tasks}
        for task in self.tasks.values():
            for dep_id in task.dependencies:
                in_degree[task.task_id] += 1

        # Kahn算法
        queue = [task_id for task_id in self.tasks if in_degree[task_id] == 0]
        order = []

        while queue:
            node = queue.pop(0)
            order.append(node)

            # 找到依赖于这个节点的所有任务
            for task_id, task in self.tasks.items():
                if node in task.dependencies:
                    in_degree[task_id] -= 1
                    if in_degree[task_id] == 0:
                        queue.append(task_id)

        self.execution_order = order
        return order

    def get_parallelizable_groups(self) -> List[List[str]]:
        """获取可并行执行的任务组"""
        # 断言:DAG必须已验证(无循环、无孤立节点)
        assert self.validate(), "DAG验证失败:存在循环或缺失依赖"

        self.topological_sort()
        groups = []
        completed = set()
        remaining = set(self.execution_order)

        while remaining:
            # 只把依赖已经在前一轮完成的任务放入同一并行层。
            # 同一层内的任务不能互相满足依赖，否则会把应串行的任务误排并行。
            ready = [
                task_id
                for task_id in self.execution_order
                if task_id in remaining
                and all(dep in completed for dep in self.tasks[task_id].dependencies)
            ]
            if not ready:
                raise ValueError("无法继续分组:存在循环依赖或缺失依赖")

            groups.append(ready)
            completed.update(ready)
            remaining.difference_update(ready)

        # 断言:所有任务都应被分配到某个组(无孤立节点)
        total_assigned = sum(len(g) for g in groups)
        assert total_assigned == len(self.tasks), \
            f"孤立节点:已分配{total_assigned}/{len(self.tasks)}个任务"

        return groups

    def visualize(self) -> str:
        """生成Mermaid图表代码"""
        lines = ["graph TD"]

        # 添加节点
        for task_id, task in self.tasks.items():
            label = f"{task.task_id}<br/>({task.task_type.value})"
            lines.append(f'    {task_id.replace("#", "_")}["{label}"]')

        # 添加边
        for task_id, task in self.tasks.items():
            for dep_id in task.dependencies:
                lines.append(f"    {dep_id.replace('#', '_')} --> {task_id.replace('#', '_')}")

        return "\n".join(lines)

# 使用示例
if __name__ == "__main__":
    dag = TaskDAG()

    # 添加任务
    dag.add_task(TaskDefinition(
        task_id="task_0",
        task_type=TaskType.LOCAL_BASH,
        description="获取论文数据"
    ))

    dag.add_task(TaskDefinition(
        task_id="task_1",
        task_type=TaskType.LOCAL_AGENT,
        description="内容分析",
        dependencies=["task_0"]
    ))

    dag.add_task(TaskDefinition(
        task_id="task_2",
        task_type=TaskType.LOCAL_AGENT,
        description="引用分析",
        dependencies=["task_0"]
    ))

    dag.add_task(TaskDefinition(
        task_id="task_3",
        task_type=TaskType.LOCAL_AGENT,
        description="综合评估",
        dependencies=["task_1", "task_2"]
    ))

    # 验证和排序
    order = dag.topological_sort()
    print("执行顺序:", order)

    # 获取并行组
    groups = dag.get_parallelizable_groups()
    print("并行组:", groups)

    # 生成可视化
    print(dag.visualize())

from typing import Dict, List, Callable, Any, Optional, Tuple
from dataclasses import dataclass, field
from enum import Enum
import json
from datetime import datetime

class ActionType(Enum):
    """Action类型"""
    TOOL_CALL = "tool_call"
    CONDITIONAL = "conditional"
    SET_VARIABLE = "set_variable"
    LOG = "log"
    PARALLEL = "parallel"

class StateType(Enum):
    """状态类型"""
    INITIAL = "initial"
    NORMAL = "normal"
    FINAL = "final"
    ERROR = "error"
    WAIT = "wait"
    PARALLEL = "parallel"

@dataclass
class Action:
    """Action定义"""
    action_id: str
    action_type: ActionType
    side_effect: bool = False
    params: Dict[str, Any] = field(default_factory=dict)
    result: Optional[Dict[str, Any]] = None
    executed: bool = False

@dataclass
class State:
    """状态定义"""
    state_id: str
    state_type: StateType
    actions: List[Action] = field(default_factory=list)
    variables: Dict[str, Any] = field(default_factory=dict)

@dataclass
class Transition:
    """转移定义"""
    from_state: str
    to_state: str
    condition: Optional[Callable[[Dict[str, Any]], bool]] = None
    on_transition: List[Callable] = field(default_factory=list)

class WorkflowExecutor:
    """工作流执行引擎"""

    def __init__(self):
        self.states: Dict[str, State] = {}
        self.transitions: List[Transition] = []
        self.current_state: Optional[str] = None
        self.execution_history: List[Dict] = []
        self.context: Dict[str, Any] = {}
        self.execution_paused = False
        self.pending_approvals: Dict[str, Action] = {}

    def add_state(self, state: State) -> None:
        """添加状态"""
        self.states[state.state_id] = state

    def add_transition(self, transition: Transition) -> None:
        """添加转移"""
        self.transitions.append(transition)

    def initialize(self, initial_state: str, context: Dict[str, Any]) -> None:
        """初始化工作流"""
        if initial_state not in self.states:
            raise ValueError(f"初始状态{initial_state}不存在")
        self.current_state = initial_state
        self.context = context
        self.execution_history = []
        self.pending_approvals = {}
        self.execution_paused = False

    def evaluate_condition(self, condition: Callable) -> bool:
        """评估转移条件"""
        try:
            return condition(self.context)
        except Exception as e:
            print(f"条件评估失败: {e}")
            return False

    def find_next_state(self) -> Optional[str]:
        """根据当前状态和条件查找下一个状态"""
        applicable_transitions = [
            t for t in self.transitions if t.from_state == self.current_state
        ]
        for transition in applicable_transitions:
            if transition.condition is None or self.evaluate_condition(transition.condition):
                return transition.to_state
        return None

from __future__ import annotations
from typing import Any, Dict


class WorkflowExecutor:
    # 省略 __init__、状态注册和转移选择方法；此处只展示执行/审批相关方法。

    def execute_action(self, action: Action) -> bool:
        """执行单个Action"""
        if action.executed:
            return True

        if action.action_type == ActionType.TOOL_CALL:
            if action.side_effect:
                self.pending_approvals[action.action_id] = action
                self.execution_paused = True
                print(f"[审批] Action {action.action_id} 需要批准")
                return False
            else:
                action.result = self._simulate_tool_call(action.params)
                action.executed = True
                return True

        elif action.action_type == ActionType.CONDITIONAL:
            result = self.evaluate_condition(action.params.get("condition"))
            action.result = {"value": result}
            action.executed = True
            return True

        elif action.action_type == ActionType.SET_VARIABLE:
            var_name = action.params.get("name")
            var_value = action.params.get("value")
            self.context[var_name] = var_value
            action.executed = True
            return True

        elif action.action_type == ActionType.LOG:
            message = action.params.get("message", "")
            print(f"[日志] {message}")
            action.executed = True
            return True
        return False

    def _simulate_tool_call(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """模拟工具调用"""
        return {"success": True, "output": f"Executed with {params}"}

    def execute_state(self) -> bool:
        """执行当前状态的所有actions"""
        if self.current_state is None:
            return False
        state = self.states[self.current_state]
        print(f"\n[进入状态] {self.current_state}")
        for action in state.actions:
            if not self.execute_action(action):
                if self.execution_paused:
                    return False
        return True

    def approve_action(self, action_id: str) -> bool:
        """批准等待的副作用Action"""
        if action_id not in self.pending_approvals:
            return False
        action = self.pending_approvals[action_id]
        action.result = self._simulate_tool_call(action.params)
        action.executed = True
        del self.pending_approvals[action_id]
        if not self.pending_approvals:
            self.execution_paused = False
        return True

    def step(self) -> bool:
        """执行一步:执行当前状态,查找下一个状态"""
        if self.current_state is None:
            return False
        if not self.execute_state():
            return False
        self.execution_history.append({
            "timestamp": datetime.now().isoformat(),
            "state": self.current_state,
            "context": dict(self.context),
        })
        next_state = self.find_next_state()
        if next_state is None:
            return False
        self.current_state = next_state
        if self.states[self.current_state].state_type == StateType.FINAL:
            return False
        return True

    def run(self) -> Dict[str, Any]:
        """执行工作流直到完成或错误"""
        max_iterations = 100
        iterations = 0
        while iterations < max_iterations:
            if self.execution_paused:
                break
            if not self.step():
                break
            iterations += 1
        return {
            "final_state": self.current_state,
            "context": self.context,
            "history": self.execution_history,
            "paused": self.execution_paused,
            "pending_approvals": list(self.pending_approvals.keys()),
        }

    def get_status(self) -> Dict[str, Any]:
        """获取当前工作流状态"""
        return {
            "current_state": self.current_state,
            "is_paused": self.execution_paused,
            "pending_approvals": list(self.pending_approvals.keys()),
            "context": self.context,
            "history_length": len(self.execution_history),
        }

# 构建审批工作流
executor = WorkflowExecutor()

# 定义状态
executor.add_state(State(
    state_id="submitted",
    state_type=StateType.INITIAL,
    actions=[Action(action_id="validate_input", action_type=ActionType.TOOL_CALL)],
))
executor.add_state(State(
    state_id="reviewing",
    state_type=StateType.NORMAL,
    actions=[Action(
        action_id="apply_changes",
        action_type=ActionType.TOOL_CALL,
        side_effect=True,  # 需要人工审批
        params={"target": "production"},
    )],
))
executor.add_state(State(state_id="approved", state_type=StateType.FINAL))
executor.add_state(State(state_id="rejected", state_type=StateType.FINAL))

# 定义转移
executor.add_transition(Transition(
    from_state="submitted",
    to_state="reviewing",
    condition=lambda ctx: ctx.get("valid", False),
))
executor.add_transition(Transition(
    from_state="submitted",
    to_state="rejected",
    condition=lambda ctx: not ctx.get("valid", False),
))
executor.add_transition(Transition(
    from_state="reviewing",
    to_state="approved",
    condition=lambda ctx: ctx.get("approved", False),
))

# 初始化并执行
executor.initialize("submitted", context={"valid": True})
result = executor.run()
print(result)
# {"final_state": "reviewing", "paused": True, "pending_approvals": ["apply_changes"], ...}

# 模拟人工审批
executor.context["approved"] = True
executor.approve_action("apply_changes")
result = executor.run()
print(result)
# {"final_state": "approved", "paused": False, ...}

def save_checkpoint(self) -> str:
    """保存检查点"""
    checkpoint = {
        "current_state": self.current_state,
        "context": self.context,
        "execution_history": self.execution_history,
        "timestamp": datetime.now().isoformat(),
    }
    checkpoint_id = hash_checkpoint(checkpoint)
    # 保存到持久化存储
    return checkpoint_id

def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
    """从检查点恢复"""
    checkpoint = load_checkpoint(checkpoint_id)
    if not checkpoint:
        return False

    self.current_state = checkpoint["current_state"]
    self.context = checkpoint["context"]
    self.execution_history = checkpoint["execution_history"]
    return True

import asyncio
import time

class OrchestrationContext:
    """编排上下文,管理多智能体间的消息"""

    def __init__(self, dag: dict = None):
        self.message_queue = asyncio.Queue()
        self.agent_channels = {}  # agent_id -> Queue
        self.state_outputs = {}   # state_name -> result
        self.dag = dag or {}      # state_name -> 前驱状态名

    async def send_message(self, from_agent: str, to_agent: str, msg: dict):
        """发送消息给特定智能体"""
        if to_agent not in self.agent_channels:
            self.agent_channels[to_agent] = asyncio.Queue()
        await self.agent_channels[to_agent].put(msg)

    async def broadcast_message(self, from_agent: str, msg: dict):
        """广播消息给所有智能体(用于讨论模式)"""
        for agent_id, queue in self.agent_channels.items():
            if agent_id != from_agent:  # 不发给自己
                await queue.put(msg)

    def store_state_output(self, state_name: str, result: dict):
        """存储状态的输出,用于后续状态访问"""
        self.state_outputs[state_name] = result

    def get_state_input(self, state_name: str) -> dict:
        """从前驱状态获取输入"""
        # 根据DAG定义查找前驱状态
        return self.state_outputs.get(self._get_predecessor(state_name), {})

    def _get_predecessor(self, state_name: str) -> str:
        """根据 DAG 定义返回前驱状态名"""
        return self.dag.get(state_name, "")

class WorkflowAbortedError(Exception):
    """工作流被错误处理器主动中止"""

class ErrorPropagation:
    """多智能体间的错误传播"""

    def __init__(self, orchestration_context: OrchestrationContext,
                 execute_state=None):
        self.ctx = orchestration_context
        self.execute_state = execute_state  # 状态执行回调,由编排引擎注入
        self.error_handlers = {}

    def register_error_handler(self, source_agent: str, handler_fn):
        """为特定智能体注册错误处理器"""
        self.error_handlers[source_agent] = handler_fn

    async def propagate_error(self, source_agent: str, error: Exception, state: str):
        """传播错误到下游智能体"""
        error_event = {
            "type": "agent_error",
            "source": source_agent,
            "state": state,
            "error_type": error.__class__.__name__,
            "error_message": str(error),
            "timestamp": time.time()
        }

        # 1. 广播错误事件
        await self.ctx.broadcast_message(source_agent, error_event)

        # 2. 触发错误处理器
        if source_agent in self.error_handlers:
            handler_result = await self.error_handlers[source_agent](error_event)

            if handler_result.action == "retry":
                # 重试该阶段
                return await self._retry_state(state)
            elif handler_result.action == "fallback":
                # 使用备选方案
                return handler_result.fallback_result
            elif handler_result.action == "abort":
                # 中止整个工作流
                raise WorkflowAbortedError(error_event)

    async def _retry_state(self, state: str, max_retries: int = 3):
        """重试某个状态"""
        for attempt in range(max_retries):
            try:
                return await self.execute_state(state)
            except Exception as e:
                if attempt == max_retries - 1:
                    raise
                await asyncio.sleep(2 ** attempt)  # 指数退避

class ExecutionContext:
    """分层执行上下文"""

    def __init__(self, workflow_id: str, parent_context: Optional['ExecutionContext'] = None):
        self.workflow_id = workflow_id
        self.parent = parent_context

        # 本层专有的变量
        self.local_vars: Dict[str, Any] = {}

        # 从父上下文继承的只读副本
        self.inherited_vars: Dict[str, Any] = parent_context.local_vars.copy() if parent_context else {}

    def get(self, key: str, default=None):
        """查询变量:先查本地,再查继承的"""
        if key in self.local_vars:
            return self.local_vars[key]
        return self.inherited_vars.get(key, default)

    def set(self, key: str, value: Any):
        """设置变量(只在本地作用域)"""
        self.local_vars[key] = value

    def commit_to_parent(self, keys: List[str]):
        """显式提交某些变量到父上下文(用于阶段间数据传递)"""
        if self.parent:
            for key in keys:
                if key in self.local_vars:
                    self.parent.set(key, self.local_vars[key])

async def execute_parallel_agents(ctx: ExecutionContext, agents: List[Agent]):
    """并发执行多个智能体,隔离上下文"""
    tasks = []

    for agent in agents:
        # 为每个智能体创建独立的子上下文
        child_ctx = ExecutionContext(f"{ctx.workflow_id}::{agent.name}", parent_context=ctx)
        task = agent.execute(child_ctx)
        tasks.append((agent.name, task, child_ctx))

    # 并发运行
    results = {}
    gathered = await asyncio.gather(
        *(task for _, task, _ in tasks),
        return_exceptions=True,
    )
    for (agent_name, _, child_ctx), result in zip(tasks, gathered):
        if isinstance(result, Exception):
            results[agent_name] = {"error": str(result)}
            continue

        results[agent_name] = result

        # 执行完毕后,父上下文可以选择性地获取某些结果
        ctx.set(f"{agent_name}_result", result)

    return results

class GANStyleEvaluationLoop:
    """GAN启发的迭代优化循环"""

    def __init__(self, generator_agent, evaluator_agent, quality_threshold: float = 4.0):
        self.generator = generator_agent
        self.evaluator = evaluator_agent
        self.quality_threshold = quality_threshold
        self.iteration_history = []

    async def iterate(self, plan: str, max_iterations: int = 3) -> dict:
        """执行GAN式的反馈循环"""
        current_output = None

        for iteration in range(max_iterations):
            # 步骤1：生成
            if iteration == 0:
                current_output = await self.generator.execute(
                    input_text=plan,
                    prompt_variant="initial_generation"
                )
            else:
                current_output = await self.generator.execute(
                    input_text=f"{plan}\n\n上一轮反馈:{previous_feedback}",
                    prompt_variant="refinement"
                )

            # 步骤2：评估
            evaluation_result = await self.evaluator.execute(
                input_text=current_output,
                evaluation_criteria={
                    "design_quality": "Does it follow best practices?",
                    "originality": "Is it distinctive and creative?",
                    "craft": "Is the implementation sound?",
                    "completeness": "Does it meet all requirements?"
                }
            )

            # 步骤3：质量判断
            avg_score = sum(evaluation_result["scores"].values()) / len(evaluation_result["scores"])

            self.iteration_history.append({
                "iteration": iteration,
                "output_length": len(current_output),
                "scores": evaluation_result["scores"],
                "avg_score": avg_score,
                "issues": evaluation_result["issues"]
            })

            # 步骤4：判断是否继续迭代
            if avg_score >= self.quality_threshold:
                return {
                    "status": "completed",
                    "final_output": current_output,
                    "final_score": avg_score,
                    "iterations": iteration + 1,
                    "history": self.iteration_history
                }
            elif iteration < max_iterations - 1:
                previous_feedback = evaluation_result["improvement_suggestions"]
            else:
                # 达到最大迭代次数
                return {
                    "status": "max_iterations_reached",
                    "final_output": current_output,
                    "final_score": avg_score,
                    "iterations": max_iterations,
                    "history": self.iteration_history
                }

        return {
            "status": "failed",
            "final_output": current_output,
            "final_score": avg_score if current_output else 0,
            "iterations": max_iterations,
            "history": self.iteration_history
        }

class ConcurrencyManager:
    """管理并发任务的执行"""

    def __init__(self, max_concurrent: int = 5):
        self.semaphore = asyncio.Semaphore(max_concurrent)
        self.active_tasks: Dict[str, asyncio.Task] = {}

    async def execute_task(self, task_id: str, coro):
        """限制并发数的任务执行"""
        async with self.semaphore:
            try:
                task = asyncio.create_task(coro)
                self.active_tasks[task_id] = task
                result = await task
                return {"status": "success", "result": result}
            except Exception as e:
                return {"status": "failed", "error": str(e)}
            finally:
                del self.active_tasks[task_id]

    async def execute_batch(self, items: List[Any], executor_fn, max_concurrent: int = 5):
        """批量执行任务,控制并发度"""
        self.semaphore = asyncio.Semaphore(max_concurrent)

        tasks = [
            self.execute_task(f"task_{i}", executor_fn(item))
            for i, item in enumerate(items)
        ]

        results = await asyncio.gather(*tasks)
        return results