""" Base class for all Gemini MCP tools """ from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Literal import os import json from google import genai from google.genai import types from mcp.types import TextContent from pydantic import BaseModel, Field from .models import ToolOutput, ClarificationRequest class ToolRequest(BaseModel): """Base request model for all tools""" model: Optional[str] = Field( None, description="Model to use (defaults to Gemini 2.5 Pro)" ) temperature: Optional[float] = Field( None, description="Temperature for response (tool-specific defaults)" ) thinking_mode: Optional[Literal["minimal", "low", "medium", "high", "max"]] = Field( None, description="Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", ) class BaseTool(ABC): """Base class for all Gemini tools""" def __init__(self): self.name = self.get_name() self.description = self.get_description() self.default_temperature = self.get_default_temperature() @abstractmethod def get_name(self) -> str: """Return the tool name""" pass @abstractmethod def get_description(self) -> str: """Return the verbose tool description for Claude""" pass @abstractmethod def get_input_schema(self) -> Dict[str, Any]: """Return the JSON schema for tool inputs""" pass @abstractmethod def get_system_prompt(self) -> str: """Return the system prompt for this tool""" pass def get_default_temperature(self) -> float: """Return default temperature for this tool""" return 0.5 def get_default_thinking_mode(self) -> str: """Return default thinking_mode for this tool""" return "medium" # Default to medium thinking for better reasoning @abstractmethod def get_request_model(self): """Return the Pydantic model for request validation""" pass def validate_file_paths(self, request) -> Optional[str]: """ Validate that all file paths in the request are absolute. Returns error message if validation fails, None if all paths are valid. """ # Check if request has 'files' attribute if hasattr(request, "files") and request.files: for file_path in request.files: if not os.path.isabs(file_path): return ( f"Error: All file paths must be absolute. " f"Received relative path: {file_path}\n" f"Please provide the full absolute path starting with '/'" ) # Check if request has 'path' attribute (for review_pending_changes) if hasattr(request, "path") and request.path: if not os.path.isabs(request.path): return ( f"Error: Path must be absolute. " f"Received relative path: {request.path}\n" f"Please provide the full absolute path starting with '/'" ) return None async def execute(self, arguments: Dict[str, Any]) -> List[TextContent]: """Execute the tool with given arguments""" try: # Validate request request_model = self.get_request_model() request = request_model(**arguments) # Validate file paths path_error = self.validate_file_paths(request) if path_error: error_output = ToolOutput( status="error", content=path_error, content_type="text", ) return [TextContent(type="text", text=error_output.model_dump_json())] # Prepare the prompt prompt = await self.prepare_prompt(request) # Get model configuration from config import DEFAULT_MODEL model_name = getattr(request, "model", None) or DEFAULT_MODEL temperature = getattr(request, "temperature", None) if temperature is None: temperature = self.get_default_temperature() thinking_mode = getattr(request, "thinking_mode", None) if thinking_mode is None: thinking_mode = self.get_default_thinking_mode() # Create and configure model model = self.create_model(model_name, temperature, thinking_mode) # Generate response response = model.generate_content(prompt) # Handle response and create standardized output if response.candidates and response.candidates[0].content.parts: raw_text = response.candidates[0].content.parts[0].text # Check if this is a clarification request tool_output = self._parse_response(raw_text, request) else: finish_reason = ( response.candidates[0].finish_reason if response.candidates else "Unknown" ) tool_output = ToolOutput( status="error", content=f"Response blocked or incomplete. Finish reason: {finish_reason}", content_type="text", ) # Serialize the standardized output as JSON return [TextContent(type="text", text=tool_output.model_dump_json())] except Exception as e: error_output = ToolOutput( status="error", content=f"Error in {self.name}: {str(e)}", content_type="text", ) return [TextContent(type="text", text=error_output.model_dump_json())] def _parse_response(self, raw_text: str, request) -> ToolOutput: """Parse the raw response and determine if it's a clarification request""" try: # Try to parse as JSON to check for clarification requests potential_json = json.loads(raw_text.strip()) if ( isinstance(potential_json, dict) and potential_json.get("status") == "requires_clarification" ): # Validate the clarification request structure clarification = ClarificationRequest(**potential_json) return ToolOutput( status="requires_clarification", content=clarification.model_dump_json(), content_type="json", metadata={ "original_request": ( request.model_dump() if hasattr(request, "model_dump") else str(request) ) }, ) except (json.JSONDecodeError, ValueError, TypeError): # Not a JSON clarification request, treat as normal response pass # Normal text response - format using tool-specific formatting formatted_content = self.format_response(raw_text, request) # Determine content type based on the formatted content content_type = ( "markdown" if any( marker in formatted_content for marker in ["##", "**", "`", "- ", "1. "] ) else "text" ) return ToolOutput( status="success", content=formatted_content, content_type=content_type, metadata={"tool_name": self.name}, ) @abstractmethod async def prepare_prompt(self, request) -> str: """Prepare the full prompt for Gemini""" pass def format_response(self, response: str, request) -> str: """Format the response for display (can be overridden)""" return response def create_model( self, model_name: str, temperature: float, thinking_mode: str = "medium" ): """Create a configured Gemini model with thinking configuration""" # Map thinking modes to budget values thinking_budgets = { "minimal": 128, # Minimum for 2.5 Pro "low": 2048, "medium": 8192, "high": 16384, "max": 32768, } thinking_budget = thinking_budgets.get(thinking_mode, 8192) # For models supporting thinking config, use the new API # Skip in test environment to allow mocking if "2.5" in model_name and not os.environ.get("PYTEST_CURRENT_TEST"): try: # Get API key api_key = os.environ.get("GEMINI_API_KEY") if not api_key: raise ValueError("GEMINI_API_KEY environment variable is required") client = genai.Client(api_key=api_key) # Create a wrapper to match the expected interface class ModelWrapper: def __init__( self, client, model_name, temperature, thinking_budget ): self.client = client self.model_name = model_name self.temperature = temperature self.thinking_budget = thinking_budget def generate_content(self, prompt): response = self.client.models.generate_content( model=self.model_name, contents=prompt, config=types.GenerateContentConfig( temperature=self.temperature, candidate_count=1, thinking_config=types.ThinkingConfig( thinking_budget=self.thinking_budget ), ), ) # Convert to match expected format class ResponseWrapper: def __init__(self, text): self.text = text self.candidates = [ type( "obj", (object,), { "content": type( "obj", (object,), { "parts": [ type( "obj", (object,), {"text": text}, ) ] }, )(), "finish_reason": "STOP", }, ) ] return ResponseWrapper(response.text) return ModelWrapper(client, model_name, temperature, thinking_budget) except Exception: # Fall back to regular genai model if new API fails pass # For non-2.5 models or if thinking not needed, use regular API # Get API key api_key = os.environ.get("GEMINI_API_KEY") if not api_key: raise ValueError("GEMINI_API_KEY environment variable is required") client = genai.Client(api_key=api_key) # Create wrapper for consistency class SimpleModelWrapper: def __init__(self, client, model_name, temperature): self.client = client self.model_name = model_name self.temperature = temperature def generate_content(self, prompt): response = self.client.models.generate_content( model=self.model_name, contents=prompt, config=types.GenerateContentConfig( temperature=self.temperature, candidate_count=1, ), ) # Convert to match expected format class ResponseWrapper: def __init__(self, text): self.text = text self.candidates = [ type( "obj", (object,), { "content": type( "obj", (object,), { "parts": [ type("obj", (object,), {"text": text}) ] }, )(), "finish_reason": "STOP", }, ) ] return ResponseWrapper(response.text) return SimpleModelWrapper(client, model_name, temperature)