MCP Client 的多工具编排逻辑

from typing import Dict, List, Any, Set, Tuple
from collections import defaultdict, deque

class DAGBuilder:
    def __init__(self):
        self.nodes = {}  # 工具节点映射: {node_id: tool_info}
        self.edges = defaultdict(list)  # 边映射: {from_node: [to_node1, to_node2, ...]}
        self.in_degree = defaultdict(int)  # 入度映射: {node_id: in_degree}
    
    def add_tool(self, tool_id: str, tool_info: Dict[str, Any]) -> None:
        """添加工具节点"""
        self.nodes[tool_id] = tool_info
        if tool_id not in self.in_degree:
            self.in_degree[tool_id] = 0
    
    def add_dependency(self, from_tool: str, to_tool: str) -> None:
        """添加工具依赖关系：from_tool 必须在 to_tool 之前执行"""
        # 验证节点是否存在
        if from_tool not in self.nodes:
            raise ValueError(f"Tool {from_tool} not found")
        if to_tool not in self.nodes:
            raise ValueError(f"Tool {to_tool} not found")
        
        # 添加边
        self.edges[from_tool].append(to_tool)
        self.in_degree[to_tool] += 1
    
    def add_conditional_dependency(self, from_tool: str, to_tool: str, condition: str) -> None:
        """添加条件依赖：只有当条件满足时，to_tool 才依赖 from_tool"""
        # 条件依赖的实现需要在执行时动态判断
        # 这里先添加基本依赖，条件判断在执行时处理
        self.add_dependency(from_tool, to_tool)
        
        # 添加条件信息到目标节点
        if "conditions" not in self.nodes[to_tool]:
            self.nodes[to_tool]["conditions"] = []
        self.nodes[to_tool]["conditions"].append({
            "type": "dependency",
            "from_tool": from_tool,
            "condition": condition
        })
    
    def validate_dag(self) -> Tuple[bool, List[str]]:
        """验证 DAG 是否有效，检查是否存在循环依赖"""
        errors = []
        
        # 检查是否存在循环依赖
        if self._has_cycle():
            errors.append("DAG contains cycle dependencies")
        
        # 检查所有节点是否可达
        unreachable_nodes = self._find_unreachable_nodes()
        if unreachable_nodes:
            errors.append(f"Unreachable nodes found: {', '.join(unreachable_nodes)}")
        
        # 检查是否存在孤立节点（没有依赖也没有被依赖）
        isolated_nodes = self._find_isolated_nodes()
        if isolated_nodes:
            errors.append(f"Isolated nodes found: {', '.join(isolated_nodes)}")
        
        return len(errors) == 0, errors
    
    def _has_cycle(self) -> bool:
        """使用 Kahn 算法检查是否存在循环依赖"""
        # 复制入度映射
        in_degree_copy = self.in_degree.copy()
        
        # 将所有入度为 0 的节点加入队列
        queue = deque()
        for node in self.nodes:
            if in_degree_copy[node] == 0:
                queue.append(node)
        
        # 执行拓扑排序
        visited = 0
        while queue:
            current = queue.popleft()
            visited += 1
            
            # 减少所有邻接节点的入度
            for neighbor in self.edges[current]:
                in_degree_copy[neighbor] -= 1
                if in_degree_copy[neighbor] == 0:
                    queue.append(neighbor)
        
        # 如果访问的节点数少于总节点数，说明存在循环
        return visited != len(self.nodes)
    
    def _find_unreachable_nodes(self) -> List[str]:
        """查找不可达节点"""
        # 执行 BFS 从所有入度为 0 的节点开始
        visited = set()
        queue = deque()
        
        for node in self.nodes:
            if self.in_degree[node] == 0:
                queue.append(node)
                visited.add(node)
        
        while queue:
            current = queue.popleft()
            for neighbor in self.edges[current]:
                if neighbor not in visited:
                    visited.add(neighbor)
                    queue.append(neighbor)
        
        # 未访问的节点就是不可达节点
        unreachable = [node for node in self.nodes if node not in visited]
        return unreachable
    
    def _find_isolated_nodes(self) -> List[str]:
        """查找孤立节点"""
        isolated = []
        for node in self.nodes:
            # 没有入边且没有出边的节点
            if self.in_degree[node] == 0 and len(self.edges[node]) == 0:
                isolated.append(node)
        return isolated
    
    def get_topological_order(self) -> List[str]:
        """获取拓扑排序结果"""
        # 使用 Kahn 算法进行拓扑排序
        in_degree_copy = self.in_degree.copy()
        queue = deque()
        order = []
        
        for node in self.nodes:
            if in_degree_copy[node] == 0:
                queue.append(node)
        
        while queue:
            current = queue.popleft()
            order.append(current)
            
            for neighbor in self.edges[current]:
                in_degree_copy[neighbor] -= 1
                if in_degree_copy[neighbor] == 0:
                    queue.append(neighbor)
        
        return order
    
    def get_dag_info(self) -> Dict[str, Any]:
        """获取 DAG 信息"""
        return {
            "nodes": self.nodes,
            "edges": dict(self.edges),
            "in_degree": dict(self.in_degree),
            "topological_order": self.get_topological_order(),
            "is_valid": self.validate_dag()[0]
        }

