blivechat/services/translate.py

# -*- coding: utf-8 -*-
import asyncio
import dataclasses
import datetime
import enum
import functools
import hashlib
import hmac
import json
import logging
import random
from typing import *

import Crypto.Cipher.AES as cry_aes  # noqa
import Crypto.Util.Padding as cry_pad  # noqa
import aiohttp
import cachetools

import config
import utils.async_io
import utils.request

logger = logging.getLogger(__name__)

_translate_providers: List['TranslateProvider'] = []
# text -> res
_translate_cache: Optional[cachetools.LRUCache] = None
# 正在翻译的Future，text -> Future
_text_future_map: Dict[str, asyncio.Future] = {}
# 正在翻译的任务队列，索引是优先级
_task_queues: List['asyncio.Queue[TranslateTask]'] = []


class Priority(enum.IntEnum):
    HIGH = 0
    NORMAL = 1


@dataclasses.dataclass
class TranslateTask:
    priority: Priority
    text: str
    future: 'asyncio.Future[Optional[str]]'
    remain_retry_count: int


def init():
    cfg = config.get_config()
    global _translate_cache, _task_queues
    _translate_cache = cachetools.LRUCache(cfg.translation_cache_size)
    # 总队列长度会超过translate_max_queue_size，不用这么严格
    _task_queues = [asyncio.Queue(cfg.translate_max_queue_size) for _ in range(len(Priority))]
    utils.async_io.create_task_with_ref(_do_init())


async def _do_init():
    cfg = config.get_config()
    if not cfg.enable_translate:
        return
    providers = []
    for trans_cfg in cfg.translator_configs:
        provider = create_translate_provider(trans_cfg)
        if provider is not None:
            providers.append(provider)
    await asyncio.gather(*(provider.init() for provider in providers))
    global _translate_providers
    _translate_providers = providers


def create_translate_provider(cfg):
    type_ = cfg['type']
    if type_ == 'TencentTranslate':
        return TencentTranslate(
            cfg['query_interval'], cfg['source_language'], cfg['target_language'],
            cfg['secret_id'], cfg['secret_key'], cfg['region']
        )
    elif type_ == 'BaiduTranslate':
        return BaiduTranslate(
            cfg['query_interval'], cfg['source_language'], cfg['target_language'],
            cfg['app_id'], cfg['secret']
        )
    elif type_ == 'GeminiTranslate':
        return GeminiTranslate(            
            cfg['query_interval'], cfg['proxy'], cfg['api_key'], cfg['model_code'],
            cfg['prompt'], cfg['temperature']
        )
    return None


def need_translate(text):
    text = text.strip()
    # 中文数，不算平时打不出的字
    zh_num = 0
    # 日文假名数
    ja_num = 0
    for c in text:
        if 0x4E00 <= ord(c) <= 0x9FFF:
            zh_num += 1
        elif 0x3040 <= ord(c) <= 0x30FF:
            ja_num += 1
        elif c == '【':
            # 弹幕同传
            return False

    # 没有中文
    if zh_num == 0:
        return False
    # 日文假名较多
    if ja_num * 2 >= zh_num:
        return False
    return True


def get_translation_from_cache(text):
    key = text.strip().lower()
    return _translate_cache.get(key, None)


def translate(text, priority=Priority.NORMAL) -> Awaitable[Optional[str]]:
    key = text.strip().lower()
    # 如果已有正在翻译的future则返回，防止重复翻译
    future = _text_future_map.get(key, None)
    if future is not None:
        return future
    # 否则创建一个翻译任务
    future = asyncio.get_running_loop().create_future()

    # 查缓存
    res = _translate_cache.get(key, None)
    if res is not None:
        future.set_result(res)
        return future

    task = TranslateTask(
        priority=priority,
        text=text,
        future=future,
        remain_retry_count=3 if priority == Priority.HIGH else 1
    )
    if not _push_task(task):
        future.set_result(None)
        return future

    _text_future_map[key] = future
    future.add_done_callback(functools.partial(_on_translate_done, key))
    return future


def _on_translate_done(key, future):
    _text_future_map.pop(key, None)
    # 缓存
    try:
        res = future.result()
    except Exception:  # noqa
        return
    if res is None:
        return
    _translate_cache[key] = res


