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_codefunction and custom MCP protocol handlers - code_grapher.py: Implements the
CodeGrapherclass 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.pyprovides a direct JSON-RPC interface compatible with MCP. -
MCP SDK Integration: The new implementation in
server.pyleverages 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.