refactor: moved registries into a separate module and code cleanup

fix: refactored dial provider to follow the same pattern
This commit is contained in:
Fahad
2025-10-07 12:59:09 +04:00
parent c27e81d6d2
commit 7c36b9255a
54 changed files with 325 additions and 282 deletions

View File

@@ -38,7 +38,7 @@ if sys.platform == "win32":
# Register providers for all tests
from providers.gemini import GeminiModelProvider # noqa: E402
from providers.openai_provider import OpenAIModelProvider # noqa: E402
from providers.openai import OpenAIModelProvider # noqa: E402
from providers.registry import ModelProviderRegistry # noqa: E402
from providers.shared import ProviderType # noqa: E402
from providers.xai import XAIModelProvider # noqa: E402

View File

@@ -9,7 +9,7 @@ import os
from unittest.mock import patch
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.shared import ProviderType
from utils.model_restrictions import ModelRestrictionService

View File

@@ -7,7 +7,7 @@ from unittest.mock import MagicMock, patch
import pytest
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.registry import ModelProviderRegistry
from providers.shared import ProviderType
from providers.xai import XAIModelProvider

View File

@@ -9,7 +9,7 @@ import pytest
import utils.env as env_config
import utils.model_restrictions as model_restrictions
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.openrouter import OpenRouterProvider
from providers.registry import ModelProviderRegistry
from providers.shared import ProviderType

View File

@@ -86,7 +86,7 @@ class TestAutoModeProviderSelection:
os.environ.pop(key, None)
# Register only OpenAI provider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -127,7 +127,7 @@ class TestAutoModeProviderSelection:
# Register both providers
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -212,7 +212,7 @@ class TestAutoModeProviderSelection:
# Register both providers
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -256,7 +256,7 @@ class TestAutoModeProviderSelection:
# Register all providers
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.xai import XAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
@@ -307,7 +307,7 @@ class TestAutoModeProviderSelection:
# Register all providers
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.xai import XAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)

View File

@@ -12,7 +12,7 @@ from unittest.mock import MagicMock, patch
import pytest
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.shared import ProviderType
from utils.model_restrictions import ModelRestrictionService

View File

@@ -105,7 +105,7 @@ async def test_chat_cross_model_continuation(monkeypatch):
ModelProviderRegistry.reset_for_testing()
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
@@ -170,7 +170,7 @@ async def test_chat_cross_model_continuation(monkeypatch):
ModelProviderRegistry.reset_for_testing()
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)

View File

@@ -54,7 +54,7 @@ async def test_chat_auto_mode_with_openai(monkeypatch):
# Reset registry and register only OpenAI provider
ModelProviderRegistry.reset_for_testing()
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -115,7 +115,7 @@ async def test_chat_openai_continuation(monkeypatch):
m.delenv(key, raising=False)
ModelProviderRegistry.reset_for_testing()
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)

View File

@@ -75,7 +75,7 @@ async def test_consensus_multi_model_consultations(monkeypatch):
# Reset providers and register only OpenAI & Gemini for deterministic behavior
ModelProviderRegistry.reset_for_testing()
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)

View File

@@ -11,7 +11,7 @@ import tempfile
from pathlib import Path
from unittest.mock import Mock, patch
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
class TestCustomOpenAITemperatureParameterFix:
@@ -79,7 +79,7 @@ class TestCustomOpenAITemperatureParameterFix:
mock_client.chat.completions.create.return_value = mock_response
# Create provider with custom config
with patch("providers.openrouter_registry.OpenRouterModelRegistry") as mock_registry_class:
with patch("providers.registries.openrouter.OpenRouterModelRegistry") as mock_registry_class:
# Mock registry to load our test config
mock_registry = Mock()
mock_registry_class.return_value = mock_registry
@@ -163,7 +163,7 @@ class TestCustomOpenAITemperatureParameterFix:
mock_client.chat.completions.create.return_value = mock_response
# Create provider with custom config
with patch("providers.openrouter_registry.OpenRouterModelRegistry") as mock_registry_class:
with patch("providers.registries.openrouter.OpenRouterModelRegistry") as mock_registry_class:
# Mock registry to load our test config
mock_registry = Mock()
mock_registry_class.return_value = mock_registry
@@ -221,7 +221,7 @@ class TestCustomOpenAITemperatureParameterFix:
mock_service.is_allowed.return_value = True
mock_restriction_service.return_value = mock_service
with patch("providers.openrouter_registry.OpenRouterModelRegistry") as mock_registry_class:
with patch("providers.registries.openrouter.OpenRouterModelRegistry") as mock_registry_class:
# Mock registry to return a custom OpenAI model
mock_registry = Mock()
mock_registry_class.return_value = mock_registry
@@ -267,7 +267,7 @@ class TestCustomOpenAITemperatureParameterFix:
mock_service.is_allowed.return_value = True
mock_restriction_service.return_value = mock_service
with patch("providers.openrouter_registry.OpenRouterModelRegistry") as mock_registry_class:
with patch("providers.registries.openrouter.OpenRouterModelRegistry") as mock_registry_class:
# Mock registry to raise an exception
mock_registry_class.side_effect = Exception("Registry not available")