def _push_task(task: TranslateTask):
    if not _has_available_translate_provider():
        return False

    queue = _task_queues[task.priority]
    if not queue.full():
        queue.put_nowait(task)
        return True

    if task.priority != Priority.HIGH:
        return False

    # 高优先级的尝试降级，挤掉低优先级的任务
    queue = _task_queues[Priority.NORMAL]
    if queue.full():
        lower_task = queue.get_nowait()
        lower_task.future.set_result(None)
    queue.put_nowait(task)
    return True


async def _pop_task() -> TranslateTask:
    # 按优先级遍历，看是否已经有任务
    for queue in _task_queues:
        if not queue.empty():
            return queue.get_nowait()

    done_future_set, pending_future_set = await asyncio.wait(
        [asyncio.create_task(queue.get()) for queue in _task_queues],
        return_when=asyncio.FIRST_COMPLETED
    )
    for future in pending_future_set:
        future.cancel()

    # 如果有多个队列都取到任务了，只返回优先级最高的那一个，剩下的放回队列
    assert len(done_future_set) != 0
    tasks = [await future for future in done_future_set]
    if len(tasks) > 1:
        tasks.sort(key=lambda task_: task_.priority)

    res = None
    for task in tasks:
        if res is None:
            res = task
            continue

        if not _push_task(task):
            task.future.set_result(None)
    return res


def _cancel_all_tasks_if_no_available_translate_provider():
    if _has_available_translate_provider():
        return

    logger.warning('No available translate provider')
    for queue in _task_queues:
        while not queue.empty():
            task = queue.get_nowait()
            task.future.set_result(None)


def _has_available_translate_provider():
    return any(provider.is_available for provider in _translate_providers)


class TranslateProvider:
    def __init__(self, query_interval):
        self._query_interval = query_interval
        self._be_available_event = asyncio.Event()
        self._be_available_event.set()

    async def init(self):
        utils.async_io.create_task_with_ref(self._translate_consumer())
        return True

    @property
    def is_available(self):
        return True

    def _on_availability_change(self):
        if self.is_available:
            self._be_available_event.set()
        else:
            self._be_available_event.clear()
            _cancel_all_tasks_if_no_available_translate_provider()

    async def _translate_consumer(self):
        cls_name = type(self).__name__
        while True:
            try:
                if not self.is_available:
                    logger.info('%s waiting to become available', cls_name)
                    await self._be_available_event.wait()
                    logger.info('%s became available', cls_name)

                task = await _pop_task()
                # 为了简化代码，约定只会在_translate_wrapper里变成不可用，所以获取task之后这里还是可用的
                assert self.is_available

                start_time = datetime.datetime.now()
                await self._translate_wrapper(task)
                cost_time = (datetime.datetime.now() - start_time).total_seconds()

                # 频率限制
                await asyncio.sleep(self._query_interval - cost_time)
            except Exception:  # noqa
                logger.exception('%s error:', cls_name)

    async def _translate_wrapper(self, task: TranslateTask) -> Optional[str]:
        try:
            exc = None
            task.remain_retry_count -= 1
            res = await self._do_translate(task.text)
        except BaseException as e:
            exc = e
            res = None
        if res is not None:
            task.future.set_result(res)
            return res

        if task.remain_retry_count > 0:
            # 还可以重试则放回队列
            if not _push_task(task):
                task.future.set_result(None)
        else:
            # 否则设置异常或None结果
            if exc is not None:
                task.future.set_exception(exc)
            else:
                task.future.set_result(None)
        return None

    async def _do_translate(self, text) -> Optional[str]:
        raise NotImplementedError


