跳转到内容
STK

添加MCP工具

本指南解释如何向 STK MCP 服务器添加新工具。

先决条件

  • 理解 Python 异步函数
  • 熟悉 MCP 协议
  • 访问 STK 源代码

添加工具的方法

有两种主要方法可以向 STK MCP 服务器添加新工具：

  1. 直接添加工具方法
  2. 通过 CLI 集成添加工具

方法一：直接添加工具方法

步骤

  1. 定位服务器文件：suan/mcp/server.py

  2. 添加一个用 @mcp.tool() 装饰的新函数：

    @mcp.tool()
async def new_tool_example(param1: str, param2: int = 0) -> str:
    """ 
    新工具的描述

    Args:
        param1: param1 的描述
        param2: param2 的描述（可选）
    """
    # 在这里实现
    try:
        # 你的代码
        return "结果"
    except Exception as e:
        return f"错误: {str(e)}"

  3. 确保安全性：

    • 验证输入
    • 将文件访问限制在工作目录内
    • 正确处理异常

方法二：通过 CLI 集成添加工具

STK MCP Server 能够自动发现并集成现有的 CLI 命令，使 AI 助手能够直接调用这些科学计算功能。这种方法的优势在于，您可以轻松地将现有的命令行工具转换为 MCP 工具，而无需编写额外的代码。

工作原理

  1. MCP 服务器启动时会自动扫描并发现项目中的 CLI 命令（通过 cli_integration.py 中的 discover_and_integrate_cli 函数）
  2. 从 Click 命令组中提取命令及其参数信息（名称、帮助信息、是否必需、默认值等）
  3. 为每个 CLI 命令生成一个 MCP 工具定义，包括工具名称、描述和参数的 JSON Schema
  4. 动态创建并注册 MCP 工具，使其可以通过 MCP 协议被调用
  5. AI 助手调用工具时，参数会被转换为命令行参数，并执行相应的命令
  6. 命令执行结果会被返回给 AI 助手

支持的 CLI 模块

MCP 服务器会自动查找以下模块中的 CLI 命令：

  • suan.cli.main - 主 CLI 模块
  • toolkits.sjob.cli - 作业管理 CLI
  • toolkits.smesh.cli - 网格处理 CLI
  • toolkits.sviz.cli - 可视化 CLI

如果您想添加新的 CLI 模块，可以在 suan/mcp/cli_integration.py 文件中的 CLIIntegration 类中修改 CLI_MODULES 列表。

工具命名规则

每个 CLI 命令会被转换为一个命名为 {模块名}_{命令名} 的 MCP 工具。例如：

  • sviz plot-scalar 命令会被转换为 sviz_plot-scalar 工具
  • smesh generate 命令会被转换为 smesh_generate 工具

对于子命令，命名格式为 {模块名}_{父命令名}_{子命令名}。

添加新的 CLI 命令

要通过 CLI 集成添加新工具，您需要：

  1. 创建一个基于 Click 的 CLI 命令：

    # 在您的模块中（例如 toolkits/your_tool/cli.py）
import click

@click.group()
def cli():
    """您的工具描述"""
    pass

@cli.command()
@click.option('--input', '-i', required=True, help='输入文件路径')
@click.option('--output', '-o', required=True, help='输出文件路径')
def process(input, output):
    """处理数据的命令"""
    # 实现您的命令逻辑
    print(f"处理 {input} 并保存到 {output}")

  2. 确保您的模块在支持的 CLI 模块列表中（如果不在，请修改 CLI_MODULES 列表）

  3. 重启 MCP 服务器，您的命令将自动被集成为 MCP 工具

配置选项

可以在 suan/mcp/config.py 文件中配置以下选项：

  • ENABLE_COMMAND_EXECUTION - 是否启用 CLI 命令执行功能（默认为 True）
  • COMMAND_TIMEOUT - 命令执行最大超时时间（默认为 60 秒）
  • ALLOWED_COMMANDS - 允许执行的命令前缀白名单（默认为 ["sviz", "smesh", "sjob"]）

这些配置选项可以帮助您控制 CLI 集成的行为和安全性。

测试新工具

测试直接添加的工具

  1. 重启 MCP 服务器：

    python -m suan.mcp

  2. 使用 MCP 客户端（如 Cursor、Claude Desktop 或带有 MCP 插件的 VSCode）连接到服务器