Gemini model rename
This commit is contained in:
@@ -58,10 +58,10 @@ DEFAULT_THINKING_MODE_THINKDEEP=high
|
|||||||
# - mini (shorthand for o4-mini)
|
# - mini (shorthand for o4-mini)
|
||||||
#
|
#
|
||||||
# Supported Google/Gemini models:
|
# Supported Google/Gemini models:
|
||||||
# - gemini-2.5-flash-preview-05-20 (1M context, fast, supports thinking)
|
# - gemini-2.5-flash (1M context, fast, supports thinking)
|
||||||
# - gemini-2.5-pro-preview-06-05 (1M context, powerful, supports thinking)
|
# - gemini-2.5-pro (1M context, powerful, supports thinking)
|
||||||
# - flash (shorthand for gemini-2.5-flash-preview-05-20)
|
# - flash (shorthand for gemini-2.5-flash)
|
||||||
# - pro (shorthand for gemini-2.5-pro-preview-06-05)
|
# - pro (shorthand for gemini-2.5-pro)
|
||||||
#
|
#
|
||||||
# Supported X.AI GROK models:
|
# Supported X.AI GROK models:
|
||||||
# - grok-3 (131K context, advanced reasoning)
|
# - grok-3 (131K context, advanced reasoning)
|
||||||
|
|||||||
@@ -89,7 +89,7 @@
|
|||||||
"description": "Google's Gemini 2.5 Pro via OpenRouter with vision"
|
"description": "Google's Gemini 2.5 Pro via OpenRouter with vision"
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"model_name": "google/gemini-2.5-flash-preview-05-20",
|
"model_name": "google/gemini-2.5-flash",
|
||||||
"aliases": ["flash","gemini-flash", "flash-openrouter", "flash-2.5"],
|
"aliases": ["flash","gemini-flash", "flash-openrouter", "flash-2.5"],
|
||||||
"context_window": 1048576,
|
"context_window": 1048576,
|
||||||
"supports_extended_thinking": false,
|
"supports_extended_thinking": false,
|
||||||
|
|||||||
10
config.py
10
config.py
@@ -14,7 +14,7 @@ import os
|
|||||||
# These values are used in server responses and for tracking releases
|
# These values are used in server responses and for tracking releases
|
||||||
# IMPORTANT: This is the single source of truth for version and author info
|
# IMPORTANT: This is the single source of truth for version and author info
|
||||||
# Semantic versioning: MAJOR.MINOR.PATCH
|
# Semantic versioning: MAJOR.MINOR.PATCH
|
||||||
__version__ = "5.1.0"
|
__version__ = "5.1.2"
|
||||||
# Last update date in ISO format
|
# Last update date in ISO format
|
||||||
__updated__ = "2025-06-18"
|
__updated__ = "2025-06-18"
|
||||||
# Primary maintainer
|
# Primary maintainer
|
||||||
@@ -60,8 +60,8 @@ MODEL_CAPABILITIES_DESC = {
|
|||||||
"grok3": "GROK-3 (131K context) - Advanced reasoning model from X.AI, excellent for complex analysis",
|
"grok3": "GROK-3 (131K context) - Advanced reasoning model from X.AI, excellent for complex analysis",
|
||||||
"grokfast": "GROK-3 Fast (131K context) - Higher performance variant, faster processing but more expensive",
|
"grokfast": "GROK-3 Fast (131K context) - Higher performance variant, faster processing but more expensive",
|
||||||
# Full model names also supported (for explicit specification)
|
# Full model names also supported (for explicit specification)
|
||||||
"gemini-2.5-flash-preview-05-20": "Ultra-fast (1M context) - Quick analysis, simple queries, rapid iterations",
|
"gemini-2.5-flash": "Ultra-fast (1M context) - Quick analysis, simple queries, rapid iterations",
|
||||||
"gemini-2.5-pro-preview-06-05": (
|
"gemini-2.5-pro": (
|
||||||
"Deep reasoning + thinking mode (1M context) - Complex problems, architecture, deep analysis"
|
"Deep reasoning + thinking mode (1M context) - Complex problems, architecture, deep analysis"
|
||||||
),
|
),
|
||||||
}
|
}
|
||||||
@@ -69,8 +69,8 @@ MODEL_CAPABILITIES_DESC = {
|
|||||||
# OpenRouter/Custom API Fallback Behavior:
|
# OpenRouter/Custom API Fallback Behavior:
|
||||||
# When only OpenRouter or Custom API is configured (no native API keys), these
|
# When only OpenRouter or Custom API is configured (no native API keys), these
|
||||||
# model aliases automatically map to equivalent models through the proxy:
|
# model aliases automatically map to equivalent models through the proxy:
|
||||||
# - "flash" → "google/gemini-2.5-flash-preview-05-20" (via OpenRouter)
|
# - "flash" → "google/gemini-2.5-flash" (via OpenRouter)
|
||||||
# - "pro" → "google/gemini-2.5-pro-preview-06-05" (via OpenRouter)
|
# - "pro" → "google/gemini-2.5-pro" (via OpenRouter)
|
||||||
# - "o3" → "openai/o3" (via OpenRouter)
|
# - "o3" → "openai/o3" (via OpenRouter)
|
||||||
# - "o3mini" → "openai/o3-mini" (via OpenRouter)
|
# - "o3mini" → "openai/o3-mini" (via OpenRouter)
|
||||||
# - "o4-mini" → "openai/o4-mini" (via OpenRouter)
|
# - "o4-mini" → "openai/o4-mini" (via OpenRouter)
|
||||||
|
|||||||
@@ -124,8 +124,8 @@ OPENROUTER_ALLOWED_MODELS=opus,sonnet,mistral
|
|||||||
- `mini` (shorthand for o4-mini)
|
- `mini` (shorthand for o4-mini)
|
||||||
|
|
||||||
**Gemini Models:**
|
**Gemini Models:**
|
||||||
- `gemini-2.5-flash-preview-05-20` (1M context, fast)
|
- `gemini-2.5-flash` (1M context, fast)
|
||||||
- `gemini-2.5-pro-preview-06-05` (1M context, powerful)
|
- `gemini-2.5-pro` (1M context, powerful)
|
||||||
- `flash` (shorthand for Flash model)
|
- `flash` (shorthand for Flash model)
|
||||||
- `pro` (shorthand for Pro model)
|
- `pro` (shorthand for Pro model)
|
||||||
|
|
||||||
|
|||||||
@@ -46,8 +46,8 @@ The server uses `conf/custom_models.json` to map convenient aliases to both Open
|
|||||||
| `haiku` | `anthropic/claude-3-haiku` |
|
| `haiku` | `anthropic/claude-3-haiku` |
|
||||||
| `gpt4o`, `4o` | `openai/gpt-4o` |
|
| `gpt4o`, `4o` | `openai/gpt-4o` |
|
||||||
| `gpt4o-mini`, `4o-mini` | `openai/gpt-4o-mini` |
|
| `gpt4o-mini`, `4o-mini` | `openai/gpt-4o-mini` |
|
||||||
| `pro`, `gemini` | `google/gemini-2.5-pro-preview-06-05` |
|
| `pro`, `gemini` | `google/gemini-2.5-pro` |
|
||||||
| `flash` | `google/gemini-2.5-flash-preview-05-20` |
|
| `flash` | `google/gemini-2.5-flash` |
|
||||||
| `mistral` | `mistral/mistral-large` |
|
| `mistral` | `mistral/mistral-large` |
|
||||||
| `deepseek`, `coder` | `deepseek/deepseek-coder` |
|
| `deepseek`, `coder` | `deepseek/deepseek-coder` |
|
||||||
| `perplexity` | `perplexity/llama-3-sonar-large-32k-online` |
|
| `perplexity` | `perplexity/llama-3-sonar-large-32k-online` |
|
||||||
@@ -153,7 +153,7 @@ CUSTOM_MODEL_NAME=your-loaded-model
|
|||||||
# OpenRouter models:
|
# OpenRouter models:
|
||||||
"Use opus for deep analysis" # → anthropic/claude-3-opus
|
"Use opus for deep analysis" # → anthropic/claude-3-opus
|
||||||
"Use sonnet to review this code" # → anthropic/claude-3-sonnet
|
"Use sonnet to review this code" # → anthropic/claude-3-sonnet
|
||||||
"Use pro via zen to analyze this" # → google/gemini-2.5-pro-preview-06-05
|
"Use pro via zen to analyze this" # → google/gemini-2.5-pro
|
||||||
"Use gpt4o via zen to analyze this" # → openai/gpt-4o
|
"Use gpt4o via zen to analyze this" # → openai/gpt-4o
|
||||||
"Use mistral via zen to optimize" # → mistral/mistral-large
|
"Use mistral via zen to optimize" # → mistral/mistral-large
|
||||||
|
|
||||||
|
|||||||
@@ -46,7 +46,7 @@ The tool displays:
|
|||||||
📋 Available Models by Provider
|
📋 Available Models by Provider
|
||||||
|
|
||||||
🔹 Google (Gemini) - ✅ Configured
|
🔹 Google (Gemini) - ✅ Configured
|
||||||
• pro (gemini-2.5-pro-preview-06-05) - 1M context, thinking modes
|
• pro (gemini-2.5-pro) - 1M context, thinking modes
|
||||||
• flash (gemini-2.0-flash-experimental) - 1M context, ultra-fast
|
• flash (gemini-2.0-flash-experimental) - 1M context, ultra-fast
|
||||||
|
|
||||||
🔹 OpenAI - ✅ Configured
|
🔹 OpenAI - ✅ Configured
|
||||||
|
|||||||
@@ -19,14 +19,14 @@ class GeminiModelProvider(ModelProvider):
|
|||||||
|
|
||||||
# Model configurations
|
# Model configurations
|
||||||
SUPPORTED_MODELS = {
|
SUPPORTED_MODELS = {
|
||||||
"gemini-2.5-flash-preview-05-20": {
|
"gemini-2.5-flash": {
|
||||||
"context_window": 1_048_576, # 1M tokens
|
"context_window": 1_048_576, # 1M tokens
|
||||||
"supports_extended_thinking": True,
|
"supports_extended_thinking": True,
|
||||||
"max_thinking_tokens": 24576, # Flash 2.5 thinking budget limit
|
"max_thinking_tokens": 24576, # Flash 2.5 thinking budget limit
|
||||||
"supports_images": True, # Vision capability
|
"supports_images": True, # Vision capability
|
||||||
"max_image_size_mb": 20.0, # Conservative 20MB limit for reliability
|
"max_image_size_mb": 20.0, # Conservative 20MB limit for reliability
|
||||||
},
|
},
|
||||||
"gemini-2.5-pro-preview-06-05": {
|
"gemini-2.5-pro": {
|
||||||
"context_window": 1_048_576, # 1M tokens
|
"context_window": 1_048_576, # 1M tokens
|
||||||
"supports_extended_thinking": True,
|
"supports_extended_thinking": True,
|
||||||
"max_thinking_tokens": 32768, # Pro 2.5 thinking budget limit
|
"max_thinking_tokens": 32768, # Pro 2.5 thinking budget limit
|
||||||
@@ -34,8 +34,8 @@ class GeminiModelProvider(ModelProvider):
|
|||||||
"max_image_size_mb": 32.0, # Higher limit for Pro model
|
"max_image_size_mb": 32.0, # Higher limit for Pro model
|
||||||
},
|
},
|
||||||
# Shorthands
|
# Shorthands
|
||||||
"flash": "gemini-2.5-flash-preview-05-20",
|
"flash": "gemini-2.5-flash",
|
||||||
"pro": "gemini-2.5-pro-preview-06-05",
|
"pro": "gemini-2.5-pro",
|
||||||
}
|
}
|
||||||
|
|
||||||
# Thinking mode configurations - percentages of model's max_thinking_tokens
|
# Thinking mode configurations - percentages of model's max_thinking_tokens
|
||||||
@@ -364,8 +364,8 @@ class GeminiModelProvider(ModelProvider):
|
|||||||
"""Check if the model supports vision (image processing)."""
|
"""Check if the model supports vision (image processing)."""
|
||||||
# Gemini 2.5 models support vision
|
# Gemini 2.5 models support vision
|
||||||
vision_models = {
|
vision_models = {
|
||||||
"gemini-2.5-flash-preview-05-20",
|
"gemini-2.5-flash",
|
||||||
"gemini-2.5-pro-preview-06-05",
|
"gemini-2.5-pro",
|
||||||
"gemini-2.0-flash",
|
"gemini-2.0-flash",
|
||||||
"gemini-1.5-pro",
|
"gemini-1.5-pro",
|
||||||
"gemini-1.5-flash",
|
"gemini-1.5-flash",
|
||||||
|
|||||||
@@ -105,7 +105,7 @@ class ModelProviderRegistry:
|
|||||||
3. OPENROUTER - Catch-all for cloud models via unified API
|
3. OPENROUTER - Catch-all for cloud models via unified API
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
model_name: Name of the model (e.g., "gemini-2.5-flash-preview-05-20", "o3-mini")
|
model_name: Name of the model (e.g., "gemini-2.5-flash", "o3-mini")
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
ModelProvider instance that supports this model
|
ModelProvider instance that supports this model
|
||||||
@@ -295,7 +295,7 @@ class ModelProviderRegistry:
|
|||||||
return custom_models[0]
|
return custom_models[0]
|
||||||
else:
|
else:
|
||||||
# Fallback to pro if nothing found
|
# Fallback to pro if nothing found
|
||||||
return "gemini-2.5-pro-preview-06-05"
|
return "gemini-2.5-pro"
|
||||||
|
|
||||||
elif tool_category == ToolModelCategory.FAST_RESPONSE:
|
elif tool_category == ToolModelCategory.FAST_RESPONSE:
|
||||||
# Prefer fast, cost-efficient models
|
# Prefer fast, cost-efficient models
|
||||||
@@ -325,7 +325,7 @@ class ModelProviderRegistry:
|
|||||||
return custom_models[0]
|
return custom_models[0]
|
||||||
else:
|
else:
|
||||||
# Default to flash
|
# Default to flash
|
||||||
return "gemini-2.5-flash-preview-05-20"
|
return "gemini-2.5-flash"
|
||||||
|
|
||||||
# BALANCED or no category specified - use existing balanced logic
|
# BALANCED or no category specified - use existing balanced logic
|
||||||
if openai_available and "o4-mini" in openai_models:
|
if openai_available and "o4-mini" in openai_models:
|
||||||
@@ -353,7 +353,7 @@ class ModelProviderRegistry:
|
|||||||
# This might happen if all models are restricted
|
# This might happen if all models are restricted
|
||||||
logging.warning("No models available due to restrictions")
|
logging.warning("No models available due to restrictions")
|
||||||
# Return a reasonable default for backward compatibility
|
# Return a reasonable default for backward compatibility
|
||||||
return "gemini-2.5-flash-preview-05-20"
|
return "gemini-2.5-flash"
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _find_extended_thinking_model(cls) -> Optional[str]:
|
def _find_extended_thinking_model(cls) -> Optional[str]:
|
||||||
@@ -383,7 +383,7 @@ class ModelProviderRegistry:
|
|||||||
preferred_models = [
|
preferred_models = [
|
||||||
"anthropic/claude-3.5-sonnet",
|
"anthropic/claude-3.5-sonnet",
|
||||||
"anthropic/claude-3-opus-20240229",
|
"anthropic/claude-3-opus-20240229",
|
||||||
"google/gemini-2.5-pro-preview-06-05",
|
"google/gemini-2.5-pro",
|
||||||
"google/gemini-pro-1.5",
|
"google/gemini-pro-1.5",
|
||||||
"meta-llama/llama-3.1-70b-instruct",
|
"meta-llama/llama-3.1-70b-instruct",
|
||||||
"mistralai/mixtral-8x7b-instruct",
|
"mistralai/mixtral-8x7b-instruct",
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ class TestModelThinkingConfig(BaseSimulatorTest):
|
|||||||
"chat",
|
"chat",
|
||||||
{
|
{
|
||||||
"prompt": "What is 2 + 2? Please think carefully and explain.",
|
"prompt": "What is 2 + 2? Please think carefully and explain.",
|
||||||
"model": "pro", # Should resolve to gemini-2.5-pro-preview-06-05
|
"model": "pro", # Should resolve to gemini-2.5-pro
|
||||||
"thinking_mode": "high", # Should use thinking_config
|
"thinking_mode": "high", # Should use thinking_config
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
@@ -55,7 +55,7 @@ class TestModelThinkingConfig(BaseSimulatorTest):
|
|||||||
"chat",
|
"chat",
|
||||||
{
|
{
|
||||||
"prompt": "What is 3 + 3? Give a quick answer.",
|
"prompt": "What is 3 + 3? Give a quick answer.",
|
||||||
"model": "flash", # Should resolve to gemini-2.5-flash-preview-05-20
|
"model": "flash", # Should resolve to gemini-2.5-flash
|
||||||
"thinking_mode": "high", # Should be ignored for Flash model
|
"thinking_mode": "high", # Should be ignored for Flash model
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
@@ -79,8 +79,8 @@ class TestModelThinkingConfig(BaseSimulatorTest):
|
|||||||
test_cases = [
|
test_cases = [
|
||||||
("pro", "should work with Pro model"),
|
("pro", "should work with Pro model"),
|
||||||
("flash", "should work with Flash model"),
|
("flash", "should work with Flash model"),
|
||||||
("gemini-2.5-pro-preview-06-05", "should work with full Pro model name"),
|
("gemini-2.5-pro", "should work with full Pro model name"),
|
||||||
("gemini-2.5-flash-preview-05-20", "should work with full Flash model name"),
|
("gemini-2.5-flash", "should work with full Flash model name"),
|
||||||
]
|
]
|
||||||
|
|
||||||
success_count = 0
|
success_count = 0
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ class OpenRouterModelsTest(BaseSimulatorTest):
|
|||||||
self.setup_test_files()
|
self.setup_test_files()
|
||||||
|
|
||||||
# Test 1: Flash alias mapping to OpenRouter
|
# Test 1: Flash alias mapping to OpenRouter
|
||||||
self.logger.info(" 1: Testing 'flash' alias (should map to google/gemini-2.5-flash-preview-05-20)")
|
self.logger.info(" 1: Testing 'flash' alias (should map to google/gemini-2.5-flash)")
|
||||||
|
|
||||||
response1, continuation_id = self.call_mcp_tool(
|
response1, continuation_id = self.call_mcp_tool(
|
||||||
"chat",
|
"chat",
|
||||||
@@ -63,7 +63,7 @@ class OpenRouterModelsTest(BaseSimulatorTest):
|
|||||||
self.logger.info(f" ✅ Got continuation_id: {continuation_id}")
|
self.logger.info(f" ✅ Got continuation_id: {continuation_id}")
|
||||||
|
|
||||||
# Test 2: Pro alias mapping to OpenRouter
|
# Test 2: Pro alias mapping to OpenRouter
|
||||||
self.logger.info(" 2: Testing 'pro' alias (should map to google/gemini-2.5-pro-preview-06-05)")
|
self.logger.info(" 2: Testing 'pro' alias (should map to google/gemini-2.5-pro)")
|
||||||
|
|
||||||
response2, _ = self.call_mcp_tool(
|
response2, _ = self.call_mcp_tool(
|
||||||
"chat",
|
"chat",
|
||||||
|
|||||||
@@ -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
|
# Set default model to a specific value for tests to avoid auto mode
|
||||||
# This prevents all tests from failing due to missing model parameter
|
# This prevents all tests from failing due to missing model parameter
|
||||||
os.environ["DEFAULT_MODEL"] = "gemini-2.5-flash-preview-05-20"
|
os.environ["DEFAULT_MODEL"] = "gemini-2.5-flash"
|
||||||
|
|
||||||
# Force reload of config module to pick up the env var
|
# Force reload of config module to pick up the env var
|
||||||
import config # noqa: E402
|
import config # noqa: E402
|
||||||
@@ -108,7 +108,7 @@ def mock_provider_availability(request, monkeypatch):
|
|||||||
if model_name in ["unavailable-model", "gpt-5-turbo", "o3"]:
|
if model_name in ["unavailable-model", "gpt-5-turbo", "o3"]:
|
||||||
return None
|
return None
|
||||||
# For common test models, return a mock provider
|
# 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
|
# Try to use the real provider first if it exists
|
||||||
real_provider = original_get_provider(model_name)
|
real_provider = original_get_provider(model_name)
|
||||||
if real_provider:
|
if real_provider:
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ from unittest.mock import Mock
|
|||||||
from providers.base import ModelCapabilities, ProviderType, RangeTemperatureConstraint
|
from providers.base import ModelCapabilities, ProviderType, RangeTemperatureConstraint
|
||||||
|
|
||||||
|
|
||||||
def create_mock_provider(model_name="gemini-2.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."""
|
"""Create a properly configured mock provider."""
|
||||||
mock_provider = Mock()
|
mock_provider = Mock()
|
||||||
|
|
||||||
|
|||||||
@@ -39,9 +39,9 @@ class TestAliasTargetRestrictions:
|
|||||||
|
|
||||||
# Should include both aliases and their targets
|
# Should include both aliases and their targets
|
||||||
assert "flash" in all_known # alias
|
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 "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
|
@patch.dict(os.environ, {"OPENAI_ALLOWED_MODELS": "o4-mini"}) # Allow target
|
||||||
def test_restriction_policy_allows_alias_when_target_allowed(self):
|
def test_restriction_policy_allows_alias_when_target_allowed(self):
|
||||||
@@ -80,7 +80,7 @@ class TestAliasTargetRestrictions:
|
|||||||
# Direct target should NOT be allowed
|
# Direct target should NOT be allowed
|
||||||
assert not provider.validate_model_name("o4-mini")
|
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):
|
def test_gemini_restriction_policy_allows_alias_when_target_allowed(self):
|
||||||
"""Test Gemini restriction policy allows alias when target is allowed."""
|
"""Test Gemini restriction policy allows alias when target is allowed."""
|
||||||
# Clear cached restriction service
|
# Clear cached restriction service
|
||||||
@@ -91,7 +91,7 @@ class TestAliasTargetRestrictions:
|
|||||||
provider = GeminiModelProvider(api_key="test-key")
|
provider = GeminiModelProvider(api_key="test-key")
|
||||||
|
|
||||||
# Both target and alias should be allowed
|
# 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")
|
assert provider.validate_model_name("flash")
|
||||||
|
|
||||||
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "flash"}) # Allow alias only
|
@patch.dict(os.environ, {"GOOGLE_ALLOWED_MODELS": "flash"}) # Allow alias only
|
||||||
@@ -107,7 +107,7 @@ class TestAliasTargetRestrictions:
|
|||||||
# Only the alias should be allowed
|
# Only the alias should be allowed
|
||||||
assert provider.validate_model_name("flash")
|
assert provider.validate_model_name("flash")
|
||||||
# Direct target should NOT be allowed
|
# 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):
|
def test_restriction_service_validation_includes_all_targets(self):
|
||||||
"""Test that restriction service validation knows about all aliases and targets."""
|
"""Test that restriction service validation knows about all aliases and targets."""
|
||||||
|
|||||||
@@ -80,9 +80,9 @@ class TestAutoModeComprehensive:
|
|||||||
"OPENROUTER_API_KEY": None,
|
"OPENROUTER_API_KEY": None,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"EXTENDED_REASONING": "gemini-2.5-pro-preview-06-05", # Pro for deep thinking
|
"EXTENDED_REASONING": "gemini-2.5-pro", # Pro for deep thinking
|
||||||
"FAST_RESPONSE": "gemini-2.5-flash-preview-05-20", # Flash for speed
|
"FAST_RESPONSE": "gemini-2.5-flash", # Flash for speed
|
||||||
"BALANCED": "gemini-2.5-flash-preview-05-20", # Flash as balanced
|
"BALANCED": "gemini-2.5-flash", # Flash as balanced
|
||||||
},
|
},
|
||||||
),
|
),
|
||||||
# Only OpenAI API available
|
# Only OpenAI API available
|
||||||
@@ -313,8 +313,8 @@ class TestAutoModeComprehensive:
|
|||||||
# Should include Gemini models
|
# Should include Gemini models
|
||||||
assert "flash" in available_models
|
assert "flash" in available_models
|
||||||
assert "pro" in available_models
|
assert "pro" in available_models
|
||||||
assert "gemini-2.5-flash-preview-05-20" in available_models
|
assert "gemini-2.5-flash" in available_models
|
||||||
assert "gemini-2.5-pro-preview-06-05" in available_models
|
assert "gemini-2.5-pro" in available_models
|
||||||
|
|
||||||
# Should also include other models (users might have OpenRouter configured)
|
# Should also include other models (users might have OpenRouter configured)
|
||||||
# The schema should show all options; validation happens at runtime
|
# The schema should show all options; validation happens at runtime
|
||||||
@@ -476,8 +476,8 @@ class TestAutoModeComprehensive:
|
|||||||
assert "o3-mini" not in available_models
|
assert "o3-mini" not in available_models
|
||||||
|
|
||||||
# Should still include all Gemini models (no restrictions)
|
# Should still include all Gemini models (no restrictions)
|
||||||
assert "gemini-2.5-flash-preview-05-20" in available_models
|
assert "gemini-2.5-flash" in available_models
|
||||||
assert "gemini-2.5-pro-preview-06-05" in available_models
|
assert "gemini-2.5-pro" in available_models
|
||||||
|
|
||||||
def test_openrouter_fallback_when_no_native_apis(self):
|
def test_openrouter_fallback_when_no_native_apis(self):
|
||||||
"""Test that OpenRouter provides fallback models when no native APIs are available."""
|
"""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 OpenRouter registry to return known models
|
||||||
mock_registry = MagicMock()
|
mock_registry = MagicMock()
|
||||||
mock_registry.list_models.return_value = [
|
mock_registry.list_models.return_value = [
|
||||||
"google/gemini-2.5-flash-preview-05-20",
|
"google/gemini-2.5-flash",
|
||||||
"google/gemini-2.5-pro-preview-06-05",
|
"google/gemini-2.5-pro",
|
||||||
"openai/o3",
|
"openai/o3",
|
||||||
"openai/o4-mini",
|
"openai/o4-mini",
|
||||||
"anthropic/claude-3-opus",
|
"anthropic/claude-3-opus",
|
||||||
@@ -564,11 +564,11 @@ class TestAutoModeComprehensive:
|
|||||||
mock_provider = MagicMock()
|
mock_provider = MagicMock()
|
||||||
mock_response = MagicMock()
|
mock_response = MagicMock()
|
||||||
mock_response.content = "test response"
|
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_response.usage = {"input_tokens": 10, "output_tokens": 5}
|
||||||
# Mock _resolve_model_name to simulate alias resolution
|
# Mock _resolve_model_name to simulate alias resolution
|
||||||
mock_provider._resolve_model_name = lambda alias: (
|
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
|
mock_provider.generate_content.return_value = mock_response
|
||||||
|
|
||||||
|
|||||||
@@ -201,7 +201,7 @@ class TestAutoModeCustomProviderOnly:
|
|||||||
|
|
||||||
# Should get a valid model name, not the hardcoded fallback
|
# Should get a valid model name, not the hardcoded fallback
|
||||||
assert (
|
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"
|
), "Should not fallback to hardcoded Gemini model when custom provider is available"
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -59,9 +59,9 @@ class TestAutoModeProviderSelection:
|
|||||||
balanced = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.BALANCED)
|
balanced = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.BALANCED)
|
||||||
|
|
||||||
# Should select appropriate Gemini models
|
# Should select appropriate Gemini models
|
||||||
assert extended_reasoning in ["gemini-2.5-pro-preview-06-05", "pro"]
|
assert extended_reasoning in ["gemini-2.5-pro", "pro"]
|
||||||
assert fast_response in ["gemini-2.5-flash-preview-05-20", "flash"]
|
assert fast_response in ["gemini-2.5-flash", "flash"]
|
||||||
assert balanced in ["gemini-2.5-flash-preview-05-20", "flash"]
|
assert balanced in ["gemini-2.5-flash", "flash"]
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
# Restore original environment
|
# Restore original environment
|
||||||
@@ -229,8 +229,8 @@ class TestAutoModeProviderSelection:
|
|||||||
assert "o3-mini" not in available_models
|
assert "o3-mini" not in available_models
|
||||||
|
|
||||||
# Should include all Gemini models (no restrictions)
|
# Should include all Gemini models (no restrictions)
|
||||||
assert "gemini-2.5-flash-preview-05-20" in available_models
|
assert "gemini-2.5-flash" in available_models
|
||||||
assert available_models["gemini-2.5-flash-preview-05-20"] == ProviderType.GOOGLE
|
assert available_models["gemini-2.5-flash"] == ProviderType.GOOGLE
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
# Restore original environment
|
# Restore original environment
|
||||||
@@ -316,8 +316,8 @@ class TestAutoModeProviderSelection:
|
|||||||
|
|
||||||
# Test that providers resolve aliases correctly
|
# Test that providers resolve aliases correctly
|
||||||
test_cases = [
|
test_cases = [
|
||||||
("flash", ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20"),
|
("flash", ProviderType.GOOGLE, "gemini-2.5-flash"),
|
||||||
("pro", ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05"),
|
("pro", ProviderType.GOOGLE, "gemini-2.5-pro"),
|
||||||
("mini", ProviderType.OPENAI, "o4-mini"),
|
("mini", ProviderType.OPENAI, "o4-mini"),
|
||||||
("o3mini", ProviderType.OPENAI, "o3-mini"),
|
("o3mini", ProviderType.OPENAI, "o3-mini"),
|
||||||
("grok", ProviderType.XAI, "grok-3"),
|
("grok", ProviderType.XAI, "grok-3"),
|
||||||
|
|||||||
@@ -85,16 +85,16 @@ class TestBuggyBehaviorPrevention:
|
|||||||
|
|
||||||
# Verify both aliases and targets are included
|
# Verify both aliases and targets are included
|
||||||
assert "flash" in all_known # alias
|
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 "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
|
# Simulate admin restricting to target model names
|
||||||
service = ModelRestrictionService()
|
service = ModelRestrictionService()
|
||||||
service.restrictions = {
|
service.restrictions = {
|
||||||
ProviderType.GOOGLE: {
|
ProviderType.GOOGLE: {
|
||||||
"gemini-2.5-flash-preview-05-20", # Target name restriction
|
"gemini-2.5-flash", # Target name restriction
|
||||||
"gemini-2.5-pro-preview-06-05", # Target name restriction
|
"gemini-2.5-pro", # Target name restriction
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -105,8 +105,8 @@ class TestBuggyBehaviorPrevention:
|
|||||||
# Should NOT warn about these valid target models
|
# Should NOT warn about these valid target models
|
||||||
all_warnings = [str(call) for call in mock_logger.warning.call_args_list]
|
all_warnings = [str(call) for call in mock_logger.warning.call_args_list]
|
||||||
for warning in all_warnings:
|
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-flash" 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-pro" not in warning or "not a recognized" not in warning
|
||||||
|
|
||||||
def test_old_bug_policy_bypass_prevention(self):
|
def test_old_bug_policy_bypass_prevention(self):
|
||||||
"""
|
"""
|
||||||
@@ -206,7 +206,7 @@ class TestBuggyBehaviorPrevention:
|
|||||||
"""
|
"""
|
||||||
providers_to_test = [
|
providers_to_test = [
|
||||||
(OpenAIModelProvider(api_key="test-key"), "mini", "o4-mini"),
|
(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:
|
for provider, alias, target in providers_to_test:
|
||||||
|
|||||||
@@ -59,7 +59,7 @@ class TestClaudeContinuationOffers:
|
|||||||
# If providers are not registered yet, tool might detect auto mode
|
# If providers are not registered yet, tool might detect auto mode
|
||||||
self.tool = ClaudeContinuationTool()
|
self.tool = ClaudeContinuationTool()
|
||||||
# Set default model to avoid effective auto mode
|
# 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("utils.conversation_memory.get_storage")
|
||||||
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
|
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
|
||||||
@@ -67,7 +67,7 @@ class TestClaudeContinuationOffers:
|
|||||||
"""Test that new conversations offer Claude continuation opportunity"""
|
"""Test that new conversations offer Claude continuation opportunity"""
|
||||||
# Create tool AFTER providers are registered (in conftest.py fixture)
|
# Create tool AFTER providers are registered (in conftest.py fixture)
|
||||||
tool = ClaudeContinuationTool()
|
tool = ClaudeContinuationTool()
|
||||||
tool.default_model = "gemini-2.5-flash-preview-05-20"
|
tool.default_model = "gemini-2.5-flash"
|
||||||
|
|
||||||
mock_client = Mock()
|
mock_client = Mock()
|
||||||
mock_storage.return_value = mock_client
|
mock_storage.return_value = mock_client
|
||||||
@@ -80,7 +80,7 @@ class TestClaudeContinuationOffers:
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Analysis complete.",
|
content="Analysis complete.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -137,7 +137,7 @@ class TestClaudeContinuationOffers:
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Continued analysis.",
|
content="Continued analysis.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -170,7 +170,7 @@ class TestClaudeContinuationOffers:
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Analysis complete. The code looks good.",
|
content="Analysis complete. The code looks good.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -212,7 +212,7 @@ I'd be happy to examine the error handling patterns in more detail if that would
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=content_with_followup,
|
content=content_with_followup,
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Continued analysis complete.",
|
content="Continued analysis complete.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Final response.",
|
content="Final response.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -336,7 +336,7 @@ class TestContinuationIntegration:
|
|||||||
def setup_method(self):
|
def setup_method(self):
|
||||||
self.tool = ClaudeContinuationTool()
|
self.tool = ClaudeContinuationTool()
|
||||||
# Set default model to avoid effective auto mode
|
# 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("utils.conversation_memory.get_storage")
|
||||||
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
|
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
|
||||||
@@ -364,7 +364,7 @@ class TestContinuationIntegration:
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Analysis result",
|
content="Analysis result",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -417,7 +417,7 @@ class TestContinuationIntegration:
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Structure analysis done.",
|
content="Structure analysis done.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -454,7 +454,7 @@ class TestContinuationIntegration:
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Performance analysis done.",
|
content="Performance analysis done.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -41,7 +41,7 @@ class TestDynamicContextRequests:
|
|||||||
mock_provider.get_provider_type.return_value = Mock(value="google")
|
mock_provider.get_provider_type.return_value = Mock(value="google")
|
||||||
mock_provider.supports_thinking_mode.return_value = False
|
mock_provider.supports_thinking_mode.return_value = False
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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.get_provider_type.return_value = Mock(value="google")
|
||||||
mock_provider.supports_thinking_mode.return_value = False
|
mock_provider.supports_thinking_mode.return_value = False
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=normal_response, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=normal_response, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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.get_provider_type.return_value = Mock(value="google")
|
||||||
mock_provider.supports_thinking_mode.return_value = False
|
mock_provider.supports_thinking_mode.return_value = False
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=malformed_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=malformed_json, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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.get_provider_type.return_value = Mock(value="google")
|
||||||
mock_provider.supports_thinking_mode.return_value = False
|
mock_provider.supports_thinking_mode.return_value = False
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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.get_provider_type.return_value = Mock(value="google")
|
||||||
mock_provider.supports_thinking_mode.return_value = False
|
mock_provider.supports_thinking_mode.return_value = False
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
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.get_provider_type.return_value = Mock(value="google")
|
||||||
mock_provider.supports_thinking_mode.return_value = False
|
mock_provider.supports_thinking_mode.return_value = False
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=clarification_json, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=clarification_json, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
|
|
||||||
@@ -340,7 +340,7 @@ class TestCollaborationWorkflow:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=final_response, usage={}, model_name="gemini-2.5-flash-preview-05-20", metadata={}
|
content=final_response, usage={}, model_name="gemini-2.5-flash", metadata={}
|
||||||
)
|
)
|
||||||
|
|
||||||
result2 = await tool.execute(
|
result2 = await tool.execute(
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ class TestConfig:
|
|||||||
def test_model_config(self):
|
def test_model_config(self):
|
||||||
"""Test model configuration"""
|
"""Test model configuration"""
|
||||||
# DEFAULT_MODEL is set in conftest.py for tests
|
# DEFAULT_MODEL is set in conftest.py for tests
|
||||||
assert DEFAULT_MODEL == "gemini-2.5-flash-preview-05-20"
|
assert DEFAULT_MODEL == "gemini-2.5-flash"
|
||||||
|
|
||||||
def test_temperature_defaults(self):
|
def test_temperature_defaults(self):
|
||||||
"""Test temperature constants"""
|
"""Test temperature constants"""
|
||||||
|
|||||||
@@ -74,7 +74,7 @@ async def test_conversation_history_field_mapping():
|
|||||||
mock_provider = MagicMock()
|
mock_provider = MagicMock()
|
||||||
mock_provider.get_capabilities.return_value = ModelCapabilities(
|
mock_provider.get_capabilities.return_value = ModelCapabilities(
|
||||||
provider=ProviderType.GOOGLE,
|
provider=ProviderType.GOOGLE,
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
friendly_name="Gemini",
|
friendly_name="Gemini",
|
||||||
context_window=200000,
|
context_window=200000,
|
||||||
supports_extended_thinking=True,
|
supports_extended_thinking=True,
|
||||||
@@ -132,7 +132,7 @@ async def test_unknown_tool_defaults_to_prompt():
|
|||||||
# Mock ModelContext to avoid calculation errors
|
# Mock ModelContext to avoid calculation errors
|
||||||
with patch("utils.model_context.ModelContext") as mock_model_context_class:
|
with patch("utils.model_context.ModelContext") as mock_model_context_class:
|
||||||
mock_model_context = MagicMock()
|
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(
|
mock_model_context.calculate_token_allocation.return_value = MagicMock(
|
||||||
total_tokens=200000,
|
total_tokens=200000,
|
||||||
content_tokens=120000,
|
content_tokens=120000,
|
||||||
|
|||||||
@@ -95,7 +95,7 @@ class TestConversationHistoryBugFix:
|
|||||||
return Mock(
|
return Mock(
|
||||||
content="Response with conversation context",
|
content="Response with conversation context",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -155,7 +155,7 @@ class TestConversationHistoryBugFix:
|
|||||||
return Mock(
|
return Mock(
|
||||||
content="Response without history",
|
content="Response without history",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -193,7 +193,7 @@ class TestConversationHistoryBugFix:
|
|||||||
return Mock(
|
return Mock(
|
||||||
content="New conversation response",
|
content="New conversation response",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -277,7 +277,7 @@ class TestConversationHistoryBugFix:
|
|||||||
return Mock(
|
return Mock(
|
||||||
content="Analysis of new files complete",
|
content="Analysis of new files complete",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -112,7 +112,7 @@ I'd be happy to review these security findings in detail if that would be helpfu
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content=content,
|
content=content,
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -159,7 +159,7 @@ I'd be happy to review these security findings in detail if that would be helpfu
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Critical security vulnerability confirmed. The authentication function always returns true, bypassing all security checks.",
|
content="Critical security vulnerability confirmed. The authentication function always returns true, bypassing all security checks.",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -284,7 +284,7 @@ I'd be happy to review these security findings in detail if that would be helpfu
|
|||||||
mock_provider.generate_content.return_value = Mock(
|
mock_provider.generate_content.return_value = Mock(
|
||||||
content="Security review of auth.py shows vulnerabilities",
|
content="Security review of auth.py shows vulnerabilities",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
|
|||||||
@@ -48,14 +48,14 @@ class TestIntelligentFallback:
|
|||||||
|
|
||||||
@patch.dict(os.environ, {"OPENAI_API_KEY": "", "GEMINI_API_KEY": "test-gemini-key"}, clear=False)
|
@patch.dict(os.environ, {"OPENAI_API_KEY": "", "GEMINI_API_KEY": "test-gemini-key"}, clear=False)
|
||||||
def test_prefers_gemini_flash_when_openai_unavailable(self):
|
def test_prefers_gemini_flash_when_openai_unavailable(self):
|
||||||
"""Test that gemini-2.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
|
# Register only Gemini provider for this test
|
||||||
from providers.gemini import GeminiModelProvider
|
from providers.gemini import GeminiModelProvider
|
||||||
|
|
||||||
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
|
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
|
||||||
|
|
||||||
fallback_model = ModelProviderRegistry.get_preferred_fallback_model()
|
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)
|
@patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-key", "GEMINI_API_KEY": "test-gemini-key"}, clear=False)
|
||||||
def test_prefers_openai_when_both_available(self):
|
def test_prefers_openai_when_both_available(self):
|
||||||
@@ -81,7 +81,7 @@ class TestIntelligentFallback:
|
|||||||
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
|
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
|
||||||
|
|
||||||
fallback_model = ModelProviderRegistry.get_preferred_fallback_model()
|
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):
|
def test_available_providers_with_keys(self):
|
||||||
"""Test the get_available_providers_with_keys method"""
|
"""Test the get_available_providers_with_keys method"""
|
||||||
@@ -186,14 +186,14 @@ class TestIntelligentFallback:
|
|||||||
|
|
||||||
history, tokens = build_conversation_history(context, model_context=None)
|
history, tokens = build_conversation_history(context, model_context=None)
|
||||||
|
|
||||||
# Should use gemini-2.5-flash-preview-05-20 when only Gemini is available
|
# Should use gemini-2.5-flash when only Gemini is available
|
||||||
mock_context_class.assert_called_once_with("gemini-2.5-flash-preview-05-20")
|
mock_context_class.assert_called_once_with("gemini-2.5-flash")
|
||||||
|
|
||||||
def test_non_auto_mode_unchanged(self):
|
def test_non_auto_mode_unchanged(self):
|
||||||
"""Test that non-auto mode behavior is unchanged"""
|
"""Test that non-auto mode behavior is unchanged"""
|
||||||
from utils.conversation_memory import ThreadContext, build_conversation_history
|
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
|
from utils.conversation_memory import ConversationTurn
|
||||||
|
|
||||||
context = ThreadContext(
|
context = ThreadContext(
|
||||||
@@ -219,7 +219,7 @@ class TestIntelligentFallback:
|
|||||||
history, tokens = build_conversation_history(context, model_context=None)
|
history, tokens = build_conversation_history(context, model_context=None)
|
||||||
|
|
||||||
# Should use the configured DEFAULT_MODEL, not the intelligent fallback
|
# 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__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
@@ -76,7 +76,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="This is a test response",
|
content="This is a test response",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -103,7 +103,7 @@ class TestLargePromptHandling:
|
|||||||
patch("utils.model_context.ModelContext") as mock_model_context_class,
|
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_provider.generate_content.return_value.content = "Processed prompt from file"
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
|
|
||||||
@@ -111,7 +111,7 @@ class TestLargePromptHandling:
|
|||||||
from utils.model_context import TokenAllocation
|
from utils.model_context import TokenAllocation
|
||||||
|
|
||||||
mock_model_context = MagicMock()
|
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(
|
mock_model_context.calculate_token_allocation.return_value = TokenAllocation(
|
||||||
total_tokens=1_048_576,
|
total_tokens=1_048_576,
|
||||||
content_tokens=838_861,
|
content_tokens=838_861,
|
||||||
@@ -293,7 +293,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="Success",
|
content="Success",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -336,7 +336,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="Response to the large prompt",
|
content="Response to the large prompt",
|
||||||
usage={"input_tokens": 12000, "output_tokens": 10, "total_tokens": 12010},
|
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"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -368,7 +368,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="Success",
|
content="Success",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -390,7 +390,7 @@ class TestLargePromptHandling:
|
|||||||
patch("utils.model_context.ModelContext") as mock_model_context_class,
|
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_provider.generate_content.return_value.content = "Success"
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
|
|
||||||
@@ -398,7 +398,7 @@ class TestLargePromptHandling:
|
|||||||
from utils.model_context import TokenAllocation
|
from utils.model_context import TokenAllocation
|
||||||
|
|
||||||
mock_model_context = MagicMock()
|
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(
|
mock_model_context.calculate_token_allocation.return_value = TokenAllocation(
|
||||||
total_tokens=1_048_576,
|
total_tokens=1_048_576,
|
||||||
content_tokens=838_861,
|
content_tokens=838_861,
|
||||||
@@ -437,7 +437,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="Weather is sunny",
|
content="Weather is sunny",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
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"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -502,7 +502,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="Hi there!",
|
content="Hi there!",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
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"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
@@ -555,7 +555,7 @@ class TestLargePromptHandling:
|
|||||||
mock_provider.generate_content.return_value = MagicMock(
|
mock_provider.generate_content.return_value = MagicMock(
|
||||||
content="Continuing our conversation...",
|
content="Continuing our conversation...",
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
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"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
mock_get_provider.return_value = mock_provider
|
mock_get_provider.return_value = mock_provider
|
||||||
|
|||||||
@@ -65,8 +65,8 @@ class TestListModelsTool:
|
|||||||
|
|
||||||
# Check Gemini shows as configured
|
# Check Gemini shows as configured
|
||||||
assert "Google Gemini ✅" in content
|
assert "Google Gemini ✅" in content
|
||||||
assert "`flash` → `gemini-2.5-flash-preview-05-20`" in content
|
assert "`flash` → `gemini-2.5-flash`" in content
|
||||||
assert "`pro` → `gemini-2.5-pro-preview-06-05`" in content
|
assert "`pro` → `gemini-2.5-pro`" in content
|
||||||
assert "1M context" in content
|
assert "1M context" in content
|
||||||
|
|
||||||
# Check summary
|
# Check summary
|
||||||
|
|||||||
@@ -96,8 +96,8 @@ class TestModelEnumeration:
|
|||||||
"grok-3-fast",
|
"grok-3-fast",
|
||||||
"grok3",
|
"grok3",
|
||||||
"grokfast", # X.AI models
|
"grokfast", # X.AI models
|
||||||
"gemini-2.5-flash-preview-05-20",
|
"gemini-2.5-flash",
|
||||||
"gemini-2.5-pro-preview-06-05", # Full Gemini names
|
"gemini-2.5-pro", # Full Gemini names
|
||||||
]
|
]
|
||||||
|
|
||||||
for model in native_models:
|
for model in native_models:
|
||||||
@@ -264,7 +264,7 @@ class TestModelEnumeration:
|
|||||||
("flash", True), # Native Gemini
|
("flash", True), # Native Gemini
|
||||||
("o3", True), # Native OpenAI
|
("o3", True), # Native OpenAI
|
||||||
("grok", True), # Native X.AI
|
("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
|
("o4-mini-high", True), # Native OpenAI variant
|
||||||
("grok-3-fast", True), # Native X.AI variant
|
("grok-3-fast", True), # Native X.AI variant
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -22,8 +22,8 @@ class TestModelRestrictionService:
|
|||||||
# Should allow all models
|
# Should allow all models
|
||||||
assert service.is_allowed(ProviderType.OPENAI, "o3")
|
assert service.is_allowed(ProviderType.OPENAI, "o3")
|
||||||
assert service.is_allowed(ProviderType.OPENAI, "o3-mini")
|
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-pro")
|
||||||
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20")
|
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, "anthropic/claude-3-opus")
|
||||||
assert service.is_allowed(ProviderType.OPENROUTER, "openai/o3")
|
assert service.is_allowed(ProviderType.OPENROUTER, "openai/o3")
|
||||||
|
|
||||||
@@ -43,7 +43,7 @@ class TestModelRestrictionService:
|
|||||||
assert not service.is_allowed(ProviderType.OPENAI, "o4-mini")
|
assert not service.is_allowed(ProviderType.OPENAI, "o4-mini")
|
||||||
|
|
||||||
# Google and OpenRouter should have no restrictions
|
# 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")
|
assert service.is_allowed(ProviderType.OPENROUTER, "anthropic/claude-3-opus")
|
||||||
|
|
||||||
def test_load_multiple_models_restriction(self):
|
def test_load_multiple_models_restriction(self):
|
||||||
@@ -59,7 +59,7 @@ class TestModelRestrictionService:
|
|||||||
# Check Google models
|
# Check Google models
|
||||||
assert service.is_allowed(ProviderType.GOOGLE, "flash")
|
assert service.is_allowed(ProviderType.GOOGLE, "flash")
|
||||||
assert service.is_allowed(ProviderType.GOOGLE, "pro")
|
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):
|
def test_case_insensitive_and_whitespace_handling(self):
|
||||||
"""Test that model names are case-insensitive and whitespace is trimmed."""
|
"""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)
|
# Google should only allow flash (and its resolved name)
|
||||||
assert service.is_allowed(ProviderType.GOOGLE, "flash")
|
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, "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):
|
def test_filter_models(self):
|
||||||
"""Test filtering a list of models based on restrictions."""
|
"""Test filtering a list of models based on restrictions."""
|
||||||
@@ -124,8 +124,8 @@ class TestModelRestrictionService:
|
|||||||
assert not service.is_allowed(ProviderType.OPENAI, "o3")
|
assert not service.is_allowed(ProviderType.OPENAI, "o3")
|
||||||
|
|
||||||
# Google should allow both models via shorthands
|
# 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-flash", "flash")
|
||||||
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05", "pro")
|
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro", "pro")
|
||||||
|
|
||||||
# Also test that full names work when specified in restrictions
|
# Also test that full names work when specified in restrictions
|
||||||
assert service.is_allowed(ProviderType.OPENAI, "o3-mini", "o3mini") # Even with shorthand
|
assert service.is_allowed(ProviderType.OPENAI, "o3-mini", "o3mini") # Even with shorthand
|
||||||
@@ -238,7 +238,7 @@ class TestProviderIntegration:
|
|||||||
provider.get_capabilities("o3")
|
provider.get_capabilities("o3")
|
||||||
assert "not allowed by restriction policy" in str(exc_info.value)
|
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):
|
def test_gemini_provider_respects_restrictions(self):
|
||||||
"""Test that Gemini provider respects restrictions."""
|
"""Test that Gemini provider respects restrictions."""
|
||||||
# Clear any cached restriction service
|
# Clear any cached restriction service
|
||||||
@@ -250,11 +250,11 @@ class TestProviderIntegration:
|
|||||||
|
|
||||||
# Should validate allowed models (both shorthand and full name allowed)
|
# Should validate allowed models (both shorthand and full name allowed)
|
||||||
assert provider.validate_model_name("flash")
|
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
|
# Should not validate disallowed model
|
||||||
assert not provider.validate_model_name("pro")
|
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
|
# get_capabilities should raise for disallowed model
|
||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
@@ -288,13 +288,13 @@ class TestProviderIntegration:
|
|||||||
|
|
||||||
# Should allow getting capabilities for "flash"
|
# Should allow getting capabilities for "flash"
|
||||||
capabilities = provider.get_capabilities("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
|
# 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
|
# 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):
|
def test_gemini_parameter_order_edge_case_full_name_only(self):
|
||||||
"""Test parameter order with only full name allowed, not alias.
|
"""Test parameter order with only full name allowed, not alias.
|
||||||
|
|
||||||
@@ -310,7 +310,7 @@ class TestProviderIntegration:
|
|||||||
provider = GeminiModelProvider(api_key="test-key")
|
provider = GeminiModelProvider(api_key="test-key")
|
||||||
|
|
||||||
# Should allow full name
|
# 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
|
# Should also allow alias that resolves to allowed full name
|
||||||
# This works because is_allowed checks both resolved_name and original_name
|
# This works because is_allowed checks both resolved_name and original_name
|
||||||
@@ -318,7 +318,7 @@ class TestProviderIntegration:
|
|||||||
|
|
||||||
# Should not allow "pro" alias
|
# Should not allow "pro" alias
|
||||||
assert not provider.validate_model_name("pro")
|
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:
|
class TestCustomProviderOpenRouterRestrictions:
|
||||||
@@ -469,8 +469,8 @@ class TestRegistryIntegration:
|
|||||||
|
|
||||||
mock_gemini = MagicMock()
|
mock_gemini = MagicMock()
|
||||||
mock_gemini.SUPPORTED_MODELS = {
|
mock_gemini.SUPPORTED_MODELS = {
|
||||||
"gemini-2.5-pro-preview-06-05": {"context_window": 1048576},
|
"gemini-2.5-pro": {"context_window": 1048576},
|
||||||
"gemini-2.5-flash-preview-05-20": {"context_window": 1048576},
|
"gemini-2.5-flash": {"context_window": 1048576},
|
||||||
}
|
}
|
||||||
mock_gemini.get_provider_type.return_value = ProviderType.GOOGLE
|
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_models = gemini_list_models
|
||||||
mock_gemini.list_all_known_models.return_value = [
|
mock_gemini.list_all_known_models.return_value = [
|
||||||
"gemini-2.5-pro-preview-06-05",
|
"gemini-2.5-pro",
|
||||||
"gemini-2.5-flash-preview-05-20",
|
"gemini-2.5-flash",
|
||||||
]
|
]
|
||||||
|
|
||||||
def get_provider_side_effect(provider_type):
|
def get_provider_side_effect(provider_type):
|
||||||
@@ -514,7 +514,7 @@ class TestRegistryIntegration:
|
|||||||
}
|
}
|
||||||
|
|
||||||
with patch.dict(
|
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
|
# Clear cached restriction service
|
||||||
import utils.model_restrictions
|
import utils.model_restrictions
|
||||||
@@ -526,8 +526,8 @@ class TestRegistryIntegration:
|
|||||||
# Should only include allowed models
|
# Should only include allowed models
|
||||||
assert "o3-mini" in available
|
assert "o3-mini" in available
|
||||||
assert "o3" not in available
|
assert "o3" not in available
|
||||||
assert "gemini-2.5-flash-preview-05-20" in available
|
assert "gemini-2.5-flash" in available
|
||||||
assert "gemini-2.5-pro-preview-06-05" not in available
|
assert "gemini-2.5-pro" not in available
|
||||||
|
|
||||||
|
|
||||||
class TestShorthandRestrictions:
|
class TestShorthandRestrictions:
|
||||||
@@ -552,7 +552,7 @@ class TestShorthandRestrictions:
|
|||||||
gemini_provider = GeminiModelProvider(api_key="test-key")
|
gemini_provider = GeminiModelProvider(api_key="test-key")
|
||||||
assert gemini_provider.validate_model_name("flash") # Should work with shorthand
|
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
|
# 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
|
assert not gemini_provider.validate_model_name("pro") # Not allowed
|
||||||
|
|
||||||
@patch.dict(os.environ, {"OPENAI_ALLOWED_MODELS": "o3mini,mini,o4-mini"})
|
@patch.dict(os.environ, {"OPENAI_ALLOWED_MODELS": "o3mini,mini,o4-mini"})
|
||||||
@@ -579,7 +579,7 @@ class TestShorthandRestrictions:
|
|||||||
|
|
||||||
@patch.dict(
|
@patch.dict(
|
||||||
os.environ,
|
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):
|
def test_both_shorthand_and_full_name_allowed(self):
|
||||||
"""Test that we can allow both shorthand and full names."""
|
"""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 - both flash and full name are allowed
|
||||||
gemini_provider = GeminiModelProvider(api_key="test-key")
|
gemini_provider = GeminiModelProvider(api_key="test-key")
|
||||||
assert gemini_provider.validate_model_name("flash")
|
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:
|
class TestAutoModeWithRestrictions:
|
||||||
@@ -688,7 +688,7 @@ class TestAutoModeWithRestrictions:
|
|||||||
|
|
||||||
# The fallback will depend on how get_available_models handles aliases
|
# The fallback will depend on how get_available_models handles aliases
|
||||||
# For now, we accept either behavior and document it
|
# 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:
|
finally:
|
||||||
# Restore original registry state
|
# Restore original registry state
|
||||||
registry = ModelProviderRegistry()
|
registry = ModelProviderRegistry()
|
||||||
|
|||||||
@@ -167,7 +167,7 @@ class TestOldBehaviorSimulation:
|
|||||||
(["mini", "o3mini"], ["mini", "o3mini", "o4-mini", "o3-mini"], "OpenAI"),
|
(["mini", "o3mini"], ["mini", "o3mini", "o4-mini", "o3-mini"], "OpenAI"),
|
||||||
(
|
(
|
||||||
["flash", "pro"],
|
["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",
|
"Gemini",
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -151,8 +151,8 @@ class TestOpenRouterAutoMode:
|
|||||||
|
|
||||||
mock_registry = Mock()
|
mock_registry = Mock()
|
||||||
mock_registry.list_models.return_value = [
|
mock_registry.list_models.return_value = [
|
||||||
"google/gemini-2.5-flash-preview-05-20",
|
"google/gemini-2.5-flash",
|
||||||
"google/gemini-2.5-pro-preview-06-05",
|
"google/gemini-2.5-pro",
|
||||||
"openai/o3",
|
"openai/o3",
|
||||||
"openai/o3-mini",
|
"openai/o3-mini",
|
||||||
"anthropic/claude-3-opus",
|
"anthropic/claude-3-opus",
|
||||||
@@ -181,7 +181,7 @@ class TestOpenRouterAutoMode:
|
|||||||
os.environ.pop("OPENAI_API_KEY", None)
|
os.environ.pop("OPENAI_API_KEY", None)
|
||||||
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
|
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
|
||||||
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
|
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"
|
os.environ["DEFAULT_MODEL"] = "auto"
|
||||||
|
|
||||||
# Force reload to pick up new environment variable
|
# Force reload to pick up new environment variable
|
||||||
@@ -191,8 +191,8 @@ class TestOpenRouterAutoMode:
|
|||||||
|
|
||||||
mock_registry = Mock()
|
mock_registry = Mock()
|
||||||
mock_registry.list_models.return_value = [
|
mock_registry.list_models.return_value = [
|
||||||
"google/gemini-2.5-flash-preview-05-20",
|
"google/gemini-2.5-flash",
|
||||||
"google/gemini-2.5-pro-preview-06-05",
|
"google/gemini-2.5-pro",
|
||||||
"anthropic/claude-3-opus",
|
"anthropic/claude-3-opus",
|
||||||
"anthropic/claude-3-sonnet",
|
"anthropic/claude-3-sonnet",
|
||||||
]
|
]
|
||||||
@@ -206,7 +206,7 @@ class TestOpenRouterAutoMode:
|
|||||||
|
|
||||||
assert len(available_models) > 0, "Should have some allowed models"
|
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 (
|
assert (
|
||||||
set(available_models.keys()) == expected_allowed
|
set(available_models.keys()) == expected_allowed
|
||||||
|
|||||||
@@ -91,13 +91,13 @@ class TestModelSelection:
|
|||||||
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
||||||
# Mock only Gemini models available
|
# Mock only Gemini models available
|
||||||
mock_get_available.return_value = {
|
mock_get_available.return_value = {
|
||||||
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
|
"gemini-2.5-pro": ProviderType.GOOGLE,
|
||||||
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
|
"gemini-2.5-flash": ProviderType.GOOGLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.EXTENDED_REASONING)
|
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.EXTENDED_REASONING)
|
||||||
# Should find the pro model for 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):
|
def test_fast_response_with_openai(self):
|
||||||
"""Test FAST_RESPONSE prefers o4-mini when OpenAI is available."""
|
"""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:
|
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
||||||
# Mock only Gemini models available
|
# Mock only Gemini models available
|
||||||
mock_get_available.return_value = {
|
mock_get_available.return_value = {
|
||||||
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
|
"gemini-2.5-pro": ProviderType.GOOGLE,
|
||||||
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
|
"gemini-2.5-flash": ProviderType.GOOGLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.FAST_RESPONSE)
|
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.FAST_RESPONSE)
|
||||||
# Should find the flash model for 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):
|
def test_balanced_category_fallback(self):
|
||||||
"""Test BALANCED category uses existing logic."""
|
"""Test BALANCED category uses existing logic."""
|
||||||
@@ -143,13 +143,13 @@ class TestModelSelection:
|
|||||||
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
||||||
# Mock only Gemini models available
|
# Mock only Gemini models available
|
||||||
mock_get_available.return_value = {
|
mock_get_available.return_value = {
|
||||||
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
|
"gemini-2.5-pro": ProviderType.GOOGLE,
|
||||||
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
|
"gemini-2.5-flash": ProviderType.GOOGLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
model = ModelProviderRegistry.get_preferred_fallback_model()
|
model = ModelProviderRegistry.get_preferred_fallback_model()
|
||||||
# Should pick a reasonable default, preferring flash for balanced use
|
# 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:
|
class TestFlexibleModelSelection:
|
||||||
@@ -168,8 +168,8 @@ class TestFlexibleModelSelection:
|
|||||||
# Case 2: Mix of Gemini shorthands and full names
|
# Case 2: Mix of Gemini shorthands and full names
|
||||||
{
|
{
|
||||||
"available": {
|
"available": {
|
||||||
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
|
"gemini-2.5-flash": ProviderType.GOOGLE,
|
||||||
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
|
"gemini-2.5-pro": ProviderType.GOOGLE,
|
||||||
},
|
},
|
||||||
"category": ToolModelCategory.FAST_RESPONSE,
|
"category": ToolModelCategory.FAST_RESPONSE,
|
||||||
"expected_contains": "flash",
|
"expected_contains": "flash",
|
||||||
@@ -220,7 +220,7 @@ class TestCustomProviderFallback:
|
|||||||
mock_find_thinking.return_value = None
|
mock_find_thinking.return_value = None
|
||||||
|
|
||||||
model = ModelProviderRegistry.get_preferred_fallback_model(ToolModelCategory.EXTENDED_REASONING)
|
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:
|
class TestAutoModeErrorMessages:
|
||||||
@@ -234,8 +234,8 @@ class TestAutoModeErrorMessages:
|
|||||||
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
with patch.object(ModelProviderRegistry, "get_available_models") as mock_get_available:
|
||||||
# Mock only Gemini models available
|
# Mock only Gemini models available
|
||||||
mock_get_available.return_value = {
|
mock_get_available.return_value = {
|
||||||
"gemini-2.5-pro-preview-06-05": ProviderType.GOOGLE,
|
"gemini-2.5-pro": ProviderType.GOOGLE,
|
||||||
"gemini-2.5-flash-preview-05-20": ProviderType.GOOGLE,
|
"gemini-2.5-flash": ProviderType.GOOGLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
tool = ThinkDeepTool()
|
tool = ThinkDeepTool()
|
||||||
@@ -245,7 +245,7 @@ class TestAutoModeErrorMessages:
|
|||||||
assert "Model parameter is required in auto mode" in result[0].text
|
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')
|
# Should suggest a model suitable for extended reasoning (either full name or with 'pro')
|
||||||
response_text = result[0].text
|
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
|
assert "(category: extended_reasoning)" in response_text
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -294,7 +294,7 @@ class TestFileContentPreparation:
|
|||||||
# Set up model context to simulate normal execution flow
|
# Set up model context to simulate normal execution flow
|
||||||
from utils.model_context import ModelContext
|
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
|
# Call the method
|
||||||
content, processed_files = tool._prepare_file_content_for_prompt(["/test/file.py"], None, "test")
|
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
|
assert len(debug_calls) > 0
|
||||||
debug_message = str(debug_calls[0])
|
debug_message = str(debug_calls[0])
|
||||||
# Should mention the model being used
|
# 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)
|
# Should mention file tokens (not content tokens)
|
||||||
assert "file tokens" in debug_message
|
assert "file tokens" in debug_message
|
||||||
|
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ class TestPromptRegression:
|
|||||||
return Mock(
|
return Mock(
|
||||||
content=text,
|
content=text,
|
||||||
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
|
||||||
model_name="gemini-2.5-flash-preview-05-20",
|
model_name="gemini-2.5-flash",
|
||||||
metadata={"finish_reason": "STOP"},
|
metadata={"finish_reason": "STOP"},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -64,7 +64,7 @@ class TestModelProviderRegistry:
|
|||||||
"""Test getting provider for a specific model"""
|
"""Test getting provider for a specific model"""
|
||||||
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
|
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
|
||||||
|
|
||||||
provider = ModelProviderRegistry.get_provider_for_model("gemini-2.5-flash-preview-05-20")
|
provider = ModelProviderRegistry.get_provider_for_model("gemini-2.5-flash")
|
||||||
|
|
||||||
assert provider is not None
|
assert provider is not None
|
||||||
assert isinstance(provider, GeminiModelProvider)
|
assert isinstance(provider, GeminiModelProvider)
|
||||||
@@ -95,10 +95,10 @@ class TestGeminiProvider:
|
|||||||
"""Test getting model capabilities"""
|
"""Test getting model capabilities"""
|
||||||
provider = GeminiModelProvider(api_key="test-key")
|
provider = GeminiModelProvider(api_key="test-key")
|
||||||
|
|
||||||
capabilities = provider.get_capabilities("gemini-2.5-flash-preview-05-20")
|
capabilities = provider.get_capabilities("gemini-2.5-flash")
|
||||||
|
|
||||||
assert capabilities.provider == ProviderType.GOOGLE
|
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.context_window == 1_048_576
|
||||||
assert capabilities.supports_extended_thinking
|
assert capabilities.supports_extended_thinking
|
||||||
|
|
||||||
@@ -106,7 +106,7 @@ class TestGeminiProvider:
|
|||||||
"""Test getting capabilities for Pro model with thinking support"""
|
"""Test getting capabilities for Pro model with thinking support"""
|
||||||
provider = GeminiModelProvider(api_key="test-key")
|
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
|
assert capabilities.supports_extended_thinking
|
||||||
|
|
||||||
@@ -118,14 +118,14 @@ class TestGeminiProvider:
|
|||||||
assert provider.validate_model_name("pro")
|
assert provider.validate_model_name("pro")
|
||||||
|
|
||||||
capabilities = provider.get_capabilities("flash")
|
capabilities = provider.get_capabilities("flash")
|
||||||
assert capabilities.model_name == "gemini-2.5-flash-preview-05-20"
|
assert capabilities.model_name == "gemini-2.5-flash"
|
||||||
|
|
||||||
def test_supports_thinking_mode(self):
|
def test_supports_thinking_mode(self):
|
||||||
"""Test thinking mode support detection"""
|
"""Test thinking mode support detection"""
|
||||||
provider = GeminiModelProvider(api_key="test-key")
|
provider = GeminiModelProvider(api_key="test-key")
|
||||||
|
|
||||||
assert provider.supports_thinking_mode("gemini-2.5-flash-preview-05-20")
|
assert provider.supports_thinking_mode("gemini-2.5-flash")
|
||||||
assert provider.supports_thinking_mode("gemini-2.5-pro-preview-06-05")
|
assert provider.supports_thinking_mode("gemini-2.5-pro")
|
||||||
|
|
||||||
@patch("google.genai.Client")
|
@patch("google.genai.Client")
|
||||||
def test_generate_content(self, mock_client_class):
|
def test_generate_content(self, mock_client_class):
|
||||||
@@ -149,12 +149,12 @@ class TestGeminiProvider:
|
|||||||
provider = GeminiModelProvider(api_key="test-key")
|
provider = GeminiModelProvider(api_key="test-key")
|
||||||
|
|
||||||
response = provider.generate_content(
|
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 isinstance(response, ModelResponse)
|
||||||
assert response.content == "Generated content"
|
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.provider == ProviderType.GOOGLE
|
||||||
assert response.usage["input_tokens"] == 10
|
assert response.usage["input_tokens"] == 10
|
||||||
assert response.usage["output_tokens"] == 20
|
assert response.usage["output_tokens"] == 20
|
||||||
|
|||||||
@@ -418,7 +418,7 @@ class TestThinkingModes:
|
|||||||
from providers.gemini import GeminiModelProvider
|
from providers.gemini import GeminiModelProvider
|
||||||
|
|
||||||
provider = GeminiModelProvider(api_key="test-key")
|
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
|
flash_max_tokens = 24576
|
||||||
|
|
||||||
expected_budgets = {
|
expected_budgets = {
|
||||||
|
|||||||
@@ -328,7 +328,7 @@ class BaseTool(ABC):
|
|||||||
logging.debug(f"Failed to add custom models to enum: {e}")
|
logging.debug(f"Failed to add custom models to enum: {e}")
|
||||||
|
|
||||||
# Note: MODEL_CAPABILITIES_DESC already includes both short aliases (e.g., "flash", "o3")
|
# Note: MODEL_CAPABILITIES_DESC already includes both short aliases (e.g., "flash", "o3")
|
||||||
# and full model names (e.g., "gemini-2.5-flash-preview-05-20") as keys
|
# and full model names (e.g., "gemini-2.5-flash") as keys
|
||||||
|
|
||||||
# Remove duplicates while preserving order
|
# Remove duplicates while preserving order
|
||||||
seen = set()
|
seen = set()
|
||||||
@@ -1097,7 +1097,7 @@ When recommending searches, be specific about what information you need and why
|
|||||||
"status": "error",
|
"status": "error",
|
||||||
"content": (
|
"content": (
|
||||||
f"Image support not available: Model '{model_name}' does not support image processing. "
|
f"Image support not available: Model '{model_name}' does not support image processing. "
|
||||||
f"Please use a vision-capable model such as 'gemini-2.5-flash-preview-05-20', 'o3', "
|
f"Please use a vision-capable model such as 'gemini-2.5-flash', 'o3', "
|
||||||
f"or 'claude-3-opus' for image analysis tasks."
|
f"or 'claude-3-opus' for image analysis tasks."
|
||||||
),
|
),
|
||||||
"content_type": "text",
|
"content_type": "text",
|
||||||
|
|||||||
@@ -80,8 +80,8 @@ class ListModelsTool(BaseTool):
|
|||||||
"name": "Google Gemini",
|
"name": "Google Gemini",
|
||||||
"env_key": "GEMINI_API_KEY",
|
"env_key": "GEMINI_API_KEY",
|
||||||
"models": {
|
"models": {
|
||||||
"flash": "gemini-2.5-flash-preview-05-20",
|
"flash": "gemini-2.5-flash",
|
||||||
"pro": "gemini-2.5-pro-preview-06-05",
|
"pro": "gemini-2.5-pro",
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
"openai": {
|
"openai": {
|
||||||
|
|||||||
@@ -159,7 +159,7 @@ class ConversationTurn(BaseModel):
|
|||||||
images: List of image paths referenced in this specific turn
|
images: List of image paths referenced in this specific turn
|
||||||
tool_name: Which tool generated this turn (for cross-tool tracking)
|
tool_name: Which tool generated this turn (for cross-tool tracking)
|
||||||
model_provider: Provider used (e.g., "google", "openai")
|
model_provider: Provider used (e.g., "google", "openai")
|
||||||
model_name: Specific model used (e.g., "gemini-2.5-flash-preview-05-20", "o3-mini")
|
model_name: Specific model used (e.g., "gemini-2.5-flash", "o3-mini")
|
||||||
model_metadata: Additional model-specific metadata (e.g., thinking mode, token usage)
|
model_metadata: Additional model-specific metadata (e.g., thinking mode, token usage)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -327,7 +327,7 @@ def add_turn(
|
|||||||
images: Optional list of images referenced in this turn
|
images: Optional list of images referenced in this turn
|
||||||
tool_name: Name of the tool adding this turn (for attribution)
|
tool_name: Name of the tool adding this turn (for attribution)
|
||||||
model_provider: Provider used (e.g., "google", "openai")
|
model_provider: Provider used (e.g., "google", "openai")
|
||||||
model_name: Specific model used (e.g., "gemini-2.5-flash-preview-05-20", "o3-mini")
|
model_name: Specific model used (e.g., "gemini-2.5-flash", "o3-mini")
|
||||||
model_metadata: Additional model info (e.g., thinking mode, token usage)
|
model_metadata: Additional model info (e.g., thinking mode, token usage)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
|
|||||||
Reference in New Issue
Block a user