Japanese-VTT-Translator/ollama_client.py

"""
Ollama API client for translation requests.
"""

import requests
from typing import Optional


class OllamaClient:
    """Client for communicating with Ollama server."""

    MASTER_PROMPT = (
        "You are a professional Japanese (日本語) to English translator. "
        "Your task is to translate the following Japanese (日本語) text into English. "
        "IMPORTANT: You MUST always provide a translation, even if the text is fragmented, unclear, or lacks context. "
        "Do not refuse to translate or provide explanations - translate everything provided, making your best interpretation. "
        "Convey the meaning and nuances of the original text while adhering to English grammar and vocabulary. "
        "Output ONLY the English translation with no explanations, commentary, or meta-discussion. "
        "Translate the following Japanese (日本語) text into English:"
    )

    def __init__(self, base_url: str, model: str):
        """
        Initialize Ollama client.

        Args:
            base_url: Base URL for Ollama server (e.g., http://ai-house:11434/)
            model: Model name (e.g., translategemma:12b)
        """
        self.base_url = base_url.rstrip('/')
        self.model = model
        self.api_endpoint = f"{self.base_url}/api/generate"

    def translate(self, text: str, timeout: int = 300) -> Optional[str]:
        """
        Translate Japanese text to English using Ollama.

        Args:
            text: Japanese text to translate
            timeout: Request timeout in seconds

        Returns:
            Translated English text or None if translation fails
        """
        prompt = f"{self.MASTER_PROMPT}\n\n{text}"

        try:
            import sys
            print(f"[Ollama] Sending request to {self.api_endpoint}...", file=sys.stderr)
            sys.stderr.flush()

            response = requests.post(
                self.api_endpoint,
                json={
                    "model": self.model,
                    "prompt": prompt,
                    "stream": False,
                },
                timeout=timeout
            )

            print(f"[Ollama] Got response, status: {response.status_code}", file=sys.stderr)
            sys.stderr.flush()

            response.raise_for_status()

            result = response.json()
            if 'response' in result:
                translated = result['response'].strip()
                print(f"[Ollama] Translation complete ({len(translated)} chars)", file=sys.stderr)
                sys.stderr.flush()
                return translated
            return None

        except requests.exceptions.Timeout as e:
            print(f"[Ollama] Timeout after {timeout}s: {e}", file=sys.stderr)
            return None
        except requests.exceptions.RequestException as e:
            print(f"[Ollama] Error: {e}", file=sys.stderr)
            return None
        except (KeyError, ValueError) as e:
            print(f"[Ollama] Parse error: {e}", file=sys.stderr)
            return None

    def is_available(self) -> bool:
        """Check if Ollama server is available."""
        try:
            response = requests.get(
                f"{self.base_url}/api/tags",
                timeout=5
            )
            return response.status_code == 200
        except requests.exceptions.RequestException:
            return False

    def get_model_info(self) -> Optional[dict]:
        """Get information about the model."""
        try:
            response = requests.post(
                f"{self.base_url}/api/show",
                json={"name": self.model},
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException:
            return None