Add encouraging message about powerful models to schema in case it's not on Opus 4 or above

OPENROUTER_ALLOWED_MODELS environment variable support to further limit the models to allow from within Claude. This will put a limit on top of even the ones listed in custom_models.json
This commit is contained in:
Fahad
2025-06-14 11:34:17 +04:00
parent 21037c2d81
commit e0a05b86f1
12 changed files with 218 additions and 17 deletions

View File

@@ -84,6 +84,9 @@ OPENAI_ALLOWED_MODELS=o4-mini,o3-mini
# Only allow specific Gemini models
GOOGLE_ALLOWED_MODELS=flash
# Only allow specific OpenRouter models
OPENROUTER_ALLOWED_MODELS=opus,sonnet,mistral
# Use shorthand names or full model names
OPENAI_ALLOWED_MODELS=mini,o3-mini # mini = o4-mini
```
@@ -99,17 +102,21 @@ OPENAI_ALLOWED_MODELS=mini,o3-mini # mini = o4-mini
# Cost control - only cheap models
OPENAI_ALLOWED_MODELS=o4-mini
GOOGLE_ALLOWED_MODELS=flash
OPENROUTER_ALLOWED_MODELS=haiku,sonnet
# Single model per provider
OPENAI_ALLOWED_MODELS=o4-mini
GOOGLE_ALLOWED_MODELS=pro
OPENROUTER_ALLOWED_MODELS=opus
```
**Notes:**
- Applies to all usage including auto mode
- Case-insensitive, whitespace tolerant
- Server warns about typos at startup
- Only affects native providers (not OpenRouter/Custom)
- `OPENAI_ALLOWED_MODELS` and `GOOGLE_ALLOWED_MODELS` only affect native providers
- `OPENROUTER_ALLOWED_MODELS` affects OpenRouter models accessed via custom provider (where `is_custom: false` in custom_models.json)
- Custom local models (`is_custom: true`) are not affected by any restrictions
## Thinking Modes

View File

@@ -84,7 +84,7 @@ isort .
## What Each Test Suite Covers
### Unit Tests (256 tests)
### Unit Tests
Test isolated components and functions:
- **Provider functionality**: Model initialization, API interactions, capability checks
- **Tool operations**: All MCP tools (chat, analyze, debug, etc.)
@@ -92,7 +92,7 @@ Test isolated components and functions:
- **File handling**: Path validation, token limits, deduplication
- **Auto mode**: Model selection logic and fallback behavior
### Simulator Tests (14 tests)
### Simulator Tests
Validate real-world usage scenarios by simulating actual Claude prompts:
- **Basic conversations**: Multi-turn chat functionality with real prompts
- **Cross-tool continuation**: Context preservation across different tools

View File

@@ -128,7 +128,20 @@ class CustomProvider(OpenAICompatibleProvider):
capabilities = self._registry.get_capabilities(model_name)
if capabilities:
# Update provider type to CUSTOM
# Check if this is an OpenRouter model and apply restrictions
config = self._registry.resolve(model_name)
if config and not config.is_custom:
# This is an OpenRouter model, check restrictions
from utils.model_restrictions import get_restriction_service
restriction_service = get_restriction_service()
if not restriction_service.is_allowed(ProviderType.OPENROUTER, config.model_name, model_name):
raise ValueError(f"OpenRouter model '{model_name}' is not allowed by restriction policy.")
# Update provider type to OPENROUTER for OpenRouter models
capabilities.provider = ProviderType.OPENROUTER
else:
# Update provider type to CUSTOM for local custom models
capabilities.provider = ProviderType.CUSTOM
return capabilities
else:
@@ -188,9 +201,20 @@ class CustomProvider(OpenAICompatibleProvider):
logging.debug(f"Model '{model_name}' -> '{model_id}' validated via registry (custom model)")
return True
else:
# This is a cloud/OpenRouter model - if OpenRouter is available, defer to it
# This is a cloud/OpenRouter model - check restrictions if available
if openrouter_available:
logging.debug(f"Model '{model_name}' -> '{model_id}' deferred to OpenRouter (cloud model)")
# Check if OpenRouter model is allowed by restrictions
from utils.model_restrictions import get_restriction_service
restriction_service = get_restriction_service()
if not restriction_service.is_allowed(ProviderType.OPENROUTER, model_id, model_name):
logging.debug(f"Model '{model_name}' -> '{model_id}' blocked by OpenRouter restrictions")
return False
logging.debug(
f"Model '{model_name}' -> '{model_id}' validated via OpenRouter (passes restrictions)"
)
return True
else:
logging.debug(f"Model '{model_name}' -> '{model_id}' rejected (cloud model, no OpenRouter)")
return False

