Skip to content

fix: validate augmented entities before finalization#197

Open
fallintoplace wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
fallintoplace:fix/validate-augmented-detections
Open

fix: validate augmented entities before finalization#197
fallintoplace wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
fallintoplace:fix/validate-augmented-detections

Conversation

@fallintoplace

Copy link
Copy Markdown

Summary

  • keep the existing seed validation step used to prepare augmentation inputs
  • add a merged validation step for the combined seed and augmented entities before finalization
  • add regression coverage for merged validation decisions across chunked validation, custom columns, and the detection workflow

Why

The detection workflow only validated seed detector entities before augmentation. After the merged entity set was built, augmented-only entities had no matching validation decisions, and finalization kept unmatched entities unchanged. That allowed hallucinated augmented entities to become final detections without validation.

Impact

Augmented entities now need a validation decision from the merged candidate pass before they can reach the final detected entity set.

Validation

  • uv run --group dev ruff check src/anonymizer/engine/constants.py src/anonymizer/engine/detection/chunked_validation.py src/anonymizer/engine/detection/custom_columns.py src/anonymizer/engine/detection/detection_workflow.py src/anonymizer/engine/workflow_columns/detection/config.py src/anonymizer/engine/workflow_columns/detection/impl.py tests/engine/test_chunked_validation.py tests/engine/test_detection_custom_columns.py tests/engine/test_detection_workflow.py
  • uv run --group dev pytest tests/engine/test_chunked_validation.py tests/engine/test_detection_custom_columns.py tests/engine/test_detection_workflow.py -q

@fallintoplace fallintoplace requested a review from a team as a code owner June 27, 2026 14:38
@github-actions

github-actions Bot commented Jun 27, 2026

Copy link
Copy Markdown
Contributor

All contributors have signed the DCO ✍️ ✅
Posted by the DCO Assistant Lite bot.

@fallintoplace

Copy link
Copy Markdown
Author

I have read the DCO document and I hereby sign the DCO.

@fallintoplace fallintoplace changed the title Validate augmented entities before finalization fix: validate augmented entities before finalization Jun 27, 2026
@greptile-apps

greptile-apps Bot commented Jun 27, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR closes a security gap where augmented entities produced by the LLM augmenter could reach the final detected-entity set without ever being validated. The previous pipeline only ran LLM validation on seed (GLiNER) entities; after the seed+augmented merge, no second validation was performed before finalization.

  • Adds two new constants (COL_MERGED_VALIDATION_DECISIONS, COL_MERGED_VALIDATED_ENTITIES) and a full second ChunkedValidationConfig pass over the merged entity set after augmentation, with enrich_merged_validation_decisions enriching the decisions before apply_validation_and_finalize consumes them.
  • Generalises ChunkedValidationParams, make_chunked_validation_generator, and ChunkedValidationConfig with configurable entities_column, candidates_column, and output_column, allowing the same chunked-validation machinery to serve both the seed and merged passes without code duplication.
  • Adds regression tests covering the new parameterisation path in chunked_validate_row, the enrich_merged_validation_decisions custom column, and the end-to-end drop behaviour in apply_validation_and_finalize.

Confidence Score: 5/5

Safe to merge; the new merged validation pass is correctly wired into the DAG and well-covered by regression tests.

The structural change is straightforward: two new columns inserted between MERGE_AND_BUILD_CANDIDATES and APPLY_VALIDATION_AND_FINALIZE, using the existing ChunkedValidation machinery generalised with configurable column parameters. The generalisation is backward-compatible (defaults match the original hardcoded values). All three changed test files have dedicated coverage for the new paths, and existing tests were updated correctly.

No files require special attention; both comments are minor documentation and test-assertion nits.

Important Files Changed

