Skip to content

feat: rewrite replace chg and LLM entity augment chg and couple bugs#198

Open
asteier2026 wants to merge 4 commits into
mainfrom
asteier2026/feature/rewrite-replace
Open

feat: rewrite replace chg and LLM entity augment chg and couple bugs#198
asteier2026 wants to merge 4 commits into
mainfrom
asteier2026/feature/rewrite-replace

Conversation

@asteier2026

@asteier2026 asteier2026 commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

Changes include:

  • Rewrite — direct identifier replacement is now a programmatic pre-step: direct identifiers are substituted from the replacement map before the rewrite LLM call, rather than relying on the LLM to apply them from a <replacement_map> block. This is a critical fix as currently our replacement approach in Rewrite misses important direct identifiers.
  • LLM entity augmentation — enhanced to handle difficult/tricky entities: added Additional extraction requirements block covering verbatim span extraction and disguised identifiers (digit-words for phone/SSN/credit card numbers, letter-by-letter name spellings). This is a critical fix if we want to do well on RAT-Bench difficulty 2.
  • MeaningUnitAspect validation removed: the enum was dropped and aspect is now an unvalidated str — the model returns open-ended values (e.g. goal, symptom) that can't be exhaustively enumerated, and validation errors were causing record loss
  • Latent entity rationale max length increased from 150 → 250: validation errors on longer rationale strings were causing record loss

UPDATE: we will be closing this PR shortly. It was broken down into:
#205
This just deals with latent reason length

#206
This one just removes the validate on meaningUnitAspect

#207
This one is the enhanced augment for tricky entities

#208
This one is just making replace programmatic

@asteier2026 asteier2026 requested a review from a team as a code owner June 29, 2026 17:00
@greptile-apps

greptile-apps Bot commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR restructures the replacement step so direct identifiers are substituted programmatically before the LLM ever sees the text, removing the <replacement_map> block from the rewrite prompt and eliminating the LLM's responsibility for applying verbatim substitutions. It also strengthens the entity augmentation prompt with verbatim-extraction rules and disguised-identifier examples, increases the latent entity rationale length limit to prevent record loss, and drops the MeaningUnitAspect enum in favour of unvalidated str to stop schema validation errors from discarding records with open-ended aspect values.

  • Pre-replace pipeline (_apply_direct_replacements): builds COL_PREREPLACE_TEXT and COL_PREREPLACE_TAGGED_TEXT by sorting replacement pairs longest-first and applying str.replace() on both the plain and tagged text; the LLM disposition block now also excludes replace entities so the rewrite prompt only instructs generalize, remove, and suppress_inference methods.
  • Prompt improvements (detection_workflow.py): new "Additional extraction requirements" block covering verbatim span extraction and examples of disguised identifiers (digit-word phone/SSN numbers, letter-by-letter name spellings).
  • Schema fixes (detection.py, rewrite.py): rationale max_length 150→250; MeaningUnitAspect enum removed; repair.py updated to reference COL_PREREPLACE_TEXT as the original-text baseline for the repair LLM.

Confidence Score: 5/5

The refactoring is well-scoped and properly tested; all changed paths produce correct outputs for the happy path and the no-replace-entities passthrough.

The pre-replace logic is correct and the longest-first sort prevents substring collisions. The warning logic across _get_replace_pairs and _apply_direct_replacements covers all missing-map and unmatched-entity cases without false positives. Schema changes are additive-relaxation only and are consistent with the stated motivation of preventing record loss.

rewrite_generation.py — the double parse of COL_SENSITIVITY_DISPOSITION on the empty-pairs path is worth cleaning up, but it has no correctness impact.

Important Files Changed