import asyncio
from typing import Dict, List, Any, Optional
from collections import defaultdict

class Orchestrator:
    def __init__(self, dag_builder, mcp_executor, result_manager):
        self.dag_builder = dag_builder
        self.mcp_executor = mcp_executor
        self.result_manager = result_manager
        
        # 执行状态
        self.execution_state = {
            "status": "idle",  # idle, running, completed, failed, canceled
            "current_step": None,
            "results": {},
            "errors": [],
            "start_time": None,
            "end_time": None,
            "execution_time": 0.0
        }
    
    async def execute_dag(self, input_data: Dict[str, Any] = None) -> Dict[str, Any]:
        """执行 DAG 编排"""
        # 1. 验证 DAG
        is_valid, errors = self.dag_builder.validate_dag()
        if not is_valid:
            self.execution_state["status"] = "failed"
            self.execution_state["errors"] = errors
            return self.execution_state
        
        # 2. 初始化执行状态
        self._reset_execution_state()
        self.execution_state["status"] = "running"
        self.execution_state["start_time"] = asyncio.get_event_loop().time()
        
        # 3. 准备执行数据
        execution_data = input_data.copy() if input_data else {}
        
        # 4. 获取拓扑排序
        topological_order = self.dag_builder.get_topological_order()
        
        # 5. 执行 DAG
        try:
            # 按拓扑顺序执行工具
            for tool_id in topological_order:
                self.execution_state["current_step"] = tool_id
                
                # 获取工具信息
                tool_info = self.dag_builder.nodes[tool_id]
                
                # 检查条件依赖
                if not await self._check_conditions(tool_id, execution_data):
                    continue
                
                # 注入依赖数据
                tool_args = await self._inject_dependencies(tool_id, execution_data)
                
                # 执行工具
                result = await self.mcp_executor.execute_tool(
                    tool_name=tool_info["name"],
                    arguments=tool_args
                )
                
                # 保存执行结果
                self.execution_state["results"][tool_id] = result
                self.result_manager.save_result(tool_id, result)
                
                # 更新执行数据
                execution_data[f"{tool_id}_result"] = result
            
            # 6. 执行完成
            self.execution_state["status"] = "completed"
            self.execution_state["end_time"] = asyncio.get_event_loop().time()
            self.execution_state["execution_time"] = self.execution_state["end_time"] - self.execution_state["start_time"]
            
        except Exception as e:
            # 7. 执行失败
            self.execution_state["status"] = "failed"
            self.execution_state["errors"].append(str(e))
            self.execution_state["end_time"] = asyncio.get_event_loop().time()
            self.execution_state["execution_time"] = self.execution_state["end_time"] - self.execution_state["start_time"]
        
        return self.execution_state
    
    async def execute_dag_parallel(self, input_data: Dict[str, Any] = None, max_workers: int = 5) -> Dict[str, Any]:
        """并行执行 DAG 编排"""
        # 1. 验证 DAG
        is_valid, errors = self.dag_builder.validate_dag()
        if not is_valid:
            self.execution_state["status"] = "failed"
            self.execution_state["errors"] = errors
            return self.execution_state
        
        # 2. 初始化执行状态
        self._reset_execution_state()
        self.execution_state["status"] = "running"
        self.execution_state["start_time"] = asyncio.get_event_loop().time()
        
        # 3. 准备执行数据
        execution_data = input_data.copy() if input_data else {}
        
        # 4. 初始化执行状态
        completed = set()
        running = set()
        results = {}
        
        # 5. 获取 DAG 信息
        dag_info = self.dag_builder.get_dag_info()
        nodes = dag_info["nodes"]
        edges = dag_info["edges"]
        in_degree = {k: v for k, v in dag_info["in_degree"].items()}
        
        # 6. 创建任务队列
        task_queue = asyncio.Queue()
        
        # 7. 将所有入度为 0 的节点加入队列
        for node_id, degree in in_degree.items():
            if degree == 0:
                await task_queue.put(node_id)
                running.add(node_id)
        
        # 8. 执行任务
        worker_semaphore = asyncio.Semaphore(max_workers)
        
        async def execute_tool_task(tool_id):
            async with worker_semaphore:
                try:
                    # 获取工具信息
                    tool_info = nodes[tool_id]
                    
                    # 检查条件依赖
                    if not await self._check_conditions(tool_id, execution_data):
                        return None, tool_id
                    
                    # 注入依赖数据
                    tool_args = await self._inject_dependencies(tool_id, execution_data)
                    
                    # 执行工具
                    result = await self.mcp_executor.execute_tool(
                        tool_name=tool_info["name"],
                        arguments=tool_args
                    )
                    
                    return result, tool_id
                except Exception as e:
                    return e, tool_id
        
        # 9. 管理并行执行
        while running:
            # 取出所有可执行的任务
            pending_tasks = []
            while not task_queue.empty():
                tool_id = await task_queue.get()
                pending_tasks.append(execute_tool_task(tool_id))
            
            if pending_tasks:
                # 等待任务完成
                done, _ = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED)
                
                # 处理完成的任务
                for task in done:
                    result, tool_id = await task
                    
                    # 从运行集合中移除
                    running.remove(tool_id)
                    
                    if isinstance(result, Exception):
                        # 任务失败
                        self.execution_state["status"] = "failed"
                        self.execution_state["errors"].append(f"Tool {tool_id} failed: {str(result)}")
                        return self.execution_state
                    
                    # 任务成功
                    completed.add(tool_id)
                    results[tool_id] = result
                    self.result_manager.save_result(tool_id, result)
                    
                    # 更新执行数据
                    execution_data[f"{tool_id}_result"] = result
                    
                    # 更新入度，并将入度为 0 的节点加入队列
                    for neighbor in edges.get(tool_id, []):
                        in_degree[neighbor] -= 1
                        if in_degree[neighbor] == 0 and neighbor not in completed and neighbor not in running:
                            await task_queue.put(neighbor)
                            running.add(neighbor)
        
        # 10. 执行完成
        self.execution_state["status"] = "completed"
        self.execution_state["results"] = results
        self.execution_state["end_time"] = asyncio.get_event_loop().time()
        self.execution_state["execution_time"] = self.execution_state["end_time"] - self.execution_state["start_time"]
        
        return self.execution_state
    
    async def _check_conditions(self, tool_id: str, execution_data: Dict[str, Any]) -> bool:
        """检查条件依赖"""
        tool_info = self.dag_builder.nodes[tool_id]
        conditions = tool_info.get("conditions", [])
        
        for condition in conditions:
            if condition["type"] == "dependency":
                # 简单条件检查，实际实现可以更复杂
                from_tool = condition["from_tool"]
                condition_expr = condition["condition"]
                
                # 检查前序工具是否执行成功
                if from_tool not in execution_data:
                    return False
                
                # 这里可以实现更复杂的条件表达式求值
                # 简单示例：检查前序工具结果是否包含某个关键字
                from_result = execution_data.get(f"{from_tool}_result", {})
                if isinstance(from_result, dict) and "content" in from_result:
                    if condition_expr not in from_result["content"]:
                        return False
        
        return True
    
    async def _inject_dependencies(self, tool_id: str, execution_data: Dict[str, Any]) -> Dict[str, Any]:
        """注入依赖数据到工具参数"""
        tool_info = self.dag_builder.nodes[tool_id]
        tool_params = tool_info.get("parameters", {})
        
        # 初始化工具参数
        tool_args = {}
        
        # 注入输入数据
        if tool_params.get("input_injection"):
            input_params = tool_params["input_injection"]
            for param_name, input_key in input_params.items():
                if input_key in execution_data:
                    tool_args[param_name] = execution_data[input_key]
        
        # 注入前序工具结果
        for from_tool in self.dag_builder.edges:
            if tool_id in self.dag_builder.edges[from_tool]:
                result_key = f"{from_tool}_result"
                if result_key in execution_data:
                    # 简单注入，实际实现可以更复杂
                    result = execution_data[result_key]
                    if isinstance(result, dict):
                        tool_args.update(result)
        
        return tool_args
    
    def _reset_execution_state(self):
        """重置执行状态"""
        self.execution_state = {
            "status": "idle",
            "current_step": None,
            "results": {},
            "errors": [],
            "start_time": None,
            "end_time": None,
            "execution_time": 0.0
        }
    
    def get_execution_state(self) -> Dict[str, Any]:
        """获取当前执行状态"""
        return self.execution_state
    
    async def cancel_execution(self) -> bool:
        """取消当前执行"""
        # 实现取消逻辑
        self.execution_state["status"] = "canceled"
        self.execution_state["end_time"] = asyncio.get_event_loop().time()
        self.execution_state["execution_time"] = self.execution_state["end_time"] - self.execution_state["start_time"]
        return True

