torbjorn/my-pal-mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'BeehiveInnovations:main' into feat/comprehensive-project-improvements

Browse Source
This commit is contained in:
PCITI
2025-06-12 20:24:26 +02:00
committed by GitHub
parent 1364475d67 22093bbf18
commit 933097813c
28 changed files with 268 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
FIX_SUMMARY.md
View File

@@ -1,40 +0,0 @@
# Fix for Conversation History Bug in Continuation Flow
## Problem
When using `continuation_id` to continue a conversation, the conversation history (with embedded files) was being lost for tools that don't have a `prompt` field. Only new file content was being passed to the tool, resulting in minimal content (e.g., 322 chars for just a NOTE about files already in history).
## Root Cause
1. `reconstruct_thread_context()` builds conversation history and stores it in `arguments["prompt"]`
2. Different tools use different field names for user input:
- `chat` → `prompt`
- `analyze` → `question`
- `debug` → `error_description`
- `codereview` → `context`
- `thinkdeep` → `current_analysis`
- `precommit` → `original_request`
3. The enhanced prompt with conversation history was being placed in the wrong field
4. Tools would only see their new input, not the conversation history
## Solution
Modified `reconstruct_thread_context()` in `server.py` to:
1. Create a mapping of tool names to their primary input fields
2. Extract the user's new input from the correct field based on the tool
3. Store the enhanced prompt (with conversation history) back into the correct field
## Changes Made
1. **server.py**:
- Added `prompt_field_mapping` to map tools to their input fields
- Modified to extract user input from the correct field
- Modified to store enhanced prompt in the correct field
2. **tests/test_conversation_field_mapping.py**:
- Added comprehensive tests to verify the fix works for all tools
- Tests ensure conversation history is properly mapped to each tool's field
## Verification
All existing tests pass, including:
- `test_conversation_memory.py` (18 tests)
- `test_cross_tool_continuation.py` (4 tests)
- New `test_conversation_field_mapping.py` (2 tests)
The fix ensures that when continuing conversations, tools receive the full conversation history with embedded files, not just new content.

12
config.py
View File

@@ -27,7 +27,15 @@ DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "auto")
# Validate DEFAULT_MODEL and set to "auto" if invalid # Validate DEFAULT_MODEL and set to "auto" if invalid
# Only include actually supported models from providers # Only include actually supported models from providers
VALID_MODELS = ["auto", "flash", "pro", "o3", "o3-mini", "gemini-2.0-flash", "gemini-2.5-pro-preview-06-05"] VALID_MODELS = [
"auto",
"flash",
"pro",
"o3",
"o3-mini",
"gemini-2.5-flash-preview-05-20",
"gemini-2.5-pro-preview-06-05",
]
if DEFAULT_MODEL not in VALID_MODELS: if DEFAULT_MODEL not in VALID_MODELS:
import logging import logging
@@ -48,7 +56,7 @@ MODEL_CAPABILITIES_DESC = {
"o3": "Strong reasoning (200K context) - Logical problems, code generation, systematic analysis", "o3": "Strong reasoning (200K context) - Logical problems, code generation, systematic analysis",
"o3-mini": "Fast O3 variant (200K context) - Balanced performance/speed, moderate complexity", "o3-mini": "Fast O3 variant (200K context) - Balanced performance/speed, moderate complexity",
# Full model names also supported # Full model names also supported
"gemini-2.0-flash": "Ultra-fast (1M context) - Quick analysis, simple queries, rapid iterations", "gemini-2.5-flash-preview-05-20": "Ultra-fast (1M context) - Quick analysis, simple queries, rapid iterations",
"gemini-2.5-pro-preview-06-05": "Deep reasoning + thinking mode (1M context) - Complex problems, architecture, deep analysis", "gemini-2.5-pro-preview-06-05": "Deep reasoning + thinking mode (1M context) - Complex problems, architecture, deep analysis",
} }

47
providers/gemini.py
View File