Filename Overview
src/anonymizer/engine/rewrite/rewrite_generation.py Core refactor: replaces _filter_replacement_map_for_prompt with _apply_direct_replacements, which performs str.replace() pre-substitution (sorted longest-first) before the LLM call; _format_rewrite_disposition_block now also excludes replace entities from the disposition block. Minor: _apply_direct_replacements re-parses COL_SENSITIVITY_DISPOSITION when _get_replace_pairs already parsed it on the empty-pairs path.
src/anonymizer/engine/detection/detection_workflow.py Adds "Additional extraction requirements" block to the LLM entity augmentation prompt, covering verbatim span extraction and disguised identifiers (digit-word phone/SSN/CC numbers, letter-by-letter name spellings).
src/anonymizer/engine/rewrite/repair.py Updates <<ORIGINAL_TEXT>> reference from COL_TEXT to COL_PREREPLACE_TEXT so the repair LLM sees the same pre-replaced text that the rewrite LLM worked from; required_columns updated to match.
src/anonymizer/engine/schemas/detection.py Increases LatentEntitySchema.rationale max_length from 150 to 250 to prevent validation errors on longer inference rationale strings that were causing record loss.
src/anonymizer/engine/schemas/rewrite.py Removes MeaningUnitAspect enum and relaxes MeaningUnitSchema.aspect to bare str, allowing open-ended LLM-returned values like "goal" or "symptom" that couldn't be enumerated without causing validation-driven record loss.
tests/engine/test_rewrite_generation.py Tests updated to cover _apply_direct_replacements (multi-occurrence, schema-instance, passthrough); _filter_replacement_map_for_prompt tests replaced. The "no replacement map but replace entities exist" warning path in _apply_direct_replacements lacks a dedicated test.
tests/engine/test_qa_generation.py Stub meaning unit fixtures updated from MeaningUnitAspect enum values to plain string literals; no logic changes.
src/anonymizer/engine/schemas/init.py Removes MeaningUnitAspect from public exports to match its removal from rewrite.py.
src/anonymizer/engine/constants.py Renames COL_REPLACEMENT_MAP_FOR_PROMPT to two new constants: COL_PREREPLACE_TEXT and COL_PREREPLACE_TAGGED_TEXT, matching the new pre-replace pipeline design.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Row as Input Row
    participant FRD as _format_rewrite_disposition_block
    participant ADR as _apply_direct_replacements
    participant LLM as LLM (Full Rewrite)
    participant ERT as _extract_rewritten_text

    Row->>FRD: COL_SENSITIVITY_DISPOSITION
    Note over FRD: Excludes leave_as_is AND replace entities
    FRD-->>Row: COL_REWRITE_DISPOSITION_BLOCK

    Row->>ADR: COL_TEXT, COL_TAGGED_TEXT, COL_REPLACEMENT_MAP, COL_SENSITIVITY_DISPOSITION
    Note over ADR: _get_replace_pairs() → sort by len desc
    Note over ADR: str.replace() on plain + tagged text
    ADR-->>Row: COL_PREREPLACE_TEXT, COL_PREREPLACE_TAGGED_TEXT

    Row->>LLM: COL_PREREPLACE_TAGGED_TEXT + COL_REWRITE_DISPOSITION_BLOCK
    Note over LLM: Only sees generalize/remove/suppress_inference entities
    Note over LLM: Direct identifiers already swapped programmatically
    LLM-->>Row: COL_FULL_REWRITE

    Row->>ERT: COL_FULL_REWRITE
    ERT-->>Row: COL_REWRITTEN_TEXT (or None on failure)
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Row as Input Row
    participant FRD as _format_rewrite_disposition_block
    participant ADR as _apply_direct_replacements
    participant LLM as LLM (Full Rewrite)
    participant ERT as _extract_rewritten_text

    Row->>FRD: COL_SENSITIVITY_DISPOSITION
    Note over FRD: Excludes leave_as_is AND replace entities
    FRD-->>Row: COL_REWRITE_DISPOSITION_BLOCK

    Row->>ADR: COL_TEXT, COL_TAGGED_TEXT, COL_REPLACEMENT_MAP, COL_SENSITIVITY_DISPOSITION
    Note over ADR: _get_replace_pairs() → sort by len desc
    Note over ADR: str.replace() on plain + tagged text
    ADR-->>Row: COL_PREREPLACE_TEXT, COL_PREREPLACE_TAGGED_TEXT

    Row->>LLM: COL_PREREPLACE_TAGGED_TEXT + COL_REWRITE_DISPOSITION_BLOCK
    Note over LLM: Only sees generalize/remove/suppress_inference entities
    Note over LLM: Direct identifiers already swapped programmatically
    LLM-->>Row: COL_FULL_REWRITE

    Row->>ERT: COL_FULL_REWRITE
    ERT-->>Row: COL_REWRITTEN_TEXT (or None on failure)
Loading

Reviews (3): Last reviewed commit: "fix: address no matching replacement val..." | Re-trigger Greptile

Comment thread src/anonymizer/engine/rewrite/rewrite_generation.py Outdated
Comment thread src/anonymizer/engine/rewrite/rewrite_generation.py
Comment thread src/anonymizer/engine/schemas/rewrite.py
asteier2026 and others added 3 commits June 29, 2026 11:21
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant