src/understanding/rule_engine.py

"""规则引擎：pyahocorasick + re2，强规则→REJECT，弱规则→REVIEW"""
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import List
import re2
import ahocorasick

class RuleLevel(str, Enum):
    STRONG = "strong"  # → REJECT
    WEAK = "weak"      # → REVIEW


@dataclass
class RuleHit:
    rule_id: str
    level: RuleLevel
    matched_text: str


_KEYWORDS_DIR = Path(__file__).parent / "keywords"

STRONG_CATEGORIES = ["porn", "violence", "drug", "gambling", "political", "illegal", "other"]
WEAK_CATEGORIES = ["weak"]


def _load_keywords(path: Path) -> List[str]:
    if not path.exists():
        return []
    return [l.strip() for l in path.read_text(encoding="utf-8").splitlines()
            if l.strip() and not l.startswith("#")]


def _build_ac(entries: list) -> ahocorasick.Automaton:
    A = ahocorasick.Automaton()
    for i, (kw, rule_id) in enumerate(entries):
        A.add_word(kw, (i, kw, rule_id))
    A.make_automaton()
    return A


_strong_ac = None
_weak_ac = None


def _get_strong_ac():
    global _strong_ac
    if _strong_ac is None:
        entries = []
        for cat in STRONG_CATEGORIES:
            path = _KEYWORDS_DIR / f"strong_{cat}.txt"
            for kw in _load_keywords(path):
                entries.append((kw, f"strong_{cat}"))
        _strong_ac = _build_ac(entries)
    return _strong_ac


def _get_weak_ac():
    global _weak_ac
    if _weak_ac is None:
        entries = [(kw, "weak_keyword") for kw in _load_keywords(_KEYWORDS_DIR / "weak.txt")]
        _weak_ac = _build_ac(entries)
    return _weak_ac


# 联系方式正则
_RE_PHONE = re2.compile(r"(?:\D|^)(1[3-9]\d{9})(?:\D|$)")
_RE_QQ = re2.compile(r"[Qq]{2}[号:：\s]*\d{5,11}")
_RE_WECHAT = re2.compile(r"微信[号:：\s]*\w{6,20}")
_RE_ID_CARD = re2.compile(r"\d{17}[\dXx]")
_RE_MARKETING = re2.compile(r"(加[我我]|领取|免费|福利|优惠|私聊|联系)")


def match(text: str) -> List[RuleHit]:
    hits = []

    for _, (_, kw, rule_id) in _get_strong_ac().iter(text):
        hits.append(RuleHit(rule_id, RuleLevel.STRONG, kw))

    has_contact = bool(_RE_PHONE.search(text) or _RE_QQ.search(text) or _RE_WECHAT.search(text))
    has_marketing = bool(_RE_MARKETING.search(text))
    if has_contact and has_marketing:
        hits.append(RuleHit("contact_marketing_combo", RuleLevel.STRONG, text[:50]))

    if has_contact and not has_marketing:
        hits.append(RuleHit("contact_only", RuleLevel.WEAK, text[:50]))

    if _RE_ID_CARD.search(text):
        hits.append(RuleHit("id_card", RuleLevel.WEAK, "id_card_format"))

    for _, (_, kw, rule_id) in _get_weak_ac().iter(text):
        hits.append(RuleHit(rule_id, RuleLevel.WEAK, kw))

    return hits