@@ -13,26 +13,29 @@ class GeminiModelProvider(ModelProvider):
# Model configurations # Model configurations
SUPPORTED_MODELS = { SUPPORTED_MODELS = {
"gemini-2.0-flash": { "gemini-2.5-flash-preview-05-20": {
"max_tokens": 1_048_576, # 1M tokens "max_tokens": 1_048_576, # 1M tokens
"supports_extended_thinking": False, "supports_extended_thinking": True,
"max_thinking_tokens": 24576, # Flash 2.5 thinking budget limit
}, },
"gemini-2.5-pro-preview-06-05": { "gemini-2.5-pro-preview-06-05": {
"max_tokens": 1_048_576, # 1M tokens "max_tokens": 1_048_576, # 1M tokens
"supports_extended_thinking": True, "supports_extended_thinking": True,
"max_thinking_tokens": 32768, # Pro 2.5 thinking budget limit
}, },
# Shorthands # Shorthands
"flash": "gemini-2.0-flash", "flash": "gemini-2.5-flash-preview-05-20",
"pro": "gemini-2.5-pro-preview-06-05", "pro": "gemini-2.5-pro-preview-06-05",
} }
# Thinking mode configurations for models that support it # Thinking mode configurations - percentages of model's max_thinking_tokens
# These percentages work across all models that support thinking
THINKING_BUDGETS = { THINKING_BUDGETS = {
"minimal": 128, # Minimum for 2.5 Pro - fast responses "minimal": 0.005, # 0.5% of max - minimal thinking for fast responses
"low": 2048, # Light reasoning tasks "low": 0.08, # 8% of max - light reasoning tasks
"medium": 8192, # Balanced reasoning (default) "medium": 0.33, # 33% of max - balanced reasoning (default)
"high": 16384, # Complex analysis "high": 0.67, # 67% of max - complex analysis
"max": 32768, # Maximum reasoning depth "max": 1.0, # 100% of max - full thinking budget
} }
def __init__(self, api_key: str, **kwargs): def __init__(self, api_key: str, **kwargs):
@@ -107,9 +110,12 @@ class GeminiModelProvider(ModelProvider):
# Add thinking configuration for models that support it # Add thinking configuration for models that support it
capabilities = self.get_capabilities(resolved_name) capabilities = self.get_capabilities(resolved_name)
if capabilities.supports_extended_thinking and thinking_mode in self.THINKING_BUDGETS: if capabilities.supports_extended_thinking and thinking_mode in self.THINKING_BUDGETS:
generation_config.thinking_config = types.ThinkingConfig( # Get model's max thinking tokens and calculate actual budget
thinking_budget=self.THINKING_BUDGETS[thinking_mode] model_config = self.SUPPORTED_MODELS.get(resolved_name)
) if model_config and "max_thinking_tokens" in model_config:
max_thinking_tokens = model_config["max_thinking_tokens"]
actual_thinking_budget = int(max_thinking_tokens * self.THINKING_BUDGETS[thinking_mode])
generation_config.thinking_config = types.ThinkingConfig(thinking_budget=actual_thinking_budget)
try: try:
# Generate content # Generate content
@@ -164,6 +170,23 @@ class GeminiModelProvider(ModelProvider):
capabilities = self.get_capabilities(model_name) capabilities = self.get_capabilities(model_name)
return capabilities.supports_extended_thinking return capabilities.supports_extended_thinking
def get_thinking_budget(self, model_name: str, thinking_mode: str) -> int:
"""Get actual thinking token budget for a model and thinking mode."""
resolved_name = self._resolve_model_name(model_name)
model_config = self.SUPPORTED_MODELS.get(resolved_name, {})
if not model_config.get("supports_extended_thinking", False):
return 0
if thinking_mode not in self.THINKING_BUDGETS:
return 0
max_thinking_tokens = model_config.get("max_thinking_tokens", 0)
if max_thinking_tokens == 0:
return 0
return int(max_thinking_tokens * self.THINKING_BUDGETS[thinking_mode])
def _resolve_model_name(self, model_name: str) -> str: def _resolve_model_name(self, model_name: str) -> str:
"""Resolve model shorthand to full name.""" """Resolve model shorthand to full name."""
# Check if it's a shorthand # Check if it's a shorthand

8
providers/registry.py
View File

@@ -67,7 +67,7 @@ class ModelProviderRegistry:
"""Get provider instance for a specific model name. """Get provider instance for a specific model name.
Args: Args:
model_name: Name of the model (e.g., "gemini-2.0-flash", "o3-mini") model_name: Name of the model (e.g., "gemini-2.5-flash-preview-05-20", "o3-mini")
Returns: Returns:
ModelProvider instance that supports this model ModelProvider instance that supports this model
@@ -137,7 +137,7 @@ class ModelProviderRegistry:
2. Gemini 2.0 Flash (fast and efficient) if Gemini API key available 2. Gemini 2.0 Flash (fast and efficient) if Gemini API key available
3. OpenAI o3 (high performance) if OpenAI API key available 3. OpenAI o3 (high performance) if OpenAI API key available
4. Gemini 2.5 Pro (deep reasoning) if Gemini API key available 4. Gemini 2.5 Pro (deep reasoning) if Gemini API key available
5. Fallback to gemini-2.0-flash (most common case) 5. Fallback to gemini-2.5-flash-preview-05-20 (most common case)
Returns: Returns:
Model name string for fallback use Model name string for fallback use
@@ -150,11 +150,11 @@ class ModelProviderRegistry:
if openai_available: if openai_available:
return "o3-mini" # Balanced performance/cost return "o3-mini" # Balanced performance/cost
elif gemini_available: elif gemini_available:
return "gemini-2.0-flash" # Fast and efficient return "gemini-2.5-flash-preview-05-20" # Fast and efficient
else: else:
# No API keys available - return a reasonable default # No API keys available - return a reasonable default
# This maintains backward compatibility for tests # This maintains backward compatibility for tests
return "gemini-2.0-flash" return "gemini-2.5-flash-preview-05-20"
@classmethod @classmethod
def get_available_providers_with_keys(cls) -> list[ProviderType]: def get_available_providers_with_keys(cls) -> list[ProviderType]:

159
setup-docker.sh
View File

