From 9703ddad2ad26a52cbf5e70fa304160121a94462 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 2 Jul 2026 08:47:09 -0700 Subject: [PATCH 1/7] feat: apply direct identifier replacements before rewrite LLM call 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 block. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: asteier2026 --- src/anonymizer/engine/constants.py | 3 +- src/anonymizer/engine/rewrite/repair.py | 6 +- .../engine/rewrite/rewrite_generation.py | 92 +++++--- tests/engine/test_rewrite_generation.py | 204 +++++++++++------- 4 files changed, 194 insertions(+), 111 deletions(-) diff --git a/src/anonymizer/engine/constants.py b/src/anonymizer/engine/constants.py index 3cb0876a..d6b577c8 100644 --- a/src/anonymizer/engine/constants.py +++ b/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" diff --git a/src/anonymizer/engine/rewrite/repair.py b/src/anonymizer/engine/rewrite/repair.py index 4a18b1f8..a438aaf9 100644 --- a/src/anonymizer/engine/rewrite/repair.py +++ b/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 = { "<>": params.privacy_goal_str, "<>": str(params.max_privacy_leak), - "<>": str(row.get(COL_TEXT, "")), + "<>": str(row.get(COL_PREREPLACE_TEXT, "")), "<>": str(row.get(COL_REWRITTEN_TEXT, "")), "<>": str(row.get(COL_LEAKAGE_MASS, 0.0)), "<>": "\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, diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 7892b1cf..9a580639 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -14,13 +14,15 @@ from anonymizer.config.rewrite import PrivacyGoal from anonymizer.engine.constants import ( 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 @@ -49,7 +51,7 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No 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 +87,8 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No Entities NOT listed above may be kept as-is. -{% if <>.replacements %} - -Synthetic replacement values for entities with protection_method "replace": -<> - -{% endif %} 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 +107,8 @@ def _get_rewrite_prompt(privacy_goal: PrivacyGoal, data_summary: str | None = No "<>": privacy_goal.to_prompt_string(), "<>": data_context_section, "<>": COL_TAG_NOTATION, - "<>": _jinja(COL_TAGGED_TEXT), + "<>": _jinja(COL_PREREPLACE_TAGGED_TEXT), "<>": COL_REWRITE_DISPOSITION_BLOCK, - "<>": COL_REPLACEMENT_MAP_FOR_PROMPT, - "<>": _jinja(COL_REPLACEMENT_MAP_FOR_PROMPT), }, ) @@ -128,12 +120,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 +146,59 @@ 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 _get_replace_pairs(row: dict[str, Any]) -> list[tuple[str, str]]: + """Return (original, synthetic) pairs for entities with protection_method='replace'.""" + 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" } 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 or not replace_values: + return [] 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} + pairs = [(r.original, r.synthetic) for r in parsed_map.replacements if r.original in replace_values] + unmatched = replace_values - {original for original, _ in pairs} + if unmatched: + logger.warning( + "Replace entities have no entry in the replacement map and will pass through unprotected: %s", + sorted(unmatched), + ) + return pairs + + +@custom_column_generator( + required_columns=[COL_SENSITIVITY_DISPOSITION, COL_REPLACEMENT_MAP, COL_TEXT, COL_TAGGED_TEXT], + 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. + + 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). Using str.replace() ensures all + occurrences are substituted, not just detected spans. + """ + pairs = _get_replace_pairs(row) + plain_text = str(row.get(COL_TEXT, "")) + tagged_text = str(row.get(COL_TAGGED_TEXT, "")) + if not pairs: + disposition = parse_sensitivity_disposition(row[COL_SENSITIVITY_DISPOSITION]) + has_replace_entities = any( + e.protection_method_suggestion == "replace" for e in disposition.sensitivity_disposition + ) + if has_replace_entities and not row.get(COL_REPLACEMENT_MAP): + logger.warning("COL_REPLACEMENT_MAP is None but entities require replacement; no replacements applied.") + row[COL_PREREPLACE_TEXT] = plain_text + row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text + return row + for original, synthetic in sorted(pairs, key=lambda p: len(p[0]), reverse=True): + plain_text = plain_text.replace(original, synthetic) + tagged_text = tagged_text.replace(original, synthetic) + row[COL_PREREPLACE_TEXT] = plain_text + row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text return row @@ -226,8 +254,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, diff --git a/tests/engine/test_rewrite_generation.py b/tests/engine/test_rewrite_generation.py index be32ad44..71c8db82 100644 --- a/tests/engine/test_rewrite_generation.py +++ b/tests/engine/test_rewrite_generation.py @@ -10,19 +10,21 @@ from anonymizer.config.rewrite import PrivacyGoal from anonymizer.engine.constants import ( 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.rewrite.rewrite_generation import ( RewriteGenerationWorkflow, + _apply_direct_replacements, _extract_rewritten_text, - _filter_replacement_map_for_prompt, _format_rewrite_disposition_block, _get_rewrite_prompt, ) @@ -70,14 +72,36 @@ def stub_replacement_map() -> dict: # --------------------------------------------------------------------------- -def test_format_rewrite_disposition_block_filters_to_protected_only( +def test_format_rewrite_disposition_block_excludes_replace_entities( stub_sensitivity_disposition: dict, ) -> None: + # replace entities are handled programmatically; they should not appear in the prompt block row = {COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition} result = _format_rewrite_disposition_block(row) + assert result[COL_REWRITE_DISPOSITION_BLOCK] == [] + + +def test_format_rewrite_disposition_block_includes_non_replace_protected_entities() -> None: + disposition = { + "sensitivity_disposition": [ + { + "id": 1, + "source": "tagged", + "category": "quasi_identifier", + "sensitivity": "medium", + "entity_label": "city", + "entity_value": "Portland", + "protection_reason": "Quasi-identifier", + "protection_method_suggestion": "generalize", + "combined_risk_level": "medium", + } + ] + } + row = {COL_SENSITIVITY_DISPOSITION: disposition} + result = _format_rewrite_disposition_block(row) block = result[COL_REWRITE_DISPOSITION_BLOCK] assert len(block) == 1 - assert block[0]["entity_value"] == "Alice" + assert block[0]["entity_value"] == "Portland" def test_format_rewrite_disposition_block_excludes_unprotected_entities() -> None: @@ -101,108 +125,134 @@ def test_format_rewrite_disposition_block_excludes_unprotected_entities() -> Non assert result[COL_REWRITE_DISPOSITION_BLOCK] == [] -def test_format_rewrite_disposition_block_serializes_required_fields( - stub_sensitivity_disposition: dict, -) -> None: - row = {COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition} - result = _format_rewrite_disposition_block(row) - block = result[COL_REWRITE_DISPOSITION_BLOCK] - entry = block[0] - assert set(entry.keys()) == { - "entity_label", - "entity_value", - "sensitivity", - "protection_method_suggestion", - "protection_reason", - } - - -def test_format_rewrite_disposition_block_empty_when_no_protected_entities() -> None: +def test_format_rewrite_disposition_block_serializes_required_fields() -> None: disposition = { "sensitivity_disposition": [ { "id": 1, "source": "tagged", "category": "quasi_identifier", - "sensitivity": "low", + "sensitivity": "medium", "entity_label": "city", "entity_value": "Portland", - "protection_reason": "Not identifying alone", - "protection_method_suggestion": "leave_as_is", - "combined_risk_level": "low", + "protection_reason": "Quasi-identifier", + "protection_method_suggestion": "generalize", + "combined_risk_level": "medium", } ] } row = {COL_SENSITIVITY_DISPOSITION: disposition} result = _format_rewrite_disposition_block(row) - assert result[COL_REWRITE_DISPOSITION_BLOCK] == [] + entry = result[COL_REWRITE_DISPOSITION_BLOCK][0] + assert set(entry.keys()) == { + "entity_label", + "entity_value", + "sensitivity", + "protection_method_suggestion", + "protection_reason", + } # --------------------------------------------------------------------------- -# Tests: _filter_replacement_map_for_prompt +# Tests: _get_replace_pairs # --------------------------------------------------------------------------- -def test_filter_replacement_map_keeps_only_replace_method_entities( +def test_get_replace_pairs_warns_on_unmatched_replace_entity( + stub_sensitivity_disposition: dict, + caplog: pytest.LogCaptureFixture, +) -> None: + import logging + + from anonymizer.engine.rewrite.rewrite_generation import _get_replace_pairs + + # Map exists but has no entry for "Alice" + row = { + COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition, + COL_REPLACEMENT_MAP: {"replacements": [{"original": "Bob", "label": "first_name", "synthetic": "Carlos"}]}, + } + with caplog.at_level(logging.WARNING, logger="anonymizer.rewrite.generation"): + pairs = _get_replace_pairs(row) + assert pairs == [] + assert "Alice" in caplog.text + assert "unprotected" in caplog.text + + +# --------------------------------------------------------------------------- +# Tests: _apply_direct_replacements +# --------------------------------------------------------------------------- + + +def test_apply_direct_replacements_substitutes_in_plain_and_tagged_text( stub_sensitivity_disposition: dict, stub_replacement_map: dict, ) -> None: - disposition_block = [ - { - "entity_label": "first_name", - "entity_value": "Alice", - "sensitivity": "high", - "protection_method_suggestion": "replace", - "protection_reason": "Direct identifier", - } - ] row = { + COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition, COL_REPLACEMENT_MAP: stub_replacement_map, - COL_REWRITE_DISPOSITION_BLOCK: disposition_block, + COL_TEXT: "Alice works at TechCorp.", + COL_TAGGED_TEXT: "[[Alice|first_name]] works at TechCorp.", } - result = _filter_replacement_map_for_prompt(row) - filtered = result[COL_REPLACEMENT_MAP_FOR_PROMPT] - assert len(filtered["replacements"]) == 1 - assert filtered["replacements"][0]["original"] == "Alice" - - -def test_filter_replacement_map_empty_when_no_replace_method() -> None: - disposition_block = [ - { - "entity_label": "city", - "entity_value": "Portland", - "sensitivity": "low", - "protection_method_suggestion": "generalize", - "protection_reason": "Quasi-identifier", - } - ] + result = _apply_direct_replacements(row) + assert result[COL_PREREPLACE_TEXT] == "Maria works at TechCorp." + assert result[COL_PREREPLACE_TAGGED_TEXT] == "[[Maria|first_name]] works at TechCorp." + + +def test_apply_direct_replacements_replaces_all_occurrences( + stub_sensitivity_disposition: dict, + stub_replacement_map: dict, +) -> None: row = { - COL_REPLACEMENT_MAP: {"replacements": [{"original": "Portland", "label": "city", "synthetic": "Seattle"}]}, - COL_REWRITE_DISPOSITION_BLOCK: disposition_block, + COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition, + COL_REPLACEMENT_MAP: stub_replacement_map, + COL_TEXT: "Alice told Alice's manager that Alice would be late.", + COL_TAGGED_TEXT: "[[Alice|first_name]] told [[Alice|first_name]]'s manager.", } - result = _filter_replacement_map_for_prompt(row) - assert result[COL_REPLACEMENT_MAP_FOR_PROMPT]["replacements"] == [] + result = _apply_direct_replacements(row) + assert "Alice" not in result[COL_PREREPLACE_TEXT] + assert "Alice" not in result[COL_PREREPLACE_TAGGED_TEXT] -def test_filter_replacement_map_accepts_schema_instance( +def test_apply_direct_replacements_passthrough_when_no_replace_entities() -> None: + disposition = { + "sensitivity_disposition": [ + { + "id": 1, + "source": "tagged", + "category": "quasi_identifier", + "sensitivity": "low", + "entity_label": "city", + "entity_value": "Portland", + "protection_reason": "Quasi-identifier", + "protection_method_suggestion": "leave_as_is", + "combined_risk_level": "low", + } + ] + } + row = { + COL_SENSITIVITY_DISPOSITION: disposition, + COL_REPLACEMENT_MAP: {"replacements": []}, + COL_TEXT: "She lives in Portland.", + COL_TAGGED_TEXT: "She lives in [[Portland|city]].", + } + result = _apply_direct_replacements(row) + assert result[COL_PREREPLACE_TEXT] == "She lives in Portland." + assert result[COL_PREREPLACE_TAGGED_TEXT] == "She lives in [[Portland|city]]." + + +def test_apply_direct_replacements_accepts_schema_instance( + stub_sensitivity_disposition: dict, stub_replacement_map: dict, ) -> None: - disposition_block = [ - { - "entity_label": "first_name", - "entity_value": "Alice", - "sensitivity": "high", - "protection_method_suggestion": "replace", - "protection_reason": "Direct identifier", - } - ] schema = EntityReplacementMapSchema.model_validate(stub_replacement_map) row = { + COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition, COL_REPLACEMENT_MAP: schema, - COL_REWRITE_DISPOSITION_BLOCK: disposition_block, + COL_TEXT: "Alice works here.", + COL_TAGGED_TEXT: "[[Alice|first_name]] works here.", } - result = _filter_replacement_map_for_prompt(row) - assert len(result[COL_REPLACEMENT_MAP_FOR_PROMPT]["replacements"]) == 1 + result = _apply_direct_replacements(row) + assert result[COL_PREREPLACE_TEXT] == "Maria works here." # --------------------------------------------------------------------------- @@ -271,10 +321,14 @@ def test_get_rewrite_prompt_no_data_context_when_none(privacy_goal: PrivacyGoal) def test_get_rewrite_prompt_references_required_columns(privacy_goal: PrivacyGoal) -> None: prompt = _get_rewrite_prompt(privacy_goal) - assert _jinja(COL_TAGGED_TEXT) in prompt + assert _jinja(COL_PREREPLACE_TAGGED_TEXT) in prompt assert COL_TAG_NOTATION in prompt assert COL_REWRITE_DISPOSITION_BLOCK in prompt - assert _jinja(COL_REPLACEMENT_MAP_FOR_PROMPT) in prompt + + +def test_get_rewrite_prompt_has_no_replacement_map_block(privacy_goal: PrivacyGoal) -> None: + prompt = _get_rewrite_prompt(privacy_goal) + assert "replacement_map" not in prompt # --------------------------------------------------------------------------- @@ -312,7 +366,7 @@ def test_columns_full_rewrite_uses_rewrite_output_schema( assert full_rewrite_col.output_format == RewriteOutputSchema.model_json_schema() -def test_columns_includes_custom_configs_for_disposition_and_text_extraction( +def test_columns_includes_custom_configs_for_disposition_prereplace_and_text_extraction( stub_rewrite_model_selection: RewriteModelSelection, privacy_goal: PrivacyGoal, ) -> None: @@ -320,5 +374,5 @@ def test_columns_includes_custom_configs_for_disposition_and_text_extraction( cols = workflow.columns(selected_models=stub_rewrite_model_selection, privacy_goal=privacy_goal) custom_names = {c.name for c in cols if isinstance(c, CustomColumnConfig)} assert COL_REWRITE_DISPOSITION_BLOCK in custom_names - assert COL_REPLACEMENT_MAP_FOR_PROMPT in custom_names + assert COL_PREREPLACE_TEXT in custom_names assert COL_REWRITTEN_TEXT in custom_names From 2f35a4bffdfdca964b91fa5b7e7de0d3ec20efa1 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 2 Jul 2026 09:02:39 -0700 Subject: [PATCH 2/7] fix: use single-pass regex to prevent cascade replacements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Signed-off-by: asteier2026 --- .../engine/rewrite/rewrite_generation.py | 9 ++-- tests/engine/test_rewrite_generation.py | 44 +++++++++++++++++++ 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 9a580639..c4031ff8 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/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 @@ -194,9 +195,11 @@ def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]: row[COL_PREREPLACE_TEXT] = plain_text row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text return row - for original, synthetic in sorted(pairs, key=lambda p: len(p[0]), reverse=True): - plain_text = plain_text.replace(original, synthetic) - tagged_text = tagged_text.replace(original, synthetic) + 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 diff --git a/tests/engine/test_rewrite_generation.py b/tests/engine/test_rewrite_generation.py index 71c8db82..206963f4 100644 --- a/tests/engine/test_rewrite_generation.py +++ b/tests/engine/test_rewrite_generation.py @@ -255,6 +255,50 @@ def test_apply_direct_replacements_accepts_schema_instance( assert result[COL_PREREPLACE_TEXT] == "Maria works here." +def test_apply_direct_replacements_no_cascade_when_synthetic_matches_another_original() -> None: + """Single-pass regex prevents Alice→Bob→Carlos cascade.""" + disposition = { + "sensitivity_disposition": [ + { + "id": 1, + "source": "tagged", + "category": "direct_identifier", + "sensitivity": "high", + "entity_label": "first_name", + "entity_value": "Alice", + "protection_reason": "Direct identifier.", + "protection_method_suggestion": "replace", + "combined_risk_level": "high", + }, + { + "id": 2, + "source": "tagged", + "category": "direct_identifier", + "sensitivity": "high", + "entity_label": "first_name", + "entity_value": "Bob", + "protection_reason": "Direct identifier.", + "protection_method_suggestion": "replace", + "combined_risk_level": "high", + }, + ] + } + replacement_map = { + "replacements": [ + {"original": "Alice", "label": "first_name", "synthetic": "Bob"}, + {"original": "Bob", "label": "first_name", "synthetic": "Carlos"}, + ] + } + row = { + COL_SENSITIVITY_DISPOSITION: disposition, + COL_REPLACEMENT_MAP: replacement_map, + COL_TEXT: "Alice and Bob met.", + COL_TAGGED_TEXT: "Alice and Bob met.", + } + result = _apply_direct_replacements(row) + assert result[COL_PREREPLACE_TEXT] == "Bob and Carlos met." + + # --------------------------------------------------------------------------- # Tests: _extract_rewritten_text # --------------------------------------------------------------------------- From 0da33e3dfc76ca7b9628d0e8d8ae26ea3eb08127 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 2 Jul 2026 09:05:34 -0700 Subject: [PATCH 3/7] fix: parse sensitivity disposition only once in _get_replace_pairs 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 Signed-off-by: asteier2026 --- .../engine/rewrite/rewrite_generation.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index c4031ff8..c9d86372 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -153,8 +153,11 @@ def _get_replace_pairs(row: dict[str, Any]) -> list[tuple[str, str]]: replace_values = { e.entity_value for e in disposition.sensitivity_disposition if e.protection_method_suggestion == "replace" } + if not replace_values: + return [] raw_map = row.get(COL_REPLACEMENT_MAP) - if not raw_map or not replace_values: + if not raw_map: + logger.warning("COL_REPLACEMENT_MAP is None but entities require replacement; no replacements applied.") return [] raw_map = normalize_payload(raw_map) if hasattr(raw_map, "model_dump"): @@ -179,27 +182,19 @@ def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]: 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). Using str.replace() ensures all - occurrences are substituted, not just detected spans. + (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). """ pairs = _get_replace_pairs(row) plain_text = str(row.get(COL_TEXT, "")) tagged_text = str(row.get(COL_TAGGED_TEXT, "")) - if not pairs: - disposition = parse_sensitivity_disposition(row[COL_SENSITIVITY_DISPOSITION]) - has_replace_entities = any( - e.protection_method_suggestion == "replace" for e in disposition.sensitivity_disposition - ) - if has_replace_entities and not row.get(COL_REPLACEMENT_MAP): - logger.warning("COL_REPLACEMENT_MAP is None but entities require replacement; no replacements applied.") - row[COL_PREREPLACE_TEXT] = plain_text - row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text - return row - 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) + 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 From c439bb2a322468f6f35ca934e7482e70d6dc1cee Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 2 Jul 2026 11:21:46 -0700 Subject: [PATCH 4/7] fix: fall back to whitespace-normalized match in _get_replace_pairs 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 Signed-off-by: asteier2026 --- .../engine/rewrite/rewrite_generation.py | 32 +++++++++++++++-- tests/engine/test_rewrite_generation.py | 34 +++++++++++++++++++ 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index c9d86372..24a13b7d 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -147,8 +147,19 @@ def _format_rewrite_disposition_block(row: dict[str, Any]) -> dict[str, Any]: return row +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]) -> list[tuple[str, str]]: - """Return (original, synthetic) pairs for entities with protection_method='replace'.""" + """Return (original, synthetic) pairs for entities with protection_method='replace'. + + 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.sensitivity_disposition if e.protection_method_suggestion == "replace" @@ -163,8 +174,23 @@ def _get_replace_pairs(row: dict[str, Any]) -> list[tuple[str, str]]: if hasattr(raw_map, "model_dump"): raw_map = raw_map.model_dump(mode="python") parsed_map = EntityReplacementMapSchema.model_validate(raw_map) - pairs = [(r.original, r.synthetic) for r in parsed_map.replacements if r.original in replace_values] - unmatched = replace_values - {original for original, _ in pairs} + + # 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]] = [] + matched: set[str] = set() + for r in parsed_map.replacements: + if r.original in replace_values: + pairs.append((r.original, r.synthetic)) + 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)) + matched.add(disposition_value) + + unmatched = replace_values - matched if unmatched: logger.warning( "Replace entities have no entry in the replacement map and will pass through unprotected: %s", diff --git a/tests/engine/test_rewrite_generation.py b/tests/engine/test_rewrite_generation.py index 206963f4..01e07472 100644 --- a/tests/engine/test_rewrite_generation.py +++ b/tests/engine/test_rewrite_generation.py @@ -178,6 +178,40 @@ def test_get_replace_pairs_warns_on_unmatched_replace_entity( assert "unprotected" in caplog.text +def test_get_replace_pairs_matches_unicode_whitespace_variant(caplog: pytest.LogCaptureFixture) -> None: + """Disposition has U+202F (narrow no-break space); map was normalized to regular space by LLM.""" + import logging + + from anonymizer.engine.rewrite.rewrite_generation import _get_replace_pairs + + narrow_nbsp = " " + disposition = { + "sensitivity_disposition": [ + { + "id": 1, + "source": "tagged", + "category": "direct_identifier", + "sensitivity": "high", + "entity_label": "street_address", + "entity_value": f"204{narrow_nbsp}Bluegrass", + "protection_reason": "Street address identifies the subject.", + "protection_method_suggestion": "replace", + "combined_risk_level": "high", + } + ] + } + row = { + COL_SENSITIVITY_DISPOSITION: disposition, + COL_REPLACEMENT_MAP: { + "replacements": [{"original": "204 Bluegrass", "label": "street_address", "synthetic": "500 Oak Lane"}] + }, + } + with caplog.at_level(logging.WARNING, logger="anonymizer.rewrite.generation"): + pairs = _get_replace_pairs(row) + assert pairs == [(f"204{narrow_nbsp}Bluegrass", "500 Oak Lane")] + assert "unprotected" not in caplog.text + + # --------------------------------------------------------------------------- # Tests: _apply_direct_replacements # --------------------------------------------------------------------------- From 48f897f4f7cd9111becc0ecaa0d0df94b4cf62a3 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 2 Jul 2026 11:23:43 -0700 Subject: [PATCH 5/7] fix: catch exceptions in _apply_direct_replacements and fall back gracefully 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 Signed-off-by: asteier2026 --- .../engine/rewrite/rewrite_generation.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 24a13b7d..90df2f59 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -211,16 +211,22 @@ def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]: (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). + + On any failure, falls back to the unmodified text so the record is not dropped. + The LLM rewrite step will still run; it just won't have pre-applied replacements. """ - pairs = _get_replace_pairs(row) plain_text = str(row.get(COL_TEXT, "")) tagged_text = str(row.get(COL_TAGGED_TEXT, "")) - 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) + 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 return row From dead64e2c0d45d5b217e2d3b82c9c2b233957ebc Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 2 Jul 2026 11:36:58 -0700 Subject: [PATCH 6/7] fix: normalize Unicode whitespace when filtering replacement map entries 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 Signed-off-by: asteier2026 --- .../engine/replace/llm_replace_workflow.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/anonymizer/engine/replace/llm_replace_workflow.py b/src/anonymizer/engine/replace/llm_replace_workflow.py index 531b6827..cf11710e 100644 --- a/src/anonymizer/engine/replace/llm_replace_workflow.py +++ b/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 From 2794d1930a33dcb67048a1b1871f5483a5e5388e Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Thu, 9 Jul 2026 09:50:08 -0700 Subject: [PATCH 7/7] fix:raise error if no replace --- .../engine/rewrite/rewrite_generation.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 90df2f59..d4d9e749 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -212,21 +212,19 @@ def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]: ensures all occurrences are replaced without cascade (a synthetic value matching another entity's original is never re-substituted). - On any failure, falls back to the unmodified text so the record is not dropped. - The LLM rewrite step will still run; it just won't have pre-applied replacements. + 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, "")) - 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) + 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