Fix auto mode when only OpenRouter is configured

The get_available_models method in ModelProviderRegistry was only checking
for providers with SUPPORTED_MODELS attribute, which OpenRouter doesn't have.
This caused auto mode to fail with "No models available" error when only
OpenRouter API key was configured.

Added special handling for OpenRouter provider to check its _registry
for available models, ensuring auto mode works correctly with OpenRouter.

Added comprehensive tests to verify:
- Auto mode works with only OpenRouter configured
- Model restrictions are respected
- Graceful handling when no providers are available
- No crashes when OpenRouter lacks _registry attribute
This commit is contained in:
Nikolai Ugelvik
2025-06-14 18:51:45 +02:00
parent 4cacd2dad9
commit be2612752a
2 changed files with 146 additions and 1 deletions

View File

@@ -183,6 +183,16 @@ class ModelProviderRegistry:
continue
models[model_name] = provider_type
elif provider_type == ProviderType.OPENROUTER:
# OpenRouter uses a registry system instead of SUPPORTED_MODELS
if hasattr(provider, "_registry") and provider._registry:
for model_name in provider._registry.list_models():
# Check restrictions if enabled
if restriction_service and not restriction_service.is_allowed(provider_type, model_name):
logging.debug(f"Model {model_name} filtered by restrictions")
continue
models[model_name] = provider_type
return models

View File

@@ -1,7 +1,9 @@
"""Tests for OpenRouter provider."""
import os
from unittest.mock import patch
from unittest.mock import Mock, patch
import pytest
from providers.base import ProviderType
from providers.openrouter import OpenRouterProvider
@@ -110,6 +112,139 @@ class TestOpenRouterProvider:
assert isinstance(provider, OpenRouterProvider)
class TestOpenRouterAutoMode:
"""Test auto mode functionality when only OpenRouter is configured."""
def setup_method(self):
"""Store original state before each test."""
self.registry = ModelProviderRegistry()
self._original_providers = self.registry._providers.copy()
self._original_initialized = self.registry._initialized_providers.copy()
self.registry._providers.clear()
self.registry._initialized_providers.clear()
self._original_env = {}
for key in ['OPENROUTER_API_KEY', 'GEMINI_API_KEY', 'OPENAI_API_KEY', 'DEFAULT_MODEL']:
self._original_env[key] = os.environ.get(key)
def teardown_method(self):
"""Restore original state after each test."""
self.registry._providers.clear()
self.registry._initialized_providers.clear()
self.registry._providers.update(self._original_providers)
self.registry._initialized_providers.update(self._original_initialized)
for key, value in self._original_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
@pytest.mark.no_mock_provider
def test_openrouter_only_auto_mode(self):
"""Test that auto mode works when only OpenRouter is configured."""
os.environ.pop('GEMINI_API_KEY', None)
os.environ.pop('OPENAI_API_KEY', None)
os.environ['OPENROUTER_API_KEY'] = 'test-openrouter-key'
os.environ['DEFAULT_MODEL'] = 'auto'
mock_registry = Mock()
mock_registry.list_models.return_value = [
'google/gemini-2.5-flash-preview-05-20',
'google/gemini-2.5-pro-preview-06-05',
'openai/o3',
'openai/o3-mini',
'anthropic/claude-3-opus',
'anthropic/claude-3-sonnet'
]
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
provider = ModelProviderRegistry.get_provider(ProviderType.OPENROUTER)
assert provider is not None, "OpenRouter provider should be available with API key"
provider._registry = mock_registry
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
assert len(available_models) > 0, "Should find OpenRouter models in auto mode"
assert all(provider_type == ProviderType.OPENROUTER for provider_type in available_models.values())
expected_models = mock_registry.list_models()
for model in expected_models:
assert model in available_models, f"Model {model} should be available"
@pytest.mark.no_mock_provider
def test_openrouter_with_restrictions(self):
"""Test that OpenRouter respects model restrictions."""
os.environ.pop('GEMINI_API_KEY', None)
os.environ.pop('OPENAI_API_KEY', None)
os.environ['OPENROUTER_API_KEY'] = 'test-openrouter-key'
os.environ.pop('OPENROUTER_ALLOWED_MODELS', None)
os.environ['OPENROUTER_ALLOWED_MODELS'] = 'anthropic/claude-3-opus,google/gemini-2.5-flash-preview-05-20'
os.environ['DEFAULT_MODEL'] = 'auto'
# Force reload to pick up new environment variable
import utils.model_restrictions
utils.model_restrictions._restriction_service = None
mock_registry = Mock()
mock_registry.list_models.return_value = [
'google/gemini-2.5-flash-preview-05-20',
'google/gemini-2.5-pro-preview-06-05',
'anthropic/claude-3-opus',
'anthropic/claude-3-sonnet'
]
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
provider = ModelProviderRegistry.get_provider(ProviderType.OPENROUTER)
provider._registry = mock_registry
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
assert len(available_models) > 0, "Should have some allowed models"
expected_allowed = {
'google/gemini-2.5-flash-preview-05-20',
'anthropic/claude-3-opus'
}
assert set(available_models.keys()) == expected_allowed, \
f"Expected {expected_allowed}, but got {set(available_models.keys())}"
@pytest.mark.no_mock_provider
def test_no_providers_fails_auto_mode(self):
"""Test that auto mode fails gracefully when no providers are available."""
os.environ.pop('GEMINI_API_KEY', None)
os.environ.pop('OPENAI_API_KEY', None)
os.environ.pop('OPENROUTER_API_KEY', None)
os.environ['DEFAULT_MODEL'] = 'auto'
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
assert len(available_models) == 0, "Should have no models when no providers are configured"
@pytest.mark.no_mock_provider
def test_openrouter_without_registry(self):
"""Test that OpenRouter without _registry attribute doesn't crash."""
os.environ.pop('GEMINI_API_KEY', None)
os.environ.pop('OPENAI_API_KEY', None)
os.environ['OPENROUTER_API_KEY'] = 'test-openrouter-key'
os.environ['DEFAULT_MODEL'] = 'auto'
mock_provider_class = Mock()
mock_provider_instance = Mock(spec=['get_provider_type'])
mock_provider_instance.get_provider_type.return_value = ProviderType.OPENROUTER
mock_provider_class.return_value = mock_provider_instance
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, mock_provider_class)
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
assert len(available_models) == 0, "Should have no models when OpenRouter has no registry"
class TestOpenRouterRegistry:
"""Test cases for OpenRouter model registry."""