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..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 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..d4d9e749 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 @@ -14,13 +15,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 +52,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 +88,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 +108,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 +121,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 +147,86 @@ 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]) -> list[tuple[str, str]]: + """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_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 [] 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: + 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"): 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]] = [] + 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", + 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). A single-pass regex substitution + ensures all occurrences are replaced without cascade (a synthetic value matching + another entity's original is never re-substituted). + + Raises on failure rather than falling back to unmodified text: replace entities are + excluded from COL_REWRITE_DISPOSITION_BLOCK, so a silent passthrough would send + PII-containing text to the LLM with no instructions to protect those entities. + """ + plain_text = str(row.get(COL_TEXT, "")) + tagged_text = str(row.get(COL_TAGGED_TEXT, "")) + pairs = _get_replace_pairs(row) + if pairs: + sorted_pairs = sorted(pairs, key=lambda p: len(p[0]), reverse=True) + pattern = re.compile("|".join(re.escape(original) for original, _ in sorted_pairs)) + lookup = dict(sorted_pairs) + plain_text = pattern.sub(lambda m: lookup[m.group(0)], plain_text) + tagged_text = pattern.sub(lambda m: lookup[m.group(0)], tagged_text) + row[COL_PREREPLACE_TEXT] = plain_text + row[COL_PREREPLACE_TAGGED_TEXT] = tagged_text return row @@ -226,8 +282,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..01e07472 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,13 +125,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 +153,184 @@ 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_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 + + +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 = _get_replace_pairs(row) + assert pairs == [(f"204{narrow_nbsp}Bluegrass", "500 Oak Lane")] + 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_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." + + +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." # --------------------------------------------------------------------------- @@ -271,10 +399,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 +444,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 +452,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