webmail-client/backend/app/translation.py

"""Email translation, on this server: the translator service (translator/:
Mozilla's Firefox Translations models and fastText language identification)
in its own container, reached only by the API. A message's text never leaves
the server.

- A message's language is detected the first time it is shown (from its own
  text: quoted replies and signatures left out) and kept on the message, so the
  reading pane can offer "Translate" when it differs from the interface language.
- A translation is of the message HTML the user sees (already sanitized; the
  result is sanitized again), cached per message, target language and whether
  remote images were shown. It is made paragraph by paragraph, as plain text
  (translation_prep.split_pieces), each paragraph from its own language.
- The languages on offer are those installed on the service (admin panel).

Off when ``TRANSLATOR_URL`` is empty.
"""

from __future__ import annotations

import html as html_lib
import logging
import re
import time
from collections.abc import Iterable
from dataclasses import dataclass

import httpx

from app import translation_prep as prep
from app.completion import own_text
from app.config import get_settings

logger = logging.getLogger(__name__)

DETECT_CHARS = 1500  # of the message's own text: plenty to tell the language
MIN_CONFIDENCE = 50.0  # below this, the detector is guessing
MAX_HTML_CHARS = 150_000  # a translation request bigger than this is refused
_TIMEOUT = httpx.Timeout(120.0, connect=5.0)  # long messages take a while on a CPU
_LANG_CACHE_S = 60


class TranslationError(Exception):
    """Translation is off, the service is unreachable, or it refused the text."""


@dataclass(frozen=True)
class Language:
    code: str
    name: str


def enabled() -> bool:
    return bool(get_settings().translator_url.strip())


def _url(path: str) -> str:
    return get_settings().translator_url.rstrip("/") + path


async def _call(method: str, path: str, payload: dict[str, object] | None = None) -> object:
    if not enabled():
        raise TranslationError("translation is not set up on this server")
    try:
        async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
            resp = await client.request(method, _url(path), json=payload)
    except httpx.HTTPError as exc:
        raise TranslationError(f"the translation service is unreachable: {exc}") from exc
    if resp.status_code not in (200, 202):
        detail = resp.text[:200]
        raise TranslationError(f"the translation service refused: {resp.status_code} {detail}")
    return resp.json()


async def _post(path: str, payload: dict[str, object]) -> object:
    return await _call("POST", path, payload)


_languages: tuple[float, list[Language]] = (0.0, [])


async def languages() -> list[Language]:
    """The installed languages, English among them (cached for a minute). The
    browser names them in the interface language; ``name`` is the code."""
    global _languages
    if not enabled():
        return []
    stamp, cached = _languages
    if cached and time.monotonic() - stamp < _LANG_CACHE_S:
        return cached
    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0, connect=3.0)) as client:
            resp = await client.get(_url("/languages"))
        resp.raise_for_status()
        codes = resp.json()["installed"]
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
        logger.info("translation languages unavailable: %s", exc)
        return []
    found = [Language(code=str(c), name=str(c)) for c in codes]
    _languages = (time.monotonic(), found)
    return found


def forget_languages() -> None:
    """After an install or removal: ask the service again."""
    global _languages
    _languages = (0.0, [])


async def status() -> dict[str, object]:
    """The service's own account (admin panel): installed, installing (bytes
    done of total), available from Mozilla, and why an install failed."""
    async with httpx.AsyncClient(timeout=httpx.Timeout(20.0, connect=3.0)) as client:
        try:
            resp = await client.get(_url("/languages"))
            resp.raise_for_status()
            data = resp.json()
        except (httpx.HTTPError, ValueError) as exc:
            raise TranslationError(f"the translation service is unreachable: {exc}") from exc
    return data if isinstance(data, dict) else {}


async def install(code: str) -> None:
    await _call("POST", f"/languages/{code}")
    forget_languages()


async def remove(code: str) -> None:
    await _call("DELETE", f"/languages/{code}")
    forget_languages()


_CJK = re.compile(r"[\u3040-\u30ff\u3400-\u9fff\uac00-\ud7af\uf900-\ufaff]")
MIN_LETTERS = 25  # fewer, and the detector guesses ("Jacqueline Lian" → Esperanto)
MAX_DETECT = 60  # pieces of one message asked about
MIN_SHARE = 0.15  # of a message's letters, for its language to count


