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

Gemini model rename

Browse Source
This commit is contained in:
Fahad
2025-06-19 05:37:40 +04:00
parent b6ad76b39a
commit d0da6ce9e4
37 changed files with 187 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/conftest.py
View File

@@ -25,7 +25,7 @@ if "XAI_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.5-flash-preview-05-20"
os.environ["DEFAULT_MODEL"] = "gemini-2.5-flash"
# Force reload of config module to pick up the env var
import config # noqa: E402
@@ -108,7 +108,7 @@ def mock_provider_availability(request, monkeypatch):
if model_name in ["unavailable-model", "gpt-5-turbo", "o3"]:
return None
# For common test models, return a mock provider
if model_name in ["gemini-2.5-flash-preview-05-20", "gemini-2.5-pro-preview-06-05", "pro", "flash"]:
if model_name in ["gemini-2.5-flash", "gemini-2.5-pro", "pro", "flash"]:
# Try to use the real provider first if it exists
real_provider = original_get_provider(model_name)
if real_provider:

2
tests/mock_helpers.py
View File

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

10
tests/test_alias_target_restrictions.py
View File

@@ -39,9 +39,9 @@ class TestAliasTargetRestrictions:
# Should include both aliases and their targets
assert "flash" in all_known # alias
assert "gemini-2.5-flash-preview-05-20" in all_known # target of 'flash'
assert "gemini-2.5-flash" in all_known # target of 'flash'
assert "pro" in all_known # alias
assert "gemini-2.5-pro-preview-06-05" in all_known # target of 'pro'
assert "gemini-2.5-pro" in all_known # target of 'pro'
@patch.dict(os.environ, {"OPENAI_ALLOWED_MODELS": "o4-mini"}) # Allow target
def test_restriction_policy_allows_alias_when_target_allowed(self):
@@ -80,7 +80,7 @@ class TestAliasTargetRestrictions:
# Direct target should NOT be allowed
assert not provider.validate_model_name("o4-mini")
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash-preview-05-20"}) # Allow target
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash"}) # Allow target
def test_gemini_restriction_policy_allows_alias_when_target_allowed(self):
"""Test Gemini restriction policy allows alias when target is allowed."""
# Clear cached restriction service
@@ -91,7 +91,7 @@ class TestAliasTargetRestrictions:
provider = GeminiModelProvider(api_key="test-key")
# Both target and alias should be allowed
assert provider.validate_model_name("gemini-2.5-flash-preview-05-20")
assert provider.validate_model_name("gemini-2.5-flash")
assert provider.validate_model_name("flash")
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "flash"}) # Allow alias only
@@ -107,7 +107,7 @@ class TestAliasTargetRestrictions:
# Only the alias should be allowed
assert provider.validate_model_name("flash")
# Direct target should NOT be allowed
assert not provider.validate_model_name("gemini-2.5-flash-preview-05-20")
assert not provider.validate_model_name("gemini-2.5-flash")
def test_restriction_service_validation_includes_all_targets(self):
"""Test that restriction service validation knows about all aliases and targets."""

22
tests/test_auto_mode_comprehensive.py
View File

