# -*- 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()