async def _detect_many(texts: list[str]) -> list[str]:
    """The language of each of ``texts`` ('' where unsure), in one request."""
    if not texts:
        return []
    result = await _post("/detect", {"q": [t[:DETECT_CHARS] for t in texts]})
    if not isinstance(result, list) or len(result) != len(texts):
        return [""] * len(texts)
    return [
        str(r.get("language", ""))
        if isinstance(r, dict) and float(r.get("confidence", 0)) >= MIN_CONFIDENCE
        else ""
        for r in result
    ]


async def _detect_one(text: str) -> str:
    return (await _detect_many([text]))[0]


def _weight(text: str) -> int:
    """How much text this is: letters, a Chinese/Japanese/Korean sign as a word."""
    return sum(3 if _CJK.match(c) else 1 for c in text if c.isalpha())


async def languages_of(parts: list[str]) -> list[str]:
    """The language of each of ``parts`` (paragraphs, in order). Parts too
    short to tell take the language of the next part that could be told (a
    greeting goes with what follows it), the last ones that of the one before;
    parts without letters get ''."""
    # Addresses say nothing about a language ("E-mail: lianshaojuan.lsj@…" → Estonian).
    parts = [_ADDRESS.sub(" ", part) for part in parts]
    told: list[str] = [""] * len(parts)
    asked = [
        i
        for i, part in enumerate(parts)
        if _CJK.search(part) or sum(c.isalpha() for c in part) >= MIN_LETTERS
    ][:MAX_DETECT]
    for i, language in zip(asked, await _detect_many([parts[i] for i in asked]), strict=True):
        told[i] = language
    out = list(told)
    following = ""
    for i in reversed(range(len(parts))):
        if told[i]:
            following = told[i]
        elif any(c.isalpha() for c in parts[i]):
            out[i] = following
    previous = ""
    for i, part in enumerate(parts):
        if told[i]:
            previous = told[i]
        elif not out[i] and any(c.isalpha() for c in part):
            out[i] = previous
    return out


async def detect(text: str) -> str:
    """The languages of ``text`` (a message's plain text), most of it first,
    comma-separated ("zh-Hans,pl,en"); '' if unsure. A message can be in
    several: each language that makes up a good part of it counts."""
    own = own_text(text)
    if len(" ".join(own.split())) < 20:
        return ""  # a line or two says too little
    lines: list[str] = []
    total = 0
    for line in (x.strip() for x in own.splitlines()):
        if line and total < DETECT_CHARS * 4:
            lines.append(line)
            total += len(line)
    found = await languages_of(lines)
    weight: dict[str, int] = {}
    for line, language in zip(lines, found, strict=True):
        if language:
            weight[language] = weight.get(language, 0) + _weight(_ADDRESS.sub(" ", line))
    if not weight:
        return await _detect_one(" ".join(own.split()))
    all_of_it = sum(weight.values())
    ranked = sorted(weight, key=lambda k: weight[k], reverse=True)
    return ",".join(k for k in ranked if weight[k] / all_of_it >= MIN_SHARE)


def _base(code: str) -> str:
    return code.lower().split("-")[0]


async def _texts(texts: list[str], source: str, target: str) -> list[str]:
    """``texts`` translated, in batches."""
    out: list[str] = []
    for start in range(0, len(texts), 100):
        batch = texts[start : start + 100]
        result = await _post("/translate", {"q": batch, "source": source, "target": target})
        got = result.get("translatedText") if isinstance(result, dict) else None
        if not isinstance(got, list) or len(got) != len(batch):
            raise TranslationError("the translation service sent no translation")
        out += [str(x) for x in got]
    return out


def _kept(name: str, translated: str) -> bool:
    """Whether ``name`` came through a translation: each of its words still
    there, give or take an ending ("Lukas" → "Lukasie": yes; "Milan" →
    "Mediolan": no)."""
    return all(
        word[: max(3, len(word) - 2)] in translated for word in name.split() if len(word) >= 2
    )


