Python MCP Server for Code Graph Extraction

This MCP (Model Context Protocol) server provides tools for extracting and analyzing Python code structures, focusing on import/export relationships between files. This is a lightweight implementation that doesn't require an agent system, making it easy to integrate into any Python application.

Features

Code Relationship Discovery: Analyze import relationships between Python files
Smart Code Extraction: Extract only the most relevant code sections to stay within token limits
Directory Context: Include files from the same directory to provide better context
Documentation Inclusion: Always include README.md files (or variants) to provide project documentation
LLM-Friendly Formatting: Format code with proper metadata for language models
MCP Protocol Support: Fully compatible with the Model Context Protocol JSON-RPC standard

The `get_python_code` Tool

The server exposes a powerful code extraction tool that:

Analyzes a target Python file and discovers all imported modules, classes, and functions
Returns the complete code of the target file
Includes code for all referenced objects from other files
Adds additional contextual files from the same directory
Respects token limits to avoid overwhelming language models

Installation

# Clone the repository
git clone https://github.com/yourusername/python-mcp-new.git
cd python-mcp-new

# Create a virtual environment
python -m venv venv
source venv/bin/activate  # On Windows, use: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

Environment Variables

Create a .env file based on the provided .env.example:

# Token limit for extraction
TOKEN_LIMIT=8000

Usage

Configuring for MCP Clients

To configure this MCP server for use in MCP-compatible clients (like Codeium Windsurf), add the following configuration to your client's MCP config file:

{
  "mcpServers": {
    "python-code-explorer": {
      "command": "python",
      "args": [
        "/path/to/python-mcp-new/server.py"
      ],
      "env": {
        "TOKEN_LIMIT": "8000"
      }
    }
  }
}

Replace /path/to/python-mcp-new/server.py with the absolute path to the server.py file on your system.

You can also customize the environment variables:

TOKEN_LIMIT: Maximum token limit for code extraction (default: 8000)

Usage Examples

Direct Function Call

from agent import get_python_code

# Get Python code structure for a specific file
result = get_python_code(
    target_file="/home/user/project/main.py",
    root_repo_path="/home/user/project"  # Optional, defaults to target file directory
)

# Process the result
target_file = result["target_file"]
print(f"Main file: {target_file['file_path']}")
print(f"Docstring: {target_file['docstring']}")

# Display related files
for ref_file in result["referenced_files"]:
    print(f"Related file: {ref_file['file_path']}")
    print(f"Object: {ref_file['object_name']}")
    print(f"Type: {ref_file['object_type']}")

# See if we're close to the token limit
print(f"Token usage: {result['token_count']}/{result['token_limit']}")

Example Response (Direct Function Call)

{
    "target_file": {
        "file_path": "main.py",
        "code": "import os\nimport sys\nfrom utils.helpers import format_output\n\ndef main():\n    args = sys.argv[1:]\n    if not args:\n        print('No arguments provided')\n        return\n    \n    result = format_output(args[0])\n    print(result)\n\nif __name__ == '__main__':\n    main()",
        "type": "target",
        "docstring": ""
    },
    "referenced_files": [
        {
            "file_path": "utils/helpers.py",
            "object_name": "format_output",
            "object_type": "function",
            "code": "def format_output(text):\n    \"\"\"Format the input text for display.\"\"\"\n    if not text:\n        return ''\n    return f'Output: {text.upper()}'\n",
            "docstring": "Format the input text for display.",
            "truncated": false
        }
    ],
    "additional_files": [
        {
            "file_path": "config.py",
            "code": "# Configuration settings\n\nDEBUG = True\nVERSION = '1.0.0'\nMAX_RETRIES = 3\n",
            "type": "related_by_directory",
            "docstring": "Configuration settings for the application."
        }
    ],
    "total_files": 3,
    "token_count": 450,
    "token_limit": 8000
}

Using the MCP Protocol

Listing Available Tools

from agent import handle_mcp_request
import json

# List available tools
list_request = {
    "jsonrpc": "2.0",
    "id": 1,
    "method": "tools/list"
}

response = handle_mcp_request(list_request)
print(json.dumps(response, indent=2))

Example Response (tools/list)

{
  "jsonrpc": "2.0",
  "id": 1,
  "result": {
    "tools": [
      {
        "name": "get_python_code",
        "description": "Return the code of a target Python file and related files based on import/export proximity.",
        "inputSchema": {
          "type": "object",
          "properties": {
            "target_file": {
              "type": "string",
              "description": "Path to the Python file to analyze."
            },
            "root_repo_path": {
              "type": "string",
              "description": "Root directory of the repository. If not provided, the directory of the target file will be used."
            }
          },
          "required": ["target_file"]
        }
      }
    ]
  }
}

Calling get_python_code Tool

from agent import handle_mcp_request
import json

