feat: apply direct identifier replacements before rewrite LLM call

src/anonymizer/engine/constants.py

@@ -94,7 +94,8 @@
 COL_SENSITIVITY_DISPOSITION = "_sensitivity_disposition"
 COL_SENSITIVITY_DISPOSITION_BLOCK = "_sensitivity_disposition_block"
 COL_REWRITE_DISPOSITION_BLOCK = "_rewrite_disposition_block"
-COL_REPLACEMENT_MAP_FOR_PROMPT = "_replacement_map_for_prompt"
+COL_PREREPLACE_TEXT = "_prereplace_text"
+COL_PREREPLACE_TAGGED_TEXT = "_prereplace_tagged_text"
 COL_FULL_REWRITE = "_full_rewrite"
 COL_MEANING_UNITS = "_meaning_units"
 COL_MEANING_UNITS_SERIALIZED = "_meaning_units_serialized"

src/anonymizer/engine/replace/llm_replace_workflow.py

@@ -164,13 +164,30 @@ def _filter_replacement_map_to_input_entities(
         for label in entity.labels
         if entity.value and label
     }
+    # Normalized form → (canonical_value, label) for fuzzy whitespace fallback.
+    # The LLM generating the map may normalise unusual Unicode whitespace to a
+    # regular space in the original field; we want to keep those entries using
+    # the canonical (detected) entity value so downstream lookups succeed.
+    _nws = lambda s: " ".join(s.split())  # noqa: E731
+    normalized_allowed: dict[tuple[str, str], tuple[str, str]] = {
+        (_nws(v), lbl): (v, lbl) for v, lbl in allowed_pairs
+    }
     protected_original_values = {value for value, _ in allowed_pairs}
     filtered: list[dict[str, str]] = []
     seen: set[tuple[str, str]] = set()
     synthetic_collision_labels: Counter[str] = Counter()
     for replacement in parsed_map.replacements:
         key = (replacement.original, replacement.label)
-        if key not in allowed_pairs or key in seen:
+        if key not in allowed_pairs:
+            # Try whitespace-normalised fallback
+            norm_key = (_nws(replacement.original), replacement.label)
+            canonical = normalized_allowed.get(norm_key)
+            if canonical is None or canonical in seen:
+                continue
+            # Rewrite original to the canonical (detected) value
+            key = canonical
+            replacement = replacement.model_copy(update={"original": canonical[0]})
+        if key in seen:
             continue
         if replacement.synthetic in protected_original_values:
             synthetic_collision_labels[replacement.label] += 1

src/anonymizer/engine/rewrite/repair.py

@@ -18,11 +18,11 @@
     COL_ANY_HIGH_LEAKED,
     COL_LEAKAGE_MASS,
     COL_LEAKED_PRIVACY_ITEMS,
+    COL_PREREPLACE_TEXT,
     COL_PRIVACY_QA,
     COL_PRIVACY_QA_REANSWER,
     COL_REWRITTEN_TEXT,
     COL_REWRITTEN_TEXT_NEXT,
-    COL_TEXT,
     COL_UTILITY_SCORE,
 )
 from anonymizer.engine.ndd.adapter import NddAdapter
