office_translator/services/translation_tm.py

"""
Translation Memory (TM) — persistent reuse layer on top of the existing
Redis-backed cache (services/translation_cache.py).

Scope is PER USER (privacy: a customer's translations are never served to
another customer) and per translation context (custom prompt / glossary /
formality are hashed into the key so a glossary change invalidates old
matches). When Redis is not configured the cache transparently falls back
to a process-local LRU — still useful within a worker's lifetime.
"""

import hashlib
from typing import Dict, List, Optional, Tuple

from core.logging import get_logger

logger = get_logger(__name__)


class TMScope:
    """Identity/context of a translation job for TM keying."""

    __slots__ = ("user_id", "context_hash")

    def __init__(self, user_id: Optional[str], context_hash: Optional[str]):
        self.user_id = user_id
        self.context_hash = context_hash

    @classmethod
    def from_prompt(cls, user_id: Optional[str], prompt: Optional[str]) -> "TMScope":
        """Build a scope; the prompt (glossary + tone + formality merged)
        is hashed so any directive change produces different TM entries."""
        context_hash = None
        if prompt:
            context_hash = hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:16]
        return cls(user_id=user_id, context_hash=context_hash)


def _cache():
    from services.translation_cache import get_cache

    return get_cache()


def lookup_tm(
    texts: List[str],
    target_language: str,
    source_language: str,
    provider_name: str,
    scope: Optional[TMScope],
) -> Tuple[Dict[int, str], List[int]]:
    """Return ({index: cached_translation}, [indices not in TM]).

    No-op (all misses) when the scope has no user_id — anonymous jobs are
    never cached or served from another user's entries.
    """
    if not scope or not scope.user_id:
        return {}, list(range(len(texts)))

    try:
        cache = _cache()
    except Exception as e:
        logger.warning("tm_init_failed", error=str(e))
        return {}, list(range(len(texts)))

    hits: Dict[int, str] = {}
    misses: List[int] = []
    for i, text in enumerate(texts):
        if not text or not text.strip():
            misses.append(i)
            continue
        cached = cache.get(
            text,
            target_language,
            source_language,
            provider_name,
            user_id=scope.user_id,
            custom_prompt_hash=scope.context_hash,
        )
        if cached is not None and cached.strip():
            hits[i] = cached
        else:
            misses.append(i)

    if hits:
        logger.info("tm_lookup_hits", hits=len(hits), misses=len(misses))
    return hits, misses


def store_tm(
    texts: List[str],
    translations: List[str],
    target_language: str,
    source_language: str,
    provider_name: str,
    scope: Optional[TMScope],
) -> int:
    """Store fresh translations in the TM. Returns the number stored."""
    if not scope or not scope.user_id:
        return 0
    try:
        cache = _cache()
    except Exception as e:
        logger.warning("tm_init_failed", error=str(e))
        return 0

    stored = 0
    for text, translation in zip(texts, translations):
        if not text or not translation:
            continue
        # Never store identity entries — they would poison future lookups
        # (an unchanged text is not a translation).
        if translation.strip() == text.strip():
            continue
        cache.set(
            text,
            target_language,
            source_language,
            provider_name,
            translation,
            user_id=scope.user_id,
            custom_prompt_hash=scope.context_hash,
        )
        stored += 1
    if stored:
        logger.info("tm_stored", entries=stored)
    return stored


def tm_stats() -> Dict:
    """Backend stats for observability."""
    try:
        return _cache().stats()
    except Exception:
        return {}


def translate_with_tm(
    texts: List[str],
    target_language: str,
    source_language: str,
    provider_name: str,
    scope: "TMScope | None",
    translate_fn,
) -> List[str]:
    """Translate a batch with TM reuse: hits come from the cache, misses go
    through ``translate_fn`` (which receives ONLY the missed texts, in
    order) and fresh results are stored back. Falls back to a plain
    ``translate_fn(texts)`` call when the TM is unavailable.
    """
    if not texts:
        return []

    tm_hits, miss_indices = lookup_tm(
        texts, target_language, source_language, provider_name, scope
    )
    miss_texts = [texts[i] for i in miss_indices]

    if not miss_texts:
        return [tm_hits.get(i, texts[i]) for i in range(len(texts))]

    try:
        translated_misses = translate_fn(miss_texts)
    except Exception:
        if tm_hits:
            # Provider failed but we have partial TM hits — keep them and
            # leave the misses untouched rather than dropping everything.
            logger.warning("tm_translate_fn_failed_partial_hits", hits=len(tm_hits))
            translated_misses = None
        else:
            raise

    if translated_misses is None:
        return [
            tm_hits.get(i, texts[i]) for i in range(len(texts))
        ]

    store_tm(
        miss_texts,
        translated_misses,
        target_language,
        source_language,
        provider_name,
        scope,
    )

    merged: List[str] = []
    miss_pos = 0
    for i in range(len(texts)):
        if i in tm_hits:
            merged.append(tm_hits[i])
        else:
            merged.append(
                translated_misses[miss_pos]
                if miss_pos < len(translated_misses)
                else texts[i]
            )
            miss_pos += 1
    return merged