@@ -200,7 +200,7 @@ else
exit 1 exit 1
fi fi
echo " - Starting Redis and MCP services... please wait" echo " - Starting Redis (needed for conversation memory)... please wait"
if $COMPOSE_CMD up -d >/dev/null 2>&1; then if $COMPOSE_CMD up -d >/dev/null 2>&1; then
echo "✅ Services started successfully!" echo "✅ Services started successfully!"
else else
@@ -223,58 +223,112 @@ echo ""
echo "📋 Service Status:" echo "📋 Service Status:"
$COMPOSE_CMD ps --format table $COMPOSE_CMD ps --format table
echo "" # Function to show configuration steps - only if CLI not already set up
echo "🔄 Next steps:" show_configuration_steps() {
NEEDS_KEY_UPDATE=false echo ""
if grep -q "your_gemini_api_key_here" .env 2>/dev/null || grep -q "your_openai_api_key_here" .env 2>/dev/null; then echo "🔄 Next steps:"
NEEDS_KEY_UPDATE=true NEEDS_KEY_UPDATE=false
if grep -q "your_gemini_api_key_here" .env 2>/dev/null || grep -q "your_openai_api_key_here" .env 2>/dev/null; then
NEEDS_KEY_UPDATE=true
fi
if [ "$NEEDS_KEY_UPDATE" = true ]; then
echo "1. Edit .env and replace placeholder API keys with actual ones"
echo " - GEMINI_API_KEY: your-gemini-api-key-here"
echo " - OPENAI_API_KEY: your-openai-api-key-here"
echo "2. Restart services: $COMPOSE_CMD restart"
echo "3. Copy the configuration below to your Claude Desktop config if required:"
else
echo "1. Copy the configuration below to your Claude Desktop config if required:"
fi
echo ""
echo "===== CLAUDE DESKTOP CONFIGURATION ====="
echo "{"
echo " \"mcpServers\": {"
echo " \"zen\": {"
echo " \"command\": \"docker\","
echo " \"args\": ["
echo " \"exec\","
echo " \"-i\","
echo " \"zen-mcp-server\","
echo " \"python\","
echo " \"server.py\""
echo " ]"
echo " }"
echo " }"
echo "}"
echo "==========================================="
}
# Function to automatically configure Claude Code CLI
# Returns: 0 if already configured, 1 if CLI not found, 2 if configured/skipped
setup_claude_code_cli() {
# Check if claude command exists
if ! command -v claude &> /dev/null; then
echo "⚠️ Claude Code CLI not found. Install it to use with CLI:"
echo " npm install -g @anthropic-ai/claude-code"
echo ""
echo "📋 Manual MCP configuration for Claude Code CLI:"
echo "claude mcp add zen -s user -- docker exec -i zen-mcp-server python server.py"
return 1
fi
echo "🔧 Configuring Claude Code CLI..."
# Get current MCP list and check if zen-mcp-server already exists
if claude mcp list 2>/dev/null | grep -q "zen-mcp-server" 2>/dev/null; then
echo "✅ Zen MCP Server already configured in Claude Code CLI"
echo ""
return 0 # Already configured
else
echo " - Zen MCP Server not found in Claude Code CLI configuration"
echo ""
echo -n "Would you like to add the Zen MCP Server to Claude Code CLI now? [Y/n]: "
read -r response
# Default to yes if empty response (just pressed enter)
if [[ -z "$response" || "$response" =~ ^[Yy]$ ]]; then
echo " - Adding Zen MCP Server to Claude Code CLI..."
if claude mcp add zen -s user -- docker exec -i zen-mcp-server python server.py >/dev/null 2>&1; then
echo "✅ Zen MCP Server added to Claude Code CLI successfully!"
echo " Use 'claude' command to start a session with the MCP server"
else
echo "⚠️ Failed to add MCP server automatically. You can add it manually:"
echo " claude mcp add zen -s user -- docker exec -i zen-mcp-server python server.py"
fi
else
echo " - Skipped adding MCP server. You can add it manually later:"
echo " claude mcp add zen -s user -- docker exec -i zen-mcp-server python server.py"
fi
echo ""
return 2 # Configured or skipped
fi
}
# Set up Claude Code CLI automatically
setup_claude_code_cli
CLI_STATUS=$?
# Only show configuration details if zen is NOT already configured
if [ $CLI_STATUS -ne 0 ]; then
# Show configuration steps
show_configuration_steps
echo ""
echo "===== CLAUDE CODE CLI CONFIGURATION ====="
echo "# Useful Claude Code CLI commands:"
echo "claude # Start interactive session"
echo "claude mcp list # List your MCP servers"
echo "claude mcp remove zen -s user # Remove if needed"
echo "==========================================="
echo ""
echo "📁 Config file locations:"
echo " macOS: ~/Library/Application Support/Claude/claude_desktop_config.json"
echo ' Windows (WSL): /mnt/c/Users/USERNAME/AppData/Roaming/Claude/claude_desktop_config.json'
echo ""
fi fi
if [ "$NEEDS_KEY_UPDATE" = true ]; then
echo "1. Edit .env and replace placeholder API keys with actual ones"
echo " - GEMINI_API_KEY: your-gemini-api-key-here"
echo " - OPENAI_API_KEY: your-openai-api-key-here"
echo "2. Restart services: $COMPOSE_CMD restart"
echo "3. Copy the configuration below to your Claude Desktop config:"
else
echo "1. Copy the configuration below to your Claude Desktop config:"
fi
echo ""
echo "===== CLAUDE DESKTOP CONFIGURATION ====="
echo "{"
echo " \"mcpServers\": {"
echo " \"zen\": {"
echo " \"command\": \"docker\","
echo " \"args\": ["
echo " \"exec\","
echo " \"-i\","
echo " \"zen-mcp-server\","
echo " \"python\","
echo " \"server.py\""
echo " ]"
echo " }"
echo " }"
echo "}"
echo "==========================================="
echo ""
echo "===== CLAUDE CODE CLI CONFIGURATION ====="
echo "# Add the MCP server via Claude Code CLI:"
echo "claude mcp add zen -s user -- docker exec -i zen-mcp-server python server.py"
echo ""
echo "# List your MCP servers to verify:"
echo "claude mcp list"
echo ""
echo "# Remove if needed:"
echo "claude mcp remove zen -s user"
echo "==========================================="
echo ""
echo "📁 Config file locations:"
echo " macOS: ~/Library/Application Support/Claude/claude_desktop_config.json"
echo ' Windows (WSL): /mnt/c/Users/USERNAME/AppData/Roaming/Claude/claude_desktop_config.json'
echo ""
echo "🔧 Useful commands:" echo "🔧 Useful commands:"
echo " Start services: $COMPOSE_CMD up -d" echo " Start services: $COMPOSE_CMD up -d"
echo " Stop services: $COMPOSE_CMD down" echo " Stop services: $COMPOSE_CMD down"
@@ -283,5 +337,4 @@ echo " Restart services: $COMPOSE_CMD restart"
echo " Service status: $COMPOSE_CMD ps" echo " Service status: $COMPOSE_CMD ps"
echo "" echo ""
echo "🗃️ Redis for conversation threading is automatically configured and running!" echo "Happy Clauding!"
echo " All AI-to-AI conversations will persist between requests."