@@ -168,7 +168,7 @@ def _render_repair_prompt(row: dict[str, Any], params: RepairParams) -> str:
     replacements = {
         "<<PRIVACY_GOAL>>": params.privacy_goal_str,
         "<<MAX_PRIVACY_LEAK>>": str(params.max_privacy_leak),
-        "<<ORIGINAL_TEXT>>": str(row.get(COL_TEXT, "")),
+        "<<ORIGINAL_TEXT>>": str(row.get(COL_PREREPLACE_TEXT, "")),
         "<<REWRITTEN_TEXT>>": str(row.get(COL_REWRITTEN_TEXT, "")),
         "<<LEAKAGE_MASS>>": str(row.get(COL_LEAKAGE_MASS, 0.0)),
         "<<HIGH_WARN>>": "\nWARNING: HIGH-SENSITIVITY LEAK DETECTED - must be fixed!"
@@ -201,7 +201,7 @@ def _make_repair_column(repairer_alias: str) -> Any:
         required_columns=[
             COL_LEAKED_PRIVACY_ITEMS,
             COL_REWRITTEN_TEXT,
-            COL_TEXT,
+            COL_PREREPLACE_TEXT,
             COL_LEAKAGE_MASS,
             COL_ANY_HIGH_LEAKED,
             COL_UTILITY_SCORE,

src/anonymizer/engine/rewrite/rewrite_generation.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import logging
+import re
 from typing import Any
 
 from data_designer.config import custom_column_generator
@@ -13,20 +14,25 @@
 from anonymizer.config.models import RewriteModelSelection
 from anonymizer.config.rewrite import PrivacyGoal
 from anonymizer.engine.constants import (
+    COL_FINAL_ENTITIES,
     COL_FULL_REWRITE,
+    COL_PREREPLACE_TAGGED_TEXT,
+    COL_PREREPLACE_TEXT,
     COL_REPLACEMENT_MAP,
-    COL_REPLACEMENT_MAP_FOR_PROMPT,
     COL_REWRITE_DISPOSITION_BLOCK,
     COL_REWRITTEN_TEXT,
     COL_SENSITIVITY_DISPOSITION,
     COL_TAG_NOTATION,
     COL_TAGGED_TEXT,
+    COL_TEXT,
     _jinja,
 )
 from anonymizer.engine.ndd.model_loader import resolve_model_alias
 from anonymizer.engine.prompt_utils import substitute_placeholders
+from anonymizer.engine.replace.strategies import ReplacementEntry, _apply_replacement_map_to_text
 from anonymizer.engine.rewrite.parsers import normalize_payload, parse_sensitivity_disposition
 from anonymizer.engine.schemas import (
+    EntitiesSchema,
     EntityReplacementMapSchema,
     RewriteOutputSchema,
 )
@@ -49,7 +55,7 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No
 
 <instructions>
 Your task is to rewrite the text below so that it protects the privacy of the entities described,
-following the entity protection rules and replacement map provided. The rewrite must read naturally as
+following the entity protection rules provided. The rewrite must read naturally as
 plain, fluent text — no tags, brackets, or annotation artifacts.
 
 Apply each protection decision consistently across ALL occurrences of the same entity value.
@@ -85,16 +91,8 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No
 Entities NOT listed above may be kept as-is.
 </sensitivity_disposition>
 
-{% if <<REPLACEMENT_MAP_COL>>.replacements %}
-<replacement_map>
-Synthetic replacement values for entities with protection_method "replace":
-<<REPLACEMENT_MAP>>
-</replacement_map>
-{% endif %}
 <output_requirements>
 Apply each protection method as follows:
-- "replace": Substitute the entity value with the corresponding synthetic value from the replacement map.
-  Use the synthetic value consistently for every occurrence.
 - "generalize": Replace with a broader category or range
   (e.g., a specific city → "a city in the Pacific Northwest", exact age → "in their late 30s").
 - "remove": Omit the detail entirely. Rewrite the surrounding sentence so it reads naturally without it.
@@ -113,10 +111,8 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No
             "<<PRIVACY_GOAL>>": privacy_goal.to_prompt_string(),
             "<<DATA_CONTEXT>>": data_context_section,
             "<<TAG_NOTATION>>": COL_TAG_NOTATION,
-            "<<TAGGED_TEXT>>": _jinja(COL_TAGGED_TEXT),
+            "<<TAGGED_TEXT>>": _jinja(COL_PREREPLACE_TAGGED_TEXT),
             "<<REWRITE_DISPOSITION_BLOCK>>": COL_REWRITE_DISPOSITION_BLOCK,
-            "<<REPLACEMENT_MAP_COL>>": COL_REPLACEMENT_MAP_FOR_PROMPT,
-            "<<REPLACEMENT_MAP>>": _jinja(COL_REPLACEMENT_MAP_FOR_PROMPT),
         },
     )
 
@@ -128,12 +124,18 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No
 
 @custom_column_generator(required_columns=[COL_SENSITIVITY_DISPOSITION])
 def _format_rewrite_disposition_block(row: dict[str, Any]) -> dict[str, Any]:
-    """Pre-filter and serialize protected entities (protection_method_suggestion != "leave_as_is") for the rewrite prompt."""
+    """Pre-filter and serialize protected entities for the rewrite prompt.
+
+    Excludes leave_as_is entities and replace entities (the latter are handled
+    programmatically by _apply_direct_replacements before the LLM sees the text).
+    """
     disposition = parse_sensitivity_disposition(row[COL_SENSITIVITY_DISPOSITION])
     block = []
     for e in disposition.sensitivity_disposition:
         if not e.needs_protection:
             continue
