translator.py

import os
import queue
import threading
import uuid
from datetime import date, datetime, timezone
from pathlib import Path
from typing import Optional, Union

import deepl
from dotenv import load_dotenv

from db import get_connection, init_db, log_translation

load_dotenv()

_client: Optional[deepl.DeepLClient] = None
_client_lock = threading.Lock()
_db_conn = None

# Bounded queue: entries dropped (not blocking) when full, so translation is never impacted.
_log_queue: queue.Queue = queue.Queue(maxsize=5000)


def _get_client() -> deepl.DeepLClient:
    global _client
    if _client is None:
        _client = deepl.DeepLClient(os.environ["DEEPL_API_KEY"])
    return _client


def _get_db():
    global _db_conn
    if _db_conn is None:
        _db_conn = get_connection()
        init_db(_db_conn)
    return _db_conn


def _log_worker() -> None:
    db = _get_db()
    while True:
        row = _log_queue.get()
        try:
            log_translation(db, row)
        except Exception as e:
            print(f"[deepl-analytics] log_translation failed: {e}", flush=True)
        finally:
            _log_queue.task_done()


# Daemon thread exits automatically when the main process ends.
_worker_thread = threading.Thread(target=_log_worker, daemon=True)
_worker_thread.start()


def translate(
    texts: Union[str, list[str]],
    target_lang: str,
    reporting_tag: Optional[str] = None,
    **kwargs,
) -> list[deepl.TextResult]:
    if isinstance(texts, str):
        texts = [texts]

    request_meta = {
        "logged_at":      datetime.now(timezone.utc),
        "logged_date":    date.today(),
        "target_lang":    target_lang,
        "api_key_alias":     os.environ.get("DEEPL_API_KEY_ALIAS", "unknown"),
        "request_id":     str(uuid.uuid4()),
        "reporting_tag":  reporting_tag,
    }

    # Lock covers header mutation + HTTP call so concurrent calls don't clobber each other's tag.
    with _client_lock:
        client = _get_client()
        if reporting_tag is not None:
            client.headers["X-DeepL-Reporting-Tag"] = reporting_tag
        elif "X-DeepL-Reporting-Tag" in client.headers:
            del client.headers["X-DeepL-Reporting-Tag"]

        try:
            result = client.translate_text(texts, target_lang=target_lang, **kwargs)
        except Exception as e:
            try:
                _log_queue.put_nowait({
                    **request_meta,
                    "source_lang":        None,
                    "billed_characters":  None,
                    "is_success":         False,
                    "error_code":         type(e).__name__,
                    "error_http_status":  getattr(e, "http_status_code", None),
                    "error_message":      str(e),
                    "translation_type":   "Text",
                })
            except queue.Full:
                pass
            raise

    for translation in result:
        try:
            _log_queue.put_nowait({
                **request_meta,
                "source_lang":       translation.detected_source_lang,
                "billed_characters": translation.billed_characters,
                "is_success":        True,
                "translation_type":  "Text",
            })
        except queue.Full:
            pass  # Drop the entry rather than block the caller.

    return result


_DOCUMENT_TYPE_MAP = {
    ".txt":   "TXT",
    ".srt":   "SRT",
    ".pdf":   "PDF",
    ".docx":  "DOCX",
    ".xlsx":  "XLSX",
    ".pptx":  "PPTX",
    ".xliff": "XLIFF",
    ".html":  "HTML",
    ".htm":   "HTML",
    ".png":   "IMG",
    ".jpg":   "IMG",
    ".jpeg":  "IMG",
}


def translate_document(
    input_path: str,
    output_path: str,
    target_lang: str,
    reporting_tag: Optional[str] = None,
    **kwargs,
) -> None:
    ext = Path(input_path).suffix.lower()
    document_type = _DOCUMENT_TYPE_MAP.get(ext, None)

    request_meta = {
        "logged_at":      datetime.now(timezone.utc),
        "logged_date":    date.today(),
        "target_lang":    target_lang,
        "api_key_alias":     os.environ.get("DEEPL_API_KEY_ALIAS", "unknown"),
        "request_id":     str(uuid.uuid4()),
        "reporting_tag":  reporting_tag,
    }

    # Lock covers header mutation + HTTP call so concurrent calls don't clobber each other's tag.
    with _client_lock:
        client = _get_client()
        if reporting_tag is not None:
            client.headers["X-DeepL-Reporting-Tag"] = reporting_tag
        elif "X-DeepL-Reporting-Tag" in client.headers:
            del client.headers["X-DeepL-Reporting-Tag"]

        try:
            with open(input_path, "rb") as in_file, open(output_path, "wb") as out_file:
                status = client.translate_document(in_file, out_file, target_lang=target_lang, **kwargs)
        except Exception as e:
            try:
                _log_queue.put_nowait({
                    **request_meta,
                    "source_lang":        None,
                    "billed_characters":  None,
                    "is_success":         False,
                    "error_code":         type(e).__name__,
                    "error_http_status":  getattr(e, "http_status_code", None),
                    "error_message":      str(e),
                    "translation_type":   "Document",
                    "document_type":      document_type,
                })
            except queue.Full:
                pass
            raise

    source_lang = getattr(status, "detected_source_language", None)
    billed_characters = status.billed_characters

    try:
        _log_queue.put_nowait({
            **request_meta,
            "source_lang":       source_lang,
            "billed_characters": billed_characters,
            "is_success":        True,
            "translation_type":  "Document",
            "document_type":     document_type,
        })
    except queue.Full:
        pass


def query_by_language_pair(
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
):
    db = _get_db()
    where, params = "", []

    if start_date and end_date:
        where, params = "WHERE logged_at BETWEEN ? AND ?", [start_date, end_date]
    elif start_date:
        where, params = "WHERE logged_at >= ?", [start_date]
    elif end_date:
        where, params = "WHERE logged_at <= ?", [end_date]

    return db.execute(
        f"""
        SELECT
            source_lang,
            target_lang,
            SUM(billed_characters) AS total_billed_chars,
            COUNT(*)               AS request_count
        FROM translation_usage
        {where}
        GROUP BY source_lang, target_lang
        ORDER BY total_billed_chars DESC
        """,
        params,
    ).fetchdf()