class TencentTranslate(TranslateProvider):
    def __init__(self, query_interval, source_language, target_language,
                 secret_id, secret_key, region):
        super().__init__(query_interval)
        self._source_language = source_language
        self._target_language = target_language
        self._secret_id = secret_id
        self._secret_key = secret_key
        self._region = region

        self._cool_down_timer_handle = None

    @property
    def is_available(self):
        return self._cool_down_timer_handle is None and super().is_available

    async def _do_translate(self, text) -> Optional[str]:
        try:
            async with self._request_tencent_cloud(
                'TextTranslate',
                '2018-03-21',
                {
                    'SourceText': text,
                    'Source': self._source_language,
                    'Target': self._target_language,
                    'ProjectId': 0
                }
            ) as r:
                if r.status != 200:
                    logger.warning('TencentTranslate request failed: status=%d %s', r.status, r.reason)
                    return None
                data = (await r.json())['Response']
        except (aiohttp.ClientConnectionError, asyncio.TimeoutError):
            return None
        error = data.get('Error', None)
        if error is not None:
            logger.warning('TencentTranslate failed: %s %s, RequestId=%s', error['Code'],
                           error['Message'], data['RequestId'])
            self._on_fail(error['Code'])
            return None
        return data['TargetText']

    def _request_tencent_cloud(self, action, version, body):
        body_bytes = json.dumps(body).encode('utf-8')

        canonical_headers = 'content-type:application/json; charset=utf-8\nhost:tmt.tencentcloudapi.com\n'
        signed_headers = 'content-type;host'
        hashed_request_payload = hashlib.sha256(body_bytes).hexdigest()
        canonical_request = f'POST\n/\n\n{canonical_headers}\n{signed_headers}\n{hashed_request_payload}'

        request_timestamp = int(datetime.datetime.now().timestamp())
        date = datetime.datetime.utcfromtimestamp(request_timestamp).strftime('%Y-%m-%d')
        credential_scope = f'{date}/tmt/tc3_request'
        hashed_canonical_request = hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()
        string_to_sign = f'TC3-HMAC-SHA256\n{request_timestamp}\n{credential_scope}\n{hashed_canonical_request}'

        def sign(key, msg):
            return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()

        secret_date = sign(('TC3' + self._secret_key).encode('utf-8'), date)
        secret_service = sign(secret_date, 'tmt')
        secret_signing = sign(secret_service, 'tc3_request')
        signature = hmac.new(secret_signing, string_to_sign.encode('utf-8'), hashlib.sha256).hexdigest()

        authorization = (
            f'TC3-HMAC-SHA256 Credential={self._secret_id}/{credential_scope}, '
            f'SignedHeaders={signed_headers}, Signature={signature}'
        )

        headers = {
            'Authorization': authorization,
            'Content-Type': 'application/json; charset=utf-8',
            'X-TC-Action': action,
            'X-TC-Version': version,
            'X-TC-Timestamp': str(request_timestamp),
            'X-TC-Region': self._region
        }

        return utils.request.http_session.post('https://tmt.tencentcloudapi.com/', headers=headers, data=body_bytes)

    def _on_fail(self, code):
        if self._cool_down_timer_handle is not None:
            return

        sleep_time = 0
        if code == 'FailedOperation.NoFreeAmount':
            # 下个月恢复免费额度
            cur_time = datetime.datetime.now()
            year = cur_time.year
            month = cur_time.month + 1
            if month > 12:
                year += 1
                month = 1
            next_month_time = datetime.datetime(year, month, 1, minute=5)
            sleep_time = (next_month_time - cur_time).total_seconds()
            # Python 3.8之前不能超过一天
            sleep_time = min(sleep_time, 24 * 60 * 60 - 1)
        elif code in ('FailedOperation.ServiceIsolate', 'LimitExceeded'):
            # 需要手动处理，等5分钟
            sleep_time = 5 * 60
        if sleep_time != 0:
            self._cool_down_timer_handle = asyncio.get_running_loop().call_later(
                sleep_time, self._on_cool_down_timeout
            )
            self._on_availability_change()

    def _on_cool_down_timeout(self):
        self._cool_down_timer_handle = None
        self._on_availability_change()


