cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/tasks/translation/providers/azure_provider.py

@@ -0,0 +1,192 @@
+from typing import Optional
+
+import aiohttp
+
+from cognee.shared.logging_utils import get_logger
+
+from .base import TranslationProvider, TranslationResult
+from ..config import get_translation_config
+from ..exceptions import TranslationProviderError
+
+logger = get_logger(__name__)
+
+
+class AzureTranslationProvider(TranslationProvider):
+    """
+    Translation provider using Azure Translator API.
+
+    Requires:
+    - AZURE_TRANSLATOR_KEY environment variable
+    - AZURE_TRANSLATOR_REGION environment variable (optional)
+    """
+
+    def __init__(self):
+        self._config = get_translation_config()
+
+    @property
+    def provider_name(self) -> str:
+        return "azure"
+
+    def is_available(self) -> bool:
+        """Check if Azure Translator is available."""
+        return self._config.azure_translator_key is not None
+
+    async def translate(
+        self,
+        text: str,
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> TranslationResult:
+        """
+        Translate text using Azure Translator API.
+
+        Args:
+            text: The text to translate
+            target_language: Target language code (default: "en")
+            source_language: Source language code (optional)
+
+        Returns:
+            TranslationResult with translated text and metadata
+        """
+        if not self.is_available():
+            raise TranslationProviderError(
+                provider=self.provider_name,
+                message="Azure Translator API key not configured. Set AZURE_TRANSLATOR_KEY environment variable.",
+            )
+
+        endpoint = f"{self._config.azure_translator_endpoint}/translate"
+
+        params = {
+            "api-version": "3.0",
+            "to": target_language,
+        }
+        if source_language:
+            params["from"] = source_language
+
+        headers = {
+            "Ocp-Apim-Subscription-Key": self._config.azure_translator_key,
+            "Content-Type": "application/json",
+        }
+        if self._config.azure_translator_region:
+            headers["Ocp-Apim-Subscription-Region"] = self._config.azure_translator_region
+
+        body = [{"text": text}]
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    endpoint,
+                    params=params,
+                    headers=headers,
+                    json=body,
+                    timeout=aiohttp.ClientTimeout(total=self._config.timeout_seconds),
+                ) as response:
+                    response.raise_for_status()
+                    result = await response.json()
+
+            translation = result[0]["translations"][0]
+            detected_language = result[0].get("detectedLanguage", {})
+
+            return TranslationResult(
+                translated_text=translation["text"],
+                source_language=source_language or detected_language.get("language", "unknown"),
+                target_language=target_language,
+                confidence_score=detected_language.get("score", 0.9),
+                provider=self.provider_name,
+                raw_response=result[0],
+            )
+
+        except Exception as e:
+            logger.error(f"Azure translation failed: {e}")
+            raise TranslationProviderError(
+                provider=self.provider_name,
+                message=f"Translation failed: {e}",
+                original_error=e,
+            )
+
+    async def translate_batch(
+        self,
+        texts: list[str],
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> list[TranslationResult]:
+        """
+        Translate multiple texts using Azure Translator API.
+
+        Azure Translator supports up to 100 texts per request.
+
+        Args:
+            texts: List of texts to translate
+            target_language: Target language code
+            source_language: Source language code (optional)
+
+        Returns:
+            List of TranslationResult objects
+        """
+        if not self.is_available():
+            raise TranslationProviderError(
+                provider=self.provider_name,
+                message="Azure Translator API key not configured. Set AZURE_TRANSLATOR_KEY environment variable.",
+            )
+
+        endpoint = f"{self._config.azure_translator_endpoint}/translate"
+
+        params = {
+            "api-version": "3.0",
+            "to": target_language,
+        }
+        if source_language:
+            params["from"] = source_language
+
+        headers = {
+            "Ocp-Apim-Subscription-Key": self._config.azure_translator_key,
+            "Content-Type": "application/json",
+        }
+        if self._config.azure_translator_region:
+            headers["Ocp-Apim-Subscription-Region"] = self._config.azure_translator_region
+
+        # Azure supports up to 100 texts per request
+        batch_size = min(100, self._config.batch_size)
+        all_results = []
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                for i in range(0, len(texts), batch_size):
+                    batch = texts[i : i + batch_size]
+                    body = [{"text": text} for text in batch]
+
+                    async with session.post(
+                        endpoint,
+                        params=params,
+                        headers=headers,
+                        json=body,
+                        timeout=aiohttp.ClientTimeout(total=self._config.timeout_seconds),
+                    ) as response:
+                        response.raise_for_status()
+                        results = await response.json()
+
+                    for result in results:
+                        translation = result["translations"][0]
+                        detected_language = result.get("detectedLanguage", {})
+
+                        all_results.append(
+                            TranslationResult(
+                                translated_text=translation["text"],
+                                source_language=source_language
+                                or detected_language.get("language", "unknown"),
+                                target_language=target_language,
+                                confidence_score=detected_language.get("score", 0.9),
+                                provider=self.provider_name,
+                                raw_response=result,
+                            )
+                        )
+
+        except Exception as e:
+            logger.error(f"Azure batch translation failed: {e}")
+            raise TranslationProviderError(
+                provider=self.provider_name,
+                message=f"Batch translation failed: {e}",
+                original_error=e,
+            )
+
+        return all_results

cognee/tasks/translation/providers/base.py

@@ -0,0 +1,85 @@
+"""
+Base classes for translation providers.
+
+This module defines the abstract interface that all translation providers must implement.
+Providers handle the actual translation of text using external services like OpenAI,
+Google Translate, or Azure Translator.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class TranslationResult:
+    """Result of a translation operation."""
+
+    translated_text: str
+    source_language: str
+    target_language: str
+    # Confidence score from the provider, or None if not available (e.g., Google Translate)
+    confidence_score: Optional[float]
+    provider: str
+    raw_response: Optional[dict] = None
+
+
+class TranslationProvider(ABC):
+    """Abstract base class for translation providers."""
+
+    @property
+    @abstractmethod
+    def provider_name(self) -> str:
+        """Return the name of this translation provider."""
+        pass
+
+    @abstractmethod
+    async def translate(
+        self,
+        text: str,
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> TranslationResult:
+        """
+        Translate text to the target language.
+
+        Args:
+            text: The text to translate
+            target_language: Target language code (default: "en")
+            source_language: Source language code (optional, will be auto-detected if not provided)
+
+        Returns:
+            TranslationResult with translated text and metadata
+        """
+        pass
+
+    @abstractmethod
+    async def translate_batch(
+        self,
+        texts: list[str],
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> list[TranslationResult]:
+        """
+        Translate multiple texts to the target language.
+
+        Args:
+            texts: List of texts to translate
+            target_language: Target language code (default: "en")
+            source_language: Source language code (optional)
+
+        Returns:
+            List of TranslationResult objects
+        """
+        pass
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Check if this provider is available (has required credentials).
+
+        All providers must implement this method to validate their credentials.
+
+        Returns:
+            True if the provider has valid credentials and is ready to use.
+        """
+        pass

cognee/tasks/translation/providers/google_provider.py

@@ -0,0 +1,158 @@
+import asyncio
+from typing import Optional
+
+from cognee.shared.logging_utils import get_logger
+
+from .base import TranslationProvider, TranslationResult
+from ..config import get_translation_config
+
+logger = get_logger(__name__)
+
+
+class GoogleTranslationProvider(TranslationProvider):
+    """
+    Translation provider using Google Cloud Translation API.
+
+    Requires:
+    - google-cloud-translate package
+    - GOOGLE_TRANSLATE_API_KEY or GOOGLE_PROJECT_ID environment variable
+    """
+
+    def __init__(self):
+        self._client = None
+        self._config = get_translation_config()
+
+    @property
+    def provider_name(self) -> str:
+        return "google"
+
+    def _get_client(self):
+        """Lazy initialization of Google Translate client."""
+        if self._client is None:
+            try:
+                from google.cloud import translate_v2 as translate
+
+                self._client = translate.Client()
+            except ImportError:
+                raise ImportError(
+                    "google-cloud-translate is required for Google translation. "
+                    "Install it with: pip install google-cloud-translate"
+                )
+            except Exception as e:
+                logger.error(f"Failed to initialize Google Translate client: {e}")
+                raise
+        return self._client
+
+    def is_available(self) -> bool:
+        """Check if Google Translate is available."""
+        try:
+            self._get_client()
+            return True
+        except Exception as e:
+            logger.debug(f"Google Translate not available: {e}")
+            return False
+
+    async def translate(
+        self,
+        text: str,
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> TranslationResult:
+        """
+        Translate text using Google Translate API.
+
+        Args:
+            text: The text to translate
+            target_language: Target language code (default: "en")
+            source_language: Source language code (optional)
+
+        Returns:
+            TranslationResult with translated text and metadata
+        """
+        try:
+            client = self._get_client()
+
+            # Run in thread pool since google-cloud-translate is synchronous
+            loop = asyncio.get_running_loop()
+
+            # Build kwargs for translate call
+            translate_kwargs = {"target_language": target_language}
+            if source_language:
+                translate_kwargs["source_language"] = source_language
+
+            result = await loop.run_in_executor(
+                None,
+                lambda: client.translate(text, **translate_kwargs),
+            )
+
+            detected_language = result.get("detectedSourceLanguage", source_language or "unknown")
+
+            return TranslationResult(
+                translated_text=result["translatedText"],
+                source_language=detected_language,
+                target_language=target_language,
+                # Google Translate API does not provide confidence scores
+                confidence_score=None,
+                provider=self.provider_name,
+                raw_response=result,
+            )
+
+        except Exception as e:
+            logger.error(f"Google translation failed: {e}")
+            raise
+
+    async def translate_batch(
+        self,
+        texts: list[str],
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> list[TranslationResult]:
+        """
+        Translate multiple texts using Google Translate API.
+
+        Google Translate supports batch translation natively.
+
+        Args:
+            texts: List of texts to translate
+            target_language: Target language code
+            source_language: Source language code (optional)
+
+        Returns:
+            List of TranslationResult objects
+        """
+        try:
+            client = self._get_client()
+            loop = asyncio.get_running_loop()
+
+            # Build kwargs for translate call
+            translate_kwargs = {"target_language": target_language}
+            if source_language:
+                translate_kwargs["source_language"] = source_language
+
+            results = await loop.run_in_executor(
+                None,
+                lambda: client.translate(texts, **translate_kwargs),
+            )
+
+            translation_results = []
+            for result in results:
+                detected_language = result.get(
+                    "detectedSourceLanguage", source_language or "unknown"
+                )
+                translation_results.append(
+                    TranslationResult(
+                        translated_text=result["translatedText"],
+                        source_language=detected_language,
+                        target_language=target_language,
+                        # Google Translate API does not provide confidence scores
+                        confidence_score=None,
+                        provider=self.provider_name,
+                        raw_response=result,
+                    )
+                )
+
+            return translation_results
+
+        except Exception as e:
+            logger.error(f"Google batch translation failed: {e}")
+            raise

cognee/tasks/translation/providers/llm_provider.py

@@ -0,0 +1,143 @@
+import asyncio
+from typing import Optional
+
+from pydantic import BaseModel
+
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.shared.logging_utils import get_logger
+
+from .base import TranslationProvider, TranslationResult
+
+logger = get_logger(__name__)
+
+
+class TranslationOutput(BaseModel):
+    """Pydantic model for structured translation output from LLM."""
+
+    translated_text: str
+    detected_source_language: str
+    translation_notes: Optional[str] = None
+
+
+class LLMTranslationProvider(TranslationProvider):
+    """
+    Translation provider using the configured LLM for translation.
+
+    This provider leverages the existing LLM infrastructure in Cognee
+    to perform translations using any LLM configured via LLM_PROVIDER
+    (OpenAI, Azure, Ollama, Anthropic, etc.).
+
+    The LLM used is determined by the cognee LLM configuration settings:
+    - LLM_PROVIDER: The LLM provider (openai, azure, ollama, etc.)
+    - LLM_MODEL: The model to use
+    - LLM_API_KEY: API key for the provider
+    """
+
+    @property
+    def provider_name(self) -> str:
+        """Return 'llm' as the provider name."""
+        return "llm"
+
+    async def translate(
+        self,
+        text: str,
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+    ) -> TranslationResult:
+        """
+        Translate text using the configured LLM.
+
+        Args:
+            text: The text to translate
+            target_language: Target language code (default: "en")
+            source_language: Source language code (optional)
+
+        Returns:
+            TranslationResult with translated text and metadata
+        """
+        try:
+            system_prompt = read_query_prompt("translate_content.txt")
+
+            # Validate system prompt was loaded successfully
+            if system_prompt is None:
+                logger.warning("translate_content.txt prompt file not found, using default prompt")
+                system_prompt = (
+                    "You are a professional translator. Translate the given text accurately "
+                    "while preserving the original meaning, tone, and style. "
+                    "Detect the source language if not provided."
+                )
+
+            # Build the input with context
+            if source_language:
+                input_text = (
+                    f"Translate the following text from {source_language} to {target_language}.\n\n"
+                    f"Text to translate:\n{text}"
+                )
+            else:
+                input_text = (
+                    f"Translate the following text to {target_language}. "
+                    f"First detect the source language.\n\n"
+                    f"Text to translate:\n{text}"
+                )
+
+            result = await LLMGateway.acreate_structured_output(
+                text_input=input_text,
+                system_prompt=system_prompt,
+                response_model=TranslationOutput,
+            )
+
+            return TranslationResult(
+                translated_text=result.translated_text,
+                source_language=source_language or result.detected_source_language,
+                target_language=target_language,
+                # TODO: Consider deriving confidence from LLM response metadata
+                # or making configurable via TranslationConfig
+                confidence_score=0.95,  # LLM translations are generally high quality
+                provider=self.provider_name,
+                raw_response={"notes": result.translation_notes},
+            )
+
+        except Exception as e:
+            logger.error(f"LLM translation failed: {e}")
+            raise
+
+    async def translate_batch(
+        self,
+        texts: list[str],
+        target_language: str = "en",
+        source_language: Optional[str] = None,
+        max_concurrent: int = 5,
+    ) -> list[TranslationResult]:
+        """
+        Translate multiple texts using the configured LLM.
+
+        Uses a semaphore to limit concurrent requests and avoid API rate limits.
+
+        Args:
+            texts: List of texts to translate
+            target_language: Target language code
+            source_language: Source language code (optional)
+            max_concurrent: Maximum concurrent translation requests (default: 5)
+
+        Returns:
+            List of TranslationResult objects
+        """
+        semaphore = asyncio.Semaphore(max_concurrent)
+
+        async def limited_translate(text: str) -> TranslationResult:
+            async with semaphore:
+                return await self.translate(text, target_language, source_language)
+
+        tasks = [limited_translate(text) for text in texts]
+        return await asyncio.gather(*tasks)
+
+    def is_available(self) -> bool:
+        """Check if LLM provider is available (has required credentials)."""
+        try:
+            llm_config = get_llm_config()
+            # Check if API key is configured (required for most providers)
+            return bool(llm_config.llm_api_key)
+        except Exception:
+            return False