cmhub/apps/moderation/providers/keyword.py

from __future__ import annotations

import threading
from dataclasses import dataclass

try:
    from ahocorapy.keywordtree import KeywordTree
except ImportError:  # pragma: no cover - exercised only when dependency is missing.
    KeywordTree = None

from apps.moderation.models import SensitiveWord
from apps.moderation.normalization import normalize_text
from apps.moderation.versioning import get_sensitive_words_version

from .base import ModerationResult, ModerationUnavailableError, Verdict


@dataclass(frozen=True)
class SensitiveWordMatch:
    id: int
    word: str
    normalized_word: str
    category: str


class KeywordMatcher:
    def __init__(self, records: list[SensitiveWordMatch]) -> None:
        self._records_by_keyword: dict[str, list[SensitiveWordMatch]] = {}
        if KeywordTree is None:
            raise ModerationUnavailableError("ahocorapy is not installed")
        self._tree = KeywordTree(case_insensitive=False)
        for record in records:
            self._records_by_keyword.setdefault(record.normalized_word, []).append(record)
        for keyword in self._records_by_keyword:
            self._tree.add(keyword)
        self._tree.finalize()

    def search(self, text: str) -> list[SensitiveWordMatch]:
        normalized = normalize_text(text)
        if not normalized:
            return []

        matches: list[SensitiveWordMatch] = []
        seen: set[int] = set()
        for keyword, _index in self._tree.search_all(normalized):
            for record in self._records_by_keyword.get(keyword, ()):
                if record.id in seen:
                    continue
                seen.add(record.id)
                matches.append(record)
        return matches


@dataclass(frozen=True)
class _MatcherState:
    version: str
    matcher: KeywordMatcher


_matcher_lock = threading.Lock()
_matcher_state: _MatcherState | None = None


def reset_keyword_matcher_cache() -> None:
    global _matcher_state
    with _matcher_lock:
        _matcher_state = None


def _load_active_records() -> list[SensitiveWordMatch]:
    rows = SensitiveWord.objects.filter(
        is_active=True,
        action=SensitiveWord.Action.BLOCK,
    ).only("id", "word", "normalized_word", "category")
    return [
        SensitiveWordMatch(
            id=row.id,
            word=row.word,
            normalized_word=row.normalized_word,
            category=row.category,
        )
        for row in rows.order_by("id")
        if row.normalized_word
    ]


def get_keyword_matcher() -> KeywordMatcher:
    global _matcher_state
    version = get_sensitive_words_version()
    state = _matcher_state
    if state is not None and state.version == version:
        return state.matcher

    with _matcher_lock:
        state = _matcher_state
        if state is not None and state.version == version:
            return state.matcher
        matcher = KeywordMatcher(_load_active_records())
        _matcher_state = _MatcherState(version=version, matcher=matcher)
        return matcher


class KeywordModerationProvider:
    def moderate_text(self, text: str, *, biz_type: str = "") -> ModerationResult:
        matches = get_keyword_matcher().search(text)
        if not matches:
            return ModerationResult(verdict=Verdict.PASS)

        labels = tuple(dict.fromkeys(match.category for match in matches))
        keywords = tuple(dict.fromkeys(match.word for match in matches))
        word_ids = tuple(match.id for match in matches)
        return ModerationResult(
            verdict=Verdict.BLOCK,
            labels=labels,
            keywords=keywords,
            raw={"word_ids": word_ids},
        )