Merge pull request #48 from pescheckit/feature_splitting-model-in-seperate-classes

Splitting model in seperate classes

Author: pescheck-bram

python_gpt_po/services/model_manager.py

@@ -5,10 +5,9 @@
 import logging
 from typing import List
 
-import requests
-
 from ..models.enums import ModelProvider
 from ..models.provider_clients import ProviderClients
+from .providers.registry import ProviderRegistry
 
 
 class ModelManager:
@@ -17,77 +16,17 @@ class ModelManager:
     @staticmethod
     def get_available_models(provider_clients: ProviderClients, provider: ModelProvider) -> List[str]:
         """Retrieve available models from a specific provider."""
-        models = []
+        provider_impl = ProviderRegistry.get_provider(provider)
 
-        try:
-            if provider == ModelProvider.OPENAI:
-                if provider_clients.openai_client:
-                    response = provider_clients.openai_client.models.list()
-                    models = [model.id for model in response.data]
-                else:
-                    logging.error("OpenAI client not initialized")
-
-            elif provider == ModelProvider.ANTHROPIC:
-                if provider_clients.anthropic_client:
-                    # Use Anthropic's models endpoint
-                    try:
-                        headers = {
-                            "x-api-key": provider_clients.anthropic_client.api_key,
-                            "anthropic-version": "2023-06-01"
-                        }
-                        response = requests.get(
-                            "https://api.anthropic.com/v1/models",
-                            headers=headers,
-                            timeout=15
-                        )
-                        response.raise_for_status()
-                        model_data = response.json().get("data", [])
-                        models = [model["id"] for model in model_data]
-                    except Exception as e:
-                        logging.error("Error fetching Anthropic models: %s", str(e))
-                        # Fallback to commonly used models if API call fails
-                        models = [
-                            "claude-3-7-sonnet-latest",
-                            "claude-3-5-haiku-latest",
-                            "claude-3-5-sonnet-latest",
-                            "claude-3-opus-20240229",
-                        ]
-                else:
-                    logging.error("Anthropic client not initialized")
-
-            elif provider == ModelProvider.DEEPSEEK:
-                if provider_clients.deepseek_api_key:
-                    headers = {
-                        "Authorization": f"Bearer {provider_clients.deepseek_api_key}",
-                        "Content-Type": "application/json"
-                    }
-                    response = requests.get(
-                        f"{provider_clients.deepseek_base_url}/models",
-                        headers=headers,
-                        timeout=15
-                    )
-                    response.raise_for_status()
-                    models = [model["id"] for model in response.json().get("data", [])]
-                else:
-                    logging.error("DeepSeek API key not set")
-
-            elif provider == ModelProvider.AZURE_OPENAI:
-                return ModelManager._get_azure_openai_models(provider_clients)
+        if not provider_impl:
+            logging.error("Provider %s not registered", provider.value)
+            return []
 
+        try:
+            return provider_impl.get_models(provider_clients)
         except Exception as e:
             logging.error("Error fetching models from %s: %s", provider.value, str(e))
-
-        return models
-
-    @staticmethod
-    def _get_azure_openai_models(provider_clients: ProviderClients) -> List[str]:
-        """Retrieve models from Azure OpenAI."""
-        if provider_clients.azure_openai_client:
-            response = provider_clients.azure_openai_client.models.list()
-            return [model.id for model in response.data]
-
-        logging.error("Azure OpenAI client not initialized")
-        return []
+            return []
 
     @staticmethod
     def validate_model(provider_clients: ProviderClients, provider: ModelProvider, model: str) -> bool:
