Ai Spark MCP Server

🚀 Generative AI MCP Protocol Tool

This is a set of generative AI tools based on the MCP protocol, primarily used to optimize the performance of PySpark code and provide a detailed comparative analysis of execution effects. Users can complete code optimization, performance evaluation, and result output through simple interface calls.

🚀 Quick Start

Initialize the Service

from spark_mcp.server import MCPServer

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

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

Client Call Example

from spark_mcp.client import SparkMCPClient

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

    # Provide the PySpark code to be optimized
    original_code = '''
    df = spark.createDataFrame(data)
    result = df.groupBy("id").agg({"value": "sum"})
    '''

    # Get the optimized code
    optimized_code = await client.optimize_spark_code(
        code=original_code,
        optimization_level="advanced",
        save_to_file=True  # Save the result to output/spark_optimized.py
    )

    # Comparative analysis of execution effects
    analysis_report = await client.analyze_performance(original_code, optimized_code)

    await client.close()

asyncio.run(optimize_code())

✨ Features

1. Code Optimization

   • Multi-level Optimization: Supports both basic and advanced optimization modes.
   • Performance Improvement: Improves execution efficiency through efficient computing strategies and resource management.
   • Result Output: Automatically saves the optimized code file.

2. Performance Analysis

   • Execution Time Comparison: Provides the running time data of the original code and the optimized code.
   • Resource Usage Report: Analyzes the consumption of resources such as memory and CPU.
   • Scalability Evaluation: Evaluates the performance under different data scales.

3. Result Output

   • Optimized Code: Saved in output/spark_optimized.py.
   • Analysis Report: Generates an output/performance_analysis.md file containing detailed comparison information.

📦 Installation

git clone [Project URL]
cd ai-mcp
pip install -r requirements.txt

💻 Usage Examples

Basic Usage

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()

Advanced Usage

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()  # Trigger the calculation
    result.show()

if __name__ == "__main__":
    optimized_main()

📚 Documentation

Project Architecture

ai-mcp/
├── spark_mcp/
│   ├── __init__.py
│   ├── server.py      # MCP server implementation
│   └── client.py      # MCP client implementation
├── examples/
│   ├── optimize_code.py           # Example usage
│   └── spark_optimized.py         # Automatically generated optimized code
├── requirements.txt
└── run_server.py                  # Service startup script

Implemented Tools

1. spark_code_optimizer

   • Function: Optimize the performance of PySpark code.
   • Features:
     • Supports multi-level optimization strategies.
     • Provides a comparative analysis of execution time.
     • Automatically generates the optimized code file.

2. performance_analyzer

   • Function: Analyze the performance differences between the original code and the optimized code.
   • Analysis content:
     • Comparison of execution time.
     • Resource usage.
     • Performance improvement rate.
     • Scalability evaluation.

Environment Variable Configuration

• ANTHROPIC_API_KEY: The API key required for the generative AI service.

Execution Effect Comparison

Detailed Comparison Report

Execution Time

• Original Code: 5.2 seconds
• Optimized Code: 3.8 seconds
• Performance Improvement: Approximately 30%

Resource Usage

Frequently Asked Questions

Q: How to configure environment variables? Create a .env file in the project root directory and add the following content:

ANTHROPIC_API_KEY=your_api_key

Q: What optimization strategies are supported? Currently, two modes are supported: basic optimization and advanced optimization. Basic optimization mainly focuses on code structure optimization, while advanced optimization includes advanced techniques such as parallel computing and result caching.

🔧 Technical Details

Project Goals

• Short-term Goal: Implement a complete MCP protocol stack and provide an efficient PySpark code optimization service.
• Medium-term Goal: Expand support for more types of AI tasks, such as image processing and natural language understanding.
• Long-term Goal: Build a general generative AI tool framework supporting multiple computing backends.

Project Advantages

1. High Performance: Achieve efficient resource management and task scheduling through the MCP protocol.
2. High Availability: Support large-scale data sets and complex computing tasks.
3. Ease of Use: Provide simple and intuitive API interfaces and detailed usage documentation.

Contribution Guidelines

Welcome contributions from the community! Please refer to CONTRIBUTING.md to learn how to participate in project development.