llm.py 5.1 KB
"""LLM-based translation backend."""

from __future__ import annotations

import logging
import time
from typing import List, Optional, Sequence, Union

from openai import OpenAI

from translation.languages import LANGUAGE_LABELS
from translation.prompts import TRANSLATION_PROMPTS
from translation.scenes import normalize_scene_name

logger = logging.getLogger(__name__)


def _build_prompt(
    text: str,
    *,
    source_lang: Optional[str],
    target_lang: str,
    scene: Optional[str],
) -> str:
    tgt = str(target_lang or "").strip().lower()
    src = str(source_lang or "auto").strip().lower() or "auto"
    normalized_scene = normalize_scene_name(scene)
    group = TRANSLATION_PROMPTS[normalized_scene]
    template = group.get(tgt) or group.get("en")
    if template is None:
        raise ValueError(f"Missing llm translation prompt for scene='{normalized_scene}' target_lang='{tgt}'")

    source_lang_label = LANGUAGE_LABELS.get(src, src)
    target_lang_label = LANGUAGE_LABELS.get(tgt, tgt)

    return template.format(
        source_lang=source_lang_label,
        src_lang_code=src,
        target_lang=target_lang_label,
        tgt_lang_code=tgt,
        text=text,
    )


class LLMTranslationBackend:
    def __init__(
        self,
        *,
        capability_name: str,
        model: str,
        timeout_sec: float,
        base_url: str,
        api_key: Optional[str],
    ) -> None:
        self.capability_name = capability_name
        self.model = model
        self.timeout_sec = float(timeout_sec)
        self.base_url = base_url
        self.api_key = api_key
        self.client = self._create_client()

    @property
    def supports_batch(self) -> bool:
        return True

    def _create_client(self) -> Optional[OpenAI]:
        if not self.api_key:
            logger.warning("DASHSCOPE_API_KEY not set; llm translation unavailable")
            return None
        try:
            return OpenAI(api_key=self.api_key, base_url=self.base_url)
        except Exception as exc:
            logger.error("Failed to initialize llm translation client: %s", exc, exc_info=True)
            return None

    def _translate_single(
        self,
        text: str,
        target_lang: str,
        source_lang: Optional[str] = None,
        scene: Optional[str] = None,
    ) -> Optional[str]:
        if not text or not str(text).strip():
            return text
        if not self.client:
            return None

        tgt = str(target_lang or "").strip().lower()
        src = str(source_lang or "auto").strip().lower() or "auto"
        if scene is None:
            raise ValueError("llm translation scene is required")
        normalized_scene = normalize_scene_name(scene)
        user_prompt = _build_prompt(
            text=text,
            source_lang=src,
            target_lang=tgt,
            scene=normalized_scene,
        )
        start = time.time()
        try:
            logger.info(
                "[llm] Request | src=%s tgt=%s model=%s prompt=%s",
                src,
                tgt,
                self.model,
                user_prompt,
            )
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": user_prompt}],
                timeout=self.timeout_sec,
            )
            content = (completion.choices[0].message.content or "").strip()
            latency_ms = (time.time() - start) * 1000
            if not content:
                logger.warning("[llm] Empty result | src=%s tgt=%s latency=%.1fms", src, tgt, latency_ms)
                return None
            logger.info(
                "[llm] Success | src=%s tgt=%s src_text=%s response=%s latency=%.1fms",
                src,
                tgt,
                text,
                content,
                latency_ms,
            )
            return content
        except Exception as exc:
            latency_ms = (time.time() - start) * 1000
            logger.warning(
                "[llm] Failed | src=%s tgt=%s latency=%.1fms error=%s",
                src,
                tgt,
                latency_ms,
                exc,
                exc_info=True,
            )
            return None

    def translate(
        self,
        text: Union[str, Sequence[str]],
        target_lang: str,
        source_lang: Optional[str] = None,
        scene: Optional[str] = None,
    ) -> Union[Optional[str], List[Optional[str]]]:
        if isinstance(text, (list, tuple)):
            results: List[Optional[str]] = []
            for item in text:
                if item is None:
                    results.append(None)
                    continue
                results.append(
                    self._translate_single(
                        text=str(item),
                        target_lang=target_lang,
                        source_lang=source_lang,
                        scene=scene,
                    )
                )
            return results

        return self._translate_single(
            text=str(text),
            target_lang=target_lang,
            source_lang=source_lang,
            scene=scene,
        )