sensitive.py 2.74 KB
"""敏感词等反垃圾相关。"""


from engine.rpc import get_current_rpc_invoker
from engine.logger import sensitive_logger
from libs.spawn_task import SpawnTask, Task


class DRIVERS(object):
    """后续可能接入多个服务"""

    YIDUN = "1"  # "易盾"


class ENDPOINT(object):
    """请求入口。"""

    # 所有请求接入到 antispam 服务中
    endpoint = {
        DRIVERS.YIDUN: "antispam/yidun",
    }

    def __getitem__(self, key):
        try:
            return self.endpoint[key]
        except KeyError:
            raise Exception("没有对应的检测平台")


class Sensitive(object):
    """敏感词等反垃圾相关入口。"""

    endpoint = ENDPOINT()

    @classmethod
    def _check_job(cls, endpoint, text):
        """内容检测"""

        rpc = get_current_rpc_invoker()
        try:
            res = rpc[endpoint](text=text).unwrap()
        except:
            res = None

        sensitive_logger.info(text)
        sensitive_logger.info(res)

        return res

    @classmethod
    def check(cls, text_list, driver=DRIVERS.YIDUN, detail=False):
        """检测"""

        endpoint = cls.endpoint[driver]

        tasks = [Task(cls._check_job, endpoint, text) for text in text_list]
        spawn = SpawnTask(tasks)

        spawn.run()

        data = []
        result = {}
        if driver == DRIVERS.YIDUN:
            for res in spawn.result:
                data.append(cls._yidun_paraser(res, False))

            for text, r in zip(text_list, data):
                result[text] = r

        return result

    @classmethod
    def _yidun_paraser(cls, data, detail=False):
        """易盾结果解析"""
        # {
        #   'code': 200,
        #   'msg': 'ok',
        #   'result': {
        #       'taskId': 'd28e036d95874361916a26d406ea7db9',
        #       'action': 2,
        #       'labels': [
        #           {
        #               'label': 600,
        #               'level': 2,
        #               'details': {
        #                   'hint': ['傻逼']
        #               }
        #           }
        #       ]
        #   }
        # }
        hints = []
        if not data or "detail" in data:  # 程序出错啦
            return

        result = data.get("result", {})
        if result:
            action = result.get("action", 0)
        else:
            action = 0

        if action == 0 or action == 1:  # "通过"
            return []

        if result:
            lables = result.get("labels", [])
        else:
            lables = []

        for lable in lables:
            hints.extend(lable.get('details', {}).get("hint"))

        # 是否需要详细敏感词列表
        if not detail:
            return True if hints else False

        return hints