diff --git a/providers/gemini.py b/providers/gemini.py index 7ade53b..5788da0 100644 --- a/providers/gemini.py +++ b/providers/gemini.py @@ -361,23 +361,6 @@ class GeminiModelProvider(ModelProvider): """Get the provider type.""" return ProviderType.GOOGLE - 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.MODEL_CAPABILITIES.get(resolved_name) - - if not model_config or not model_config.supports_extended_thinking: - return 0 - - if thinking_mode not in self.THINKING_BUDGETS: - return 0 - - max_thinking_tokens = model_config.max_thinking_tokens - if max_thinking_tokens == 0: - return 0 - - return int(max_thinking_tokens * self.THINKING_BUDGETS[thinking_mode]) - def _extract_usage(self, response) -> dict[str, int]: """Extract token usage from Gemini response.""" usage = {} diff --git a/tests/test_thinking_modes.py b/tests/test_thinking_modes.py index b2e8a61..a0bc839 100644 --- a/tests/test_thinking_modes.py +++ b/tests/test_thinking_modes.py @@ -389,47 +389,3 @@ class TestThinkingModes: # Reload config and clear registry importlib.reload(config) ModelProviderRegistry._instance = None - - def test_thinking_budget_mapping(self): - """Test that thinking modes map to correct budget values""" - from tools.shared.base_tool import BaseTool - - # Create a simple test tool - class TestTool(BaseTool): - def get_name(self): - return "test" - - def get_description(self): - return "test" - - def get_input_schema(self): - return {} - - def get_system_prompt(self): - return "test" - - def get_request_model(self): - return None - - async def prepare_prompt(self, request): - return "test" - - # 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" - flash_max_tokens = 24576 - - expected_budgets = { - "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 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}"