View File

@@ -46,10 +46,15 @@ class TestCustomProvider:
"""Test get_capabilities returns registry capabilities when available."""
provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1")
# Test with a model that should be in the registry
# Test with a model that should be in the registry (OpenRouter model)
capabilities = provider.get_capabilities("llama")
assert capabilities.provider == ProviderType.CUSTOM
assert capabilities.provider == ProviderType.OPENROUTER # llama is an OpenRouter model (is_custom=false)
assert capabilities.context_window > 0
# Test with a custom model (is_custom=true)
capabilities = provider.get_capabilities("local-llama")
assert capabilities.provider == ProviderType.CUSTOM # local-llama has is_custom=true
assert capabilities.context_window > 0
def test_get_capabilities_generic_fallback(self):

View File

@@ -24,10 +24,13 @@ class TestModelRestrictionService:
assert service.is_allowed(ProviderType.OPENAI, "o3-mini")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-pro-preview-06-05")
assert service.is_allowed(ProviderType.GOOGLE, "gemini-2.5-flash-preview-05-20")
assert service.is_allowed(ProviderType.OPENROUTER, "anthropic/claude-3-opus")
assert service.is_allowed(ProviderType.OPENROUTER, "openai/o3")
# Should have no restrictions
assert not service.has_restrictions(ProviderType.OPENAI)
assert not service.has_restrictions(ProviderType.GOOGLE)
assert not service.has_restrictions(ProviderType.OPENROUTER)
def test_load_single_model_restriction(self):
"""Test loading a single allowed model."""
@@ -39,8 +42,9 @@ class TestModelRestrictionService:
assert not service.is_allowed(ProviderType.OPENAI, "o3")
assert not service.is_allowed(ProviderType.OPENAI, "o4-mini")
# Google 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.OPENROUTER, "anthropic/claude-3-opus")
def test_load_multiple_models_restriction(self):
"""Test loading multiple allowed models."""
@@ -146,6 +150,68 @@ class TestModelRestrictionService:
assert "o4-mimi" in caplog.text
assert "not a recognized" in caplog.text
def test_openrouter_model_restrictions(self):
"""Test OpenRouter model restrictions functionality."""
with patch.dict(os.environ, {"OPENROUTER_ALLOWED_MODELS": "opus,sonnet"}):
service = ModelRestrictionService()
# Should only allow specified OpenRouter models
assert service.is_allowed(ProviderType.OPENROUTER, "opus")
assert service.is_allowed(ProviderType.OPENROUTER, "sonnet")
assert service.is_allowed(ProviderType.OPENROUTER, "anthropic/claude-3-opus", "opus") # With original name
assert not service.is_allowed(ProviderType.OPENROUTER, "haiku")
assert not service.is_allowed(ProviderType.OPENROUTER, "anthropic/claude-3-haiku")
assert not service.is_allowed(ProviderType.OPENROUTER, "mistral-large")
# Other providers should have no restrictions
assert service.is_allowed(ProviderType.OPENAI, "o3")
assert service.is_allowed(ProviderType.GOOGLE, "pro")
# Should have restrictions for OpenRouter
assert service.has_restrictions(ProviderType.OPENROUTER)
assert not service.has_restrictions(ProviderType.OPENAI)
assert not service.has_restrictions(ProviderType.GOOGLE)
def test_openrouter_filter_models(self):
"""Test filtering OpenRouter models based on restrictions."""
with patch.dict(os.environ, {"OPENROUTER_ALLOWED_MODELS": "opus,mistral"}):
service = ModelRestrictionService()
models = ["opus", "sonnet", "haiku", "mistral", "llama"]
filtered = service.filter_models(ProviderType.OPENROUTER, models)
assert filtered == ["opus", "mistral"]
def test_combined_provider_restrictions(self):
"""Test that restrictions work correctly when set for multiple providers."""
with patch.dict(
os.environ,
{
"OPENAI_ALLOWED_MODELS": "o3-mini",
"GOOGLE_ALLOWED_MODELS": "flash",
"OPENROUTER_ALLOWED_MODELS": "opus,sonnet",
},
):
service = ModelRestrictionService()
# OpenAI restrictions
assert service.is_allowed(ProviderType.OPENAI, "o3-mini")
assert not service.is_allowed(ProviderType.OPENAI, "o3")
# Google restrictions
assert service.is_allowed(ProviderType.GOOGLE, "flash")
assert not service.is_allowed(ProviderType.GOOGLE, "pro")
# OpenRouter restrictions
assert service.is_allowed(ProviderType.OPENROUTER, "opus")
assert service.is_allowed(ProviderType.OPENROUTER, "sonnet")
assert not service.is_allowed(ProviderType.OPENROUTER, "haiku")
# All providers should have restrictions
assert service.has_restrictions(ProviderType.OPENAI)
assert service.has_restrictions(ProviderType.GOOGLE)
assert service.has_restrictions(ProviderType.OPENROUTER)
class TestProviderIntegration:
"""Test integration with actual providers."""
@@ -195,6 +261,96 @@ class TestProviderIntegration:
assert "not allowed by restriction policy" in str(exc_info.value)
class TestCustomProviderOpenRouterRestrictions:
"""Test custom provider integration with OpenRouter restrictions."""
@patch.dict(os.environ, {"OPENROUTER_ALLOWED_MODELS": "opus,sonnet", "OPENROUTER_API_KEY": "test-key"})
def test_custom_provider_respects_openrouter_restrictions(self):
"""Test that custom provider respects OpenRouter restrictions for cloud models."""
# Clear any cached restriction service
import utils.model_restrictions
utils.model_restrictions._restriction_service = None
from providers.custom import CustomProvider
provider = CustomProvider(base_url="http://test.com/v1")
# Should validate allowed OpenRouter models (is_custom=false)
assert provider.validate_model_name("opus")
assert provider.validate_model_name("sonnet")
# Should not validate disallowed OpenRouter models
assert not provider.validate_model_name("haiku")
# Should still validate custom models (is_custom=true) regardless of restrictions
assert provider.validate_model_name("local-llama") # This has is_custom=true
@patch.dict(os.environ, {"OPENROUTER_ALLOWED_MODELS": "opus", "OPENROUTER_API_KEY": "test-key"})
def test_custom_provider_openrouter_capabilities_restrictions(self):
"""Test that custom provider's get_capabilities respects OpenRouter restrictions."""
# Clear any cached restriction service
import utils.model_restrictions
utils.model_restrictions._restriction_service = None
from providers.custom import CustomProvider
provider = CustomProvider(base_url="http://test.com/v1")
# Should work for allowed OpenRouter model
capabilities = provider.get_capabilities("opus")
assert capabilities.provider == ProviderType.OPENROUTER
# Should raise for disallowed OpenRouter model
with pytest.raises(ValueError) as exc_info:
provider.get_capabilities("haiku")
assert "not allowed by restriction policy" in str(exc_info.value)
# Should still work for custom models (is_custom=true)
capabilities = provider.get_capabilities("local-llama")
assert capabilities.provider == ProviderType.CUSTOM
@patch.dict(os.environ, {"OPENROUTER_ALLOWED_MODELS": "opus"}, clear=False)
def test_custom_provider_no_openrouter_key_ignores_restrictions(self):
"""Test that when OpenRouter key is not set, cloud models are rejected regardless of restrictions."""
# Make sure OPENROUTER_API_KEY is not set
if "OPENROUTER_API_KEY" in os.environ:
del os.environ["OPENROUTER_API_KEY"]
# Clear any cached restriction service
import utils.model_restrictions
utils.model_restrictions._restriction_service = None
from providers.custom import CustomProvider
provider = CustomProvider(base_url="http://test.com/v1")
# Should not validate OpenRouter models when key is not available
assert not provider.validate_model_name("opus") # Even though it's in allowed list
assert not provider.validate_model_name("haiku")
# Should still validate custom models
assert provider.validate_model_name("local-llama")
@patch.dict(os.environ, {"OPENROUTER_ALLOWED_MODELS": "", "OPENROUTER_API_KEY": "test-key"})
def test_custom_provider_empty_restrictions_allows_all_openrouter(self):
"""Test that empty OPENROUTER_ALLOWED_MODELS allows all OpenRouter models."""
# Clear any cached restriction service
import utils.model_restrictions
utils.model_restrictions._restriction_service = None
from providers.custom import CustomProvider
provider = CustomProvider(base_url="http://test.com/v1")
# Should validate all OpenRouter models when restrictions are empty
assert provider.validate_model_name("opus")
assert provider.validate_model_name("sonnet")
assert provider.validate_model_name("haiku")
class TestRegistryIntegration:
"""Test integration with ModelProviderRegistry."""

