From e97d4bd48f38466d30e728d58b859e038e008d15 Mon Sep 17 00:00:00 2001 From: Fahad Date: Thu, 12 Jun 2025 18:54:44 +0400 Subject: [PATCH 1/6] Removed --- FIX_SUMMARY.md | 40 ---------------------------------------- 1 file changed, 40 deletions(-) delete mode 100644 FIX_SUMMARY.md diff --git a/FIX_SUMMARY.md b/FIX_SUMMARY.md deleted file mode 100644 index d5e1bad..0000000 --- a/FIX_SUMMARY.md +++ /dev/null @@ -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. \ No newline at end of file From b64f5095de508b374ed2644d9b9bb35dd5019892 Mon Sep 17 00:00:00 2001 From: Fahad Date: Thu, 12 Jun 2025 19:20:41 +0400 Subject: [PATCH 2/6] Make setup easier, ask user if they'd like the mcp set up automatically --- setup-docker.sh | 159 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 106 insertions(+), 53 deletions(-) diff --git a/setup-docker.sh b/setup-docker.sh index b7acc61..523820a 100755 --- a/setup-docker.sh +++ b/setup-docker.sh @@ -200,7 +200,7 @@ else exit 1 fi -echo " - Starting Redis and MCP services... please wait" +echo " - Starting Redis (needed for conversational context persistence)... please wait" if $COMPOSE_CMD up -d >/dev/null 2>&1; then echo "✅ Services started successfully!" else @@ -223,58 +223,112 @@ echo "" echo "📋 Service Status:" $COMPOSE_CMD ps --format table -echo "" -echo "🔄 Next steps:" -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 +# Function to show configuration steps - only if CLI not already set up +show_configuration_steps() { + echo "" + echo "🔄 Next steps:" + 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 -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 " Start services: $COMPOSE_CMD up -d" echo " Stop services: $COMPOSE_CMD down" @@ -283,5 +337,4 @@ echo " Restart services: $COMPOSE_CMD restart" echo " Service status: $COMPOSE_CMD ps" echo "" -echo "🗃️ Redis for conversation threading is automatically configured and running!" -echo " All AI-to-AI conversations will persist between requests." \ No newline at end of file +echo "Happy Clauding!" \ No newline at end of file From 4f2763689bf85768c889a4830db035b2fac5bce5 Mon Sep 17 00:00:00 2001 From: Fahad Date: Thu, 12 Jun 2025 19:22:32 +0400 Subject: [PATCH 3/6] Make setup easier, ask user if they'd like the mcp set up automatically --- README.md | 2 +- setup-docker.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index c1bf0e9..15adefe 100644 --- a/README.md +++ b/README.md @@ -117,7 +117,7 @@ cd zen-mcp-server - **Creates .env file** (automatically uses `$GEMINI_API_KEY` and `$OPENAI_API_KEY` if set in environment) - **Starts Redis service** for AI-to-AI conversation memory - **Starts MCP server** with providers based on available API keys -- **Shows exact Claude Desktop configuration** to copy (optional when only using claude code) +- **Adds Zen to Claude Code automatically** ### 3. Add Your API Keys diff --git a/setup-docker.sh b/setup-docker.sh index 523820a..4f489c4 100755 --- a/setup-docker.sh +++ b/setup-docker.sh @@ -200,7 +200,7 @@ else exit 1 fi -echo " - Starting Redis (needed for conversational context persistence)... please wait" +echo " - Starting Redis (needed for conversation memory)... please wait" if $COMPOSE_CMD up -d >/dev/null 2>&1; then echo "✅ Services started successfully!" else From b34c63d710c93bd125f26c4b827671eb4bb91738 Mon Sep 17 00:00:00 2001 From: Fahad Date: Thu, 12 Jun 2025 20:05:55 +0400 Subject: [PATCH 4/6] Fixed web-search prompt, models can prompt claude to perform web searches on their behalf if and as needed --- README.md | 25 ++++++------------------- tools/base.py | 41 +++++++++++++++++++++++++---------------- 2 files changed, 31 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 15adefe..1c5869a 100644 --- a/README.md +++ b/README.md @@ -390,7 +390,7 @@ Use zen and perform a thorough precommit ensuring there aren't any new regressio - Supports specialized analysis types: architecture, performance, security, quality - Uses file paths (not content) for clean terminal output - Can identify patterns, anti-patterns, and refactoring opportunities -- **Web search capability**: When enabled with `use_websearch`, can look up framework documentation, design patterns, and best practices relevant to the code being analyzed +- **Web search capability**: When enabled with `use_websearch`, the model can request Claude to perform web searches and share results back to enhance analysis with current documentation, design patterns, and best practices ### 7. `get_version` - Server Information ``` "Get zen to show its version" @@ -409,7 +409,7 @@ All tools that work with files support **both individual files and entire direct - `analysis_type`: architecture|performance|security|quality|general - `output_format`: summary|detailed|actionable - `thinking_mode`: minimal|low|medium|high|max (default: medium, Gemini only) -- `use_websearch`: Enable web search for documentation and best practices (default: false) +- `use_websearch`: Enable web search for documentation and best practices - allows model to request Claude perform searches (default: true) ``` "Analyze the src/ directory for architectural patterns" (auto mode picks best model) @@ -442,7 +442,7 @@ All tools that work with files support **both individual files and entire direct - `runtime_info`: Environment details - `previous_attempts`: What you've tried - `thinking_mode`: minimal|low|medium|high|max (default: medium, Gemini only) -- `use_websearch`: Enable web search for error messages and solutions (default: false) +- `use_websearch`: Enable web search for error messages and solutions - allows model to request Claude perform searches (default: true) ``` "Debug this logic error with context from backend/" (auto mode picks best model) @@ -457,7 +457,7 @@ All tools that work with files support **both individual files and entire direct - `focus_areas`: Specific aspects to focus on - `files`: Files or directories for context - `thinking_mode`: minimal|low|medium|high|max (default: max, Gemini only) -- `use_websearch`: Enable web search for documentation and insights (default: false) +- `use_websearch`: Enable web search for documentation and insights - allows model to request Claude perform searches (default: true) ``` "Think deeper about my design with reference to src/models/" (auto mode picks best model) @@ -705,25 +705,12 @@ Claude can then search for these specific topics and provide you with the most c - More collaborative approach between the two AI assistants - Reduces hallucination by encouraging verification of assumptions -**Disabling web search:** -If you prefer Gemini to work only with its training data, you can disable web search: +**Web search control:** +Web search is enabled by default, allowing models to request Claude perform searches for current documentation and solutions. If you prefer the model to work only with its training data, you can disable web search: ``` "Use gemini to review this code with use_websearch false" ``` -### Standardized Response Format -All tools now return structured JSON responses for consistent handling: -```json -{ - "status": "success|error|requires_clarification", - "content": "The actual response content", - "content_type": "text|markdown|json", - "metadata": {"tool_name": "analyze", ...} -} -``` - -This enables better integration, error handling, and support for the dynamic context request feature. - ## Configuration The server includes several configurable properties that control its behavior: diff --git a/tools/base.py b/tools/base.py index 940bf22..9e6d73b 100644 --- a/tools/base.py +++ b/tools/base.py @@ -60,8 +60,8 @@ class ToolRequest(BaseModel): description="Thinking depth: minimal (128), low (2048), medium (8192), high (16384), max (32768)", ) use_websearch: Optional[bool] = Field( - False, - 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.", + True, + 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( None, @@ -473,7 +473,7 @@ class BaseTool(ABC): 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: use_websearch: Whether web search is enabled @@ -487,30 +487,39 @@ class BaseTool(ABC): base_instruction = """ -WEB SEARCH REASONING: As you analyze this request, consider whether web searches would enhance your response. -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.""" +WEB SEARCH CAPABILITY: You can request Claude to perform web searches to enhance your analysis with current information! + +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: return f"""{base_instruction} {tool_specific} -In your response, if web searches would be beneficial, include a section like: -**Recommended Web Searches for Claude:** -- [Specific topic/framework/library] - to verify/understand/confirm [specific aspect] -- [Another search topic] - to explore [specific concern or feature]""" +When recommending searches, be specific about what information you need and why it would improve your analysis.""" # Default instruction for all tools return f"""{base_instruction} -Consider searches for: -- Current documentation and best practices -- Similar issues and community solutions -- API references and usage examples -- Recent developments and updates +Consider requesting searches for: +- Current documentation and API references +- Recent best practices and patterns +- Known issues and community solutions +- 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 def get_request_model(self): From 3aedb161011be90bc5594b77dfa4bd0dd4cc354f Mon Sep 17 00:00:00 2001 From: Fahad Date: Thu, 12 Jun 2025 20:46:54 +0400 Subject: [PATCH 5/6] Use the new Gemini 2.5 Flash Updated to support Thinking Tokens as a ratio of the max allowed Updated tests Updated README --- README.md | 2 +- config.py | 12 ++++- providers/gemini.py | 47 ++++++++++++++----- providers/registry.py | 8 ++-- simulator_tests/test_model_thinking_config.py | 4 +- tests/conftest.py | 2 +- tests/mock_helpers.py | 2 +- tests/test_claude_continuation.py | 18 +++---- tests/test_collaboration.py | 14 +++--- tests/test_config.py | 2 +- tests/test_conversation_field_mapping.py | 2 +- tests/test_conversation_history_bug.py | 8 ++-- tests/test_cross_tool_continuation.py | 6 +-- tests/test_intelligent_fallback.py | 10 ++-- tests/test_large_prompt_handling.py | 12 ++--- tests/test_prompt_regression.py | 2 +- tests/test_providers.py | 16 ++++--- tests/test_server.py | 2 +- tests/test_thinking_modes.py | 38 ++++++++------- tests/test_tools.py | 10 ++-- tools/analyze.py | 2 +- tools/base.py | 2 +- tools/chat.py | 2 +- tools/codereview.py | 2 +- tools/debug.py | 2 +- tools/thinkdeep.py | 2 +- utils/conversation_memory.py | 4 +- 27 files changed, 135 insertions(+), 98 deletions(-) diff --git a/README.md b/README.md index 1c5869a..0abd47e 100644 --- a/README.md +++ b/README.md @@ -390,7 +390,7 @@ Use zen and perform a thorough precommit ensuring there aren't any new regressio - Supports specialized analysis types: architecture, performance, security, quality - Uses file paths (not content) for clean terminal output - Can identify patterns, anti-patterns, and refactoring opportunities -- **Web search capability**: When enabled with `use_websearch`, the model can request Claude to perform web searches and share results back to enhance analysis with current documentation, design patterns, and best practices +- **Web search capability**: When enabled with `use_websearch` (default: true), the model can request Claude to perform web searches and share results back to enhance analysis with current documentation, design patterns, and best practices ### 7. `get_version` - Server Information ``` "Get zen to show its version" diff --git a/config.py b/config.py index 9e213f9..c1762bb 100644 --- a/config.py +++ b/config.py @@ -26,7 +26,15 @@ DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "auto") # Validate DEFAULT_MODEL and set to "auto" if invalid # 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: import logging @@ -47,7 +55,7 @@ MODEL_CAPABILITIES_DESC = { "o3": "Strong reasoning (200K context) - Logical problems, code generation, systematic analysis", "o3-mini": "Fast O3 variant (200K context) - Balanced performance/speed, moderate complexity", # 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", } diff --git a/providers/gemini.py b/providers/gemini.py index a80b4e4..5fe435e 100644 --- a/providers/gemini.py +++ b/providers/gemini.py @@ -13,26 +13,29 @@ class GeminiModelProvider(ModelProvider): # Model configurations SUPPORTED_MODELS = { - "gemini-2.0-flash": { + "gemini-2.5-flash-preview-05-20": { "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": { "max_tokens": 1_048_576, # 1M tokens "supports_extended_thinking": True, + "max_thinking_tokens": 32768, # Pro 2.5 thinking budget limit }, # Shorthands - "flash": "gemini-2.0-flash", + "flash": "gemini-2.5-flash-preview-05-20", "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 = { - "minimal": 128, # Minimum for 2.5 Pro - fast responses - "low": 2048, # Light reasoning tasks - "medium": 8192, # Balanced reasoning (default) - "high": 16384, # Complex analysis - "max": 32768, # Maximum reasoning depth + "minimal": 0.005, # 0.5% of max - minimal thinking for fast responses + "low": 0.08, # 8% of max - light reasoning tasks + "medium": 0.33, # 33% of max - balanced reasoning (default) + "high": 0.67, # 67% of max - complex analysis + "max": 1.0, # 100% of max - full thinking budget } def __init__(self, api_key: str, **kwargs): @@ -107,9 +110,12 @@ class GeminiModelProvider(ModelProvider): # Add thinking configuration for models that support it capabilities = self.get_capabilities(resolved_name) if capabilities.supports_extended_thinking and thinking_mode in self.THINKING_BUDGETS: - generation_config.thinking_config = types.ThinkingConfig( - thinking_budget=self.THINKING_BUDGETS[thinking_mode] - ) + # Get model's max thinking tokens and calculate actual budget + 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: # Generate content @@ -164,6 +170,23 @@ class GeminiModelProvider(ModelProvider): capabilities = self.get_capabilities(model_name) 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: """Resolve model shorthand to full name.""" # Check if it's a shorthand diff --git a/providers/registry.py b/providers/registry.py index 057821c..8d126b2 100644 --- a/providers/registry.py +++ b/providers/registry.py @@ -67,7 +67,7 @@ class ModelProviderRegistry: """Get provider instance for a specific model name. 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: 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 3. OpenAI o3 (high performance) if OpenAI 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: Model name string for fallback use @@ -150,11 +150,11 @@ class ModelProviderRegistry: if openai_available: return "o3-mini" # Balanced performance/cost elif gemini_available: - return "gemini-2.0-flash" # Fast and efficient + return "gemini-2.5-flash-preview-05-20" # Fast and efficient else: # No API keys available - return a reasonable default # This maintains backward compatibility for tests - return "gemini-2.0-flash" + return "gemini-2.5-flash-preview-05-20" @classmethod def get_available_providers_with_keys(cls) -> list[ProviderType]: diff --git a/simulator_tests/test_model_thinking_config.py b/simulator_tests/test_model_thinking_config.py index b1b096f..94105c8 100644 --- a/simulator_tests/test_model_thinking_config.py +++ b/simulator_tests/test_model_thinking_config.py @@ -55,7 +55,7 @@ class TestModelThinkingConfig(BaseSimulatorTest): "chat", { "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 }, ) @@ -80,7 +80,7 @@ class TestModelThinkingConfig(BaseSimulatorTest): ("pro", "should work with Pro model"), ("flash", "should work with Flash model"), ("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 diff --git a/tests/conftest.py b/tests/conftest.py index 57718c2..bae4921 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 # 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 import config # noqa: E402 diff --git a/tests/mock_helpers.py b/tests/mock_helpers.py index 0aa4c5c..447dd5b 100644 --- a/tests/mock_helpers.py +++ b/tests/mock_helpers.py @@ -5,7 +5,7 @@ from unittest.mock import Mock 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.""" mock_provider = Mock() diff --git a/tests/test_claude_continuation.py b/tests/test_claude_continuation.py index bed5408..1de4b45 100644 --- a/tests/test_claude_continuation.py +++ b/tests/test_claude_continuation.py @@ -72,7 +72,7 @@ class TestClaudeContinuationOffers: mock_provider.generate_content.return_value = Mock( content="Analysis complete.", 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"}, ) mock_get_provider.return_value = mock_provider @@ -129,7 +129,7 @@ class TestClaudeContinuationOffers: mock_provider.generate_content.return_value = Mock( content="Continued analysis.", 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"}, ) mock_get_provider.return_value = mock_provider @@ -162,7 +162,7 @@ class TestClaudeContinuationOffers: mock_provider.generate_content.return_value = Mock( content="Analysis complete. The code looks good.", 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"}, ) 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( content=content_with_followup, 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"}, ) 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( content="Continued analysis complete.", 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"}, ) 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( content="Final response.", 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"}, ) mock_get_provider.return_value = mock_provider @@ -358,7 +358,7 @@ class TestContinuationIntegration: mock_provider.generate_content.return_value = Mock( content="Analysis result", 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"}, ) mock_get_provider.return_value = mock_provider @@ -411,7 +411,7 @@ class TestContinuationIntegration: mock_provider.generate_content.return_value = Mock( content="Structure analysis done.", 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"}, ) mock_get_provider.return_value = mock_provider @@ -448,7 +448,7 @@ class TestContinuationIntegration: mock_provider.generate_content.return_value = Mock( content="Performance analysis done.", 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"}, ) diff --git a/tests/test_collaboration.py b/tests/test_collaboration.py index 966cc39..e2e5254 100644 --- a/tests/test_collaboration.py +++ b/tests/test_collaboration.py @@ -41,7 +41,7 @@ class TestDynamicContextRequests: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -82,7 +82,7 @@ class TestDynamicContextRequests: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -106,7 +106,7 @@ class TestDynamicContextRequests: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -146,7 +146,7 @@ class TestDynamicContextRequests: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -233,7 +233,7 @@ class TestCollaborationWorkflow: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -272,7 +272,7 @@ class TestCollaborationWorkflow: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -299,7 +299,7 @@ class TestCollaborationWorkflow: """ 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( diff --git a/tests/test_config.py b/tests/test_config.py index 0ac6368..6220226 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -32,7 +32,7 @@ class TestConfig: def test_model_config(self): """Test model configuration""" # 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 def test_temperature_defaults(self): diff --git a/tests/test_conversation_field_mapping.py b/tests/test_conversation_field_mapping.py index a26f3b8..28cd82e 100644 --- a/tests/test_conversation_field_mapping.py +++ b/tests/test_conversation_field_mapping.py @@ -75,7 +75,7 @@ async def test_conversation_history_field_mapping(): mock_provider = MagicMock() mock_provider.get_capabilities.return_value = ModelCapabilities( provider=ProviderType.GOOGLE, - model_name="gemini-2.0-flash", + model_name="gemini-2.5-flash-preview-05-20", friendly_name="Gemini", max_tokens=200000, supports_extended_thinking=True, diff --git a/tests/test_conversation_history_bug.py b/tests/test_conversation_history_bug.py index e73bb8b..c5c6386 100644 --- a/tests/test_conversation_history_bug.py +++ b/tests/test_conversation_history_bug.py @@ -95,7 +95,7 @@ class TestConversationHistoryBugFix: return Mock( content="Response with conversation context", 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"}, ) @@ -155,7 +155,7 @@ class TestConversationHistoryBugFix: return Mock( content="Response without history", 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"}, ) @@ -193,7 +193,7 @@ class TestConversationHistoryBugFix: return Mock( content="New conversation response", 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"}, ) @@ -277,7 +277,7 @@ class TestConversationHistoryBugFix: return Mock( content="Analysis of new files complete", 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"}, ) diff --git a/tests/test_cross_tool_continuation.py b/tests/test_cross_tool_continuation.py index f3f3af5..ac4a95a 100644 --- a/tests/test_cross_tool_continuation.py +++ b/tests/test_cross_tool_continuation.py @@ -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( content=content, 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"}, ) 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( 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}, - model_name="gemini-2.0-flash", + model_name="gemini-2.5-flash-preview-05-20", metadata={"finish_reason": "STOP"}, ) 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( content="Security review of auth.py shows vulnerabilities", 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"}, ) mock_get_provider.return_value = mock_provider diff --git a/tests/test_intelligent_fallback.py b/tests/test_intelligent_fallback.py index 112f5bb..2d9bbca 100644 --- a/tests/test_intelligent_fallback.py +++ b/tests/test_intelligent_fallback.py @@ -33,10 +33,10 @@ class TestIntelligentFallback: @patch.dict(os.environ, {"OPENAI_API_KEY": "", "GEMINI_API_KEY": "test-gemini-key"}, clear=False) 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() 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) def test_prefers_openai_when_both_available(self): @@ -50,7 +50,7 @@ class TestIntelligentFallback: """Test fallback behavior when no API keys are available""" ModelProviderRegistry.clear_cache() 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): """Test the get_available_providers_with_keys method""" @@ -140,8 +140,8 @@ class TestIntelligentFallback: history, tokens = build_conversation_history(context, model_context=None) - # Should use gemini-2.0-flash when only Gemini is available - mock_context_class.assert_called_once_with("gemini-2.0-flash") + # Should use gemini-2.5-flash-preview-05-20 when only Gemini is available + mock_context_class.assert_called_once_with("gemini-2.5-flash-preview-05-20") def test_non_auto_mode_unchanged(self): """Test that non-auto mode behavior is unchanged""" diff --git a/tests/test_large_prompt_handling.py b/tests/test_large_prompt_handling.py index 33573aa..943de76 100644 --- a/tests/test_large_prompt_handling.py +++ b/tests/test_large_prompt_handling.py @@ -75,7 +75,7 @@ class TestLargePromptHandling: mock_provider.generate_content.return_value = MagicMock( content="This is a test response", 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"}, ) mock_get_provider.return_value = mock_provider @@ -100,7 +100,7 @@ class TestLargePromptHandling: mock_provider.generate_content.return_value = MagicMock( content="Processed large prompt", 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"}, ) mock_get_provider.return_value = mock_provider @@ -212,7 +212,7 @@ class TestLargePromptHandling: mock_provider.generate_content.return_value = MagicMock( content="Success", 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"}, ) mock_get_provider.return_value = mock_provider @@ -245,7 +245,7 @@ class TestLargePromptHandling: mock_provider.generate_content.return_value = MagicMock( content="Success", 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"}, ) mock_get_provider.return_value = mock_provider @@ -276,7 +276,7 @@ class TestLargePromptHandling: mock_provider.generate_content.return_value = MagicMock( content="Success", 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"}, ) mock_get_provider.return_value = mock_provider @@ -298,7 +298,7 @@ class TestLargePromptHandling: mock_provider.generate_content.return_value = MagicMock( content="Success", 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"}, ) mock_get_provider.return_value = mock_provider diff --git a/tests/test_prompt_regression.py b/tests/test_prompt_regression.py index cd5cedc..92291cf 100644 --- a/tests/test_prompt_regression.py +++ b/tests/test_prompt_regression.py @@ -31,7 +31,7 @@ class TestPromptRegression: return Mock( content=text, 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"}, ) diff --git a/tests/test_providers.py b/tests/test_providers.py index e7370de..ea35301 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -49,7 +49,7 @@ class TestModelProviderRegistry: """Test getting provider for a specific model""" 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 isinstance(provider, GeminiModelProvider) @@ -80,10 +80,10 @@ class TestGeminiProvider: """Test getting model capabilities""" 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.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 not capabilities.supports_extended_thinking @@ -103,13 +103,13 @@ class TestGeminiProvider: assert provider.validate_model_name("pro") 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): """Test thinking mode support detection""" 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") @patch("google.genai.Client") @@ -133,11 +133,13 @@ class TestGeminiProvider: 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 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.usage["input_tokens"] == 10 assert response.usage["output_tokens"] == 20 diff --git a/tests/test_server.py b/tests/test_server.py index 4d81015..78caf1f 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -56,7 +56,7 @@ class TestServerTools: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 diff --git a/tests/test_thinking_modes.py b/tests/test_thinking_modes.py index 8df8137..0fb4a68 100644 --- a/tests/test_thinking_modes.py +++ b/tests/test_thinking_modes.py @@ -45,7 +45,7 @@ class TestThinkingModes: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = True 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 @@ -82,7 +82,7 @@ class TestThinkingModes: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = True 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 @@ -114,7 +114,7 @@ class TestThinkingModes: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = True 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 @@ -145,7 +145,7 @@ class TestThinkingModes: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = True 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 @@ -175,7 +175,7 @@ class TestThinkingModes: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = True 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 @@ -222,18 +222,22 @@ class TestThinkingModes: async def prepare_prompt(self, request): 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 = { - "minimal": 128, - "low": 2048, - "medium": 8192, - "high": 16384, - "max": 32768, + "minimal": int(flash_max_tokens * 0.005), # 123 + "low": int(flash_max_tokens * 0.08), # 1966 + "medium": int(flash_max_tokens * 0.33), # 8110 + "high": int(flash_max_tokens * 0.67), # 16465 + "max": int(flash_max_tokens * 1.0), # 24576 } - # Check each mode in create_model - for _mode, _expected_budget in expected_budgets.items(): - # The budget mapping is inside create_model - # We can't easily test it without calling the method - # But we've verified the values are correct in the code - pass + # Check each mode using the helper method + for mode, expected_budget in expected_budgets.items(): + actual_budget = provider.get_thinking_budget(flash_model, mode) + assert actual_budget == expected_budget, f"Mode {mode}: expected {expected_budget}, got {actual_budget}" diff --git a/tests/test_tools.py b/tests/test_tools.py index 73aba51..bff688c 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -37,7 +37,7 @@ class TestThinkDeepTool: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = True 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 @@ -88,7 +88,7 @@ class TestCodeReviewTool: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -133,7 +133,7 @@ class TestDebugIssueTool: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -181,7 +181,7 @@ class TestAnalyzeTool: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 @@ -295,7 +295,7 @@ class TestAbsolutePathValidation: mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.supports_thinking_mode.return_value = False 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 diff --git a/tools/analyze.py b/tools/analyze.py index a3638b1..bd1f597 100644 --- a/tools/analyze.py +++ b/tools/analyze.py @@ -83,7 +83,7 @@ class AnalyzeTool(BaseTool): "thinking_mode": { "type": "string", "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": { "type": "boolean", diff --git a/tools/base.py b/tools/base.py index 9e6d73b..b7a64e9 100644 --- a/tools/base.py +++ b/tools/base.py @@ -57,7 +57,7 @@ class ToolRequest(BaseModel): # Higher values allow for more complex reasoning but increase latency and cost thinking_mode: Optional[Literal["minimal", "low", "medium", "high", "max"]] = Field( 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( True, diff --git a/tools/chat.py b/tools/chat.py index b44ce31..704f71f 100644 --- a/tools/chat.py +++ b/tools/chat.py @@ -68,7 +68,7 @@ class ChatTool(BaseTool): "thinking_mode": { "type": "string", "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": { "type": "boolean", diff --git a/tools/codereview.py b/tools/codereview.py index bd32777..e6889b3 100644 --- a/tools/codereview.py +++ b/tools/codereview.py @@ -126,7 +126,7 @@ class CodeReviewTool(BaseTool): "thinking_mode": { "type": "string", "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": { "type": "boolean", diff --git a/tools/debug.py b/tools/debug.py index 62d66e7..a58758e 100644 --- a/tools/debug.py +++ b/tools/debug.py @@ -86,7 +86,7 @@ class DebugIssueTool(BaseTool): "thinking_mode": { "type": "string", "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": { "type": "boolean", diff --git a/tools/thinkdeep.py b/tools/thinkdeep.py index 85a1388..e2d5f86 100644 --- a/tools/thinkdeep.py +++ b/tools/thinkdeep.py @@ -81,7 +81,7 @@ class ThinkDeepTool(BaseTool): "thinking_mode": { "type": "string", "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": { "type": "boolean", diff --git a/utils/conversation_memory.py b/utils/conversation_memory.py index cdef754..8b450a5 100644 --- a/utils/conversation_memory.py +++ b/utils/conversation_memory.py @@ -74,7 +74,7 @@ class ConversationTurn(BaseModel): files: List of file paths referenced in this specific turn tool_name: Which tool generated this turn (for cross-tool tracking) 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) """ @@ -249,7 +249,7 @@ def add_turn( files: Optional list of files referenced in this turn tool_name: Name of the tool adding this turn (for attribution) 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) Returns: From 22093bbf183724b8f04a1712f9002deb9cee0b97 Mon Sep 17 00:00:00 2001 From: Fahad Date: Thu, 12 Jun 2025 21:00:53 +0400 Subject: [PATCH 6/6] Fixed tests --- tests/test_providers.py | 2 +- tests/test_thinking_modes.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_providers.py b/tests/test_providers.py index ea35301..f1827f5 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -85,7 +85,7 @@ class TestGeminiProvider: assert capabilities.provider == ProviderType.GOOGLE assert capabilities.model_name == "gemini-2.5-flash-preview-05-20" 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): """Test getting capabilities for Pro model with thinking support""" diff --git a/tests/test_thinking_modes.py b/tests/test_thinking_modes.py index 0fb4a68..4a623e6 100644 --- a/tests/test_thinking_modes.py +++ b/tests/test_thinking_modes.py @@ -23,6 +23,7 @@ def setup_test_env(): class TestThinkingModes: """Test thinking modes across all tools""" + @patch("config.DEFAULT_THINKING_MODE_THINKDEEP", "high") def test_default_thinking_modes(self): """Test that tools have correct default thinking modes""" tools = [ @@ -169,6 +170,7 @@ class TestThinkingModes: @pytest.mark.asyncio @patch("tools.base.BaseTool.get_model_provider") + @patch("config.DEFAULT_THINKING_MODE_THINKDEEP", "high") async def test_thinking_mode_max(self, mock_get_provider): """Test max thinking mode (default for thinkdeep)""" mock_provider = create_mock_provider()