class BaiduTranslate(TranslateProvider):
    def __init__(self, query_interval, source_language, target_language,
                 app_id, secret):
        super().__init__(query_interval)
        self._source_language = source_language
        self._target_language = target_language
        self._app_id = app_id
        self._secret = secret

        self._cool_down_timer_handle = None

    @property
    def is_available(self):
        return self._cool_down_timer_handle is None and super().is_available

    async def _do_translate(self, text) -> Optional[str]:
        try:
            async with utils.request.http_session.post(
                'https://fanyi-api.baidu.com/api/trans/vip/translate',
                data=self._add_sign({
                    'q': text,
                    'from': self._source_language,
                    'to': self._target_language,
                    'appid': self._app_id,
                    'salt': random.randint(1, 999999999)
                })
            ) as r:
                if r.status != 200:
                    logger.warning('BaiduTranslate request failed: status=%d %s', r.status, r.reason)
                    return None
                data = await r.json()
        except (aiohttp.ClientConnectionError, asyncio.TimeoutError):
            return None
        error_code = data.get('error_code', None)
        if error_code is not None:
            logger.warning('BaiduTranslate failed: %s %s', error_code, data['error_msg'])
            self._on_fail(error_code)
            return None
        return ''.join(result['dst'] for result in data['trans_result'])

    def _add_sign(self, data):
        str_to_sign = f"{self._app_id}{data['q']}{data['salt']}{self._secret}"
        sign = hashlib.md5(str_to_sign.encode('utf-8')).hexdigest()
        return {**data, 'sign': sign}

    def _on_fail(self, code):
        if self._cool_down_timer_handle is not None:
            return

        sleep_time = 0
        if code == '54004':
            # 账户余额不足，需要手动处理，等5分钟
            sleep_time = 5 * 60
        if sleep_time != 0:
            self._cool_down_timer_handle = asyncio.get_running_loop().call_later(
                sleep_time, self._on_cool_down_timeout
            )
            self._on_availability_change()

    def _on_cool_down_timeout(self):
        self._cool_down_timer_handle = None
        self._on_availability_change()


class GeminiTranslate(TranslateProvider):
    def __init__(self, query_interval, proxy, api_key, model_code, prompt, temperature):
        super().__init__(query_interval)
        self._proxy = proxy or None
        self._api_key = api_key
        self._url = f'https://generativelanguage.googleapis.com/v1beta/{model_code}:generateContent'
        self._prompt = prompt
        self._body = {
            'contents': [
                {
                    # 'role': 'user',
                    'parts': [{'text': ''}]
                }
            ],
            'safetySettings': [
                {
                    'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
                    'threshold': 'BLOCK_NONE'
                },
                {
                    'category': 'HARM_CATEGORY_HATE_SPEECH',
                    'threshold': 'BLOCK_NONE'
                },
                {
                    'category': 'HARM_CATEGORY_HARASSMENT',
                    'threshold': 'BLOCK_NONE'
                },
                {
                    'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',
                    'threshold': 'BLOCK_NONE'
                }
            ],
            'generationConfig': {
                'temperature': temperature,
                'topP': 1,
                'topK': 32,
                'candidateCount': 1,
                'maxOutputTokens': 8192
            }
        }

        self._cool_down_timer_handle = None

    @property
    def is_available(self):
        return self._cool_down_timer_handle is None and super().is_available

    async def _do_translate(self, text) -> Optional[str]:
        input_text = self._prompt.format(original_text=text)
        self._body['contents'][0]['parts'][0]['text'] = input_text
        try:
            async with utils.request.http_session.post(
                self._url,
                params={'key': self._api_key},
                json=self._body,
                proxy=self._proxy,
            ) as r:
                if r.status != 200:
                    logger.warning('GeminiTranslate request failed: status=%d %s', r.status, r.reason)
                    self._on_fail(r.status)
                    return None
                data = await r.json()
        except (aiohttp.ClientConnectionError, asyncio.TimeoutError) as e:
            logger.warning('GeminiTranslate request failed, %s: %s', type(e).__name__, e)
            return None

        try:
            candidates = data['candidates']
            if not candidates:
                block_reason = data['promptFeedback'].get('blockReason', '')
                logger.warning('GeminiTranslate no candidate, block_reason=%s, text=%s', block_reason, text)
                return None

            first_content_parts = candidates[0]['content']['parts']
            return ''.join(part.get('text', '') for part in first_content_parts)
        except (KeyError, IndexError):
            logger.warning('GeminiTranslate failed to parse response: %s', data)
            return None

    def _on_fail(self, code):
        if self._cool_down_timer_handle is not None:
            return

        if code in (401, 403):
            # API密钥无效，没有权限，或者不在有效地区。需要手动处理，等5分钟
            self._cool_down_timer_handle = asyncio.get_running_loop().call_later(
                5 * 60, self._on_cool_down_timeout
            )
            self._on_availability_change()

    def _on_cool_down_timeout(self):
        self._cool_down_timer_handle = None
        self._on_availability_change()