View File

@@ -39,7 +39,7 @@ class TestIntelligentFallback:
def test_prefers_openai_o3_mini_when_available(self):
"""Test that gpt-5 is preferred when OpenAI API key is available (based on new preference order)"""
# Register only OpenAI provider for this test
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -62,7 +62,7 @@ class TestIntelligentFallback:
"""Test that OpenAI is preferred when both API keys are available"""
# Register both OpenAI and Gemini providers
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
@@ -75,7 +75,7 @@ class TestIntelligentFallback:
"""Test fallback behavior when no API keys are available"""
# Register providers but with no API keys available
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)
@@ -86,7 +86,7 @@ class TestIntelligentFallback:
def test_available_providers_with_keys(self):
"""Test the get_available_providers_with_keys method"""
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-key", "GEMINI_API_KEY": ""}, clear=False):
# Clear and register providers
@@ -119,7 +119,7 @@ class TestIntelligentFallback:
patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-key", "GEMINI_API_KEY": ""}, clear=False),
):
# Register only OpenAI provider for this test
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)

View File

@@ -6,7 +6,7 @@ Issue: Custom OpenAI models (gpt-5, o3) use temperature despite the config havin
from unittest.mock import Mock, patch
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
def test_issue_245_custom_openai_temperature_ignored():
@@ -14,7 +14,7 @@ def test_issue_245_custom_openai_temperature_ignored():
with patch("utils.model_restrictions.get_restriction_service") as mock_restriction:
with patch("providers.openai_compatible.OpenAI") as mock_openai:
with patch("providers.openrouter_registry.OpenRouterModelRegistry") as mock_registry_class:
with patch("providers.registries.openrouter.OpenRouterModelRegistry") as mock_registry_class:
# Mock restriction service
mock_service = Mock()

View File

@@ -97,7 +97,7 @@ class TestListModelsRestrictions(unittest.TestCase):
},
)
@patch("utils.model_restrictions.get_restriction_service")
@patch("providers.openrouter_registry.OpenRouterModelRegistry")
@patch("providers.registries.openrouter.OpenRouterModelRegistry")
@patch.object(ModelProviderRegistry, "get_available_models")
@patch.object(ModelProviderRegistry, "get_provider")
def test_listmodels_respects_openrouter_restrictions(
@@ -239,7 +239,7 @@ class TestListModelsRestrictions(unittest.TestCase):
self.assertIn("OpenRouter models restricted by", result)
@patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key", "GEMINI_API_KEY": "gemini-test-key"}, clear=True)
@patch("providers.openrouter_registry.OpenRouterModelRegistry")
@patch("providers.registries.openrouter.OpenRouterModelRegistry")
@patch.object(ModelProviderRegistry, "get_provider")
def test_listmodels_shows_all_models_without_restrictions(self, mock_get_provider, mock_registry_class):
"""Test that listmodels shows all models when no restrictions are set."""

View File

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock, patch
import pytest
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.shared import ProviderType
from utils.model_restrictions import ModelRestrictionService
@@ -767,7 +767,7 @@ class TestAutoModeWithRestrictions:
# Clear registry and register only OpenAI and Gemini providers
ModelProviderRegistry._instance = None
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)

View File

@@ -7,7 +7,7 @@ for O3 models while maintaining them for regular models.
from unittest.mock import Mock, patch
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
class TestO3TemperatureParameterFixSimple:
@@ -175,7 +175,7 @@ class TestO3TemperatureParameterFixSimple:
@patch("utils.model_restrictions.get_restriction_service")
def test_all_o3_models_have_correct_temperature_capability(self, mock_restriction_service):
"""Test that all O3/O4 models have supports_temperature=False in their capabilities."""
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
# Mock restriction service to allow all models
mock_service = Mock()
@@ -211,7 +211,7 @@ class TestO3TemperatureParameterFixSimple:
@patch("utils.model_restrictions.get_restriction_service")
def test_openai_provider_temperature_constraints(self, mock_restriction_service):
"""Test that OpenAI provider has correct temperature constraints for O3 models."""
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
# Mock restriction service to allow all models
mock_service = Mock()

