feat: apply direct identifier replacements before rewrite LLM call#208
feat: apply direct identifier replacements before rewrite LLM call#208asteier2026 wants to merge 7 commits into
Conversation
Direct identifiers are now substituted programmatically from the replacement map before the rewrite LLM sees the text, ensuring all occurrences are replaced consistently without relying on the LLM to apply a <replacement_map> block. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: asteier2026 <asteier@nvidia.com>
Greptile SummaryThis PR moves the substitution of
Confidence Score: 4/5The core pre-replacement logic is correct, but a documented safety contract in The single-pass regex, cascade-prevention, and Unicode whitespace fallback work as designed. The gap is that
Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Sequential str.replace() calls could incorrectly replace a synthetic value that happened to match another entity's original string (e.g. Alice→Bob then Bob→Carlos making Alice appear as Carlos). A combined regex alternation matches all originals simultaneously, eliminating the cascade. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: asteier2026 <asteier@nvidia.com>
Moved the missing-map warning into _get_replace_pairs so the disposition is parsed in one place. _apply_direct_replacements no longer re-parses COL_SENSITIVITY_DISPOSITION when _get_replace_pairs returns an empty list. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: asteier2026 <asteier@nvidia.com>
The LLM generating the replacement map occasionally normalises unusual Unicode whitespace (e.g. U+202F narrow no-break space) to a regular space in the original field. The exact-match lookup then misses the entity, triggering the unprotected warning. Add _normalize_ws and a second-pass lookup so that if the exact match fails, a whitespace-normalised comparison is tried. When a normalised match is found the disposition entity value (which reflects what is actually in the text) is used as the substitution key. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: asteier2026 <asteier@nvidia.com>
…cefully Any failure (malformed disposition, bad replacement map, regex error) previously caused the entire record to be skipped. Now the error is logged and the original text is passed through unchanged so the LLM rewrite step can still run. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: asteier2026 <asteier@nvidia.com>
The LLM generating the replacement map normalises unusual Unicode whitespace (e.g. U+202F narrow no-break space) to a regular space in the original field. _filter_replacement_map_to_input_entities was using an exact match against the detected entity values, so these entries were silently dropped from the map. Add a whitespace-normalized fallback: when the exact (original, label) pair is not in allowed_pairs, try a normalized comparison and, if it matches, rewrite original to the canonical detected value so all downstream lookups succeed. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: asteier2026 <asteier@nvidia.com>
| try: | ||
| pairs = _get_replace_pairs(row) | ||
| if pairs: | ||
| 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 = dict(sorted_pairs) | ||
| plain_text = pattern.sub(lambda m: lookup[m.group(0)], plain_text) | ||
| tagged_text = pattern.sub(lambda m: lookup[m.group(0)], tagged_text) | ||
| except Exception: | ||
| logger.warning("Direct replacement failed; passing text through unchanged.", exc_info=True) | ||
| row[COL_PREREPLACE_TEXT] = plain_text | ||
| row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text |
There was a problem hiding this comment.
replace entities completely unprotected
When the try block throws, plain_text and tagged_text fall back to the originals (with PII), but _format_rewrite_disposition_block has already run and excluded all replace entities from COL_REWRITE_DISPOSITION_BLOCK. The LLM prompt also no longer contains a replacement map. The net effect is that on any exception the LLM sees PII-containing text with zero instructions to anonymize the affected entities — they flow through to COL_REWRITTEN_TEXT completely unprotected.
In the old code the LLM had the replacement map as a fallback; that safety net no longer exists. Consider either re-raising (let the record fail explicitly) or, at minimum, logging at ERROR level and setting a failure flag so downstream callers can skip or quarantine the row rather than emit unprotected PII.
| def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]: | ||
| """Programmatically replace direct identifier entities before the rewrite LLM call. | ||
|
|
||
| Applies each (original → synthetic) pair as a global string replacement on both | ||
| the plain text (COL_TEXT → COL_PREREPLACE_TEXT) and the tagged text | ||
| (COL_TAGGED_TEXT → COL_PREREPLACE_TAGGED_TEXT). A single-pass regex substitution | ||
| ensures all occurrences are replaced without cascade (a synthetic value matching | ||
| another entity's original is never re-substituted). | ||
|
|
||
| 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 = _get_replace_pairs(row) | ||
| if pairs: | ||
| 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 = dict(sorted_pairs) | ||
| plain_text = pattern.sub(lambda m: lookup[m.group(0)], plain_text) | ||
| tagged_text = pattern.sub(lambda m: lookup[m.group(0)], tagged_text) | ||
| row[COL_PREREPLACE_TEXT] = plain_text | ||
| row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text | ||
| return row |
There was a problem hiding this comment.
_apply_direct_replacements documents that it "Raises on failure rather than falling back to unmodified text", citing exactly the risk scenario: replace entities are excluded from COL_REWRITE_DISPOSITION_BLOCK, so a silent passthrough would emit PII with no LLM instructions. But _get_replace_pairs has two paths that return [] without raising:
if not raw_map— logs a warning and returns[]- unmatched entities — logs a warning and returns partial pairs (possibly
[])
When either path fires, pairs is empty or incomplete, the if pairs: branch is skipped, plain_text/tagged_text remain as the raw PII-containing originals, and both columns are written straight to COL_PREREPLACE_TEXT/COL_PREREPLACE_TAGGED_TEXT. The LLM then sees unredacted PII with zero instructions to protect it. The test test_get_replace_pairs_warns_on_unmatched_replace_entity explicitly asserts pairs == [] for this path, confirming the silent-passthrough survives.
Consider having _apply_direct_replacements compute replace_values (or receive it from _get_replace_pairs) and raise when required pairs are missing, consistent with the documented contract.
1 participant
Summary
the replacement map before the rewrite LLM call, using a single-pass regex to prevent cascade
replacements (e.g. Alice→Bob→Carlos)
replacements
_get_replace_pairs to handle LLM-normalised Unicode whitespace (e.g. U+202F → U+0020) in entity values
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.