@@ -119,13 +58,13 @@ def get_default_model(provider: ModelProvider) -> str:
         Returns:
             str: The default model ID
         """
-        default_models = {
-            ModelProvider.OPENAI: "gpt-4o-mini",
-            ModelProvider.ANTHROPIC: "claude-3-5-haiku-latest",
-            ModelProvider.DEEPSEEK: "deepseek-chat",
-            ModelProvider.AZURE_OPENAI: "gpt-35-turbo",
-        }
-        return default_models.get(provider, "")
+        provider_impl = ProviderRegistry.get_provider(provider)
+
+        if not provider_impl:
+            logging.warning("Provider %s not registered, returning empty default", provider.value)
+            return ""
+
+        return provider_impl.get_default_model()
 
     @staticmethod
     def verify_model_capabilities(
@@ -170,25 +109,23 @@ def suggest_model(provider_clients: ProviderClients, provider: ModelProvider,
         Returns:
             str: The suggested model ID
         """
-        # For translation tasks, prefer more capable models when available
-        if task == "translation":
-            preferred_models = {
-                ModelProvider.OPENAI: ["gpt-4", "gpt-4o", "gpt-3.5-turbo"],
-                ModelProvider.ANTHROPIC: ["claude-3-opus", "claude-3-5-sonnet", "claude-3-5-haiku"],
-                ModelProvider.DEEPSEEK: ["deepseek-chat"]
-            }
-
-            available_models = ModelManager.get_available_models(provider_clients, provider)
-
-            # Try to find a match from the preferred models list
-            for preferred in preferred_models.get(provider, []):
-                for available in available_models:
-                    if preferred in available.lower():
-                        return available
-
-            # Fall back to the first available model or the default
-            if available_models:
-                return available_models[0]
-
-        # Default to the standard default model
-        return ModelManager.get_default_model(provider)
+        provider_impl = ProviderRegistry.get_provider(provider)
+
+        if not provider_impl:
+            logging.warning("Provider %s not registered", provider.value)
+            return ""
+
+        preferred_models = provider_impl.get_preferred_models(task)
+        available_models = ModelManager.get_available_models(provider_clients, provider)
+
+        # Try to find a match from the preferred models list
+        for preferred in preferred_models:
+            for available in available_models:
+                if preferred in available.lower():
+                    return available
+
+        # Fall back to the first available model or the default
+        if available_models:
+            return available_models[0]
+
+        return provider_impl.get_default_model()

python_gpt_po/services/providers/__init__.py

@@ -0,0 +1,8 @@
+"""
+Provider implementations for model management.
+"""
+from . import provider_init  # noqa: F401 - Auto-registers providers
+from .base import ModelProviderInterface
+from .registry import ProviderRegistry
+
+__all__ = ["ModelProviderInterface", "ProviderRegistry"]

python_gpt_po/services/providers/anthropic_provider.py

@@ -0,0 +1,64 @@
+"""
+Anthropic provider implementation.
+"""
+import logging
+from typing import List
+
+import requests
+
+from ...models.provider_clients import ProviderClients
+from .base import ModelProviderInterface
+
+
+class AnthropicProvider(ModelProviderInterface):
+    """Anthropic model provider implementation."""
+
+    def get_models(self, provider_clients: ProviderClients) -> List[str]:
+        """Retrieve available models from Anthropic."""
+        models = []
+
+        if not self.is_client_initialized(provider_clients):
+            logging.error("Anthropic client not initialized")
+            return models
+
+        try:
+            headers = {
+                "x-api-key": provider_clients.anthropic_client.api_key,
+                "anthropic-version": "2023-06-01"
+            }
+            response = requests.get(
+                "https://api.anthropic.com/v1/models",
+                headers=headers,
+                timeout=15
+            )
+            response.raise_for_status()
+            model_data = response.json().get("data", [])
+            models = [model["id"] for model in model_data]
+        except Exception as e:
+            logging.error("Error fetching Anthropic models: %s", str(e))
+            models = self.get_fallback_models()
+
+        return models
+
+    def get_default_model(self) -> str:
+        """Get the default Anthropic model."""
+        return "claude-3-5-haiku-latest"
+
+    def get_preferred_models(self, task: str = "translation") -> List[str]:
+        """Get preferred Anthropic models for a task."""
+        if task == "translation":
+            return ["claude-3-opus", "claude-3-5-sonnet", "claude-3-5-haiku"]
+        return ["claude-3-5-haiku-latest"]
+
+    def is_client_initialized(self, provider_clients: ProviderClients) -> bool:
+        """Check if Anthropic client is initialized."""
+        return provider_clients.anthropic_client is not None
+
+    def get_fallback_models(self) -> List[str]:
+        """Get fallback models for Anthropic."""
+        return [
+            "claude-3-7-sonnet-latest",
+            "claude-3-5-haiku-latest",
+            "claude-3-5-sonnet-latest",
+            "claude-3-opus-20240229",
+        ]