# Call the get_python_code tool
tool_request = {
    "jsonrpc": "2.0",
    "id": 2,
    "method": "tools/call",
    "params": {
        "name": "get_python_code",
        "arguments": {
            "target_file": "/home/user/project/main.py",
            "root_repo_path": "/home/user/project"  # Optional
        }
    }
}

response = handle_mcp_request(tool_request)
print(json.dumps(response, indent=2))

Example Response (tools/call)

{
  "jsonrpc": "2.0",
  "id": 2,
  "result": {
    "content": [
      {
        "type": "text",
        "text": "Python code analysis for /home/user/project/main.py"
      },
      {
        "type": "resource",
        "resource": {
          "uri": "resource://python-code/main.py",
          "mimeType": "application/json",
          "data": {
            "target_file": {
              "file_path": "main.py",
              "code": "import os\nimport sys\nfrom utils.helpers import format_output\n\ndef main():\n    args = sys.argv[1:]\n    if not args:\n        print('No arguments provided')\n        return\n    \n    result = format_output(args[0])\n    print(result)\n\nif __name__ == '__main__':\n    main()",
              "type": "target",
              "docstring": ""
            },
            "referenced_files": [
              {
                "file_path": "utils/helpers.py",
                "object_name": "format_output",
                "object_type": "function",
                "code": "def format_output(text):\n    \"\"\"Format the input text for display.\"\"\"\n    if not text:\n        return ''\n    return f'Output: {text.upper()}'\n",
                "docstring": "Format the input text for display.",
                "truncated": false
              }
            ],
            "additional_files": [
              {
                "file_path": "config.py",
                "code": "# Configuration settings\n\nDEBUG = True\nVERSION = '1.0.0'\nMAX_RETRIES = 3\n",
                "type": "related_by_directory",
                "docstring": "Configuration settings for the application."
              }
            ],
            "total_files": 3,
            "token_count": 450,
            "token_limit": 8000
          }
        }
      }
    ],
    "isError": false
  }
}

Handling Errors

from agent import handle_mcp_request

# Call with invalid file path
faulty_request = {
    "jsonrpc": "2.0",
    "id": 3,
    "method": "tools/call",
    "params": {
        "name": "get_python_code",
        "arguments": {
            "target_file": "/path/to/nonexistent.py"
        }
    }
}

response = handle_mcp_request(faulty_request)
print(json.dumps(response, indent=2))

Example Error Response

{
  "jsonrpc": "2.0",
  "id": 3,
  "result": {
    "content": [
      {
        "type": "text",
        "text": "Error processing Python code: No such file or directory: '/path/to/nonexistent.py'"
      }
    ],
    "isError": true
  }
}

Testing

Run the tests to verify functionality:

python -m unittest discover tests

Key Components

agent.py: Contains the get_python_code function and custom MCP protocol handlers
code_grapher.py: Implements the CodeGrapher class for Python code analysis
server.py: Full MCP server implementation using the MCP Python SDK
run_server.py: CLI tool for running the MCP server
examples/: Example scripts showing how to use the MCP server and client
tests/: Comprehensive test cases for all functionality

Response Format Details

The get_python_code tool returns a structured JSON object with the following fields:

Field	Type	Description
`target_file`	Object	Information about the target Python file
`referenced_files`	Array	List of objects imported by the target file
`additional_files`	Array	Additional context files from the same directory
`total_files`	Number	Total number of files included in the response
`token_count`	Number	Approximate count of tokens in all included code
`token_limit`	Number	Maximum token limit configured for extraction

Target File Object

Field	Type	Description
`file_path`	String	Relative path to the file from the repository root
`code`	String	Complete source code of the file
`type`	String	Always "target"
`docstring`	String	Module-level docstring if available

Referenced File Object

Field	Type	Description
`file_path`	String	Relative path to the file
`object_name`	String	Name of the imported object (class, function, etc.)
`object_type`	String	Type of the object ("class", "function", etc.)
`code`	String	Source code of the specific object
`docstring`	String	Docstring of the object if available
`truncated`	Boolean	Whether the code was truncated due to token limits

Additional File Object

Field	Type	Description
`file_path`	String	Relative path to the file
`code`	String	Complete source code of the file
`type`	String	Type of relation (e.g., "related_by_directory")
`docstring`	String	Module-level docstring if available

Using the MCP SDK Server

This project now includes a full-featured Model Context Protocol (MCP) server built with the official Python MCP SDK. The server exposes our code extraction functionality in a standardized way that can be used with any MCP client, including Claude Desktop.

Starting the Server

# Start the server with default settings
python run_server.py

# Specify a custom name
python run_server.py --name "My Code Explorer"

# Use a specific .env file
python run_server.py --env-file .env.production

Using the MCP Development Mode

With the MCP SDK installed, you can run the server in development mode using the MCP CLI:

# Install the MCP CLI
pip install "mcp[cli]"

# Start the server in development mode with the Inspector UI
mcp dev server.py