View File

@@ -42,6 +42,7 @@ class AnalyzeTool(BaseTool):
"Use this when you need to analyze files, examine code, or understand specific aspects of a codebase. "
"Perfect for: codebase exploration, dependency analysis, pattern detection. "
"Always uses file paths for clean terminal output. "
"Note: If you're not currently using a top-tier model such as Opus 4 or above, these tools can provide enhanced capabilities."
)
def get_input_schema(self) -> dict[str, Any]:

View File

@@ -44,6 +44,7 @@ class ChatTool(BaseTool):
"Also great for: explanations, comparisons, general development questions. "
"Use this when you want to ask questions, brainstorm ideas, get opinions, discuss topics, "
"share your thinking, or need explanations about concepts and approaches. "
"Note: If you're not currently using a top-tier model such as Opus 4 or above, these tools can provide enhanced capabilities."
)
def get_input_schema(self) -> dict[str, Any]:

View File

@@ -79,6 +79,7 @@ class CodeReviewTool(BaseTool):
"Choose thinking_mode based on review scope: 'low' for small code snippets, "
"'medium' for standard files/modules (default), 'high' for complex systems/architectures, "
"'max' for critical security audits or large codebases requiring deepest analysis. "
"Note: If you're not currently using a top-tier model such as Opus 4 or above, these tools can provide enhanced capabilities."
)
def get_input_schema(self) -> dict[str, Any]:

