refactor: cleanup and comprehensive documentation

Major changes:
- Add comprehensive documentation to all modules with detailed docstrings
- Remove unused THINKING_MODEL config (use single GEMINI_MODEL with thinking_mode param)
- Remove list_models functionality (simplified to single model configuration)
- Rename DEFAULT_MODEL to GEMINI_MODEL for clarity
- Remove unused python-dotenv dependency
- Fix missing pydantic in setup.py dependencies

Documentation improvements:
- Document security measures in file_utils.py (path validation, sandboxing)
- Add detailed comments to critical logic sections
- Document tool creation process in BaseTool
- Explain configuration values and their impact
- Add comprehensive function-level documentation

Code quality:
- Apply black formatting to all files
- Fix all ruff linting issues
- Update tests to match refactored code
- All 63 tests passing

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Fahad
2025-06-09 19:04:24 +04:00
parent fd6e2f9b64
commit 783ba73181
12 changed files with 639 additions and 260 deletions

197
server.py
View File

@@ -1,5 +1,21 @@
"""
Gemini MCP Server - Main server implementation
This module implements the core MCP (Model Context Protocol) server that provides
AI-powered tools for code analysis, review, and assistance using Google's Gemini models.
The server follows the MCP specification to expose various AI tools as callable functions
that can be used by MCP clients (like Claude). Each tool provides specialized functionality
such as code review, debugging, deep thinking, and general chat capabilities.
Key Components:
- MCP Server: Handles protocol communication and tool discovery
- Tool Registry: Maps tool names to their implementations
- Request Handler: Processes incoming tool calls and returns formatted responses
- Configuration: Manages API keys and model settings
The server runs on stdio (standard input/output) and communicates using JSON-RPC messages
as defined by the MCP protocol.
"""
import asyncio
@@ -9,14 +25,13 @@ import sys
from datetime import datetime
from typing import Any, Dict, List
from google import genai
from mcp.server import Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool
from config import (
DEFAULT_MODEL,
GEMINI_MODEL,
MAX_CONTEXT_TOKENS,
__author__,
__updated__,
@@ -31,41 +46,67 @@ from tools import (
ThinkDeeperTool,
)
# Configure logging
# Configure logging for server operations
# Set to INFO level to capture important operational messages without being too verbose
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Create the MCP server instance
# Create the MCP server instance with a unique name identifier
# This name is used by MCP clients to identify and connect to this specific server
server: Server = Server("gemini-server")
# Initialize tools
# Initialize the tool registry with all available AI-powered tools
# Each tool provides specialized functionality for different development tasks
# Tools are instantiated once and reused across requests (stateless design)
TOOLS = {
"think_deeper": ThinkDeeperTool(),
"review_code": ReviewCodeTool(),
"debug_issue": DebugIssueTool(),
"analyze": AnalyzeTool(),
"chat": ChatTool(),
"review_changes": ReviewChanges(),
"think_deeper": ThinkDeeperTool(), # Extended reasoning for complex problems
"review_code": ReviewCodeTool(), # Comprehensive code review and quality analysis
"debug_issue": DebugIssueTool(), # Root cause analysis and debugging assistance
"analyze": AnalyzeTool(), # General-purpose file and code analysis
"chat": ChatTool(), # Interactive development chat and brainstorming
"review_changes": ReviewChanges(), # Pre-commit review of git changes
}
def configure_gemini():
"""Configure Gemini API with the provided API key"""
"""
Configure Gemini API with the provided API key.
This function validates that the GEMINI_API_KEY environment variable is set.
The actual API key is used when creating Gemini clients within individual tools
to ensure proper isolation and error handling.
Raises:
ValueError: If GEMINI_API_KEY environment variable is not set
"""
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise ValueError(
"GEMINI_API_KEY environment variable is required. "
"Please set it with your Gemini API key."
)
# API key is used when creating clients in tools
# Note: We don't store the API key globally for security reasons
# Each tool creates its own Gemini client with the API key when needed
logger.info("Gemini API key found")
@server.list_tools()
async def handle_list_tools() -> List[Tool]:
"""List all available tools with verbose descriptions"""
"""
List all available tools with their descriptions and input schemas.
This handler is called by MCP clients during initialization to discover
what tools are available. Each tool provides:
- name: Unique identifier for the tool
- description: Detailed explanation of what the tool does
- inputSchema: JSON Schema defining the expected parameters
Returns:
List of Tool objects representing all available tools
"""
tools = []
# Add all registered AI-powered tools from the TOOLS registry
for tool in TOOLS.values():
tools.append(
Tool(
@@ -75,17 +116,10 @@ async def handle_list_tools() -> List[Tool]:
)
)
# Add utility tools
# Add utility tools that provide server metadata and configuration info
# These tools don't require AI processing but are useful for clients
tools.extend(
[
Tool(
name="list_models",
description=(
"LIST AVAILABLE MODELS - Show all Gemini models you can use. "
"Lists model names, descriptions, and which one is the default."
),
inputSchema={"type": "object", "properties": {}},
),
Tool(
name="get_version",
description=(
@@ -102,100 +136,65 @@ async def handle_list_tools() -> List[Tool]:
@server.call_tool()
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
"""Handle tool execution requests"""
"""
Handle incoming tool execution requests from MCP clients.
# Handle dynamic tools
This is the main request dispatcher that routes tool calls to their
appropriate handlers. It supports both AI-powered tools (from TOOLS registry)
and utility tools (implemented as static functions).
Args:
name: The name of the tool to execute
arguments: Dictionary of arguments to pass to the tool
Returns:
List of TextContent objects containing the tool's response
"""
# Route to AI-powered tools that require Gemini API calls
if name in TOOLS:
tool = TOOLS[name]
return await tool.execute(arguments)
# Handle static tools
elif name == "list_models":
return await handle_list_models()
# Route to utility tools that provide server information
elif name == "get_version":
return await handle_get_version()
# Handle unknown tool requests gracefully
else:
return [TextContent(type="text", text=f"Unknown tool: {name}")]
async def handle_list_models() -> List[TextContent]:
"""List available Gemini models"""
try:
import json
# Get API key
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
return [TextContent(type="text", text="Error: GEMINI_API_KEY not set")]
client = genai.Client(api_key=api_key)
models = []
# List models using the new API
try:
model_list = client.models.list()
for model_info in model_list:
models.append(
{
"name": getattr(model_info, "id", "Unknown"),
"display_name": getattr(
model_info,
"display_name",
getattr(model_info, "id", "Unknown"),
),
"description": getattr(
model_info, "description", "No description"
),
"is_default": getattr(model_info, "id", "").endswith(
DEFAULT_MODEL
),
}
)
except Exception:
# Fallback: return some known models
models = [
{
"name": "gemini-2.5-pro-preview-06-05",
"display_name": "Gemini 2.5 Pro",
"description": "Latest Gemini 2.5 Pro model",
"is_default": True,
},
{
"name": "gemini-2.0-flash-thinking-exp",
"display_name": "Gemini 2.0 Flash Thinking",
"description": "Enhanced reasoning model",
"is_default": False,
},
]
return [TextContent(type="text", text=json.dumps(models, indent=2))]
except Exception as e:
return [TextContent(type="text", text=f"Error listing models: {str(e)}")]
async def handle_get_version() -> List[TextContent]:
"""Get version and configuration information"""
"""
Get comprehensive version and configuration information about the server.
Provides details about the server version, configuration settings,
available tools, and runtime environment. Useful for debugging and
understanding the server's capabilities.
Returns:
Formatted text with version and configuration details
"""
# Gather comprehensive server information
version_info = {
"version": __version__,
"updated": __updated__,
"author": __author__,
"default_model": DEFAULT_MODEL,
"gemini_model": GEMINI_MODEL,
"max_context_tokens": f"{MAX_CONTEXT_TOKENS:,}",
"python_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
"server_started": datetime.now().isoformat(),
"available_tools": list(TOOLS.keys()) + ["chat", "list_models", "get_version"],
"available_tools": list(TOOLS.keys()) + ["get_version"],
}
# Format the information in a human-readable way
text = f"""Gemini MCP Server v{__version__}
Updated: {__updated__}
Author: {__author__}
Configuration:
- Default Model: {DEFAULT_MODEL}
- Gemini Model: {GEMINI_MODEL}
- Max Context: {MAX_CONTEXT_TOKENS:,} tokens
- Python: {version_info['python_version']}
- Started: {version_info['server_started']}
@@ -209,11 +208,21 @@ For updates, visit: https://github.com/BeehiveInnovations/gemini-mcp-server"""
async def main():
"""Main entry point for the server"""
# Configure Gemini API
"""
Main entry point for the MCP server.
Initializes the Gemini API configuration and starts the server using
stdio transport. The server will continue running until the client
disconnects or an error occurs.
The server communicates via standard input/output streams using the
MCP protocol's JSON-RPC message format.
"""
# Validate that Gemini API key is available before starting
configure_gemini()
# Run the server using stdio transport
# Run the server using stdio transport (standard input/output)
# This allows the server to be launched by MCP clients as a subprocess
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
@@ -221,7 +230,7 @@ async def main():
InitializationOptions(
server_name="gemini",
server_version=__version__,
capabilities={"tools": {}},
capabilities={"tools": {}}, # Advertise tool support capability
),
)