diff --git a/tests/test_mapping.py b/tests/test_mapping.py index 33ca964..24eb77d 100644 --- a/tests/test_mapping.py +++ b/tests/test_mapping.py @@ -2,7 +2,7 @@ from vaultengine import config as cfg from vaultengine.mapping import Vault, is_token_like -from vaultengine.spans import (CAT_DATE, CAT_ID, CAT_ORG, CAT_PERSON, Span) +from vaultengine.spans import (CAT_DATE, CAT_ORG, CAT_PERSON, Span) class MappingTests(unittest.TestCase): @@ -94,5 +94,24 @@ def test_save_load(self): self.assertEqual(v2.rehydrate_text(v.apply("王五在此")), "王五在此") + def test_ascii_word_boundary_guards(self): + v = Vault() + pairs = v.assign([ + Span("Jack", CAT_PERSON, source="llm"), + Span("张三", CAT_PERSON, source="llm"), + ]) + text = "Jack went to jackpot with 张三, hijack, and 张三李四." + out = v.apply(text, pairs) + # Jack -> P-n1 (assuming it's P-n1) + # jackpot -> unchanged (partial hit prevented) + # hijack -> unchanged (partial hit prevented) + # 张三 -> P-n2 + # 张三李四 -> P-n2李四 (CJK has no word boundary check) + self.assertNotIn("Jack went", out) + self.assertIn("jackpot", out) + self.assertIn("hijack", out) + self.assertIn("P-n2李四", out) + + if __name__ == "__main__": unittest.main() diff --git a/vaultengine/mapping.py b/vaultengine/mapping.py index 93a6967..dd7a1ef 100644 --- a/vaultengine/mapping.py +++ b/vaultengine/mapping.py @@ -136,7 +136,19 @@ def apply(self, text: str, pairs: Optional[Dict[str, str]] = None) -> str: pairs = pairs if pairs is not None else self.replacement_pairs() for surface in sorted(pairs, key=len, reverse=True): if surface and surface in text: - text = text.replace(surface, pairs[surface]) + token = pairs[surface] + # Check if first and last characters are ASCII alphanumeric or underscore + first = surface[0] + last = surface[-1] + is_ascii_word = ( + ("a" <= first <= "z" or "A" <= first <= "Z" or "0" <= first <= "9" or first == "_") and + ("a" <= last <= "z" or "A" <= last <= "Z" or "0" <= last <= "9" or last == "_") + ) + if is_ascii_word: + pattern = re.compile(r"\b" + re.escape(surface) + r"\b") + text = pattern.sub(token, text) + else: + text = text.replace(surface, token) return text def replacement_pairs(self) -> Dict[str, str]: