Ai Spark MCP Server

🚀 生成式AI MCP协议工具

这是一套基于MCP协议的生成式AI工具，能优化PySpark代码性能，还可提供详细的执行效果对比分析。用户通过简单接口调用，就能完成代码优化、性能评估和结果输出。

🚀 快速开始

初始化服务

from spark_mcp.server import MCPServer

async def main():
    server = MCPServer()
    await server.start()

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

客户端调用示例

from spark_mcp.client import SparkMCPClient

async def optimize_code():
    client = SparkMCPClient()
    await client.connect()

    # 提供需要优化的PySpark代码
    original_code = '''
    df = spark.createDataFrame(data)
    result = df.groupBy("id").agg({"value": "sum"})
    '''

    # 获取优化后的代码
    optimized_code = await client.optimize_spark_code(
        code=original_code,
        optimization_level="advanced",
        save_to_file=True  # 结果保存至output/spark_optimized.py
    )

    # 执行效果对比分析
    analysis_report = await client.analyze_performance(original_code, optimized_code)

    await client.close()

asyncio.run(optimize_code())

✨ 主要特性

代码优化

• 多级优化：支持基础和高级两种优化模式。
• 性能提升：通过高效的计算策略和资源管理提高执行效率。
• 结果输出：自动保存优化后的代码文件。

性能分析

• 执行时间对比：提供原始代码与优化后代码的运行时长数据。
• 资源使用报告：分析内存、CPU等资源消耗情况。
• 可扩展性评估：评估不同数据规模下的性能表现。

结果输出

• 优化后代码：保存在output/spark_optimized.py。
• 分析报告：生成output/performance_analysis.md文件，包含详细对比信息。

📦 安装指南

git clone [项目地址]
cd ai-mcp
pip install -r requirements.txt

💻 使用示例

基础用法

# 初始化服务代码
from spark_mcp.server import MCPServer

async def main():
    server = MCPServer()
    await server.start()

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

高级用法

# 客户端调用示例代码
from spark_mcp.client import SparkMCPClient

async def optimize_code():
    client = SparkMCPClient()
    await client.connect()

    # 提供需要优化的PySpark代码
    original_code = '''
    df = spark.createDataFrame(data)
    result = df.groupBy("id").agg({"value": "sum"})
    '''

    # 获取优化后的代码
    optimized_code = await client.optimize_spark_code(
        code=original_code,
        optimization_level="advanced",
        save_to_file=True  # 结果保存至output/spark_optimized.py
    )

    # 执行效果对比分析
    analysis_report = await client.analyze_performance(original_code, optimized_code)

    await client.close()

asyncio.run(optimize_code())

📚 详细文档

项目架构

ai-mcp/
├── spark_mcp/
│   ├── __init__.py
│   ├── server.py      # MCP服务端实现
│   └── client.py      # MCP客户端实现
├── examples/
│   ├── optimize_code.py           # 示例用法
│   └── spark_optimized.py         # 自动生成的优化代码
├── requirements.txt
└── run_server.py                  # 服务启动脚本

已实现工具

1. spark_code_optimizer
   • 功能：对PySpark代码进行性能优化。
   • 特性：
     • 支持多级优化策略。
     • 提供执行时间对比分析。
     • 自动生成优化后代码文件。
2. performance_analyzer
   • 功能：分析原始代码与优化后代码的性能差异。
   • 分析内容：
     • 执行时长对比。
     • 资源使用情况。
     • 性能提升幅度。
     • 可扩展性评估。

环境变量配置

• ANTHROPIC_API_KEY：生成式AI服务所需的API密钥。

示例代码说明

优化前的PySpark代码

from pyspark.sql import SparkSession

def main():
    spark = SparkSession.builder.appName("example").getOrCreate()
    data = [("a", 1), ("b", 2)]
    df = spark.createDataFrame(data)
    result = df.groupBy("id").agg({"value": "sum"})
    result.show()

if __name__ == "__main__":
    main()

优化后的PySpark代码

from pyspark.sql import SparkSession

def optimized_main():
    spark = SparkSession.builder.appName("optimized_example").getOrCreate()
    data = [("a", 1), ("b", 2)]
    df = spark.createDataFrame(data)
    result = df.groupBy("id").agg({"value": "sum"}).cache()
    result.count()  # 触发计算
    result.show()

if __name__ == "__main__":
    optimized_main()

执行效果对比

详细对比报告

执行时间

• 原始代码：5.2秒
• 优化后代码：3.8秒
• 性能提升：约30%

资源使用情况

常见问题解答

Q：如何配置环境变量？ 在项目根目录下创建.env文件，添加以下内容：

ANTHROPIC_API_KEY=your_api_key

Q：支持哪些优化策略？ 目前支持基础优化和高级优化两种模式。基础优化主要进行代码结构优化，而高级优化则包括并行计算、结果缓存等高级技术。

项目目标

• 短期目标：实现完整的MCP协议栈，提供高效的PySpark代码优化服务。
• 中期目标：扩展支持更多AI任务类型，如图像处理、自然语言理解等。
• 长期目标：构建一个通用的生成式AI工具框架，支持多种计算后端。

项目优势

1. 高性能：通过MCP协议实现高效的资源管理和任务调度。
2. 高可用性：支持大规模数据集和复杂计算任务。
3. 易用性：提供简洁直观的API接口和详细的使用文档。

贡献指南

欢迎社区贡献代码！请参考CONTRIBUTING.md了解如何参与项目开发。