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瞭解如何參與項目開發。