import json
from typing import Dict, Any, Optional
from datetime import datetime

class ResultManager:
    def __init__(self, cache_enabled: bool = True, cache_ttl: int = 3600):
        self.results = {}
        self.cache_enabled = cache_enabled
        self.cache_ttl = cache_ttl  # 秒
    
    def save_result(self, tool_id: str, result: Any, metadata: Dict[str, Any] = None):
        """保存工具执行结果"""
        result_entry = {
            "result": result,
            "timestamp": datetime.now().isoformat(),
            "metadata": metadata or {}
        }
        
        self.results[tool_id] = result_entry
    
    def get_result(self, tool_id: str, use_cache: bool = True) -> Optional[Dict[str, Any]]:
        """获取工具执行结果"""
        if tool_id not in self.results:
            return None
        
        result_entry = self.results[tool_id]
        
        # 检查缓存是否过期
        if not use_cache or not self.cache_enabled:
            return result_entry
        
        # 检查时间是否过期
        timestamp = datetime.fromisoformat(result_entry["timestamp"])
        time_diff = datetime.now() - timestamp
        
        if time_diff.total_seconds() > self.cache_ttl:
            # 缓存过期，移除结果
            del self.results[tool_id]
            return None
        
        return result_entry
    
    def get_result_value(self, tool_id: str, use_cache: bool = True) -> Optional[Any]:
        """获取工具执行结果的值"""
        result_entry = self.get_result(tool_id, use_cache)
        if result_entry:
            return result_entry["result"]
        return None
    
    def get_all_results(self, use_cache: bool = True) -> Dict[str, Any]:
        """获取所有结果"""
        all_results = {}
        for tool_id in list(self.results.keys()):
            result_entry = self.get_result(tool_id, use_cache)
            if result_entry:
                all_results[tool_id] = result_entry
        
        return all_results
    
    def clear_result(self, tool_id: str):
        """清除指定工具的结果"""
        if tool_id in self.results:
            del self.results[tool_id]
    
    def clear_all_results(self):
        """清除所有结果"""
        self.results.clear()
    
    def save_results_to_file(self, file_path: str):
        """将结果保存到文件"""
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(self.results, f, ensure_ascii=False, indent=2, default=str)
    
    def load_results_from_file(self, file_path: str):
        """从文件加载结果"""
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                self.results = json.load(f)
        except Exception as e:
            print(f"Failed to load results from file: {e}")
    
    def get_result_stats(self) -> Dict[str, Any]:
        """获取结果统计信息"""
        stats = {
            "total_results": len(self.results),
            "cached_results": 0,
            "expired_results": 0
        }
        
        # 统计缓存状态
        for tool_id in list(self.results.keys()):
            result_entry = self.get_result(tool_id, use_cache=True)
            if result_entry:
                stats["cached_results"] += 1
            else:
                stats["expired_results"] += 1
        
        return stats

