299f7d3897
- Add collaboration demo showing dynamic context requests - Implement chat tool for general conversations and brainstorming - Add tool selection guide with clear boundaries - Introduce models configuration system - Update prompts for better tool descriptions - Refactor server to remove redundant functionality - Add comprehensive tests for collaboration features - Enhance base tool with collaborative features This enables Claude to request additional context from Gemini during tool execution, improving analysis quality and accuracy. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
323 lines
12 KiB
Python
323 lines
12 KiB
Python
"""
|
|
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
|
|
|
|
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)
|
|
|
|
# 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)
|