feat: apply direct identifier replacements before rewrite LLM call

[asteier2026]

Update 8/13: I think PR #246 supersedes this PR now

Summary

• Adds COL_PREREPLACE_TEXT and COL_PREREPLACE_TAGGED_TEXT constants
• Adds _apply_direct_replacements generator that programmatically substitutes direct identifiers from
  the replacement map before the rewrite LLM call, using a single-pass regex to prevent cascade
  replacements (e.g. Alice→Bob→Carlos)
• Removes <replacement_map> block from the rewrite prompt — the LLM no longer needs to apply
  replacements
• Updates repair.py to use pre-replaced text as the repair baseline
• Removes COL_REPLACEMENT_MAP_FOR_PROMPT (no longer needed)
• Adds whitespace-normalized fallback matching in both _filter_replacement_map_to_input_entities and
  _get_replace_pairs to handle LLM-normalised Unicode whitespace (e.g. U+202F → U+0020) in entity values
• _apply_direct_replacements falls back to passing text through unchanged on error rather than
  dropping the record

Motivation

The rewrite LLM was inconsistently applying replacement map entries, especially for entities that
appear multiple times. Programmatic pre-replacement guarantees all occurrences are substituted before
the LLM sees the text. The whitespace fixes handle cases where the LLM normalises unusual Unicode
whitespace in entity values, which caused map lookups to silently miss.

[greptile-apps]

Greptile Summary

This PR moves direct-identifier replacement from the rewrite LLM to a deterministic pre-processing step. Before the LLM sees the text, _apply_direct_replacements substitutes all replace-method entities using span-aware (character offset) replacement for plain text and a single-pass tagged-span regex for tagged text, then writes the results to COL_PREREPLACE_TEXT / COL_PREREPLACE_TAGGED_TEXT. The replacement_map block is removed from the rewrite prompt, replace entities are excluded from COL_REWRITE_DISPOSITION_BLOCK, and the repair baseline switches from COL_TEXT to COL_PREREPLACE_TEXT.

• Adds _apply_direct_replacements with a hard RuntimeError when any required replacement is unmatched, preventing unprotected PII from reaching the LLM.
• Adds whitespace-normalized fallback matching in both _filter_replacement_map_to_input_entities and _get_replace_pairs to handle LLM-normalised Unicode whitespace (e.g. U+202F → U+0020) in entity values.
• Replaces _filter_replacement_map_for_prompt / COL_REPLACEMENT_MAP_FOR_PROMPT entirely; the repair step now uses COL_PREREPLACE_TEXT as its baseline.

Confidence Score: 5/5

• The change is safe to merge. The core correctness guarantees — single-pass cascade prevention, hard failure on unmatched replacements, and correct repair baseline — are all working and tested.
• All previously discussed issues (silent passthrough on missing map, cascade replacement, double disposition parse) have been resolved. Remaining observations are edge-case testing gaps and a log-clarity nit, none of which affect the correctness of the production pipeline.
• The fallback regex branch in _apply_direct_replacements (lines 284-288 of rewrite_generation.py) is worth a second look before adding word-boundary-sensitive entity names to production workloads.

Important Files Changed

Filename	Overview
src/anonymizer/engine/rewrite/rewrite_generation.py	Core change: replaces _filter_replacement_map_for_prompt with _apply_direct_replacements. Uses span-aware + single-pass regex substitution. The fallback regex path (when COL_FINAL_ENTITIES has no valid spans) performs bare substring matching without word boundaries, which could corrupt partial matches like "Ann" inside "Anna" in production.
src/anonymizer/engine/replace/llm_replace_workflow.py	Adds whitespace-normalized fallback matching in _filter_replacement_map_to_input_entities; correctly rewrites the replacement entry to the canonical entity value so downstream lookups succeed.
src/anonymizer/engine/rewrite/repair.py	Correctly switches from COL_TEXT to COL_PREREPLACE_TEXT as the repair baseline, so the LLM sees pre-replaced text (synthetics already substituted) rather than the original PII-containing text.
tests/engine/test_rewrite_generation.py	Good coverage of replacement, error cases, and single-pass cascade prevention. The cascade test (test_apply_direct_replacements_no_cascade_when_synthetic_matches_another_original) omits COL_FINAL_ENTITIES, so it exercises the regex fallback path rather than the production span-aware _apply_replacement_map_to_text path.
src/anonymizer/engine/constants.py	Straightforward constant rename: COL_REPLACEMENT_MAP_FOR_PROMPT replaced by COL_PREREPLACE_TEXT and COL_PREREPLACE_TAGGED_TEXT.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Row with COL_TEXT and COL_TAGGED_TEXT] --> B[_format_rewrite_disposition_block]
    B --> C{protection_method?}
    C -->|replace| D[Excluded from REWRITE_DISPOSITION_BLOCK]
    C -->|generalize or remove| E[Included in REWRITE_DISPOSITION_BLOCK]

    A --> F[_apply_direct_replacements]
    F --> G[_get_replace_pairs]
    G --> H{raw_map present?}
    H -->|No| I[pairs is empty, replace_values filled]
    H -->|Yes| J[Match map entries with whitespace fallback]
    J --> K[pairs and replace_values returned]
    I --> L{unmatched entries?}
    K --> L
    L -->|yes| M[raise RuntimeError - PII protection refused]
    L -->|no| N{pairs non-empty?}
    N -->|No| O[Pass text through unchanged]
    N -->|Yes| P{COL_FINAL_ENTITIES has valid spans?}
    P -->|Yes| Q[_apply_replacement_map_to_text - span-aware]
    P -->|No| R[Fallback - single-pass regex without word boundaries]
    Q --> S[COL_PREREPLACE_TEXT]
    R --> S
    A --> T[_apply_tagged_text_replacements - single-pass regex on full tagged spans]
    T --> U[COL_PREREPLACE_TAGGED_TEXT]

    S --> V[LLM Rewrite - replacement_map block removed]
    U --> V
    E --> V
    V --> W[COL_REWRITTEN_TEXT]
    S --> X[Repair Prompt - baseline is COL_PREREPLACE_TEXT]
    W --> X
    X --> Y[COL_REWRITTEN_TEXT_NEXT]