+        if e.protection_method_suggestion == "replace":
+            continue
         d = e.model_dump(mode="json")
         block.append(
             {
@@ -148,29 +150,158 @@ def _format_rewrite_disposition_block(row: dict[str, Any]) -> dict[str, Any]:
     return row
 
 
-@custom_column_generator(required_columns=[COL_REPLACEMENT_MAP, COL_REWRITE_DISPOSITION_BLOCK])
-def _filter_replacement_map_for_prompt(row: dict[str, Any]) -> dict[str, Any]:
-    """Keep only replacement entries for entities with protection_method_suggestion='replace'."""
-    disposition_block: list[dict] = row.get(COL_REWRITE_DISPOSITION_BLOCK, [])
+def _normalize_ws(s: str) -> str:
+    """Collapse all Unicode whitespace variants to a single ASCII space."""
+    return " ".join(s.split())
+
+
+def _get_replace_pairs(row: dict[str, Any]) -> tuple[list[tuple[str, str, str]], set[str]]:
+    """Return (pairs, replace_values) for entities with protection_method='replace'.
+
+    ``pairs`` contains (original, synthetic, label) tuples ready for substitution.
+    ``replace_values`` is the full set of entity values that required replacement,
+    returned so the caller can detect and raise on any unmatched entries.
+
+    Falls back to whitespace-normalized matching when the map's ``original`` field
+    differs only in Unicode whitespace from the disposition entity value (e.g. the
+    LLM normalised U+202F → U+0020). In that case the disposition value is used as
+    the substitution key because it reflects what is actually present in the text.
+    """
+    disposition = parse_sensitivity_disposition(row[COL_SENSITIVITY_DISPOSITION])
     replace_values = {
-        e["entity_value"] for e in disposition_block if e.get("protection_method_suggestion") == "replace"
+        e.entity_value for e in disposition.sensitivity_disposition if e.protection_method_suggestion == "replace"
     }
+    if not replace_values:
+        return [], set()
     raw_map = row.get(COL_REPLACEMENT_MAP)
-    if raw_map is None:
-        if replace_values:
-            logger.warning(
-                "COL_REPLACEMENT_MAP is None but entities require replacement; prompt will have no replacements."
-            )
-        row[COL_REPLACEMENT_MAP_FOR_PROMPT] = {"replacements": []}
-        return row
+    if not raw_map:
+        return [], replace_values
     raw_map = normalize_payload(raw_map)
     if hasattr(raw_map, "model_dump"):
         raw_map = raw_map.model_dump(mode="python")
     parsed_map = EntityReplacementMapSchema.model_validate(raw_map)
-    filtered = [
-        replacement.model_dump() for replacement in parsed_map.replacements if replacement.original in replace_values
-    ]
-    row[COL_REPLACEMENT_MAP_FOR_PROMPT] = {"replacements": filtered}
+
+    # normalized form → original disposition value (for fuzzy fallback)
+    normalized_to_disposition: dict[str, str] = {_normalize_ws(v): v for v in replace_values}
+
+    pairs: list[tuple[str, str, str]] = []  # (original, synthetic, label)
+    matched: set[str] = set()
+    for r in parsed_map.replacements:
+        if r.original in replace_values:
+            pairs.append((r.original, r.synthetic, r.label))
+            matched.add(r.original)
+        else:
+            disposition_value = normalized_to_disposition.get(_normalize_ws(r.original))
+            if disposition_value is not None and disposition_value not in matched:
+                pairs.append((disposition_value, r.synthetic, r.label))
+                matched.add(disposition_value)
+
+    return pairs, replace_values
+
+
+def _apply_tagged_text_replacements(
+    tagged_text: str, pairs: list[tuple[str, str, str]], tag_notation: str
+) -> str:
+    """Replace entity values in tagged text using tag-boundary-aware matching.
+
+    Matches each entity value only when it appears as the text content of its
+    corresponding tag wrapper, preventing substring corruption (e.g. 'Ann' inside
+    'Anna' is safe because the tagged form '<first_name>Ann</first_name>' is bounded
+    by tag delimiters that 'Anna' does not share).
+    """
+    for original, synthetic, label in sorted(pairs, key=lambda p: len(p[0]), reverse=True):
+        esc_o = re.escape(original)
+        esc_l = re.escape(label)
+        if tag_notation == "xml":
+            tagged_text = re.sub(
+                r"(<" + esc_l + r">)" + esc_o + r"(</" + esc_l + r">)",
+                lambda m, s=synthetic: m.group(1) + s + m.group(2),
+                tagged_text,
+            )
+        elif tag_notation == "bracket":
+            tagged_text = re.sub(
+                r"\[\[" + esc_o + r"\|" + esc_l + r"\]\]",
+                lambda m, s=synthetic, l=label: "[[" + s + "|" + l + "]]",
+                tagged_text,
+            )
+        elif tag_notation == "paren":
+            tagged_text = re.sub(
+                r"\(\(SENSITIVE:" + esc_l + r"\|" + esc_o + r"\)\)",
+                lambda m, s=synthetic, l=label: "((SENSITIVE:" + l + "|" + s + "))",
+                tagged_text,
+            )
+        else:  # sentinel
+            tagged_text = re.sub(
+                r"(<<SENSITIVE:" + esc_l + r">>)" + esc_o + r"(<</SENSITIVE:" + esc_l + r">>)",
+                lambda m, s=synthetic: m.group(1) + s + m.group(2),
+                tagged_text,
+            )
+    return tagged_text
+
+
+@custom_column_generator(
+    required_columns=[
+        COL_SENSITIVITY_DISPOSITION,
+        COL_REPLACEMENT_MAP,
+        COL_TEXT,
+        COL_TAGGED_TEXT,
+        COL_FINAL_ENTITIES,
+        COL_TAG_NOTATION,
+    ],
+    side_effect_columns=[COL_PREREPLACE_TAGGED_TEXT],
+)
+def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]:
+    """Programmatically replace direct identifier entities before the rewrite LLM call.
+
+    Uses span-aware replacement for plain text (character offsets from COL_FINAL_ENTITIES)
+    and tag-boundary-aware replacement for tagged text, preventing substring corruption
+    (e.g. replacing 'Ann' must not modify 'Anna').
+
+    Raises on failure rather than falling back to unmodified text: replace entities are
+    excluded from COL_REWRITE_DISPOSITION_BLOCK, so a silent passthrough would send
+    PII-containing text to the LLM with no instructions to protect those entities.
+    """
+    plain_text = str(row.get(COL_TEXT, ""))
+    tagged_text = str(row.get(COL_TAGGED_TEXT, ""))
+    pairs, replace_values = _get_replace_pairs(row)
+    matched = {original for original, _, _ in pairs}
+    unmatched = replace_values - matched
+    if unmatched:
+        disposition = parse_sensitivity_disposition(row[COL_SENSITIVITY_DISPOSITION])
+        value_to_label: dict[str, str] = {e.entity_value: e.entity_label for e in disposition.sensitivity_disposition}
+        unmatched_labels = sorted({value_to_label.get(v, "unknown") for v in unmatched})
+        raise RuntimeError(
+            f"Replace entities have no entry in the replacement map; refusing to pass PII-containing text "
+            f"to the rewrite LLM without protection instructions. "
+            f"{len(unmatched)} entities unmatched (labels: {unmatched_labels})"
+        )
+    if pairs:
+        tag_notation = str(row.get(COL_TAG_NOTATION, "xml"))
+        # Plain text: span-aware using character offsets from COL_FINAL_ENTITIES.
+        # Falls back to sorted-pairs regex when entities are absent (e.g. in unit tests
+        # that call this function directly without a full detection pipeline).
+        all_entities = EntitiesSchema.from_raw(row.get(COL_FINAL_ENTITIES, {}))
+        replace_value_set = {original for original, _, _ in pairs}
+        replace_entities = EntitiesSchema(
+            entities=[
+                e for e in all_entities.entities
+                if e.value in replace_value_set and e.end_position > e.start_position
+            ]
+        )
+        replacement_entries = [
+            ReplacementEntry(original=original, label=label, synthetic=synthetic)
+            for original, synthetic, label in pairs
+        ]
+        if replace_entities.entities:
+            plain_text = _apply_replacement_map_to_text(plain_text, replace_entities, replacement_entries)
+        else:
+            sorted_pairs = sorted(pairs, key=lambda p: len(p[0]), reverse=True)
+            pattern = re.compile("|".join(re.escape(original) for original, _, _ in sorted_pairs))
+            lookup = {original: synthetic for original, synthetic, _ in sorted_pairs}
+            plain_text = pattern.sub(lambda m: lookup[m.group(0)], plain_text)
+        tagged_text = _apply_tagged_text_replacements(tagged_text, pairs, tag_notation)
+    row[COL_PREREPLACE_TEXT] = plain_text
+    row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text
     return row
 
 
@@ -226,8 +357,8 @@ def columns(
                 generator_function=_format_rewrite_disposition_block,
             ),
             CustomColumnConfig(
-                name=COL_REPLACEMENT_MAP_FOR_PROMPT,
-                generator_function=_filter_replacement_map_for_prompt,
+                name=COL_PREREPLACE_TEXT,
+                generator_function=_apply_direct_replacements,
             ),
             LLMStructuredColumnConfig(
                 name=COL_FULL_REWRITE,