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

…