# 示例：多步数据分析任务编排
import asyncio
from mcp_client.dag_builder import DAGBuilder
from mcp_client.orchestrator import Orchestrator
from mcp_client.mcp_executor import MCPExecutor
from mcp_client.result_manager import ResultManager

async def main():
    # 1. 创建 DAG 构建器
    dag_builder = DAGBuilder()
    
    # 2. 添加工具节点
    # 工具 1：数据获取
    dag_builder.add_tool("fetch_data", {
        "name": "web_scraper",
        "description": "从网页获取数据",
        "parameters": {
            "input_injection": {
                "url": "target_url"
            }
        }
    })
    
    # 工具 2：数据清洗
    dag_builder.add_tool("clean_data", {
        "name": "data_cleaner",
        "description": "清洗获取的数据",
        "parameters": {
            "input_injection": {
                "raw_data": "fetch_data_result"
            }
        }
    })
    
    # 工具 3：数据分析
    dag_builder.add_tool("analyze_data", {
        "name": "data_analyzer",
        "description": "分析清洗后的数据",
        "parameters": {
            "input_injection": {
                "clean_data": "clean_data_result"
            }
        }
    })
    
    # 工具 4：生成报告
    dag_builder.add_tool("generate_report", {
        "name": "report_generator",
        "description": "生成数据分析报告",
        "parameters": {
            "input_injection": {
                "analysis_result": "analyze_data_result"
            }
        }
    })
    
    # 工具 5：发送报告
    dag_builder.add_tool("send_report", {
        "name": "email_sender",
        "description": "发送数据分析报告",
        "parameters": {
            "input_injection": {
                "report": "generate_report_result",
                "recipient": "email_recipient"
            }
        }
    })
    
    # 3. 添加依赖关系
    dag_builder.add_dependency("fetch_data", "clean_data")
    dag_builder.add_dependency("clean_data", "analyze_data")
    dag_builder.add_dependency("analyze_data", "generate_report")
    dag_builder.add_dependency("generate_report", "send_report")
    
    # 4. 验证 DAG
    is_valid, errors = dag_builder.validate_dag()
    print(f"DAG is valid: {is_valid}")
    if not is_valid:
        print(f"DAG errors: {errors}")
        return
    
    # 5. 获取拓扑排序
    topological_order = dag_builder.get_topological_order()
    print(f"Topological order: {topological_order}")
    
    # 6. 创建 MCP 执行器
    mcp_executor = MCPExecutor(server_url="http://localhost:8000/mcp")
    
    # 7. 创建结果管理器
    result_manager = ResultManager()
    
    # 8. 创建编排执行器
    orchestrator = Orchestrator(dag_builder, mcp_executor, result_manager)
    
    # 9. 执行 DAG
    input_data = {
        "target_url": "https://example.com/data",
        "email_recipient": "user@example.com"
    }
    
    print("Executing DAG sequentially...")
    result = await orchestrator.execute_dag(input_data)
    print(f"Execution result: {result['status']}")
    print(f"Execution time: {result['execution_time']:.2f} seconds")
    print(f"Results: {result['results']}")
    
    # 10. 并行执行 DAG
    print("\nExecuting DAG in parallel...")
    result = await orchestrator.execute_dag_parallel(input_data, max_workers=3)
    print(f"Execution result: {result['status']}")
    print(f"Execution time: {result['execution_time']:.2f} seconds")
    print(f"Results: {result['results']}")

