base_agent/mcp/mcp_protocol.py

"""
mcp/mcp_protocol.py
MCP 协议数据结构定义
新增: ToolStep / ChainPlan / StepResult / ChainResult 支持多工具串行调用
"""

import uuid
from dataclasses import dataclass, field
from typing import Any


# ── MCP 方法常量 ───────────────────────────────────────────────
class MCPMethod:
    TOOLS_LIST     = "tools/list"
    TOOLS_CALL     = "tools/call"
    RESOURCES_READ = "resources/read"


# ════════════════════════════════════════════════════════════════
# JSON-RPC 基础消息
# ════════════════════════════════════════════════════════════════

@dataclass
class MCPRequest:
    """
    MCP 工具调用请求（JSON-RPC 2.0）

    示例:
        {
            "jsonrpc": "2.0",
            "id": "abc-123",
            "method": "tools/call",
            "params": {"name": "calculator", "arguments": {"expression": "1+1"}}
        }
    """
    method:  str
    params:  dict[str, Any] = field(default_factory=dict)
    jsonrpc: str            = "2.0"
    id:      str            = field(default_factory=lambda: str(uuid.uuid4())[:8])

    def to_dict(self) -> dict:
        return {
            "jsonrpc": self.jsonrpc,
            "id":      self.id,
            "method":  self.method,
            "params":  self.params,
        }


@dataclass
class MCPResponse:
    """
    MCP 工具调用响应（JSON-RPC 2.0）
    """
    id:      str
    result:  dict[str, Any] | None = None
    error:   dict[str, Any] | None = None
    jsonrpc: str                   = "2.0"

    @property
    def success(self) -> bool:
        return self.error is None

    @property
    def content(self) -> str:
        """提取响应中的文本内容"""
        if not self.success or not self.result:
            return self.error.get("message", "Unknown error") if self.error else ""
        items = self.result.get("content", [])
        return "\n".join(
            item.get("text", "") for item in items if item.get("type") == "text"
        )

    def to_dict(self) -> dict:
        base = {"jsonrpc": self.jsonrpc, "id": self.id}
        if self.success:
            base["result"] = self.result
        else:
            base["error"] = self.error
        return base


@dataclass
class ToolSchema:
    """工具元数据描述，供 LLM 识别和选择工具"""
    name:        str
    description: str
    parameters:  dict[str, Any]

    def to_dict(self) -> dict:
        return {
            "name":        self.name,
            "description": self.description,
            "inputSchema": {
                "type":       "object",
                "properties": self.parameters,
            },
        }


# ════════════════════════════════════════════════════════════════
# 串行调用链数据结构（新增）
# ════════════════════════════════════════════════════════════════

@dataclass
class ToolStep:
    """
    调用链中的单个执行步骤

    Attributes:
        step_id:      步骤编号（从 1 开始）
        tool_name:    要调用的工具名称
        arguments:    工具参数（支持 {{STEP_N_OUTPUT}} 占位符引用前步结果）
        description:  该步骤的自然语言说明
        depends_on:   依赖的前置步骤编号列表（用于上下文注入）

    占位符示例:
        arguments = {"query": "{{STEP_1_OUTPUT}}"}
        → 执行时自动替换为第 1 步的输出内容
    """
    step_id:     int
    tool_name:   str
    arguments:   dict[str, Any]
    description: str            = ""
    depends_on:  list[int]      = field(default_factory=list)

    def to_mcp_request(self) -> MCPRequest:
        """将步骤转换为 MCP 请求"""
        return MCPRequest(
            method=MCPMethod.TOOLS_CALL,
            params={"name": self.tool_name, "arguments": self.arguments},
        )

    def inject_context(self, context: dict[str, str]) -> "ToolStep":
        """
        将前步输出注入当前步骤的参数占位符

        Args:
            context: {"STEP_1_OUTPUT": "...", "STEP_2_OUTPUT": "..."}

        Returns:
            注入后的新 ToolStep（不修改原对象）
        """
        resolved_args = {}
        for key, value in self.arguments.items():
            if isinstance(value, str):
                for placeholder, replacement in context.items():
                    # 截取前 500 字符避免参数过长
                    value = value.replace(f"{{{{{placeholder}}}}}", replacement[:500])
            resolved_args[key] = value
        return ToolStep(
            step_id=self.step_id,
            tool_name=self.tool_name,
            arguments=resolved_args,
            description=self.description,
            depends_on=self.depends_on,
        )


@dataclass
class ChainPlan:
    """
    完整的工具调用链执行计划

    Attributes:
        steps:       有序的步骤列表（按执行顺序排列）
        goal:        整体目标描述
        is_single:   是否为单步调用（优化路径）
    """
    steps:     list[ToolStep]
    goal:      str  = ""
    is_single: bool = False

    @property
    def step_count(self) -> int:
        return len(self.steps)

    def __repr__(self) -> str:
        steps_desc = " → ".join(
            f"[{s.step_id}]{s.tool_name}" for s in self.steps
        )
        return f"ChainPlan(goal={self.goal!r}, chain={steps_desc})"


@dataclass
class StepResult:
    """
    单个步骤的执行结果

    Attributes:
        step_id:    对应的步骤编号
        tool_name:  执行的工具名称
        success:    是否执行成功
        output:     工具输出内容
        error:      失败时的错误信息
    """
    step_id:   int
    tool_name: str
    success:   bool
    output:    str
    error:     str | None = None

    @property
    def context_key(self) -> str:
        """生成占位符 key，供后续步骤引用"""
        return f"STEP_{self.step_id}_OUTPUT"


@dataclass
class ChainResult:
    """
    完整调用链的汇总结果

    Attributes:
        goal:          原始目标
        step_results:  每步的执行结果列表
        final_reply:   LLM 整合后的最终回复
        success:       整体是否成功
        failed_step:   首个失败步骤编号（成功时为 None）
    """
    goal:         str
    step_results: list[StepResult]
    final_reply:  str  = ""
    success:      bool = True
    failed_step:  int | None = None

    @property
    def completed_steps(self) -> int:
        return sum(1 for r in self.step_results if r.success)

    @property
    def total_steps(self) -> int:
        return len(self.step_results)

    def get_summary(self) -> str:
        """生成执行摘要字符串"""
        lines = [f"📋 执行计划: {self.goal}", f"📊 完成步骤: {self.completed_steps}/{self.total_steps}"]
        for r in self.step_results:
            icon = "✅" if r.success else "❌"
            lines.append(f"  {icon} Step {r.step_id} [{r.tool_name}]: {r.output[:60]}...")
        return "\n".join(lines)