4
simulator_tests/test_model_thinking_config.py
View File

@@ -55,7 +55,7 @@ class TestModelThinkingConfig(BaseSimulatorTest):
"chat", "chat",
{ {
"prompt": "What is 3 + 3? Give a quick answer.", "prompt": "What is 3 + 3? Give a quick answer.",
"model": "flash", # Should resolve to gemini-2.0-flash "model": "flash", # Should resolve to gemini-2.5-flash-preview-05-20
"thinking_mode": "high", # Should be ignored for Flash model "thinking_mode": "high", # Should be ignored for Flash model
}, },
) )
@@ -80,7 +80,7 @@ class TestModelThinkingConfig(BaseSimulatorTest):
("pro", "should work with Pro model"), ("pro", "should work with Pro model"),
("flash", "should work with Flash model"), ("flash", "should work with Flash model"),
("gemini-2.5-pro-preview-06-05", "should work with full Pro model name"), ("gemini-2.5-pro-preview-06-05", "should work with full Pro model name"),
("gemini-2.0-flash", "should work with full Flash model name"), ("gemini-2.5-flash-preview-05-20", "should work with full Flash model name"),
] ]
success_count = 0 success_count = 0

2
tests/conftest.py
View File

@@ -24,7 +24,7 @@ if "OPENAI_API_KEY" not in os.environ:
# Set default model to a specific value for tests to avoid auto mode # Set default model to a specific value for tests to avoid auto mode
# This prevents all tests from failing due to missing model parameter # This prevents all tests from failing due to missing model parameter
os.environ["DEFAULT_MODEL"] = "gemini-2.0-flash" os.environ["DEFAULT_MODEL"] = "gemini-2.5-flash-preview-05-20"
# Force reload of config module to pick up the env var # Force reload of config module to pick up the env var
import config # noqa: E402 import config # noqa: E402

2
tests/mock_helpers.py
View File

@@ -5,7 +5,7 @@ from unittest.mock import Mock
from providers.base import ModelCapabilities, ProviderType, RangeTemperatureConstraint from providers.base import ModelCapabilities, ProviderType, RangeTemperatureConstraint
def create_mock_provider(model_name="gemini-2.0-flash", max_tokens=1_048_576): def create_mock_provider(model_name="gemini-2.5-flash-preview-05-20", max_tokens=1_048_576):
"""Create a properly configured mock provider.""" """Create a properly configured mock provider."""
mock_provider = Mock() mock_provider = Mock()

18
tests/test_claude_continuation.py
View File

@@ -72,7 +72,7 @@ class TestClaudeContinuationOffers:
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Analysis complete.", content="Analysis complete.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -129,7 +129,7 @@ class TestClaudeContinuationOffers:
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Continued analysis.", content="Continued analysis.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -162,7 +162,7 @@ class TestClaudeContinuationOffers:
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Analysis complete. The code looks good.", content="Analysis complete. The code looks good.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -208,7 +208,7 @@ I'd be happy to examine the error handling patterns in more detail if that would
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=content_with_followup, content=content_with_followup,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -253,7 +253,7 @@ I'd be happy to examine the error handling patterns in more detail if that would
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Continued analysis complete.", content="Continued analysis complete.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -309,7 +309,7 @@ I'd be happy to examine the error handling patterns in more detail if that would
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Final response.", content="Final response.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -358,7 +358,7 @@ class TestContinuationIntegration:
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Analysis result", content="Analysis result",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -411,7 +411,7 @@ class TestContinuationIntegration:
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Structure analysis done.", content="Structure analysis done.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -448,7 +448,7 @@ class TestContinuationIntegration:
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Performance analysis done.", content="Performance analysis done.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )

14
tests/test_collaboration.py
View File

