refactor: cleanup, remove unused method
This commit is contained in:
@@ -361,23 +361,6 @@ class GeminiModelProvider(ModelProvider):
|
|||||||
"""Get the provider type."""
|
"""Get the provider type."""
|
||||||
return ProviderType.GOOGLE
|
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]:
|
def _extract_usage(self, response) -> dict[str, int]:
|
||||||
"""Extract token usage from Gemini response."""
|
"""Extract token usage from Gemini response."""
|
||||||
usage = {}
|
usage = {}
|
||||||
|
|||||||
@@ -389,47 +389,3 @@ class TestThinkingModes:
|
|||||||
# Reload config and clear registry
|
# Reload config and clear registry
|
||||||
importlib.reload(config)
|
importlib.reload(config)
|
||||||
ModelProviderRegistry._instance = None
|
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}"
|
|
||||||
|
|||||||
Reference in New Issue
Block a user