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:
197
server.py
197
server.py
@@ -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
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user