@@ -41,7 +41,7 @@ class TestDynamicContextRequests:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=clarification_json, usage={}, model_name="gemini-2.0-flash", metadata={} content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -82,7 +82,7 @@ class TestDynamicContextRequests:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=normal_response, usage={}, model_name="gemini-2.0-flash", metadata={} content=normal_response, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -106,7 +106,7 @@ class TestDynamicContextRequests:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=malformed_json, usage={}, model_name="gemini-2.0-flash", metadata={} content=malformed_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -146,7 +146,7 @@ class TestDynamicContextRequests:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=clarification_json, usage={}, model_name="gemini-2.0-flash", metadata={} content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -233,7 +233,7 @@ class TestCollaborationWorkflow:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=clarification_json, usage={}, model_name="gemini-2.0-flash", metadata={} content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -272,7 +272,7 @@ class TestCollaborationWorkflow:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=clarification_json, usage={}, model_name="gemini-2.0-flash", metadata={} content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -299,7 +299,7 @@ class TestCollaborationWorkflow:
""" """
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=final_response, usage={}, model_name="gemini-2.0-flash", metadata={} content=final_response, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
result2 = await tool.execute( result2 = await tool.execute(

2
tests/test_config.py
View File

@@ -32,7 +32,7 @@ class TestConfig:
def test_model_config(self): def test_model_config(self):
"""Test model configuration""" """Test model configuration"""
# DEFAULT_MODEL is set in conftest.py for tests # DEFAULT_MODEL is set in conftest.py for tests
assert DEFAULT_MODEL == "gemini-2.0-flash" assert DEFAULT_MODEL == "gemini-2.5-flash-preview-05-20"
assert MAX_CONTEXT_TOKENS == 1_000_000 assert MAX_CONTEXT_TOKENS == 1_000_000
def test_temperature_defaults(self): def test_temperature_defaults(self):

2
tests/test_conversation_field_mapping.py
View File

@@ -75,7 +75,7 @@ async def test_conversation_history_field_mapping():
mock_provider = MagicMock() mock_provider = MagicMock()
mock_provider.get_capabilities.return_value = ModelCapabilities( mock_provider.get_capabilities.return_value = ModelCapabilities(
provider=ProviderType.GOOGLE, provider=ProviderType.GOOGLE,
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
friendly_name="Gemini", friendly_name="Gemini",
max_tokens=200000, max_tokens=200000,
supports_extended_thinking=True, supports_extended_thinking=True,

8
tests/test_conversation_history_bug.py
View File

@@ -95,7 +95,7 @@ class TestConversationHistoryBugFix:
return Mock( return Mock(
content="Response with conversation context", content="Response with conversation context",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
@@ -155,7 +155,7 @@ class TestConversationHistoryBugFix:
return Mock( return Mock(
content="Response without history", content="Response without history",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
@@ -193,7 +193,7 @@ class TestConversationHistoryBugFix:
return Mock( return Mock(
content="New conversation response", content="New conversation response",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
@@ -277,7 +277,7 @@ class TestConversationHistoryBugFix:
return Mock( return Mock(
content="Analysis of new files complete", content="Analysis of new files complete",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )

6
tests/test_cross_tool_continuation.py
View File

@@ -112,7 +112,7 @@ I'd be happy to review these security findings in detail if that would be helpfu
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=content, content=content,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -159,7 +159,7 @@ I'd be happy to review these security findings in detail if that would be helpfu
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Critical security vulnerability confirmed. The authentication function always returns true, bypassing all security checks.", content="Critical security vulnerability confirmed. The authentication function always returns true, bypassing all security checks.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -284,7 +284,7 @@ I'd be happy to review these security findings in detail if that would be helpfu
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Security review of auth.py shows vulnerabilities", content="Security review of auth.py shows vulnerabilities",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider

10
tests/test_intelligent_fallback.py
View File

@@ -33,10 +33,10 @@ class TestIntelligentFallback:
@patch.dict(os.environ, {"OPENAI_API_KEY": "", "GEMINI_API_KEY": "test-gemini-key"}, clear=False) @patch.dict(os.environ, {"OPENAI_API_KEY": "", "GEMINI_API_KEY": "test-gemini-key"}, clear=False)
def test_prefers_gemini_flash_when_openai_unavailable(self): def test_prefers_gemini_flash_when_openai_unavailable(self):
"""Test that gemini-2.0-flash is used when only Gemini API key is available""" """Test that gemini-2.5-flash-preview-05-20 is used when only Gemini API key is available"""
ModelProviderRegistry.clear_cache() ModelProviderRegistry.clear_cache()
fallback_model = ModelProviderRegistry.get_preferred_fallback_model() fallback_model = ModelProviderRegistry.get_preferred_fallback_model()
assert fallback_model == "gemini-2.0-flash" assert fallback_model == "gemini-2.5-flash-preview-05-20"
@patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-key", "GEMINI_API_KEY": "test-gemini-key"}, clear=False) @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-key", "GEMINI_API_KEY": "test-gemini-key"}, clear=False)
def test_prefers_openai_when_both_available(self): def test_prefers_openai_when_both_available(self):
@@ -50,7 +50,7 @@ class TestIntelligentFallback:
"""Test fallback behavior when no API keys are available""" """Test fallback behavior when no API keys are available"""
ModelProviderRegistry.clear_cache() ModelProviderRegistry.clear_cache()
fallback_model = ModelProviderRegistry.get_preferred_fallback_model() fallback_model = ModelProviderRegistry.get_preferred_fallback_model()
assert fallback_model == "gemini-2.0-flash" # Default fallback assert fallback_model == "gemini-2.5-flash-preview-05-20" # Default fallback
def test_available_providers_with_keys(self): def test_available_providers_with_keys(self):
"""Test the get_available_providers_with_keys method""" """Test the get_available_providers_with_keys method"""
@@ -140,8 +140,8 @@ class TestIntelligentFallback:
history, tokens = build_conversation_history(context, model_context=None) history, tokens = build_conversation_history(context, model_context=None)
# Should use gemini-2.0-flash when only Gemini is available # Should use gemini-2.5-flash-preview-05-20 when only Gemini is available
mock_context_class.assert_called_once_with("gemini-2.0-flash") mock_context_class.assert_called_once_with("gemini-2.5-flash-preview-05-20")
def test_non_auto_mode_unchanged(self): def test_non_auto_mode_unchanged(self):
"""Test that non-auto mode behavior is unchanged""" """Test that non-auto mode behavior is unchanged"""

12
tests/test_large_prompt_handling.py
View File

@@ -75,7 +75,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock( mock_provider.generate_content.return_value = MagicMock(
content="This is a test response", content="This is a test response",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -100,7 +100,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock( mock_provider.generate_content.return_value = MagicMock(
content="Processed large prompt", content="Processed large prompt",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -212,7 +212,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock( mock_provider.generate_content.return_value = MagicMock(
content="Success", content="Success",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -245,7 +245,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock( mock_provider.generate_content.return_value = MagicMock(
content="Success", content="Success",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -276,7 +276,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock( mock_provider.generate_content.return_value = MagicMock(
content="Success", content="Success",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -298,7 +298,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock( mock_provider.generate_content.return_value = MagicMock(
content="Success", content="Success",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider

2
tests/test_prompt_regression.py
View File

@@ -31,7 +31,7 @@ class TestPromptRegression:
return Mock( return Mock(
content=text, content=text,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash", model_name="gemini-2.5-flash-preview-05-20",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
) )

18
tests/test_providers.py
View File

@@ -49,7 +49,7 @@ class TestModelProviderRegistry:
"""Test getting provider for a specific model""" """Test getting provider for a specific model"""
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider) ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
provider = ModelProviderRegistry.get_provider_for_model("gemini-2.0-flash") provider = ModelProviderRegistry.get_provider_for_model("gemini-2.5-flash-preview-05-20")
assert provider is not None assert provider is not None
assert isinstance(provider, GeminiModelProvider) assert isinstance(provider, GeminiModelProvider)
@@ -80,12 +80,12 @@ class TestGeminiProvider:
"""Test getting model capabilities""" """Test getting model capabilities"""
provider = GeminiModelProvider(api_key="test-key") provider = GeminiModelProvider(api_key="test-key")
capabilities = provider.get_capabilities("gemini-2.0-flash") capabilities = provider.get_capabilities("gemini-2.5-flash-preview-05-20")
assert capabilities.provider == ProviderType.GOOGLE assert capabilities.provider == ProviderType.GOOGLE
assert capabilities.model_name == "gemini-2.0-flash" assert capabilities.model_name == "gemini-2.5-flash-preview-05-20"
assert capabilities.max_tokens == 1_048_576 assert capabilities.max_tokens == 1_048_576
assert not capabilities.supports_extended_thinking assert capabilities.supports_extended_thinking
def test_get_capabilities_pro_model(self): def test_get_capabilities_pro_model(self):
"""Test getting capabilities for Pro model with thinking support""" """Test getting capabilities for Pro model with thinking support"""
@@ -103,13 +103,13 @@ class TestGeminiProvider:
assert provider.validate_model_name("pro") assert provider.validate_model_name("pro")
capabilities = provider.get_capabilities("flash") capabilities = provider.get_capabilities("flash")
assert capabilities.model_name == "gemini-2.0-flash" assert capabilities.model_name == "gemini-2.5-flash-preview-05-20"
def test_supports_thinking_mode(self): def test_supports_thinking_mode(self):
"""Test thinking mode support detection""" """Test thinking mode support detection"""
provider = GeminiModelProvider(api_key="test-key") provider = GeminiModelProvider(api_key="test-key")
assert not provider.supports_thinking_mode("gemini-2.0-flash") assert provider.supports_thinking_mode("gemini-2.5-flash-preview-05-20")
assert provider.supports_thinking_mode("gemini-2.5-pro-preview-06-05") assert provider.supports_thinking_mode("gemini-2.5-pro-preview-06-05")
@patch("google.genai.Client") @patch("google.genai.Client")
@@ -133,11 +133,13 @@ class TestGeminiProvider:
provider = GeminiModelProvider(api_key="test-key") provider = GeminiModelProvider(api_key="test-key")
response = provider.generate_content(prompt="Test prompt", model_name="gemini-2.0-flash", temperature=0.7) response = provider.generate_content(
prompt="Test prompt", model_name="gemini-2.5-flash-preview-05-20", temperature=0.7
)
assert isinstance(response, ModelResponse) assert isinstance(response, ModelResponse)
assert response.content == "Generated content" assert response.content == "Generated content"
assert response.model_name == "gemini-2.0-flash" assert response.model_name == "gemini-2.5-flash-preview-05-20"
assert response.provider == ProviderType.GOOGLE assert response.provider == ProviderType.GOOGLE
assert response.usage["input_tokens"] == 10 assert response.usage["input_tokens"] == 10
assert response.usage["output_tokens"] == 20 assert response.usage["output_tokens"] == 20

2
tests/test_server.py
View File

@@ -56,7 +56,7 @@ class TestServerTools:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Chat response", usage={}, model_name="gemini-2.0-flash", metadata={} content="Chat response", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider

40
tests/test_thinking_modes.py
View File

@@ -23,6 +23,7 @@ def setup_test_env():
class TestThinkingModes: class TestThinkingModes:
"""Test thinking modes across all tools""" """Test thinking modes across all tools"""
@patch("config.DEFAULT_THINKING_MODE_THINKDEEP", "high")
def test_default_thinking_modes(self): def test_default_thinking_modes(self):
"""Test that tools have correct default thinking modes""" """Test that tools have correct default thinking modes"""
tools = [ tools = [
@@ -45,7 +46,7 @@ class TestThinkingModes:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = True mock_provider.supports_thinking_mode.return_value = True
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Minimal thinking response", usage={}, model_name="gemini-2.0-flash", metadata={} content="Minimal thinking response", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -82,7 +83,7 @@ class TestThinkingModes:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = True mock_provider.supports_thinking_mode.return_value = True
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Low thinking response", usage={}, model_name="gemini-2.0-flash", metadata={} content="Low thinking response", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -114,7 +115,7 @@ class TestThinkingModes:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = True mock_provider.supports_thinking_mode.return_value = True
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Medium thinking response", usage={}, model_name="gemini-2.0-flash", metadata={} content="Medium thinking response", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -145,7 +146,7 @@ class TestThinkingModes:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = True mock_provider.supports_thinking_mode.return_value = True
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="High thinking response", usage={}, model_name="gemini-2.0-flash", metadata={} content="High thinking response", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -169,13 +170,14 @@ class TestThinkingModes:
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("tools.base.BaseTool.get_model_provider") @patch("tools.base.BaseTool.get_model_provider")
@patch("config.DEFAULT_THINKING_MODE_THINKDEEP", "high")
async def test_thinking_mode_max(self, mock_get_provider): async def test_thinking_mode_max(self, mock_get_provider):
"""Test max thinking mode (default for thinkdeep)""" """Test max thinking mode (default for thinkdeep)"""
mock_provider = create_mock_provider() mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = True mock_provider.supports_thinking_mode.return_value = True
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Max thinking response", usage={}, model_name="gemini-2.0-flash", metadata={} content="Max thinking response", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -222,18 +224,22 @@ class TestThinkingModes:
async def prepare_prompt(self, request): async def prepare_prompt(self, request):
return "test" return "test"
# Expected mappings # Test dynamic budget calculation for Flash 2.5
from providers.gemini import GeminiModelProvider
provider = GeminiModelProvider(api_key="test-key")
flash_model = "gemini-2.5-flash-preview-05-20"
flash_max_tokens = 24576
expected_budgets = { expected_budgets = {
"minimal": 128, "minimal": int(flash_max_tokens * 0.005), # 123
"low": 2048, "low": int(flash_max_tokens * 0.08), # 1966
"medium": 8192, "medium": int(flash_max_tokens * 0.33), # 8110
"high": 16384, "high": int(flash_max_tokens * 0.67), # 16465
"max": 32768, "max": int(flash_max_tokens * 1.0), # 24576
} }
# Check each mode in create_model # Check each mode using the helper method
for _mode, _expected_budget in expected_budgets.items(): for mode, expected_budget in expected_budgets.items():
# The budget mapping is inside create_model actual_budget = provider.get_thinking_budget(flash_model, mode)
# We can't easily test it without calling the method assert actual_budget == expected_budget, f"Mode {mode}: expected {expected_budget}, got {actual_budget}"
# But we've verified the values are correct in the code
pass

10
tests/test_tools.py
View File

@@ -37,7 +37,7 @@ class TestThinkDeepTool:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = True mock_provider.supports_thinking_mode.return_value = True
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Extended analysis", usage={}, model_name="gemini-2.0-flash", metadata={} content="Extended analysis", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -88,7 +88,7 @@ class TestCodeReviewTool:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Security issues found", usage={}, model_name="gemini-2.0-flash", metadata={} content="Security issues found", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -133,7 +133,7 @@ class TestDebugIssueTool:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Root cause: race condition", usage={}, model_name="gemini-2.0-flash", metadata={} content="Root cause: race condition", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -181,7 +181,7 @@ class TestAnalyzeTool:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Architecture analysis", usage={}, model_name="gemini-2.0-flash", metadata={} content="Architecture analysis", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider
@@ -295,7 +295,7 @@ class TestAbsolutePathValidation:
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content="Analysis complete", usage={}, model_name="gemini-2.0-flash", metadata={} content="Analysis complete", usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
) )
mock_get_provider.return_value = mock_provider mock_get_provider.return_value = mock_provider

2
tools/analyze.py
View File

@@ -83,7 +83,7 @@ class AnalyzeTool(BaseTool):
"thinking_mode": { "thinking_mode": {
"type": "string", "type": "string",
"enum": ["minimal", "low", "medium", "high", "max"], "enum": ["minimal", "low", "medium", "high", "max"],
"description": "Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", "description": "Thinking depth: minimal (0.5% of model max), low (8%), medium (33%), high (67%), max (100% of model max)",
}, },
"use_websearch": { "use_websearch": {
"type": "boolean", "type": "boolean",

43
tools/base.py
View File

@@ -57,11 +57,11 @@ class ToolRequest(BaseModel):
# Higher values allow for more complex reasoning but increase latency and cost # Higher values allow for more complex reasoning but increase latency and cost
thinking_mode: Optional[Literal["minimal", "low", "medium", "high", "max"]] = Field( thinking_mode: Optional[Literal["minimal", "low", "medium", "high", "max"]] = Field(
None, None,
description="Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", description="Thinking depth: minimal (0.5% of model max), low (8%), medium (33%), high (67%), max (100% of model max)",
) )
use_websearch: Optional[bool] = Field( use_websearch: Optional[bool] = Field(
False, True,
description="Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.", description="Enable web search for documentation, best practices, and current information. When enabled, the model can request Claude to perform web searches and share results back during conversations. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
) )
continuation_id: Optional[str] = Field( continuation_id: Optional[str] = Field(
None, None,
@@ -473,7 +473,7 @@ class BaseTool(ABC):
def get_websearch_instruction(self, use_websearch: bool, tool_specific: Optional[str] = None) -> str: def get_websearch_instruction(self, use_websearch: bool, tool_specific: Optional[str] = None) -> str:
""" """
Generate web search instruction based on the use_websearch parameter. Generate standardized web search instruction based on the use_websearch parameter.
Args: Args:
use_websearch: Whether web search is enabled use_websearch: Whether web search is enabled
@@ -487,30 +487,39 @@ class BaseTool(ABC):
base_instruction = """ base_instruction = """
WEB SEARCH REASONING: As you analyze this request, consider whether web searches would enhance your response. WEB SEARCH CAPABILITY: You can request Claude to perform web searches to enhance your analysis with current information!
If you identify areas where current documentation, API references, or community solutions would be valuable,
please note in your response what specific searches Claude should perform and why they would be helpful.""" IMPORTANT: When you identify areas where web searches would significantly improve your response (such as checking current documentation, finding recent solutions, verifying best practices, or gathering community insights), you MUST explicitly instruct Claude to perform specific web searches and then respond back using the continuation_id from this response to continue the analysis.
Use clear, direct language based on the value of the search:
For valuable supplementary information: "Please perform a web search on '[specific topic/query]' and then continue this analysis using the continuation_id from this response if you find relevant information."
For important missing information: "Please search for '[specific topic/query]' and respond back with the findings using the continuation_id from this response - this information is needed to provide a complete analysis."
For critical/essential information: "SEARCH REQUIRED: Please immediately perform a web search on '[specific topic/query]' and respond back with the results using the continuation_id from this response. Cannot provide accurate analysis without this current information."
This ensures you get the most current and comprehensive information while maintaining conversation context through the continuation_id."""
if tool_specific: if tool_specific:
return f"""{base_instruction} return f"""{base_instruction}
{tool_specific} {tool_specific}
In your response, if web searches would be beneficial, include a section like: When recommending searches, be specific about what information you need and why it would improve your analysis."""
**Recommended Web Searches for Claude:**
- [Specific topic/framework/library] - to verify/understand/confirm [specific aspect]
- [Another search topic] - to explore [specific concern or feature]"""
# Default instruction for all tools # Default instruction for all tools
return f"""{base_instruction} return f"""{base_instruction}
Consider searches for: Consider requesting searches for:
- Current documentation and best practices - Current documentation and API references
- Similar issues and community solutions - Recent best practices and patterns
- API references and usage examples - Known issues and community solutions
- Recent developments and updates - Framework updates and compatibility
- Security advisories and patches
- Performance benchmarks and optimizations
If any of these would strengthen your analysis, specify what Claude should search for and why.""" When recommending searches, be specific about what information you need and why it would improve your analysis. Always remember to instruct Claude to use the continuation_id from this response when providing search results."""
@abstractmethod @abstractmethod
def get_request_model(self): def get_request_model(self):

2
tools/chat.py
View File

@@ -68,7 +68,7 @@ class ChatTool(BaseTool):
"thinking_mode": { "thinking_mode": {
"type": "string", "type": "string",
"enum": ["minimal", "low", "medium", "high", "max"], "enum": ["minimal", "low", "medium", "high", "max"],
"description": "Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", "description": "Thinking depth: minimal (0.5% of model max), low (8%), medium (33%), high (67%), max (100% of model max)",
}, },
"use_websearch": { "use_websearch": {
"type": "boolean", "type": "boolean",

2
tools/codereview.py
View File

@@ -126,7 +126,7 @@ class CodeReviewTool(BaseTool):
"thinking_mode": { "thinking_mode": {
"type": "string", "type": "string",
"enum": ["minimal", "low", "medium", "high", "max"], "enum": ["minimal", "low", "medium", "high", "max"],
"description": "Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", "description": "Thinking depth: minimal (0.5% of model max), low (8%), medium (33%), high (67%), max (100% of model max)",
}, },
"use_websearch": { "use_websearch": {
"type": "boolean", "type": "boolean",

2
tools/debug.py
View File

@@ -86,7 +86,7 @@ class DebugIssueTool(BaseTool):
"thinking_mode": { "thinking_mode": {
"type": "string", "type": "string",
"enum": ["minimal", "low", "medium", "high", "max"], "enum": ["minimal", "low", "medium", "high", "max"],
"description": "Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", "description": "Thinking depth: minimal (0.5% of model max), low (8%), medium (33%), high (67%), max (100% of model max)",
}, },
"use_websearch": { "use_websearch": {
"type": "boolean", "type": "boolean",

2
tools/thinkdeep.py
View File

@@ -81,7 +81,7 @@ class ThinkDeepTool(BaseTool):
"thinking_mode": { "thinking_mode": {
"type": "string", "type": "string",
"enum": ["minimal", "low", "medium", "high", "max"], "enum": ["minimal", "low", "medium", "high", "max"],
"description": f"Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768). Defaults to '{self.get_default_thinking_mode()}' if not specified.", "description": f"Thinking depth: minimal (0.5% of model max), low (8%), medium (33%), high (67%), max (100% of model max). Defaults to '{self.get_default_thinking_mode()}' if not specified.",
}, },
"use_websearch": { "use_websearch": {
"type": "boolean", "type": "boolean",

4
utils/conversation_memory.py
View File

@@ -74,7 +74,7 @@ class ConversationTurn(BaseModel):
files: List of file paths referenced in this specific turn files: List of file paths referenced in this specific turn
tool_name: Which tool generated this turn (for cross-tool tracking) tool_name: Which tool generated this turn (for cross-tool tracking)
model_provider: Provider used (e.g., "google", "openai") model_provider: Provider used (e.g., "google", "openai")
model_name: Specific model used (e.g., "gemini-2.0-flash", "o3-mini") model_name: Specific model used (e.g., "gemini-2.5-flash-preview-05-20", "o3-mini")
model_metadata: Additional model-specific metadata (e.g., thinking mode, token usage) model_metadata: Additional model-specific metadata (e.g., thinking mode, token usage)
""" """
@@ -249,7 +249,7 @@ def add_turn(
files: Optional list of files referenced in this turn files: Optional list of files referenced in this turn
tool_name: Name of the tool adding this turn (for attribution) tool_name: Name of the tool adding this turn (for attribution)
model_provider: Provider used (e.g., "google", "openai") model_provider: Provider used (e.g., "google", "openai")
model_name: Specific model used (e.g., "gemini-2.0-flash", "o3-mini") model_name: Specific model used (e.g., "gemini-2.5-flash-preview-05-20", "o3-mini")
model_metadata: Additional model info (e.g., thinking mode, token usage) model_metadata: Additional model info (e.g., thinking mode, token usage)
Returns: Returns: