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/replace/llm_replace_workflow.py b/src/anonymizer/engine/replace/llm_replace_workflow.py index 531b6827..784cbbab 100644 --- a/src/anonymizer/engine/replace/llm_replace_workflow.py +++ b/src/anonymizer/engine/replace/llm_replace_workflow.py @@ -164,13 +164,28 @@ 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 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..e229320b 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 @@ -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 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. -{% 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 +111,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 +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,145 @@ 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 a single-pass, tag-boundary-aware substitution. + + Builds a lookup from full tagged span → replacement tagged span, then applies all + substitutions in one regex pass so a synthetic value that matches another entity's + original is never re-replaced (same cascade-prevention guarantee as the plain-text path). + """ + if not pairs: + return tagged_text + lookup: dict[str, str] = {} + for original, synthetic, label in sorted(pairs, key=lambda p: len(p[0]), reverse=True): + if tag_notation == "xml": + tagged_original = f"<{label}>{original}" + tagged_synthetic = f"<{label}>{synthetic}" + elif tag_notation == "bracket": + tagged_original = f"[[{original}|{label}]]" + tagged_synthetic = f"[[{synthetic}|{label}]]" + elif tag_notation == "paren": + tagged_original = f"((SENSITIVE:{label}|{original}))" + tagged_synthetic = f"((SENSITIVE:{label}|{synthetic}))" + else: # sentinel + tagged_original = f"<>{original}<>" + tagged_synthetic = f"<>{synthetic}<>" + lookup[tagged_original] = tagged_synthetic + pattern = re.compile("|".join(re.escape(k) for k in lookup)) + return pattern.sub(lambda m: lookup[m.group(0)], 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 +344,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..e6fd4b6e 100644 --- a/tests/engine/test_rewrite_generation.py +++ b/tests/engine/test_rewrite_generation.py @@ -9,20 +9,23 @@ 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.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 +73,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,13 +126,25 @@ 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} +def test_format_rewrite_disposition_block_serializes_required_fields() -> 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] - entry = block[0] + entry = result[COL_REWRITE_DISPOSITION_BLOCK][0] assert set(entry.keys()) == { "entity_label", "entity_value", @@ -117,92 +154,267 @@ def test_format_rewrite_disposition_block_serializes_required_fields( } -def test_format_rewrite_disposition_block_empty_when_no_protected_entities() -> None: +# --------------------------------------------------------------------------- +# Tests: _get_replace_pairs +# --------------------------------------------------------------------------- + + +def test_apply_direct_replacements_raises_on_unmatched_replace_entity( + stub_sensitivity_disposition: dict, +) -> None: + # Map exists but has no entry for "Alice" — must raise, not silently pass through PII + row = { + COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition, + COL_REPLACEMENT_MAP: {"replacements": [{"original": "Bob", "label": "first_name", "synthetic": "Carlos"}]}, + COL_TEXT: "Alice works here.", + COL_TAGGED_TEXT: "[[Alice|first_name]] works here.", + COL_TAG_NOTATION: "bracket", + } + with pytest.raises(RuntimeError, match="unmatched"): + _apply_direct_replacements(row) + + +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": "quasi_identifier", - "sensitivity": "low", - "entity_label": "city", - "entity_value": "Portland", - "protection_reason": "Not identifying alone", - "protection_method_suggestion": "leave_as_is", - "combined_risk_level": "low", + "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} - result = _format_rewrite_disposition_block(row) - assert result[COL_REWRITE_DISPOSITION_BLOCK] == [] + 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, replace_values = _get_replace_pairs(row) + assert pairs == [(f"204{narrow_nbsp}Bluegrass", "500 Oak Lane", "street_address")] + assert replace_values == {f"204{narrow_nbsp}Bluegrass"} + assert "unprotected" not in caplog.text # --------------------------------------------------------------------------- -# 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_raises_when_replacement_map_missing( + stub_sensitivity_disposition: dict, +) -> None: + # No replacement map at all — must raise rather than pass PII through + row = { + COL_SENSITIVITY_DISPOSITION: stub_sensitivity_disposition, + COL_REPLACEMENT_MAP: None, + COL_TEXT: "Alice works here.", + COL_TAGGED_TEXT: "[[Alice|first_name]] works here.", + COL_TAG_NOTATION: "bracket", + } + with pytest.raises(RuntimeError, match="unmatched"): + _apply_direct_replacements(row) + + +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.", + COL_TAG_NOTATION: "bracket", } - 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.", + COL_TAG_NOTATION: "bracket", } - 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.", + COL_TAG_NOTATION: "bracket", } - 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." + + +def test_apply_direct_replacements_span_aware_prevents_substring_corruption() -> None: + """Replacing 'Ann' must not corrupt 'Anna' — only the detected span is substituted.""" + disposition = { + "sensitivity_disposition": [ + { + "id": 1, + "source": "tagged", + "category": "direct_identifier", + "sensitivity": "high", + "entity_label": "first_name", + "entity_value": "Ann", + "protection_reason": "Direct identifier.", + "protection_method_suggestion": "replace", + "combined_risk_level": "high", + } + ] + } + row = { + COL_SENSITIVITY_DISPOSITION: disposition, + COL_REPLACEMENT_MAP: {"replacements": [{"original": "Ann", "label": "first_name", "synthetic": "Maria"}]}, + COL_TEXT: "Ann met Anna.", + COL_TAGGED_TEXT: "Ann met Anna.", + COL_FINAL_ENTITIES: { + "entities": [{"label": "first_name", "value": "Ann", "start_position": 0, "end_position": 3}] + }, + COL_TAG_NOTATION: "xml", + } + result = _apply_direct_replacements(row) + assert result[COL_PREREPLACE_TEXT] == "Maria met Anna." + assert result[COL_PREREPLACE_TAGGED_TEXT] == "Maria met Anna." + + +def test_apply_direct_replacements_all_replace_spans_transformed() -> None: + """Every detected span selected for replace must be transformed in the output.""" + 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", + } + ] + } + row = { + COL_SENSITIVITY_DISPOSITION: disposition, + COL_REPLACEMENT_MAP: {"replacements": [{"original": "Alice", "label": "first_name", "synthetic": "Maria"}]}, + COL_TEXT: "Alice works. Alice lives here.", + COL_TAGGED_TEXT: "Alice works. Alice lives here.", + COL_FINAL_ENTITIES: { + "entities": [ + {"label": "first_name", "value": "Alice", "start_position": 0, "end_position": 5}, + {"label": "first_name", "value": "Alice", "start_position": 13, "end_position": 18}, + ] + }, + COL_TAG_NOTATION: "xml", + } + result = _apply_direct_replacements(row) + assert "Alice" not in result[COL_PREREPLACE_TEXT] + assert "Alice" not in result[COL_PREREPLACE_TAGGED_TEXT] + + +def test_apply_direct_replacements_no_cascade_when_synthetic_matches_another_original() -> None: + """Single-pass regex prevents Alice→Bob→Carlos cascade in plain text.""" + 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.", + COL_TAG_NOTATION: "xml", + } + result = _apply_direct_replacements(row) + assert result[COL_PREREPLACE_TEXT] == "Bob and Carlos met." + assert result[COL_PREREPLACE_TAGGED_TEXT] == "Bob and Carlos met." # --------------------------------------------------------------------------- @@ -271,10 +483,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 +528,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 +536,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