Loading

_{Reviews (8): Last reviewed commit: "style: apply ruff format to rewrite_gene..." | Re-trigger Greptile}

[lipikaramaswamy]

This global regex does not preserve the span-aware matching semantics used by the Substitute workflow. For example, Ann → Maria transforms Ann met Anna into Maria met Mariaa, whereas _apply_replacement_map_to_text() replaces only detected entity spans.

Can the rewrite path reuse or adapt the existing span-aware replacement logic—including for tagged text—and add this case as a regression test?

[lipikaramaswamy]

One additional thought for the span-aware approach: could the regression test also confirm that every detected span selected for replace was actually transformed? That would give us an explicit coverage check while preserving the existing boundary-aware semantics.

[lipikaramaswamy]

Small logging-safety suggestion: sorted(unmatched) includes raw detected entity values, so the PII being protected can be persisted in operational logs. If this becomes an exception, including the values there would have the same issue. Could we report only safe metadata such as the record ID, entity IDs or labels, and counts?

[lipikaramaswamy]

Thanks for moving direct replacement into deterministic preprocessing—the overall direction makes sense. I’m requesting changes for the two safety and correctness gaps discussed inline: incomplete replacement maps currently fail open after replace entities are omitted from the prompt, and the global regex can modify text outside detected entity spans. There is also a smaller logging-safety concern around including raw entity values. Fail-safe handling, span-aware replacement, and focused regression tests would put this in good shape for another look.

[asteier2026]

Made these changes:

• Plain text: uses _apply_replacement_map_to_text with entity spans from COL_FINAL_ENTITIES — only replaces at exact detected span boundaries, fixing the Ann/Anna problem
• Tagged text: new _apply_tagged_text_replacements matches entity values only within their tag wrappers for all four notations (xml/bracket/paren/sentinel)
• Error logging: raw PII values removed; reports count + entity labels only
• _get_replace_pairs: now returns (original, synthetic, label) triples to support both fixes
• Two new regression tests: substring guard (Ann/Anna) and full-coverage check; existing tests updated with COL_TAG_NOTATION

[greptile-apps]

Sequential loop in _apply_tagged_text_replacements can still cascade

The function iterates over pairs sequentially, so a synthetic value that matches another entity's original is re-replaced in a later iteration. For example, with Alice → "Bob" and Bob → "Carlos", iteration 1 rewrites [[Alice|first_name]] to [[Bob|first_name]], and iteration 2 then matches the freshly written [[Bob|first_name]] and replaces it with [[Carlos|first_name]] — Alice ends up as Carlos in COL_PREREPLACE_TAGGED_TEXT while plain text correctly has "Bob". The LLM then sees wrong synthetic values in the tagged text it is asked to rewrite.

The existing cascade test (test_apply_direct_replacements_no_cascade_when_synthetic_matches_another_original) passes silently because COL_TAGGED_TEXT is set to "Alice and Bob met." (plain, untagged text). The xml-mode regex (<first_name>)Alice(</first_name>) finds no matches in that string, so _apply_tagged_text_replacements makes no substitutions and the cascade is never exercised. There is also no assertion on COL_PREREPLACE_TAGGED_TEXT in that test.

The plain-text path uses a single-pass regex (re.compile("|".join(...)).sub(...)) to avoid this exact problem. The tagged-text path needs the same treatment — build a single combined regex per notation format that matches all tagged originals simultaneously, then look up the replacement in one pass.

[asteier2026]

Just fixed this too