@@ -80,9 +80,9 @@ class TestAutoModeComprehensive:
"OPENROUTER_API_KEY": None,
},
{
"EXTENDED_REASONING": "gemini-2.5-pro-preview-06-05", # Pro for deep thinking
"FAST_RESPONSE": "gemini-2.5-flash-preview-05-20", # Flash for speed
"BALANCED": "gemini-2.5-flash-preview-05-20", # Flash as balanced
"EXTENDED_REASONING": "gemini-2.5-pro", # Pro for deep thinking
"FAST_RESPONSE": "gemini-2.5-flash", # Flash for speed
"BALANCED": "gemini-2.5-flash", # Flash as balanced
},
),
# Only OpenAI API available
@@ -313,8 +313,8 @@ class TestAutoModeComprehensive:
# Should include Gemini models
assert "flash" in available_models
assert "pro" in available_models
assert "gemini-2.5-flash-preview-05-20" in available_models
assert "gemini-2.5-pro-preview-06-05" in available_models
assert "gemini-2.5-flash" in available_models
assert "gemini-2.5-pro" in available_models
# Should also include other models (users might have OpenRouter configured)
# The schema should show all options; validation happens at runtime
@@ -476,8 +476,8 @@ class TestAutoModeComprehensive:
assert "o3-mini" not in available_models
# Should still include all Gemini models (no restrictions)
assert "gemini-2.5-flash-preview-05-20" in available_models
assert "gemini-2.5-pro-preview-06-05" in available_models
assert "gemini-2.5-flash" in available_models
assert "gemini-2.5-pro" in available_models
def test_openrouter_fallback_when_no_native_apis(self):
"""Test that OpenRouter provides fallback models when no native APIs are available."""
@@ -511,8 +511,8 @@ class TestAutoModeComprehensive:
# Mock OpenRouter registry to return known models
mock_registry = MagicMock()
mock_registry.list_models.return_value = [
"google/gemini-2.5-flash-preview-05-20",
"google/gemini-2.5-pro-preview-06-05",
"google/gemini-2.5-flash",
"google/gemini-2.5-pro",
"openai/o3",
"openai/o4-mini",
"anthropic/claude-3-opus",
@@ -564,11 +564,11 @@ class TestAutoModeComprehensive:
mock_provider = MagicMock()
mock_response = MagicMock()
mock_response.content = "test response"
mock_response.model_name = "gemini-2.5-flash-preview-05-20" # The resolved name
mock_response.model_name = "gemini-2.5-flash" # The resolved name
mock_response.usage = {"input_tokens": 10, "output_tokens": 5}
# Mock _resolve_model_name to simulate alias resolution
mock_provider._resolve_model_name = lambda alias: (
"gemini-2.5-flash-preview-05-20" if alias == "flash" else alias
"gemini-2.5-flash" if alias == "flash" else alias
)
mock_provider.generate_content.return_value = mock_response

2
tests/test_auto_mode_custom_provider_only.py
View File

@@ -201,7 +201,7 @@ class TestAutoModeCustomProviderOnly:
# Should get a valid model name, not the hardcoded fallback
assert (
fallback_model != "gemini-2.5-flash-preview-05-20"
fallback_model != "gemini-2.5-flash"
), "Should not fallback to hardcoded Gemini model when custom provider is available"
except Exception as e:

14
tests/test_auto_mode_provider_selection.py
View File

@@ -59,9 +59,9 @@ class TestAutoModeProviderSelection:
balanced = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.BALANCED)
# Should select appropriate Gemini models
assert extended_reasoning in ["gemini-2.5-pro-preview-06-05", "pro"]
assert fast_response in ["gemini-2.5-flash-preview-05-20", "flash"]
assert balanced in ["gemini-2.5-flash-preview-05-20", "flash"]
assert extended_reasoning in ["gemini-2.5-pro", "pro"]
assert fast_response in ["gemini-2.5-flash", "flash"]
assert balanced in ["gemini-2.5-flash", "flash"]
finally:
# Restore original environment
@@ -229,8 +229,8 @@ class TestAutoModeProviderSelection:
assert "o3-mini" not in available_models
# Should include all Gemini models (no restrictions)
assert "gemini-2.5-flash-preview-05-20" in available_models
assert available_models["gemini-2.5-flash-preview-05-20"] == ProviderType.GOOGLE
assert "gemini-2.5-flash" in available_models
assert available_models["gemini-2.5-flash"] == ProviderType.GOOGLE
finally:
# Restore original environment
@@ -316,8 +316,8 @@ class TestAutoModeProviderSelection:
# Test that providers resolve aliases correctly
test_cases = [
("flash", ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20"),
("pro", ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05"),
("flash", ProviderType.GOOGLE, "gemini-2.5-flash"),
("pro", ProviderType.GOOGLE, "gemini-2.5-pro"),
("mini", ProviderType.OPENAI, "o4-mini"),
("o3mini", ProviderType.OPENAI, "o3-mini"),
("grok", ProviderType.XAI, "grok-3"),

14
tests/test_buggy_behavior_prevention.py
View File

@@ -85,16 +85,16 @@ class TestBuggyBehaviorPrevention:
# Verify both aliases and targets are included
assert "flash" in all_known # alias
assert "gemini-2.5-flash-preview-05-20" in all_known # target
assert "gemini-2.5-flash" in all_known # target
assert "pro" in all_known # alias
assert "gemini-2.5-pro-preview-06-05" in all_known # target
assert "gemini-2.5-pro" in all_known # target
# Simulate admin restricting to target model names
service = ModelRestrictionService()
service.restrictions = {
ProviderType.GOOGLE: {
"gemini-2.5-flash-preview-05-20", # Target name restriction
"gemini-2.5-pro-preview-06-05", # Target name restriction
"gemini-2.5-flash", # Target name restriction
"gemini-2.5-pro", # Target name restriction
}
}
@@ -105,8 +105,8 @@ class TestBuggyBehaviorPrevention:
# Should NOT warn about these valid target models
all_warnings = [str(call) for call in mock_logger.warning.call_args_list]
for warning in all_warnings:
assert "gemini-2.5-flash-preview-05-20" not in warning or "not a recognized" not in warning
assert "gemini-2.5-pro-preview-06-05" not in warning or "not a recognized" not in warning
assert "gemini-2.5-flash" not in warning or "not a recognized" not in warning
assert "gemini-2.5-pro" not in warning or "not a recognized" not in warning
def test_old_bug_policy_bypass_prevention(self):
"""
@@ -206,7 +206,7 @@ class TestBuggyBehaviorPrevention:
"""
providers_to_test = [
(OpenAIModelProvider(api_key="test-key"), "mini", "o4-mini"),
(GeminiModelProvider(api_key="test-key"), "flash", "gemini-2.5-flash-preview-05-20"),
(GeminiModelProvider(api_key="test-key"), "flash", "gemini-2.5-flash"),
]
for provider, alias, target in providers_to_test:

24
tests/test_claude_continuation.py
View File

@@ -59,7 +59,7 @@ class TestClaudeContinuationOffers:
# If providers are not registered yet, tool might detect auto mode
self.tool = ClaudeContinuationTool()
# Set default model to avoid effective auto mode
self.tool.default_model = "gemini-2.5-flash-preview-05-20"
self.tool.default_model = "gemini-2.5-flash"
@patch("utils.conversation_memory.get_storage")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
@@ -67,7 +67,7 @@ class TestClaudeContinuationOffers:
"""Test that new conversations offer Claude continuation opportunity"""
# Create tool AFTER providers are registered (in conftest.py fixture)
tool = ClaudeContinuationTool()
tool.default_model = "gemini-2.5-flash-preview-05-20"
tool.default_model = "gemini-2.5-flash"
mock_client = Mock()
mock_storage.return_value = mock_client
@@ -80,7 +80,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -137,7 +137,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -170,7 +170,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -212,7 +212,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -257,7 +257,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -313,7 +313,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -336,7 +336,7 @@ class TestContinuationIntegration:
def setup_method(self):
self.tool = ClaudeContinuationTool()
# Set default model to avoid effective auto mode
self.tool.default_model = "gemini-2.5-flash-preview-05-20"
self.tool.default_model = "gemini-2.5-flash"
@patch("utils.conversation_memory.get_storage")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
@@ -364,7 +364,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -417,7 +417,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -454,7 +454,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock(
content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
)
mock_get_provider.return_value = mock_provider
@@ -86,7 +86,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.5-flash-preview-05-20", metadata={}
content=normal_response, usage={}, model_name="gemini-2.5-flash", metadata={}
)
mock_get_provider.return_value = mock_provider
@@ -110,7 +110,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.5-flash-preview-05-20", metadata={}
content=malformed_json, usage={}, model_name="gemini-2.5-flash", metadata={}
)
mock_get_provider.return_value = mock_provider
@@ -150,7 +150,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.5-flash-preview-05-20", metadata={}
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
)
mock_get_provider.return_value = mock_provider
@@ -274,7 +274,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.5-flash-preview-05-20", metadata={}
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
)
mock_get_provider.return_value = mock_provider
@@ -313,7 +313,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.5-flash-preview-05-20", metadata={}
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
)
mock_get_provider.return_value = mock_provider
@@ -340,7 +340,7 @@ class TestCollaborationWorkflow:
"""
mock_provider.generate_content.return_value = Mock(
content=final_response, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
content=final_response, usage={}, model_name="gemini-2.5-flash", metadata={}
)
result2 = await tool.execute(

2
tests/test_config.py
View File

@@ -31,7 +31,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.5-flash-preview-05-20"
assert DEFAULT_MODEL == "gemini-2.5-flash"
def test_temperature_defaults(self):
"""Test temperature constants"""

4
tests/test_conversation_field_mapping.py
View File

@@ -74,7 +74,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
friendly_name="Gemini",
context_window=200000,
supports_extended_thinking=True,
@@ -132,7 +132,7 @@ async def test_unknown_tool_defaults_to_prompt():
# Mock ModelContext to avoid calculation errors
with patch("utils.model_context.ModelContext") as mock_model_context_class:
mock_model_context = MagicMock()
mock_model_context.model_name = "gemini-2.5-flash-preview-05-20"
mock_model_context.model_name = "gemini-2.5-flash"
mock_model_context.calculate_token_allocation.return_value = MagicMock(
total_tokens=200000,
content_tokens=120000,

8
tests/test_conversation_history_bug.py
View File

@@ -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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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(
content=content,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider

14
tests/test_intelligent_fallback.py
View File

@@ -48,14 +48,14 @@ 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.5-flash-preview-05-20 is used when only Gemini API key is available"""
"""Test that gemini-2.5-flash is used when only Gemini API key is available"""
# Register only Gemini provider for this test
from providers.gemini import GeminiModelProvider
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
fallback_model = ModelProviderRegistry.get_preferred_fallback_model()
assert fallback_model == "gemini-2.5-flash-preview-05-20"
assert fallback_model == "gemini-2.5-flash"
@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):
@@ -81,7 +81,7 @@ class TestIntelligentFallback:
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
fallback_model = ModelProviderRegistry.get_preferred_fallback_model()
assert fallback_model == "gemini-2.5-flash-preview-05-20" # Default fallback
assert fallback_model == "gemini-2.5-flash" # Default fallback
def test_available_providers_with_keys(self):
"""Test the get_available_providers_with_keys method"""
@@ -186,14 +186,14 @@ class TestIntelligentFallback:
history, tokens = build_conversation_history(context, model_context=None)
# 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")
# Should use gemini-2.5-flash when only Gemini is available
mock_context_class.assert_called_once_with("gemini-2.5-flash")
def test_non_auto_mode_unchanged(self):
"""Test that non-auto mode behavior is unchanged"""
from utils.conversation_memory import ThreadContext, build_conversation_history
with patch("config.IS_AUTO_MODE", False), patch("config.DEFAULT_MODEL", "gemini-2.5-pro-preview-06-05"):
with patch("config.IS_AUTO_MODE", False), patch("config.DEFAULT_MODEL", "gemini-2.5-pro"):
from utils.conversation_memory import ConversationTurn
context = ThreadContext(
@@ -219,7 +219,7 @@ class TestIntelligentFallback:
history, tokens = build_conversation_history(context, model_context=None)
# Should use the configured DEFAULT_MODEL, not the intelligent fallback
mock_context_class.assert_called_once_with("gemini-2.5-pro-preview-06-05")
mock_context_class.assert_called_once_with("gemini-2.5-pro")
if __name__ == "__main__":

22
tests/test_large_prompt_handling.py
View File

@@ -76,7 +76,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -103,7 +103,7 @@ class TestLargePromptHandling:
patch("utils.model_context.ModelContext") as mock_model_context_class,
):
mock_provider = create_mock_provider(model_name="gemini-2.5-flash-preview-05-20", context_window=1_048_576)
mock_provider = create_mock_provider(model_name="gemini-2.5-flash", context_window=1_048_576)
mock_provider.generate_content.return_value.content = "Processed prompt from file"
mock_get_provider.return_value = mock_provider
@@ -111,7 +111,7 @@ class TestLargePromptHandling:
from utils.model_context import TokenAllocation
mock_model_context = MagicMock()
mock_model_context.model_name = "gemini-2.5-flash-preview-05-20"
mock_model_context.model_name = "gemini-2.5-flash"
mock_model_context.calculate_token_allocation.return_value = TokenAllocation(
total_tokens=1_048_576,
content_tokens=838_861,
@@ -293,7 +293,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -336,7 +336,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock(
content="Response to the large prompt",
usage={"input_tokens": 12000, "output_tokens": 10, "total_tokens": 12010},
model_name="gemini-2.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -368,7 +368,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.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -390,7 +390,7 @@ class TestLargePromptHandling:
patch("utils.model_context.ModelContext") as mock_model_context_class,
):
mock_provider = create_mock_provider(model_name="gemini-2.5-flash-preview-05-20", context_window=1_048_576)
mock_provider = create_mock_provider(model_name="gemini-2.5-flash", context_window=1_048_576)
mock_provider.generate_content.return_value.content = "Success"
mock_get_provider.return_value = mock_provider
@@ -398,7 +398,7 @@ class TestLargePromptHandling:
from utils.model_context import TokenAllocation
mock_model_context = MagicMock()
mock_model_context.model_name = "gemini-2.5-flash-preview-05-20"
mock_model_context.model_name = "gemini-2.5-flash"
mock_model_context.calculate_token_allocation.return_value = TokenAllocation(
total_tokens=1_048_576,
content_tokens=838_861,
@@ -437,7 +437,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock(
content="Weather is sunny",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -502,7 +502,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock(
content="Hi there!",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
@@ -555,7 +555,7 @@ class TestLargePromptHandling:
mock_provider.generate_content.return_value = MagicMock(
content="Continuing our conversation...",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.5-flash-preview-05-20",
model_name="gemini-2.5-flash",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider

4
tests/test_listmodels.py
View File

@@ -65,8 +65,8 @@ class TestListModelsTool:
# Check Gemini shows as configured
assert "Google Gemini ✅" in content
assert "`flash` → `gemini-2.5-flash-preview-05-20`" in content
assert "`pro` → `gemini-2.5-pro-preview-06-05`" in content
assert "`flash` → `gemini-2.5-flash`" in content
assert "`pro` → `gemini-2.5-pro`" in content
assert "1M context" in content
# Check summary

6
tests/test_model_enumeration.py
View File

@@ -96,8 +96,8 @@ class TestModelEnumeration:
"grok-3-fast",
"grok3",
"grokfast", # X.AI models
"gemini-2.5-flash-preview-05-20",
"gemini-2.5-pro-preview-06-05", # Full Gemini names
"gemini-2.5-flash",
"gemini-2.5-pro", # Full Gemini names
]
for model in native_models:
@@ -264,7 +264,7 @@ class TestModelEnumeration:
("flash", True), # Native Gemini
("o3", True), # Native OpenAI
("grok", True), # Native X.AI
("gemini-2.5-flash-preview-05-20", True), # Full native name
("gemini-2.5-flash", True), # Full native name
("o4-mini-high", True), # Native OpenAI variant
("grok-3-fast", True), # Native X.AI variant
],

54
tests/test_model_restrictions.py
View File

@@ -22,8 +22,8 @@ class TestModelRestrictionService:
# Should allow all models
assert service.is_allowed(ProviderType.OPENAI, "o3")
assert service.is_allowed(ProviderType.OPENAI, "o3-mini")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash")
assert service.is_allowed(ProviderType.OPENROUTER, "anthropic/claude-3-opus")
assert service.is_allowed(ProviderType.OPENROUTER, "openai/o3")
@@ -43,7 +43,7 @@ class TestModelRestrictionService:
assert not service.is_allowed(ProviderType.OPENAI, "o4-mini")
# Google and OpenRouter should have no restrictions
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro")
assert service.is_allowed(ProviderType.OPENROUTER, "anthropic/claude-3-opus")
def test_load_multiple_models_restriction(self):
@@ -59,7 +59,7 @@ class TestModelRestrictionService:
# Check Google models
assert service.is_allowed(ProviderType.GOOGLE, "flash")
assert service.is_allowed(ProviderType.GOOGLE, "pro")
assert not service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05")
assert not service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro")
def test_case_insensitive_and_whitespace_handling(self):
"""Test that model names are case-insensitive and whitespace is trimmed."""
@@ -84,9 +84,9 @@ class TestModelRestrictionService:
# Google should only allow flash (and its resolved name)
assert service.is_allowed(ProviderType.GOOGLE, "flash")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20", "flash")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash", "flash")
assert not service.is_allowed(ProviderType.GOOGLE, "pro")
assert not service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05", "pro")
assert not service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro", "pro")
def test_filter_models(self):
"""Test filtering a list of models based on restrictions."""
@@ -124,8 +124,8 @@ class TestModelRestrictionService:
assert not service.is_allowed(ProviderType.OPENAI, "o3")
# Google should allow both models via shorthands
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20", "flash")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05", "pro")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash", "flash")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro", "pro")
# Also test that full names work when specified in restrictions
assert service.is_allowed(ProviderType.OPENAI, "o3-mini", "o3mini") # Even with shorthand
@@ -238,7 +238,7 @@ class TestProviderIntegration:
provider.get_capabilities("o3")
assert "not allowed by restriction policy" in str(exc_info.value)
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash-preview-05-20,flash"})
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash,flash"})
def test_gemini_provider_respects_restrictions(self):
"""Test that Gemini provider respects restrictions."""
# Clear any cached restriction service
@@ -250,11 +250,11 @@ class TestProviderIntegration:
# Should validate allowed models (both shorthand and full name allowed)
assert provider.validate_model_name("flash")
assert provider.validate_model_name("gemini-2.5-flash-preview-05-20")
assert provider.validate_model_name("gemini-2.5-flash")
# Should not validate disallowed model
assert not provider.validate_model_name("pro")
assert not provider.validate_model_name("gemini-2.5-pro-preview-06-05")
assert not provider.validate_model_name("gemini-2.5-pro")
# get_capabilities should raise for disallowed model
with pytest.raises(ValueError) as exc_info:
@@ -288,13 +288,13 @@ class TestProviderIntegration:
# Should allow getting capabilities for "flash"
capabilities = provider.get_capabilities("flash")
assert capabilities.model_name == "gemini-2.5-flash-preview-05-20"
assert capabilities.model_name == "gemini-2.5-flash"
# Test the edge case: Try to use full model name when only alias is allowed
# This should NOT be allowed - only the alias "flash" is in the restriction list
assert not provider.validate_model_name("gemini-2.5-flash-preview-05-20")
assert not provider.validate_model_name("gemini-2.5-flash")
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash-preview-05-20"})
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash"})
def test_gemini_parameter_order_edge_case_full_name_only(self):
"""Test parameter order with only full name allowed, not alias.
@@ -310,7 +310,7 @@ class TestProviderIntegration:
provider = GeminiModelProvider(api_key="test-key")
# Should allow full name
assert provider.validate_model_name("gemini-2.5-flash-preview-05-20")
assert provider.validate_model_name("gemini-2.5-flash")
# Should also allow alias that resolves to allowed full name
# This works because is_allowed checks both resolved_name and original_name
@@ -318,7 +318,7 @@ class TestProviderIntegration:
# Should not allow "pro" alias
assert not provider.validate_model_name("pro")
assert not provider.validate_model_name("gemini-2.5-pro-preview-06-05")
assert not provider.validate_model_name("gemini-2.5-pro")
class TestCustomProviderOpenRouterRestrictions:
@@ -469,8 +469,8 @@ class TestRegistryIntegration:
mock_gemini = MagicMock()
mock_gemini.SUPPORTED_MODELS = {
"gemini-2.5-pro-preview-06-05": {"context_window": 1048576},
"gemini-2.5-flash-preview-05-20": {"context_window": 1048576},
"gemini-2.5-pro": {"context_window": 1048576},
"gemini-2.5-flash": {"context_window": 1048576},
}
mock_gemini.get_provider_type.return_value = ProviderType.GOOGLE
@@ -493,8 +493,8 @@ class TestRegistryIntegration:
mock_gemini.list_models = gemini_list_models
mock_gemini.list_all_known_models.return_value = [
"gemini-2.5-pro-preview-06-05",
"gemini-2.5-flash-preview-05-20",
"gemini-2.5-pro",
"gemini-2.5-flash",
]
def get_provider_side_effect(provider_type):
@@ -514,7 +514,7 @@ class TestRegistryIntegration:
}
with patch.dict(
os.environ, {"OPENAI_ALLOWED_MODELS": "o3-mini", "GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash-preview-05-20"}
os.environ, {"OPENAI_ALLOWED_MODELS": "o3-mini", "GOOGLE_ALLOWED_MODELS": "gemini-2.5-flash"}
):
# Clear cached restriction service
import utils.model_restrictions
@@ -526,8 +526,8 @@ class TestRegistryIntegration:
# Should only include allowed models
assert "o3-mini" in available
assert "o3" not in available
assert "gemini-2.5-flash-preview-05-20" in available
assert "gemini-2.5-pro-preview-06-05" not in available
assert "gemini-2.5-flash" in available
assert "gemini-2.5-pro" not in available
class TestShorthandRestrictions:
@@ -552,7 +552,7 @@ class TestShorthandRestrictions:
gemini_provider = GeminiModelProvider(api_key="test-key")
assert gemini_provider.validate_model_name("flash") # Should work with shorthand
# Same for Gemini - if you restrict to "flash", you can't use the full name
assert not gemini_provider.validate_model_name("gemini-2.5-flash-preview-05-20") # Not allowed
assert not gemini_provider.validate_model_name("gemini-2.5-flash") # Not allowed
assert not gemini_provider.validate_model_name("pro") # Not allowed
@patch.dict(os.environ, {"OPENAI_ALLOWED_MODELS": "o3mini,mini,o4-mini"})
@@ -579,7 +579,7 @@ class TestShorthandRestrictions:
@patch.dict(
os.environ,
{"OPENAI_ALLOWED_MODELS": "mini,o4-mini", "GOOGLE_ALLOWED_MODELS": "flash,gemini-2.5-flash-preview-05-20"},
{"OPENAI_ALLOWED_MODELS": "mini,o4-mini", "GOOGLE_ALLOWED_MODELS": "flash,gemini-2.5-flash"},
)
def test_both_shorthand_and_full_name_allowed(self):
"""Test that we can allow both shorthand and full names."""
@@ -596,7 +596,7 @@ class TestShorthandRestrictions:
# Gemini - both flash and full name are allowed
gemini_provider = GeminiModelProvider(api_key="test-key")
assert gemini_provider.validate_model_name("flash")
assert gemini_provider.validate_model_name("gemini-2.5-flash-preview-05-20")
assert gemini_provider.validate_model_name("gemini-2.5-flash")
class TestAutoModeWithRestrictions:
@@ -688,7 +688,7 @@ class TestAutoModeWithRestrictions:
# The fallback will depend on how get_available_models handles aliases
# For now, we accept either behavior and document it
assert model in ["o4-mini", "gemini-2.5-flash-preview-05-20"]
assert model in ["o4-mini", "gemini-2.5-flash"]
finally:
# Restore original registry state
registry = ModelProviderRegistry()

2
tests/test_old_behavior_simulation.py
View File

@@ -167,7 +167,7 @@ class TestOldBehaviorSimulation:
(["mini", "o3mini"], ["mini", "o3mini", "o4-mini", "o3-mini"], "OpenAI"),
(
["flash", "pro"],
["flash", "pro", "gemini-2.5-flash-preview-05-20", "gemini-2.5-pro-preview-06-05"],
["flash", "pro", "gemini-2.5-flash", "gemini-2.5-pro"],
"Gemini",
),
]

12
tests/test_openrouter_provider.py
View File

@@ -151,8 +151,8 @@ class TestOpenRouterAutoMode:
mock_registry = Mock()
mock_registry.list_models.return_value = [
"google/gemini-2.5-flash-preview-05-20",
"google/gemini-2.5-pro-preview-06-05",
"google/gemini-2.5-flash",
"google/gemini-2.5-pro",
"openai/o3",
"openai/o3-mini",
"anthropic/claude-3-opus",
@@ -181,7 +181,7 @@ class TestOpenRouterAutoMode:
os.environ.pop("OPENAI_API_KEY", None)
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
os.environ["OPENROUTER_ALLOWED_MODELS"] = "anthropic/claude-3-opus,google/gemini-2.5-flash-preview-05-20"
os.environ["OPENROUTER_ALLOWED_MODELS"] = "anthropic/claude-3-opus,google/gemini-2.5-flash"
os.environ["DEFAULT_MODEL"] = "auto"
# Force reload to pick up new environment variable
@@ -191,8 +191,8 @@ class TestOpenRouterAutoMode:
mock_registry = Mock()
mock_registry.list_models.return_value = [
"google/gemini-2.5-flash-preview-05-20",
"google/gemini-2.5-pro-preview-06-05",
"google/gemini-2.5-flash",
"google/gemini-2.5-pro",
"anthropic/claude-3-opus",
"anthropic/claude-3-sonnet",
]
@@ -206,7 +206,7 @@ class TestOpenRouterAutoMode:
assert len(available_models) > 0, "Should have some allowed models"
expected_allowed = {"google/gemini-2.5-flash-preview-05-20", "anthropic/claude-3-opus"}
expected_allowed = {"google/gemini-2.5-flash", "anthropic/claude-3-opus"}
assert (
set(available_models.keys()) == expected_allowed

34
tests/test_per_tool_model_defaults.py
View File

@@ -91,13 +91,13 @@ class TestModelSelection:
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
# Mock only Gemini models available
mock_get_available.return_value = {
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
"gemini-2.5-pro": ProviderType.GOOGLE,
"gemini-2.5-flash": ProviderType.GOOGLE,
}
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.EXTENDED_REASONING)
# Should find the pro model for extended reasoning
assert "pro" in model or model == "gemini-2.5-pro-preview-06-05"
assert "pro" in model or model == "gemini-2.5-pro"
def test_fast_response_with_openai(self):
"""Test FAST_RESPONSE prefers o4-mini when OpenAI is available."""
@@ -117,13 +117,13 @@ class TestModelSelection:
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
# Mock only Gemini models available
mock_get_available.return_value = {
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
"gemini-2.5-pro": ProviderType.GOOGLE,
"gemini-2.5-flash": ProviderType.GOOGLE,
}
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.FAST_RESPONSE)
# Should find the flash model for fast response
assert "flash" in model or model == "gemini-2.5-flash-preview-05-20"
assert "flash" in model or model == "gemini-2.5-flash"
def test_balanced_category_fallback(self):
"""Test BALANCED category uses existing logic."""
@@ -143,13 +143,13 @@ class TestModelSelection:
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
# Mock only Gemini models available
mock_get_available.return_value = {
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
"gemini-2.5-pro": ProviderType.GOOGLE,
"gemini-2.5-flash": ProviderType.GOOGLE,
}
model = ModelProviderRegistry.get_preferred_fallback_model()
# Should pick a reasonable default, preferring flash for balanced use
assert "flash" in model or model == "gemini-2.5-flash-preview-05-20"
assert "flash" in model or model == "gemini-2.5-flash"
class TestFlexibleModelSelection:
@@ -168,8 +168,8 @@ class TestFlexibleModelSelection:
# Case 2: Mix of Gemini shorthands and full names
{
"available": {
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
"gemini-2.5-flash": ProviderType.GOOGLE,
"gemini-2.5-pro": ProviderType.GOOGLE,
},
"category": ToolModelCategory.FAST_RESPONSE,
"expected_contains": "flash",
@@ -220,7 +220,7 @@ class TestCustomProviderFallback:
mock_find_thinking.return_value = None
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.EXTENDED_REASONING)
assert model == "gemini-2.5-pro-preview-06-05"
assert model == "gemini-2.5-pro"
class TestAutoModeErrorMessages:
@@ -234,8 +234,8 @@ class TestAutoModeErrorMessages:
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
# Mock only Gemini models available
mock_get_available.return_value = {
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
"gemini-2.5-pro": ProviderType.GOOGLE,
"gemini-2.5-flash": ProviderType.GOOGLE,
}
tool = ThinkDeepTool()
@@ -245,7 +245,7 @@ class TestAutoModeErrorMessages:
assert "Model parameter is required in auto mode" in result[0].text
# Should suggest a model suitable for extended reasoning (either full name or with 'pro')
response_text = result[0].text
assert "gemini-2.5-pro-preview-06-05" in response_text or "pro" in response_text
assert "gemini-2.5-pro" in response_text or "pro" in response_text
assert "(category: extended_reasoning)" in response_text
@pytest.mark.asyncio
@@ -294,7 +294,7 @@ class TestFileContentPreparation:
# Set up model context to simulate normal execution flow
from utils.model_context import ModelContext
tool._model_context = ModelContext("gemini-2.5-pro-preview-06-05")
tool._model_context = ModelContext("gemini-2.5-pro")
# Call the method
content, processed_files = tool._prepare_file_content_for_prompt(["/test/file.py"], None, "test")
@@ -304,7 +304,7 @@ class TestFileContentPreparation:
assert len(debug_calls) > 0
debug_message = str(debug_calls[0])
# Should mention the model being used
assert "gemini-2.5-pro-preview-06-05" in debug_message
assert "gemini-2.5-pro" in debug_message
# Should mention file tokens (not content tokens)
assert "file tokens" in debug_message

2
tests/test_prompt_regression.py
View File

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

18
tests/test_providers.py
View File

@@ -64,7 +64,7 @@ class TestModelProviderRegistry:
"""Test getting provider for a specific model"""
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
provider = ModelProviderRegistry.get_provider_for_model("gemini-2.5-flash-preview-05-20")
provider = ModelProviderRegistry.get_provider_for_model("gemini-2.5-flash")
assert provider is not None
assert isinstance(provider, GeminiModelProvider)
@@ -95,10 +95,10 @@ class TestGeminiProvider:
"""Test getting model capabilities"""
provider = GeminiModelProvider(api_key="test-key")
capabilities = provider.get_capabilities("gemini-2.5-flash-preview-05-20")
capabilities = provider.get_capabilities("gemini-2.5-flash")
assert capabilities.provider == ProviderType.GOOGLE
assert capabilities.model_name == "gemini-2.5-flash-preview-05-20"
assert capabilities.model_name == "gemini-2.5-flash"
assert capabilities.context_window == 1_048_576
assert capabilities.supports_extended_thinking
@@ -106,7 +106,7 @@ class TestGeminiProvider:
"""Test getting capabilities for Pro model with thinking support"""
provider = GeminiModelProvider(api_key="test-key")
capabilities = provider.get_capabilities("gemini-2.5-pro-preview-06-05")
capabilities = provider.get_capabilities("gemini-2.5-pro")
assert capabilities.supports_extended_thinking
@@ -118,14 +118,14 @@ class TestGeminiProvider:
assert provider.validate_model_name("pro")
capabilities = provider.get_capabilities("flash")
assert capabilities.model_name == "gemini-2.5-flash-preview-05-20"
assert capabilities.model_name == "gemini-2.5-flash"
def test_supports_thinking_mode(self):
"""Test thinking mode support detection"""
provider = GeminiModelProvider(api_key="test-key")
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-flash")
assert provider.supports_thinking_mode("gemini-2.5-pro")
@patch("google.genai.Client")
def test_generate_content(self, mock_client_class):
@@ -149,12 +149,12 @@ class TestGeminiProvider:
provider = GeminiModelProvider(api_key="test-key")
response = provider.generate_content(
prompt="Test prompt", model_name="gemini-2.5-flash-preview-05-20", temperature=0.7
prompt="Test prompt", model_name="gemini-2.5-flash", temperature=0.7
)
assert isinstance(response, ModelResponse)
assert response.content == "Generated content"
assert response.model_name == "gemini-2.5-flash-preview-05-20"
assert response.model_name == "gemini-2.5-flash"
assert response.provider == ProviderType.GOOGLE
assert response.usage["input_tokens"] == 10
assert response.usage["output_tokens"] == 20

2
tests/test_thinking_modes.py
View File

@@ -418,7 +418,7 @@ class TestThinkingModes:
from providers.gemini import GeminiModelProvider
provider = GeminiModelProvider(api_key="test-key")
flash_model = "gemini-2.5-flash-preview-05-20"
flash_model = "gemini-2.5-flash"
flash_max_tokens = 24576
expected_budgets = {