if __name__ == "__main__":
    asyncio.run(main())

# MCP Client 多工具编排配置
orchestration:
  # DAG 配置
  dag:
    validate_on_start: true
    auto_optimize: true
    max_parallelism: 10
    default_timeout: 3600
  
  # 执行器配置
  executor:
    type: "asyncio"  # asyncio, threading, multiprocessing
    max_workers: 5
    retry_policy:
      enabled: true
      max_retries: 3
      retry_delay: 1.0  # 秒
      backoff_factor: 2.0
  
  # 结果管理配置
  result_manager:
    cache_enabled: true
    cache_ttl: 3600  # 秒
    cache_max_size: 1000
    persist_results: true
    persistence_format: "json"  # json, pickle
    persistence_path: "./results"
  
  # 错误处理配置
  error_handling:
    strategy: "fail_fast"  # fail_fast, continue_on_error, skip_failed
    error_threshold: 0.1
    fallback_strategy: "null"
  
  # 监控配置
  monitoring:
    enabled: true
    metrics_port: 9090
    tracing_enabled: false
    logging_level: "INFO"
  
  # 可视化配置
  visualization:
    enabled: true
    port: 8080
    allowed_origins: ["*"]
  
  # 安全配置
  security:
    enable_auth: false
    allowed_ips: ["127.0.0.1"]
    encrypt_results: false
  
  # 集成配置
  integrations:
    airflow:
      enabled: false
      airflow_url: "http://localhost:8080"
    prefect:
      enabled: false
      prefect_url: "http://localhost:4200"