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

from __future__ import annotations

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

from openai import OpenAI

from config.env_config import DASHSCOPE_API_KEY
from config.services_config import get_translation_config
from config.translate_prompts import TRANSLATION_PROMPTS
from config.tenant_config_loader import SOURCE_LANG_CODE_MAP

logger = logging.getLogger(__name__)

DEFAULT_QWEN_BASE_URL = "https://dashscope-us.aliyuncs.com/compatible-mode/v1"
DEFAULT_LLM_MODEL = "qwen-flash"


def _build_prompt(
    text: str,
    *,
    source_lang: Optional[str],
    target_lang: str,
    scene: Optional[str],
) -> str:
    tgt = (target_lang or "").lower() or "en"
    src = (source_lang or "auto").lower()
    normalized_scene = (scene or "").strip() or "general"
    if normalized_scene in {"query", "ecommerce_search", "ecommerce_search_query"}:
        group_key = "ecommerce_search_query"
    elif normalized_scene in {"product_title", "sku_name"}:
        group_key = "sku_name"
    else:
        group_key = normalized_scene
    group = TRANSLATION_PROMPTS.get(group_key) or TRANSLATION_PROMPTS["general"]
    template = group.get(tgt) or group.get("en")
    if not template:
        template = (
            "You are a professional {source_lang} ({src_lang_code}) to "
            "{target_lang} ({tgt_lang_code}) translator, output only the translation: {text}"
        )

    source_lang_label = SOURCE_LANG_CODE_MAP.get(src, src)
    target_lang_label = SOURCE_LANG_CODE_MAP.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,
        *,
        model: Optional[str] = None,
        timeout_sec: float = 30.0,
        base_url: Optional[str] = None,
    ) -> None:
        cfg = get_translation_config()
        llm_cfg = cfg.get_capability_cfg("llm")
        self.model = model or llm_cfg.get("model") or DEFAULT_LLM_MODEL
        self.timeout_sec = float(llm_cfg.get("timeout_sec") or timeout_sec or 30.0)
        self.base_url = (
            (base_url or "").strip()
            or (llm_cfg.get("base_url") or "").strip()
            or os.getenv("DASHSCOPE_BASE_URL")
            or DEFAULT_QWEN_BASE_URL
        )
        self.client = self._create_client()

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

    def _create_client(self) -> Optional[OpenAI]:
        api_key = DASHSCOPE_API_KEY or os.getenv("DASHSCOPE_API_KEY")
        if not api_key:
            logger.warning("DASHSCOPE_API_KEY not set; llm translation unavailable")
            return None
        try:
            return OpenAI(api_key=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,
        context: Optional[str] = None,
        prompt: Optional[str] = None,
    ) -> Optional[str]:
        if not text or not str(text).strip():
            return text
        if not self.client:
            return None

        tgt = (target_lang or "").lower() or "en"
        src = (source_lang or "auto").lower()
        scene = context or "default"
        user_prompt = prompt or _build_prompt(
            text=text,
            source_lang=src,
            target_lang=tgt,
            scene=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,
        context: Optional[str] = None,
        prompt: 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,
                        context=context,
                        prompt=prompt,
                    )
                )
            return results

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


LLMTranslatorProvider = LLMTranslationBackend


def llm_translate(
    text: Union[str, Sequence[str]],
    target_lang: str,
    *,
    source_lang: Optional[str] = None,
    source_lang_label: Optional[str] = None,
    target_lang_label: Optional[str] = None,
    timeout_sec: Optional[float] = None,
) -> Union[Optional[str], List[Optional[str]]]:
    del source_lang_label, target_lang_label
    provider = LLMTranslationBackend(timeout_sec=timeout_sec or 30.0)
    return provider.translate(
        text=text,
        target_lang=target_lang,
        source_lang=source_lang,
        context=None,
    )