python_gpt_po/services/providers/azure_openai_provider.py

@@ -0,0 +1,47 @@
+"""
+Azure OpenAI provider implementation.
+"""
+import logging
+from typing import List
+
+from ...models.provider_clients import ProviderClients
+from .base import ModelProviderInterface
+
+
+class AzureOpenAIProvider(ModelProviderInterface):
+    """Azure OpenAI model provider implementation."""
+
+    def get_models(self, provider_clients: ProviderClients) -> List[str]:
+        """Retrieve available models from Azure OpenAI."""
+        models = []
+
+        if not self.is_client_initialized(provider_clients):
+            logging.error("Azure OpenAI client not initialized")
+            return models
+
+        try:
+            response = provider_clients.azure_openai_client.models.list()
+            models = [model.id for model in response.data]
+        except Exception as e:
+            logging.error("Error fetching Azure OpenAI models: %s", str(e))
+            models = self.get_fallback_models()
+
+        return models
+
+    def get_default_model(self) -> str:
+        """Get the default Azure OpenAI model."""
+        return "gpt-35-turbo"
+
+    def get_preferred_models(self, task: str = "translation") -> List[str]:
+        """Get preferred Azure OpenAI models for a task."""
+        if task == "translation":
+            return ["gpt-4", "gpt-35-turbo"]
+        return ["gpt-35-turbo"]
+
+    def is_client_initialized(self, provider_clients: ProviderClients) -> bool:
+        """Check if Azure OpenAI client is initialized."""
+        return provider_clients.azure_openai_client is not None
+
+    def get_fallback_models(self) -> List[str]:
+        """Get fallback models for Azure OpenAI."""
+        return ["gpt-35-turbo", "gpt-4"]

python_gpt_po/services/providers/base.py

@@ -0,0 +1,60 @@
+"""
+Base interface for model providers.
+"""
+from abc import ABC, abstractmethod
+from typing import List
+
+from ...models.provider_clients import ProviderClients
+
+
+class ModelProviderInterface(ABC):
+    """Abstract base class for model providers."""
+
+    @abstractmethod
+    def get_models(self, provider_clients: ProviderClients) -> List[str]:
+        """Retrieve available models from the provider.
+
+        Args:
+            provider_clients: Initialized provider clients
+
+        Returns:
+            List of available model IDs
+        """
+
+    @abstractmethod
+    def get_default_model(self) -> str:
+        """Get the default model for this provider.
+
+        Returns:
+            Default model ID
+        """
+
+    @abstractmethod
+    def get_preferred_models(self, task: str = "translation") -> List[str]:
+        """Get preferred models for a specific task.
+
+        Args:
+            task: The task type (default: "translation")
+
+        Returns:
+            List of preferred model IDs in order of preference
+        """
+
+    @abstractmethod
+    def is_client_initialized(self, provider_clients: ProviderClients) -> bool:
+        """Check if the provider client is initialized.
+
+        Args:
+            provider_clients: Provider clients instance
+
+        Returns:
+            True if client is initialized, False otherwise
+        """
+
+    def get_fallback_models(self) -> List[str]:
+        """Get fallback models when API calls fail.
+
+        Returns:
+            List of fallback model IDs
+        """
+        return []