# A capitalised word that starts a sentence ("Enclosed is a summary…").
_SENTENCE_START = re.compile(r"(?:^|(?<=[.!?…]\s))(?P<word>[^\W\d_][^\W\d_'’-]{2,})\b")


def _copied_starts(core: str, text: str, names: list[str]) -> list[str]:
    """Words that start a sentence of ``core`` and came through ``text``
    untranslated, and are no known name: the models take a capitalised
    sentence start for a name ("Enclosed is…" → "Enclosed jest…")."""
    found: list[str] = []
    for m in _SENTENCE_START.finditer(core):
        word = m.group("word")
        if not word[0].isupper() or word.isupper() or any(word in n for n in names):
            continue
        if re.search(rf"(?<!\w){re.escape(word)}(?!\w)", text):
            found.append(word)
    return found


def _lowered(core: str, words: list[str]) -> str:
    """``core`` with ``words`` lower-cased where they start a sentence."""
    return _SENTENCE_START.sub(
        lambda m: m.group(0).lower() if m.group("word") in words else m.group(0), core
    )


def _capitalised(text: str) -> str:
    """Sentence starts capitalised again (after ``_lowered``): the first
    letter, and one after a sentence's end — a full stop after a word of four
    letters or more, so not after "r.", "np.", "tzw." (abbreviations)."""
    text = re.sub(r"^(\W*)([^\W\d_])", lambda m: m.group(1) + m.group(2).upper(), text)
    return re.sub(
        r"((?:[^\W\d_]{4,}[.…]|[!?])\s+)([^\W\d_])",
        lambda m: m.group(1) + m.group(2).upper(),
        text,
    )


def _as_before(text: str, first: str) -> str:
    """A word after a comma capitalised in ``text`` but not in ``first`` (the
    first try) is lower-cased: lower-cased input makes the models capitalise
    elsewhere ("…2026, od 13 lipca" → "…2026 r., Od 13 lipca")."""

    def repl(m: re.Match[str]) -> str:
        word = m.group(2)
        lower = word[0].lower() + word[1:]
        return m.group(1) + (lower if f", {lower}" in first else word)

    return re.sub(r"(,\s+)([^\W\d_][\w'’-]*)", repl, text)


async def _with_names_kept(core: str, names: list[str], source: str, target: str) -> str:
    """``core`` translated around ``names``, which are left as they are."""
    pattern = re.compile("(" + "|".join(re.escape(n) for n in names) + ")")
    parts = pattern.split(core)
    todo = [
        (i, m.groups())
        for i, part in enumerate(parts)
        if i % 2 == 0 and (m := _EDGE_SPACE.match(part)) and any(c.isalnum() for c in part)
    ]
    done = await _texts([groups[1] for _, groups in todo], source, target)
    for (i, (lead, _, trail)), text in zip(todo, done, strict=True):
        parts[i] = f"{lead}{text}{trail}"
    return "".join(parts)