This will start the MCP Inspector, a web interface for testing and debugging your server.

Claude Desktop Integration

You can install the server into Claude Desktop to access your code exploration tools directly from Claude:

# Install the server in Claude Desktop
mcp install server.py

# With custom configuration
mcp install server.py --name "Python Code Explorer" -f .env

Custom Server Deployment

For custom deployments, you can use the MCP server directly:

from server import mcp

# Configure the server
mcp.name = "Custom Code Explorer"

# Run the server
mcp.run()

Using the MCP Client

You can use the MCP Python SDK to connect to the server programmatically. See the provided example in examples/mcp_client_example.py:

from mcp.client import Client, Transport

# Connect to the server
client = Client(Transport.subprocess(["python", "server.py"]))
client.initialize()

# List available tools
for tool in client.tools:
    print(f"Tool: {tool.name}")

# Use the get_code tool
result = client.tools.get_code(target_file="path/to/your/file.py")
print(f"Found {len(result['referenced_files'])} referenced files")

# Clean up
client.shutdown()

Run the example:

python examples/mcp_client_example.py [optional_target_file.py]

Adding Additional Tools

You can add additional tools to the MCP server by decorating functions with the @mcp.tool() decorator in server.py:

@mcp.tool()
def analyze_imports(target_file: str) -> Dict[str, Any]:
    """Analyze all imports in a Python file."""
    # Implementation code here
    return {
        "file": target_file,
        "imports": [],  # List of imports found
        "analysis": ""  # Analysis of the imports
    }
    
@mcp.tool()
def find_python_files(directory: str, pattern: str = "*.py") -> list[str]:
    """Find Python files matching a pattern in a directory."""
    from pathlib import Path
    return [str(p) for p in Path(directory).glob(pattern) if p.is_file()]

You can also add resource endpoints to provide data directly:

@mcp.resource("python_stats://{directory}")
def get_stats(directory: str) -> Dict[str, Any]:
    """Get statistics about Python files in a directory."""
    from pathlib import Path
    stats = {
        "directory": directory,
        "file_count": 0,
        "total_lines": 0,
        "average_lines": 0
    }
    
    files = list(Path(directory).glob("**/*.py"))
    stats["file_count"] = len(files)
    
    if files:
        total_lines = 0
        for file in files:
            with open(file, "r") as f:
                total_lines += len(f.readlines())
        stats["total_lines"] = total_lines
        stats["average_lines"] = total_lines / len(files)
    
    return stats

Model Context Protocol Integration

This project fully embraces the Model Context Protocol (MCP) standard, providing two implementation options:

Native MCP Integration: The original implementation in agent.py provides a direct JSON-RPC interface compatible with MCP.
MCP SDK Integration: The new implementation in server.py leverages the official MCP Python SDK for a more robust and feature-rich experience.

Benefits of MCP Integration

Standardized Interface: Makes your tools available to any MCP-compatible client
Enhanced Security: Built-in permissions model and resource controls
Better LLM Integration: Seamless integration with Claude Desktop and other LLM platforms
Improved Developer Experience: Comprehensive tooling like the MCP Inspector

MCP Protocol Version

This implementation supports MCP Protocol version 0.7.0.

For more information about MCP, refer to the official documentation.

Python MCP Server

Integrations

Python MCP Server for Code Graph Extraction

Features

The `get_python_code` Tool

Installation

Environment Variables

Usage

Configuring for MCP Clients

Usage Examples

Direct Function Call

Example Response (Direct Function Call)

Using the MCP Protocol

Listing Available Tools

Example Response (tools/list)

Calling get_python_code Tool

Example Response (tools/call)

Handling Errors

Example Error Response

Testing

Key Components

Response Format Details

Target File Object

Referenced File Object

Additional File Object

Using the MCP SDK Server

Starting the Server

Using the MCP Development Mode

Claude Desktop Integration

Custom Server Deployment

Using the MCP Client

Adding Additional Tools

Model Context Protocol Integration

Benefits of MCP Integration

MCP Protocol Version

Related Resources

Integrations

Python MCP Server for Code Graph Extraction

Features

The get_python_code Tool

Installation

Environment Variables

Usage

Configuring for MCP Clients

Usage Examples

Direct Function Call

Example Response (Direct Function Call)

Using the MCP Protocol

Listing Available Tools

Example Response (tools/list)

Calling get_python_code Tool

Example Response (tools/call)

Handling Errors

Example Error Response

Testing

Key Components

Response Format Details

Target File Object

Referenced File Object

Additional File Object

Using the MCP SDK Server

Starting the Server

Using the MCP Development Mode

Claude Desktop Integration

Custom Server Deployment

Using the MCP Client

Adding Additional Tools

Model Context Protocol Integration

Benefits of MCP Integration

MCP Protocol Version

Related Resources

The `get_python_code` Tool