Merge pull request #52 from vbossica/move_translation_method

Move the _get_*_response methods inside the provider classes

Author: brammittendorff

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
     "anthropic==0.48.0",
     "requests==2.32.3",
     "responses==0.25.6",
+    "isort==6.0.1",
 ]
 classifiers = [
     "Development Status :: 5 - Production/Stable",

python_gpt_po/services/providers/anthropic_provider.py

@@ -62,3 +62,16 @@ def get_fallback_models(self) -> List[str]:
             "claude-3-5-sonnet-latest",
             "claude-3-opus-20240229",
         ]
+
+    def translate(self, provider_clients: ProviderClients, model: str, content: str) -> str:
+        """Get response from Anthropic API."""
+        if not self.is_client_initialized(provider_clients):
+            raise ValueError("Anthropic client not initialized")
+
+        message = {"role": "user", "content": content}
+        completion = provider_clients.anthropic_client.messages.create(
+            model=model,
+            max_tokens=4000,
+            messages=[message]
+        )
+        return completion.content[0].text.strip()

python_gpt_po/services/providers/azure_openai_provider.py

@@ -45,3 +45,16 @@ def is_client_initialized(self, provider_clients: ProviderClients) -> bool:
     def get_fallback_models(self) -> List[str]:
         """Get fallback models for Azure OpenAI."""
         return ["gpt-35-turbo", "gpt-4"]
+
+    def translate(self, provider_clients: ProviderClients, model: str, content: str) -> str:
+        """Get response from OpenAI API."""
+        if not self.is_client_initialized(provider_clients):
+            raise ValueError("OpenAI client not initialized")
+
+        message = {"role": "user", "content": content}
+        completion = provider_clients.azure_openai_client.chat.completions.create(
+            model=model,
+            max_tokens=4000,
+            messages=[message]
+        )
+        return completion.choices[0].message.content.strip()

python_gpt_po/services/providers/base.py