Filename Overview
src/anonymizer/engine/constants.py Adds two new pipeline column constants for the merged validation pass; no issues.
src/anonymizer/engine/detection/chunked_validation.py Generalized to accept configurable entity/candidate/output columns; chunked_validate_row docstring is now stale (still references COL_VALIDATION_DECISIONS).
src/anonymizer/engine/detection/custom_columns.py Refactors enrich_validation_decisions into a shared helper; adds enrich_merged_validation_decisions; updates apply_validation_and_finalize to use merged validated entities. Clean and correct.
src/anonymizer/engine/detection/detection_workflow.py Inserts merged ChunkedValidationConfig and ENRICH_MERGED_VALIDATION_DECISIONS steps into the pipeline column list; wiring and ordering look correct.
src/anonymizer/engine/workflow_columns/detection/config.py Adds ENRICH_MERGED_VALIDATION_DECISIONS to required/side-effect column maps and entities_column/candidates_column fields to ChunkedValidationConfig; consistent with the rest of the config layer.
src/anonymizer/engine/workflow_columns/detection/impl.py Adds enrich_merged_validation_decisions to the transform dispatch table and threads entities_column/candidates_column/output_column into ChunkedValidationParams; straightforward.
tests/engine/test_chunked_validation.py Adds a well-constructed test for non-default entity/candidate/output columns, including an assertion on COL_MERGED_VALIDATION_DECISIONS.
tests/engine/test_detection_custom_columns.py New tests for enrich_merged_validation_decisions and the drop-augmented-entity end-to-end path; coverage is solid.
tests/engine/test_detection_workflow.py Adds assertions for the new merged-validation column config; missing drop=True check on merged_validation_col that is present for the seed validation col.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[GLiNER NER\nCOL_RAW_DETECTED] --> B[parse_detected_entities\nCOL_SEED_ENTITIES]
    B --> C[prepare_validation_inputs\nCOL_SEED_VALIDATION_CANDIDATES]
    C --> D[ChunkedValidation seed pass\nCOL_VALIDATION_DECISIONS]
    D --> E[enrich_validation_decisions\nCOL_VALIDATED_ENTITIES]
    E --> F[apply_validation_to_seed_entities\nCOL_VALIDATED_SEED_ENTITIES]
    F --> G[LLM Augmenter\nCOL_AUGMENTED_ENTITIES]
    G --> H[merge_and_build_candidates\nCOL_MERGED_ENTITIES + COL_VALIDATION_CANDIDATES]
    H --> I[ChunkedValidation merged pass NEW\nCOL_MERGED_VALIDATION_DECISIONS]
    I --> J[enrich_merged_validation_decisions NEW\nCOL_MERGED_VALIDATED_ENTITIES]
    J --> K[apply_validation_and_finalize\nCOL_DETECTED_ENTITIES]

    style I fill:#d4edda,stroke:#28a745
    style J fill:#d4edda,stroke:#28a745
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"}}}%%
flowchart TD
    A[GLiNER NER\nCOL_RAW_DETECTED] --> B[parse_detected_entities\nCOL_SEED_ENTITIES]
    B --> C[prepare_validation_inputs\nCOL_SEED_VALIDATION_CANDIDATES]
    C --> D[ChunkedValidation seed pass\nCOL_VALIDATION_DECISIONS]
    D --> E[enrich_validation_decisions\nCOL_VALIDATED_ENTITIES]
    E --> F[apply_validation_to_seed_entities\nCOL_VALIDATED_SEED_ENTITIES]
    F --> G[LLM Augmenter\nCOL_AUGMENTED_ENTITIES]
    G --> H[merge_and_build_candidates\nCOL_MERGED_ENTITIES + COL_VALIDATION_CANDIDATES]
    H --> I[ChunkedValidation merged pass NEW\nCOL_MERGED_VALIDATION_DECISIONS]
    I --> J[enrich_merged_validation_decisions NEW\nCOL_MERGED_VALIDATED_ENTITIES]
    J --> K[apply_validation_and_finalize\nCOL_DETECTED_ENTITIES]

    style I fill:#d4edda,stroke:#28a745
    style J fill:#d4edda,stroke:#28a745
Loading

Reviews (2): Last reviewed commit: "test: cover merged validation output col..." | Re-trigger Greptile

Comment on lines +509 to +521
params = ChunkedValidationParams(
pool=["v0"],
max_entities_per_call=10,
excerpt_window_chars=100,
entities_column=COL_MERGED_ENTITIES,
candidates_column=COL_VALIDATION_CANDIDATES,
prompt_template=_MINIMAL_TEMPLATE,
)

out = chunked_validate_row(row, params, {"v0": facade})

decisions = {d["id"]: d["decision"] for d in out[COL_VALIDATION_DECISIONS]["decisions"]}
assert decisions == {"first_name_0_5": "keep", "api_key_13_29": "drop"}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 output_column not exercised in the new parameterization test

The test creates ChunkedValidationParams without setting output_column, so the result lands in COL_VALIDATION_DECISIONS (the default). In production, ChunkedValidationGenerator._params() sets output_column=self.config.name, which for the merged pass resolves to COL_MERGED_VALIDATION_DECISIONS. The new test therefore doesn't cover the path where the output is redirected to a non-default column — if the output_column wiring broke, this test would still pass. Adding output_column=COL_MERGED_VALIDATION_DECISIONS to the params and asserting on out[COL_MERGED_VALIDATION_DECISIONS] would close that gap.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

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