From 48c716c4c5291df599ddfb52a80c4dd4e11f1cf2 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Mon, 29 Jun 2026 09:49:52 -0700 Subject: [PATCH 1/4] feature: rewrite replacement chg and entity augment chg --- src/anonymizer/engine/constants.py | 3 +- .../engine/detection/detection_workflow.py | 18 ++ src/anonymizer/engine/rewrite/repair.py | 6 +- .../engine/rewrite/rewrite_generation.py | 81 ++++---- src/anonymizer/engine/schemas/__init__.py | 2 - src/anonymizer/engine/schemas/detection.py | 2 +- src/anonymizer/engine/schemas/rewrite.py | 23 +-- tests/engine/test_qa_generation.py | 5 +- tests/engine/test_rewrite_generation.py | 179 ++++++++++-------- 9 files changed, 181 insertions(+), 138 deletions(-) diff --git a/src/anonymizer/engine/constants.py b/src/anonymizer/engine/constants.py index d9e1f079..d32c5a66 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/detection/detection_workflow.py b/src/anonymizer/engine/detection/detection_workflow.py index 1c762844..947fe667 100644 --- a/src/anonymizer/engine/detection/detection_workflow.py +++ b/src/anonymizer/engine/detection/detection_workflow.py @@ -682,6 +682,24 @@ def _get_augment_prompt(*, data_summary: str | None, labels: list[str], strict_l - Filename Exclusions: The "filename" label should be reserved for user-created documents or data exports (e.g., .pdf, .xlsx, .csv, .txt). - Executable Distinction: Do not tag executable binaries (ending in .exe, .dll, or .sys) as filename. Treat these extensions as non-sensitive system identifiers. +Additional extraction requirements: +- The "value" field must be the EXACT verbatim span from the input text. + Copy the text character-for-character exactly as it appears. + Do NOT normalize, correct spelling, expand abbreviations, decode encodings, + infer hidden values, translate text, reformat numbers, or otherwise modify + the extracted span. +- Extract only text that is explicitly present in the input. + Never reconstruct, guess, or generate a value that does not appear verbatim. +- Identifiers may be disguised, fragmented, hyphenated, misspelled, obfuscated, + spaced out, mixed with punctuation, or written in words instead of digits. + Detect the identifier and extract the exact text as written. + Examples of disguised identifiers to detect: + - Phone numbers, SSNs, and credit card numbers spoken as digit words, + including "o" or "oh" used in place of zero: + "nine o two, five five five, one two three four" + - Names spelled out letter by letter with hyphens or commas: + "J-O-H-N", "M, A, R, Y" + <> --- 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..17322ab4 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,48 @@ 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} + return [(r.original, r.synthetic) for r in parsed_map.replacements if r.original in replace_values] + + +@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: + if 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 pairs: + 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 +243,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/src/anonymizer/engine/schemas/__init__.py b/src/anonymizer/engine/schemas/__init__.py index 96607d66..bc99215d 100644 --- a/src/anonymizer/engine/schemas/__init__.py +++ b/src/anonymizer/engine/schemas/__init__.py @@ -34,7 +34,6 @@ EntityCategory, EntityDispositionSchema, EntitySource, - MeaningUnitAspect, MeaningUnitImportance, MeaningUnitSchema, MeaningUnitsSchema, @@ -90,7 +89,6 @@ "EntityCategory", "EntityDispositionSchema", "EntitySource", - "MeaningUnitAspect", "MeaningUnitImportance", "MeaningUnitSchema", "MeaningUnitsSchema", diff --git a/src/anonymizer/engine/schemas/detection.py b/src/anonymizer/engine/schemas/detection.py index 6c6c3dea..cd5eed6d 100644 --- a/src/anonymizer/engine/schemas/detection.py +++ b/src/anonymizer/engine/schemas/detection.py @@ -157,7 +157,7 @@ class LatentEntitySchema(BaseModel): ) rationale: str = Field( min_length=20, - max_length=150, + max_length=250, description="One sentence explaining the inference without adding new facts", ) diff --git a/src/anonymizer/engine/schemas/rewrite.py b/src/anonymizer/engine/schemas/rewrite.py index 912281dd..646a630f 100644 --- a/src/anonymizer/engine/schemas/rewrite.py +++ b/src/anonymizer/engine/schemas/rewrite.py @@ -33,7 +33,7 @@ Uses LLMJudgeColumnConfig with Score rubrics (no custom schema needed) Supporting enums: Domain, EntitySource, EntityCategory, SensitivityLevel, - ProtectionMethod, CombinedRiskLevel, MeaningUnitAspect, PrivacyAnswer + ProtectionMethod, CombinedRiskLevel, PrivacyAnswer """ from __future__ import annotations @@ -249,25 +249,6 @@ class StrictSensitivityDispositionSchema(SensitivityDispositionSchema): # --------------------------------------------------------------------------- -class MeaningUnitAspect(str, Enum): - ROLE = "role" - PROCESS = "process" - RELATIONSHIP = "relationship" - ENVIRONMENT = "environment" - ROUTINE = "routine" - CREATIVE_OUTPUT = "creative_output" - VALUE = "value" - MOTIVATION = "motivation" - INFLUENCE = "influence" - AUDIENCE = "audience" - LEGAL_BASIS = "legal_basis" - INSTITUTION = "institution" - JUSTIFICATION = "justification" - PROCEDURAL_STATUS = "procedural_status" - TEMPORAL_SEQUENCE = "temporal_sequence" - RIGHTS_IMPACT = "rights_impact" - - class MeaningUnitImportance(str, Enum): critical = "critical" important = "important" @@ -275,7 +256,7 @@ class MeaningUnitImportance(str, Enum): class MeaningUnitSchema(BaseModel): id: int = Field(ge=1) - aspect: MeaningUnitAspect + aspect: str unit: str = Field(min_length=1) importance: MeaningUnitImportance diff --git a/tests/engine/test_qa_generation.py b/tests/engine/test_qa_generation.py index 74831f9d..da2cc5a5 100644 --- a/tests/engine/test_qa_generation.py +++ b/tests/engine/test_qa_generation.py @@ -34,7 +34,6 @@ EntityCategory, EntityDispositionSchema, EntitySource, - MeaningUnitAspect, MeaningUnitImportance, MeaningUnitSchema, MeaningUnitsSchema, @@ -75,13 +74,13 @@ units=[ MeaningUnitSchema( id=1, - aspect=MeaningUnitAspect.ROLE, + aspect="role", unit="An individual works as a software engineer.", importance=MeaningUnitImportance.critical, ), MeaningUnitSchema( id=2, - aspect=MeaningUnitAspect.ENVIRONMENT, + aspect="environment", unit="The individual works remotely.", importance=MeaningUnitImportance.important, ), diff --git a/tests/engine/test_rewrite_generation.py b/tests/engine/test_rewrite_generation.py index be32ad44..0f4423b7 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,109 @@ 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: _apply_direct_replacements # --------------------------------------------------------------------------- -def test_filter_replacement_map_keeps_only_replace_method_entities( +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 +296,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 +341,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 +349,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 7dc6962d323838fc56df2547e9938d91bafe04c7 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Mon, 29 Jun 2026 11:21:03 -0700 Subject: [PATCH 2/4] Update src/anonymizer/engine/rewrite/rewrite_generation.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- src/anonymizer/engine/rewrite/rewrite_generation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 17322ab4..533ad3a5 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -183,7 +183,7 @@ 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 pairs: + 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 From 76f114f60691c154196b93f983d031f250f8e2e4 Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Mon, 29 Jun 2026 11:22:19 -0700 Subject: [PATCH 3/4] Update src/anonymizer/engine/rewrite/rewrite_generation.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- src/anonymizer/engine/rewrite/rewrite_generation.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 533ad3a5..67633a94 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -178,7 +178,11 @@ def _apply_direct_replacements(row: dict[str, Any]) -> dict[str, Any]: plain_text = str(row.get(COL_TEXT, "")) tagged_text = str(row.get(COL_TAGGED_TEXT, "")) if not pairs: - if not row.get(COL_REPLACEMENT_MAP): + 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 From 183624db9db6b19fa91b83a0f4ab4dbbe206851b Mon Sep 17 00:00:00 2001 From: asteier2026 Date: Mon, 29 Jun 2026 13:26:45 -0700 Subject: [PATCH 4/4] fix: address no matching replacement value --- .../engine/rewrite/rewrite_generation.py | 9 ++++++- tests/engine/test_rewrite_generation.py | 25 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/src/anonymizer/engine/rewrite/rewrite_generation.py b/src/anonymizer/engine/rewrite/rewrite_generation.py index 67633a94..9a580639 100644 --- a/src/anonymizer/engine/rewrite/rewrite_generation.py +++ b/src/anonymizer/engine/rewrite/rewrite_generation.py @@ -159,7 +159,14 @@ 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) - return [(r.original, r.synthetic) for r in parsed_map.replacements if r.original in replace_values] + 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( diff --git a/tests/engine/test_rewrite_generation.py b/tests/engine/test_rewrite_generation.py index 0f4423b7..71c8db82 100644 --- a/tests/engine/test_rewrite_generation.py +++ b/tests/engine/test_rewrite_generation.py @@ -153,6 +153,31 @@ def test_format_rewrite_disposition_block_serializes_required_fields() -> None: } +# --------------------------------------------------------------------------- +# Tests: _get_replace_pairs +# --------------------------------------------------------------------------- + + +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 # ---------------------------------------------------------------------------