@@ -58,3 +58,16 @@ def get_fallback_models(self) -> List[str]:
             List of fallback model IDs
         """
         return []
+
+    @abstractmethod
+    def translate(self, provider_clients: ProviderClients, model: str, content: str) -> str:
+        """Translate content using the specified model.
+
+        Args:
+            provider_clients: Provider clients instance
+            model: Model to use for translation
+            content: Content to translate
+
+        Returns:
+            Translated content
+        """

python_gpt_po/services/providers/deepseek_provider.py

@@ -54,3 +54,26 @@ def is_client_initialized(self, provider_clients: ProviderClients) -> bool:
     def get_fallback_models(self) -> List[str]:
         """Get fallback models for DeepSeek."""
         return ["deepseek-chat", "deepseek-coder"]
+
+    def translate(self, provider_clients: ProviderClients, model: str, content: str) -> str:
+        """Get response from DeepSeek API."""
+        if not self.is_client_initialized(provider_clients):
+            raise ValueError("DeepSeek client not initialized")
+
+        headers = {
+            "Authorization": f"Bearer {provider_clients.deepseek_api_key}",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "model": model,
+            "messages": [{"role": "user", "content": content}],
+            "max_tokens": 4000
+        }
+        response = requests.post(
+            f"{provider_clients.deepseek_base_url}/chat/completions",
+            headers=headers,
+            json=payload,
+            timeout=30
+        )
+        response.raise_for_status()
+        return response.json()["choices"][0]["message"]["content"].strip()

python_gpt_po/services/providers/openai_provider.py

@@ -51,3 +51,15 @@ def get_fallback_models(self) -> List[str]:
             "gpt-4",
             "gpt-3.5-turbo"
         ]
+
+    def translate(self, provider_clients: ProviderClients, model: str, content: str) -> str:
+        """Get response from OpenAI API."""
+        if not self.is_client_initialized(provider_clients):
+            raise ValueError("OpenAI client not initialized")
+
+        message = {"role": "user", "content": content}
+        completion = provider_clients.openai_client.chat.completions.create(
+            model=model,
+            messages=[message]
+        )
+        return completion.choices[0].message.content.strip()

python_gpt_po/services/translation_service.py

@@ -10,13 +10,12 @@
 from typing import Any, Dict, List, Optional
 
 import polib
-import requests
 from tenacity import retry, stop_after_attempt, wait_fixed
 
 from ..models.config import TranslationConfig
-from ..models.enums import ModelProvider
 from .model_manager import ModelManager
 from .po_file_handler import POFileHandler
+from .providers.registry import ProviderRegistry
 
 
 class TranslationService:
@@ -35,67 +34,6 @@ def __init__(self, config: TranslationConfig, batch_size: int = 40):
         self.po_file_handler = POFileHandler()
         self.model_manager = ModelManager()
 
-    def _get_openai_response(self, content: str) -> str:
-        """Get response from OpenAI API."""
-        if not self.config.provider_clients.openai_client:
-            raise ValueError("OpenAI client not initialized")
-
-        message = {"role": "user", "content": content}
-        completion = self.config.provider_clients.openai_client.chat.completions.create(
-            model=self.config.model,
-            messages=[message]
-        )
-        return completion.choices[0].message.content.strip()
-
-    def _get_anthropic_response(self, content: str) -> str:
-        """Get response from Anthropic API."""
-        if not self.config.provider_clients.anthropic_client:
-            raise ValueError("Anthropic client not initialized")
-
-        message = {"role": "user", "content": content}
-        completion = self.config.provider_clients.anthropic_client.messages.create(
-            model=self.config.model,
-            max_tokens=4000,
-            messages=[message]
-        )
-        return completion.content[0].text.strip()
-
-    def _get_deepseek_response(self, content: str) -> str:
-        """Get response from DeepSeek API."""
-        if not self.config.provider_clients.deepseek_api_key:
-            raise ValueError("DeepSeek API key not set")
-
-        headers = {
-            "Authorization": f"Bearer {self.config.provider_clients.deepseek_api_key}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "model": self.config.model,
-            "messages": [{"role": "user", "content": content}],
-            "max_tokens": 4000
-        }
-        response = requests.post(
-            f"{self.config.provider_clients.deepseek_base_url}/chat/completions",
-            headers=headers,
-            json=payload,
-            timeout=30
-        )
-        response.raise_for_status()
-        return response.json()["choices"][0]["message"]["content"].strip()
-
-    def _get_azure_openai_response(self, content: str) -> str:
-        """Get response from OpenAI API."""
-        if not self.config.provider_clients.azure_openai_client:
-            raise ValueError("OpenAI client not initialized")
-
-        message = {"role": "user", "content": content}
-        completion = self.config.provider_clients.azure_openai_client.chat.completions.create(
-            model=self.config.model,
-            max_tokens=4000,
-            messages=[message]
-        )
-        return completion.choices[0].message.content.strip()
-
     def validate_provider_connection(self) -> bool:
         """Validates the connection to the selected provider by making a test API call."""
         provider = self.config.provider
@@ -238,15 +176,13 @@ def _get_provider_response(self, content: str) -> str:
         """Get translation response from the selected provider."""
         provider = self.config.provider
 
-        if provider == ModelProvider.OPENAI:
-            return self._get_openai_response(content)
-        if provider == ModelProvider.ANTHROPIC:
-            return self._get_anthropic_response(content)
-        if provider == ModelProvider.DEEPSEEK:
-            return self._get_deepseek_response(content)
-        if provider == ModelProvider.AZURE_OPENAI:
-            return self._get_azure_openai_response(content)
-        return ""
+        if not provider:
+            return ""
+
+        provider_instance = ProviderRegistry.get_provider(provider)
+        if not provider_instance:
+            return ""
+        return provider_instance.translate(self.config.provider_clients, self.config.model, content)
 
     def _process_bulk_response(self, response_text: str, original_texts: List[str]) -> List[str]:
         """Process a bulk translation response."""

python_gpt_po/tests/providers/__init__.py

@@ -0,0 +1,32 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from python_gpt_po.models.provider_clients import ProviderClients
+from python_gpt_po.services.providers.anthropic_provider import AnthropicProvider
+
+
+@pytest.fixture
+def mock_provider_clients() -> ProviderClients:
+    """Mock provider clients for testing."""
+    clients = ProviderClients()
+    clients.anthropic_client = MagicMock()
+    clients.anthropic_client.api_key = "sk-ant-mock-key"
+    return clients
+
+def test_translate(mock_provider_clients: ProviderClients) -> None:
+    """Test bulk translation with Anthropic."""
+    # Setup mock response
+    mock_chatcompletion = MagicMock()
+    mock_chatcompletion.content = [MagicMock()]
+    mock_chatcompletion.content[0].text = '["Bonjour", "Monde", "Bienvenue dans notre application", "Au revoir"]'
+    mock_provider_clients.anthropic_client.messages.create.return_value = mock_chatcompletion
+
+    provider = AnthropicProvider()
+    translations = provider.translate(
+        provider_clients=mock_provider_clients,
+        model="gpt-4",
+        content="['Hello', 'World', 'Welcome to our application', 'Goodbye']"
+    )
+
+    assert translations == '["Bonjour", "Monde", "Bienvenue dans notre application", "Au revoir"]'

python_gpt_po/tests/providers/test_anthropic_provider.py

@@ -0,0 +1,32 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from python_gpt_po.models.provider_clients import ProviderClients
+from python_gpt_po.services.providers.azure_openai_provider import AzureOpenAIProvider
+
+
+@pytest.fixture
+def mock_provider_clients() -> ProviderClients:
+    """Mock provider clients for testing."""
+    clients = ProviderClients()
+    clients.azure_openai_client = MagicMock()
+    clients.azure_openai_client.api_key = "sk-aoi-mock-key"
+    return clients
+
+def test_translate(mock_provider_clients: ProviderClients) -> None:
+    """Test bulk translation with Azure OpenAI."""
+    # Setup mock response
+    mock_chatcompletion = MagicMock()
+    mock_chatcompletion.choices = [MagicMock()]
+    mock_chatcompletion.choices[0].message.content = '["Bonjour", "Monde", "Bienvenue dans notre application", "Au revoir"]'
+    mock_provider_clients.azure_openai_client.chat.completions.create.return_value = mock_chatcompletion
+
+    provider = AzureOpenAIProvider()
+    translations = provider.translate(
+        provider_clients=mock_provider_clients,
+        model="gpt-4",
+        content="['Hello', 'World', 'Welcome to our application', 'Goodbye']"
+    )
+
+    assert translations == '["Bonjour", "Monde", "Bienvenue dans notre application", "Au revoir"]'