View File

@@ -50,6 +50,7 @@ class DebugIssueTool(BaseTool):
"Choose thinking_mode based on issue complexity: 'low' for simple errors, "
"'medium' for standard debugging (default), 'high' for complex system issues, "
"'max' for extremely challenging bugs requiring deepest analysis. "
"Note: If you're not currently using a top-tier model such as Opus 4 or above, these tools can provide enhanced capabilities."
)
def get_input_schema(self) -> dict[str, Any]:

View File

@@ -100,6 +100,7 @@ class Precommit(BaseTool):
"Choose thinking_mode based on changeset size: 'low' for small focused changes, "
"'medium' for standard commits (default), 'high' for large feature branches or complex refactoring, "
"'max' for critical releases or when reviewing extensive changes across multiple systems. "
"Note: If you're not currently using a top-tier model such as Opus 4 or above, these tools can provide enhanced capabilities."
)
def get_input_schema(self) -> dict[str, Any]:

View File

@@ -48,6 +48,7 @@ class ThinkDeepTool(BaseTool):
"'low' for quick analysis, 'medium' for standard problems, 'high' for complex issues (default), "
"'max' for extremely complex challenges requiring deepest analysis. "
"When in doubt, err on the side of a higher mode for truly deep thought and evaluation. "
"Note: If you're not currently using a top-tier model such as Opus 4 or above, these tools can provide enhanced capabilities."
)
def get_input_schema(self) -> dict[str, Any]:

View File

@@ -9,10 +9,12 @@ standardization purposes.
Environment Variables:
- OPENAI_ALLOWED_MODELS: Comma-separated list of allowed OpenAI models
- GOOGLE_ALLOWED_MODELS: Comma-separated list of allowed Gemini models
- OPENROUTER_ALLOWED_MODELS: Comma-separated list of allowed OpenRouter models
Example:
OPENAI_ALLOWED_MODELS=o3-mini,o4-mini
GOOGLE_ALLOWED_MODELS=flash
OPENROUTER_ALLOWED_MODELS=opus,sonnet,mistral
"""
import logging
@@ -38,6 +40,7 @@ class ModelRestrictionService:
ENV_VARS = {
ProviderType.OPENAI: "OPENAI_ALLOWED_MODELS",
ProviderType.GOOGLE: "GOOGLE_ALLOWED_MODELS",
ProviderType.OPENROUTER: "OPENROUTER_ALLOWED_MODELS",
}
def __init__(self):