fix(guardrail): fold fullwidth, accented, lookalike and invisible spellings before to_ascii erases them

cosmos_framework/auxiliary/guardrail/blocklist/blocklist.py

@@ -10,6 +10,7 @@
 import nltk
 from better_profanity import profanity
 
+from cosmos_framework.auxiliary.guardrail.blocklist.lookalikes import fold_lookalikes
 from cosmos_framework.auxiliary.guardrail.blocklist.utils import read_keyword_list_from_dir, to_ascii
 from cosmos_framework.auxiliary.guardrail.common.core import (
     GUARDRAIL1_CHECKPOINT,
@@ -37,6 +38,14 @@
 INVISIBLE_CHARS = re.compile(r"[­​-‏‪-‮⁠-⁤]")
 
 
+class _NoLemmatizer:
+    """Leaves every token as it is. The default for Blocklist.from_lists."""
+
+    @staticmethod
+    def lemmatize(token: str) -> str:
+        return token
+
+
 class Blocklist(ContentSafetyGuardrail):
     def __init__(
         self,
@@ -57,10 +66,49 @@ def __init__(
         self.guardrail_partial_match_min_chars = guardrail_partial_match_min_chars
         self.guardrail_partial_match_letter_count = guardrail_partial_match_letter_count
 
-        # Load blocklist and whitelist keywords
-        self.blocklist_words = read_keyword_list_from_dir(os.path.join(self.checkpoint_dir, "custom"))
-        self.whitelist_words = read_keyword_list_from_dir(os.path.join(self.checkpoint_dir, "whitelist"))
-        self.exact_match_words = read_keyword_list_from_dir(os.path.join(self.checkpoint_dir, "exact_match"))
+        self.load_lists(
+            blocklist_words=read_keyword_list_from_dir(os.path.join(self.checkpoint_dir, "custom")),
+            whitelist_words=read_keyword_list_from_dir(os.path.join(self.checkpoint_dir, "whitelist")),
+            exact_match_words=read_keyword_list_from_dir(os.path.join(self.checkpoint_dir, "exact_match")),
+        )
+
+    @classmethod
+    def from_lists(
+        cls,
+        blocklist_words: list[str],
+        whitelist_words: list[str] | None = None,
+        exact_match_words: list[str] | None = None,
+        lemmatizer=None,
+        guardrail_partial_match_min_chars: int = 6,
+        guardrail_partial_match_letter_count: float = 0.4,
+    ) -> "Blocklist":
+        """Build a Blocklist from keyword lists in memory, without the checkpoint.
+
+        For the tests and the list-maintenance tools. It goes through load_lists
+        like __init__ does, so callers exercise the production bookkeeping rather
+        than a copy of it. The matcher is a fresh instance, not the module-level
+        singleton production shares. ``lemmatizer`` defaults to leaving tokens
+        unchanged; pass nltk's WordNetLemmatizer to match production exactly.
+        """
+        from better_profanity.better_profanity import Profanity
+
+        blocklist = cls.__new__(cls)
+        blocklist.profanity = Profanity()
+        blocklist.lemmatizer = lemmatizer if lemmatizer is not None else _NoLemmatizer()
+        blocklist.guardrail_partial_match_min_chars = guardrail_partial_match_min_chars
+        blocklist.guardrail_partial_match_letter_count = guardrail_partial_match_letter_count
+        blocklist.load_lists(blocklist_words, whitelist_words or [], exact_match_words or [])
+        return blocklist
+
+    def load_lists(self, blocklist_words: list[str], whitelist_words: list[str], exact_match_words: list[str]) -> None:
+        """Load the keyword lists into the matcher.
+
+        The single place that turns lists into matcher state, shared by __init__
+        and from_lists.
+        """
+        self.blocklist_words = blocklist_words
+        self.whitelist_words = whitelist_words
+        self.exact_match_words = exact_match_words
 
         self.profanity.load_censor_words(custom_words=self.blocklist_words, whitelist_words=self.whitelist_words)
         log.debug(f"Loaded {len(self.blocklist_words)} words/phrases from blocklist")
@@ -86,9 +134,22 @@ def normalize_for_matching(prompt: str) -> str:
 
         The collapse also closes a plainer hole: a multi-word entry was defeated
         by typing two spaces between its words.
+
+        The code-point folding is a separate method because is_safe also needs it
+        on its own: see alternate_spellings.
         """
+        return Blocklist.fold_invisible_characters(Blocklist.normalize_code_points(prompt))
+
+    @staticmethod
+    def normalize_code_points(prompt: str) -> str:
+        """NFKC, combining marks removed, runs of whitespace collapsed."""
         prompt = unicodedata.normalize("NFKC", prompt)
         prompt = "".join(c for c in unicodedata.normalize("NFKD", prompt) if not unicodedata.combining(c))
+        return " ".join(prompt.split())
+
+    @staticmethod
+    def fold_invisible_characters(prompt: str) -> str:
+        """Rewrite zero-width and bidirectional characters to a space."""
         return " ".join(INVISIBLE_CHARS.sub(" ", prompt).split())
 
     def censor_prompt(self, input_prompt: str) -> tuple[bool, str]:
@@ -226,9 +287,66 @@ def is_safe(self, input_prompt: str = "") -> tuple[bool, str]:
         # Check if the input is empty
         if not input_prompt:
             return False, "Input is empty"
-        input_prompt = to_ascii(input_prompt)
 
-        # Check full sentence for censored words
+        safe, message = self.check_passes(input_prompt)
+        if not safe or input_prompt.isascii():
+            return safe, message
+
+        # The passes run once more on each other way the prompt can be read.
+        # Like the raw pass, these only add checks: a prompt blocked as written
+        # is still blocked.
+        for spelling in self.alternate_spellings(input_prompt):
+            safe, message = self.check_passes(spelling)
+            if not safe:
+                return safe, message
+        return safe, message
+
+    @classmethod
+    def alternate_spellings(cls, prompt: str) -> list[str]:
+        """Other readings of a non-ASCII prompt that the passes must also see.
+
+        * The code points folded (NFKC, combining marks removed). censor_prompt
+          folds them itself, but the lemmatized pass and the exact-match check
+          only see the prompt after to_ascii, which deletes a fullwidth or
+          accented letter rather than folding it. Without this reading an entry
+          that is only on the exact-match list is missed when spelled that way.
+        * Invisible characters deleted rather than turned into a space. The
+          space reading catches "snow<U+200B>white" as the entry "snow white",
+          but splits "sn<U+200B>ow white" into "sn ow white", which matches
+          nothing. Deleting them reads it as the entry.
+        * Letters from other scripts drawn like Latin ones ("sn<U+043E>w", with
+          a Cyrillic o) folded to Latin. They survive NFKC, and to_ascii deletes
+          them, so every pass reads the word with a letter missing. NFKC runs
+          first, so a word mixing fullwidth and lookalike letters is recognized
+          as one disguised word.
+
+        Each reading is returned only if it differs from the prompt.
+        """
+        normalized = cls.normalize_code_points(prompt)
+        joined = " ".join(INVISIBLE_CHARS.sub("", normalized).split())
+        spellings: list[str] = []
+        for spelling in (normalized, joined, fold_lookalikes(normalized), fold_lookalikes(joined)):
+            if spelling != prompt and spelling not in spellings:
+                spellings.append(spelling)
+        return spellings
+
+    def check_passes(self, input_prompt: str) -> tuple[bool, str]:
+        """Run every blocklist pass on one spelling of the prompt."""
+        # An extra first pass sees the prompt as written. to_ascii rewrites every
+        # non-ASCII run to a space, which erases exactly the spellings
+        # censor_prompt was taught to fold: a fullwidth or accented letter is
+        # gone rather than normalized, so the word it spells never reaches the
+        # matcher and the prompt reads as clean. censor_prompt does its own
+        # code-point folding, so it needs the original text to fold.
+        censored, message = self.censor_prompt(input_prompt)
+        if censored:
+            return False, message
+
+        # The ASCII pass is kept rather than replaced. Some separators survive the
+        # folding -- "boston·dynamics" stays one token on the raw pass -- and
+        # only to_ascii turns them into the space the entry needs. Adding a pass
+        # means nothing that was blocked before can pass now.
+        input_prompt = to_ascii(input_prompt)
         censored, message = self.censor_prompt(input_prompt)
         if censored:
             return False, message