async def _translate_paragraphs(
    html: str, target: str, sources: list[str], names: Iterable[str] = ()
) -> str:
    """``html`` in ``target``, as plain text paragraph by paragraph: from the
    one language in ``sources``, or (several, or none known) each paragraph
    from its own. Paragraphs already in ``target``, or in a language that is
    not installed, are left as they are.

    ``names`` go along with their sentence, which then translates as a whole
    (kept apart, "I came across <Imaginary Pixels>" lost its verb: "Przyszedłem
    Imaginary Pixels"). Only a piece where one came out changed ("Milan" →
    "Mediolan") is translated again with it left as it is."""
    watch = sorted({n.strip() for n in names if len(n.strip()) >= 2}, key=len, reverse=True)
    pieces, paragraphs = prep.split_pieces(html)
    if len(sources) == 1:
        found = [sources[0]] * len(paragraphs)
    else:
        found = await languages_of([p.text for p in paragraphs])
    installed = {x.code for x in await languages()}
    # Chinese and Japanese go without spaces: two pieces of one sentence
    # ("…奎琳，</span>我从…") need one between them once translated.
    followed = {i for p in paragraphs for i in p.pieces[:-1]}
    by_language: dict[str, list[int]] = {}
    for paragraph, language in zip(paragraphs, found, strict=True):
        if language in installed and _base(language) != _base(target):
            by_language.setdefault(language, []).extend(paragraph.pieces)
    for language, where in by_language.items():
        todo: list[tuple[int, str, str, str]] = []
        for i in where:
            m = _EDGE_SPACE.match(html_lib.unescape(pieces[i]))
            lead, core, trail = m.groups() if m else ("", "", "")
            # A list marker stays out: the models repeat it ("1. Thursday" →
            # "1. 1. Czwartek").
            marker = _LIST_MARKER.match(core)
            if marker:
                lead, core = lead + marker.group(0), core[marker.end() :]
            if not trail and i in followed and _base(language) in _NO_SPACES:
                trail = " "
            if any(c.isalnum() for c in core):
                todo.append((i, lead, core, trail))
        done = await _texts([core for _, _, core, _ in todo], language, target)
        # Sentence starts that came through untranslated: once more, lower-case.
        again = [
            (k, words)
            for k, ((_, _, core, _), text) in enumerate(zip(todo, done, strict=True))
            if (words := _copied_starts(core, text, watch))
        ]
        if again:
            retried = await _texts(
                [_lowered(todo[k][2], words) for k, words in again], language, target
            )
            for (k, words), text in zip(again, retried, strict=True):
                # Still there: a name after all, and the first try stands.
                if len(_copied_starts(todo[k][2], text, watch)) < len(words):
                    done[k] = _as_before(_capitalised(text), done[k])
        for (i, lead, core, trail), text in zip(todo, done, strict=True):
            lost = [n for n in watch if n in core and not _kept(n, text)]
            if lost:
                text = await _with_names_kept(core, lost, language, target)
            pieces[i] = html_lib.escape(f"{lead}{text}{trail}", quote=False)
    return "".join(pieces)


_NO_SPACES = {"zh", "ja", "th"}
_LIST_MARKER = re.compile(r"(?:\d{1,3}[.)]|[a-zA-Z][.)]|[-•*–])\s+(?=\S)")
_EDGE_SPACE = re.compile(r"^(\s*)(.*?)(\s*)$", re.DOTALL)


# Web and mail addresses in the text: the models break them up ("https: / /
# vimeo.com"), so they are marked translate="no" and not sent.
_ADDRESS = re.compile(r"(?:https?://|www\.)[^\s<>\"']+|[\w.+-]+@[\w-]+(?:\.[\w-]+)+", re.IGNORECASE)
_TAG = re.compile(r"(<[^>]*>)")


def protect_addresses(html: str) -> str:
    """Mark the addresses in ``html``'s text as not to be translated (style
    blocks and tag attributes are left alone)."""
    out: list[str] = []
    in_style = False
    for piece in _TAG.split(html):
        if piece.startswith("<"):
            lowered = piece.lower()
            if lowered.startswith("<style"):
                in_style = True
            elif lowered.startswith("</style"):
                in_style = False
            out.append(piece)
        elif in_style or not piece:
            out.append(piece)
        else:
            out.append(_ADDRESS.sub(lambda m: f'<span translate="no">{m.group(0)}</span>', piece))
    return "".join(out)


# Bumped when what goes to the service, or what is done with the result,
# changes: translations kept from before are made again.
PIPELINE_VERSION = 8


async def translate_html(
    html: str, source: str, target: str, keep: Iterable[str] = (), names: Iterable[str] = ()
) -> str:
    """``html`` (a sanitized message body) in ``target``; source "auto", a code,
    or several ("zh-Hans,pl,en": then each paragraph from its own language).
    ``keep``: words always left as they are; ``names``: left as they are if
    the model would change them (translation_prep)."""
    if len(html) > MAX_HTML_CHARS:
        raise TranslationError("this message is too long to translate")
    sources = [x for x in (source or "").split(",") if x and x != "auto"]
    prepared = prep.join_wrapped_lines(html)
    for language in sources:
        prepared = prep.translate_phrases(prepared, language, target)
    candidates = {*names, *prep.find_names(prepared, sources[0] if sources else "")}
    prepared = prep.keep_symbols(protect_addresses(prep.keep_terms(prepared, keep)))
    return prep.tidy(await _translate_paragraphs(prepared, target, sources, candidates))