View File

@@ -3,7 +3,7 @@
import os
from unittest.mock import MagicMock, patch
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.shared import ProviderType

View File

@@ -282,7 +282,7 @@ class TestOpenRouterRegistry:
def test_registry_loading(self):
"""Test registry loads models from config."""
from providers.openrouter_registry import OpenRouterModelRegistry
from providers.registries.openrouter import OpenRouterModelRegistry
registry = OpenRouterModelRegistry()
@@ -301,7 +301,7 @@ class TestOpenRouterRegistry:
def test_registry_capabilities(self):
"""Test registry provides correct capabilities."""
from providers.openrouter_registry import OpenRouterModelRegistry
from providers.registries.openrouter import OpenRouterModelRegistry
registry = OpenRouterModelRegistry()
@@ -322,7 +322,7 @@ class TestOpenRouterRegistry:
def test_multiple_aliases_same_model(self):
"""Test multiple aliases pointing to same model."""
from providers.openrouter_registry import OpenRouterModelRegistry
from providers.registries.openrouter import OpenRouterModelRegistry
registry = OpenRouterModelRegistry()

View File

@@ -7,7 +7,7 @@ from unittest.mock import patch
import pytest
from providers.openrouter_registry import OpenRouterModelRegistry
from providers.registries.openrouter import OpenRouterModelRegistry
from providers.shared import ModelCapabilities, ProviderType

View File

@@ -90,7 +90,7 @@ class TestModelSelection:
ModelProviderRegistry.unregister_provider(provider_type)
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=False):
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -125,7 +125,7 @@ class TestModelSelection:
ModelProviderRegistry.unregister_provider(provider_type)
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=False):
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -159,7 +159,7 @@ class TestModelSelection:
ModelProviderRegistry.unregister_provider(provider_type)
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=False):
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
@@ -220,7 +220,7 @@ class TestFlexibleModelSelection:
with patch.dict(os.environ, case["env"], clear=False):
# Register the appropriate provider
if case["provider_type"] == ProviderType.OPENAI:
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
ModelProviderRegistry.register_provider(ProviderType.OPENAI, OpenAIModelProvider)
elif case["provider_type"] == ProviderType.GOOGLE:

View File

@@ -4,7 +4,7 @@ from types import SimpleNamespace
import pytest
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
def _mock_chat_response(content: str = "retry success") -> SimpleNamespace:

View File

@@ -189,7 +189,7 @@ class TestProviderRoutingBugs:
# Register providers in priority order (like server.py)
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.openrouter import OpenRouterProvider
ModelProviderRegistry.register_provider(ProviderType.GOOGLE, GeminiModelProvider)

View File

@@ -11,7 +11,7 @@ from unittest.mock import Mock, patch
import pytest
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.shared import ProviderType

View File

@@ -7,7 +7,7 @@ import pytest
from providers import ModelProviderRegistry, ModelResponse
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.shared import ProviderType

View File

@@ -3,7 +3,7 @@ Test to verify structured error code-based retry logic.
"""
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
def test_openai_structured_error_retry_logic():

View File

@@ -2,7 +2,7 @@
from providers.dial import DIALModelProvider
from providers.gemini import GeminiModelProvider
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.xai import XAIModelProvider

View File

@@ -5,7 +5,7 @@ import tempfile
from pathlib import Path
from unittest.mock import patch
from providers.openrouter_registry import OpenRouterModelRegistry
from providers.registries.openrouter import OpenRouterModelRegistry
class TestUvxPathResolution:
@@ -55,7 +55,7 @@ class TestUvxPathResolution:
assert registry.config_path == config_path
assert len(registry.list_models()) > 0
@patch("providers.model_registry_base.importlib.resources.files")
@patch("providers.registries.base.importlib.resources.files")
def test_multiple_path_fallback(self, mock_files):
"""Test that file-system fallback works when resource loading fails."""
mock_files.side_effect = Exception("Resource loading failed")

View File

@@ -22,7 +22,7 @@ def inject_transport(monkeypatch, cassette_path: str):
transport = inject_transport(monkeypatch, "path/to/cassette.json")
"""
# Ensure OpenAI provider is registered - always needed for transport injection
from providers.openai_provider import OpenAIModelProvider
from providers.openai import OpenAIModelProvider
from providers.registry import ModelProviderRegistry
from providers.shared import ProviderType