fix: validate augmented entities before finalization

src/anonymizer/engine/constants.py

@@ -38,6 +38,8 @@
 COL_VALIDATION_SKELETON = "_validation_skeleton"
 COL_VALIDATION_DECISIONS = "_validation_decisions"
 COL_VALIDATED_ENTITIES = "_validated_entities"
+COL_MERGED_VALIDATION_DECISIONS = "_merged_validation_decisions"
+COL_MERGED_VALIDATED_ENTITIES = "_merged_validated_entities"
 
 # Step 6: apply_validation_and_finalize
 COL_DETECTED_ENTITIES = "_detected_entities"

src/anonymizer/engine/detection/chunked_validation.py

@@ -107,6 +107,12 @@ class ChunkedValidationParams(BaseModel):
             excerpt window. The default preserves production parity with the
             pre-chunking validation path; benchmarks may disable it to probe
             compact validation prompts.
+        entities_column: Row column containing the entity spans whose
+            positions anchor chunk ordering and excerpt tagging.
+        candidates_column: Row column containing the validation candidate
+            payload to send to the validator.
+        output_column: Row column where merged validation decisions should be
+            written.
         prompt_template: Jinja2 source for the validation prompt (with
             ``_seed_tagged_text``, ``_validation_skeleton``, ``_tag_notation``
             placeholders). Typically produced by ``_get_validation_prompt``.
@@ -126,6 +132,9 @@ class ChunkedValidationParams(BaseModel):
     excerpt_window_chars: int = Field(gt=0)
     max_parallel_chunks: int | None = Field(default=None, gt=0)
     single_chunk_full_text: bool = True
+    entities_column: str = Field(default=COL_SEED_ENTITIES, min_length=1)
+    candidates_column: str = Field(default=COL_SEED_VALIDATION_CANDIDATES, min_length=1)
+    output_column: str = Field(default=COL_VALIDATION_DECISIONS, min_length=1)
     prompt_template: str = Field(repr=False)
     system_prompt: str | None = Field(default=None, repr=False)
 
@@ -141,10 +150,9 @@ def order_candidates_by_position(
 ) -> list[tuple[Any, EntitySpan]]:
     """Pair each candidate with its matching seed entity and sort by text position.
 
-    Every candidate id must resolve to a seed entity. A missing id indicates an
-    upstream bug in ``merge_and_build_candidates`` or ``prepare_validation_inputs``
-    (both produce candidates whose ids come from ``EntitySpan.entity_id``). We
-    raise early with the offending id so the failure is easy to triage.
+    Every candidate id must resolve to a source entity. A missing id indicates
+    an upstream bug in the candidate-building transform. We raise early with
+    the offending id so the failure is easy to triage.
     """
     seed_by_id = {span.entity_id: span for span in seed_entities}
     paired: list[tuple[Any, EntitySpan]] = []
@@ -153,10 +161,9 @@ def order_candidates_by_position(
         if seed is None:
             raise ValueError(
                 f"Validation candidate id {candidate.id!r} has no matching seed entity. "
-                "Every candidate produced by merge_and_build_candidates or "
-                "prepare_validation_inputs must correspond to a seed entity with "
+                "Every candidate produced by the upstream candidate builder must correspond to an entity with "
                 "start_position and end_position populated; this inconsistency "
-                "indicates a bug in one of those upstream generators."
+                "indicates a bug in that upstream generator."
             )
         paired.append((candidate, seed))
     paired.sort(key=lambda pair: (pair[1].start_position, pair[1].end_position, pair[1].entity_id))
@@ -460,8 +467,8 @@ def _build_dispatch_kwargs_per_chunk(
         )
 
     text = str(row.get(COL_TEXT, ""))
-    candidates = ValidationCandidatesSchema.from_raw(row.get(COL_SEED_VALIDATION_CANDIDATES, {}))
-    seed_entities_schema = EntitiesSchema.from_raw(row.get(COL_SEED_ENTITIES, {}))
+    candidates = ValidationCandidatesSchema.from_raw(row.get(params.candidates_column, {}))
+    seed_entities_schema = EntitiesSchema.from_raw(row.get(params.entities_column, {}))
     notation_raw = row.get(COL_TAG_NOTATION) or TagNotation.sentinel.value
     notation = TagNotation(str(notation_raw))
 
@@ -594,7 +601,7 @@ def chunked_validate_row(
             futures = [executor.submit(_dispatch_chunk, **kwargs) for kwargs in dispatch_kwargs_per_chunk]
             chunk_results = [f.result() for f in futures]
 
-    row[COL_VALIDATION_DECISIONS] = merge_chunk_decisions(chunk_results, candidates)
+    row[params.output_column] = merge_chunk_decisions(chunk_results, candidates)
     return row
 
 
@@ -613,7 +620,7 @@ async def chunked_validate_row_async(
             *[_bounded_dispatch_chunk_async(semaphore, kwargs) for kwargs in dispatch_kwargs_per_chunk]
         )
 
-    row[COL_VALIDATION_DECISIONS] = merge_chunk_decisions(chunk_results, candidates)
+    row[params.output_column] = merge_chunk_decisions(chunk_results, candidates)
     return row
 
 
@@ -622,24 +629,30 @@ async def chunked_validate_row_async(
 # ---------------------------------------------------------------------------
 
 
-def make_chunked_validation_generator(pool: list[str]) -> Any:
+def make_chunked_validation_generator(
+    pool: list[str],
+    *,
+    entities_column: str = COL_SEED_ENTITIES,
+    candidates_column: str = COL_SEED_VALIDATION_CANDIDATES,
+) -> Any:
     """Build a ``@custom_column_generator``-decorated function bound to ``pool``.
 
     ``model_aliases`` must be declared statically on the decorator so
     DataDesigner knows which facades to materialise for the generator. Since
     the pool is config-driven (per-run), we generate the function dynamically.
     The required_columns are exhaustive for DataDesigner's DAG ordering: the
-    generator reads the raw text, seed entities (for positions), the candidate
-    list (what to decide), and the tag notation (for excerpt tagging).
+    generator reads the raw text, the caller-configured entity column (seed
+    entities by default, for positions), the caller-configured candidate list
+    (what to decide), and the tag notation (for excerpt tagging).
     """
     if not pool:
         raise ValueError("Cannot build chunked validation generator: pool is empty.")
 
     @custom_column_generator(
         required_columns=[
             COL_TEXT,
-            COL_SEED_ENTITIES,
-            COL_SEED_VALIDATION_CANDIDATES,
+            entities_column,
+            candidates_column,
             COL_TAG_NOTATION,
         ],
         model_aliases=list(pool),

src/anonymizer/engine/detection/custom_columns.py

@@ -22,6 +22,8 @@
     COL_INITIAL_TAGGED_TEXT,
     COL_MERGED_ENTITIES,
     COL_MERGED_TAGGED_TEXT,
+    COL_MERGED_VALIDATED_ENTITIES,
+    COL_MERGED_VALIDATION_DECISIONS,
     COL_RAW_DETECTED,
     COL_SEED_ENTITIES,
     COL_SEED_ENTITIES_JSON,
@@ -137,40 +139,40 @@ def prepare_validation_inputs(row: dict[str, Any]) -> dict[str, Any]:
 )
 def enrich_validation_decisions(row: dict[str, Any]) -> dict[str, Any]:
     """Enrich validation decisions with entity value and filter to known candidate IDs only."""
-    raw_decisions = RawValidationDecisionsSchema.from_raw(row.get(COL_VALIDATION_DECISIONS, {}))
-    candidates = ValidationCandidatesSchema.from_raw(row.get(COL_SEED_VALIDATION_CANDIDATES, {}))
-
-    candidate_lookup = {c.id: c for c in candidates.candidates}
-    valid_ids = set(candidate_lookup)
+    _enrich_validation_decisions(
+        row,
+        decisions_column=COL_VALIDATION_DECISIONS,
+        candidates_column=COL_SEED_VALIDATION_CANDIDATES,
+        output_column=COL_VALIDATED_ENTITIES,
+    )
+    return row
 
-    enriched = [
-        ValidatedDecisionSchema(
-            id=d.id,
-            decision=d.decision,
-            proposed_label=d.proposed_label,
-            reason=d.reason,
-            value=candidate_lookup[d.id].value,
-            label=candidate_lookup[d.id].label,
-        )
-        for d in raw_decisions.decisions
-        if d.id in valid_ids
-    ]
 
-    row[COL_VALIDATED_ENTITIES] = ValidatedDecisionsSchema(decisions=enriched).model_dump(mode="json")
+@custom_column_generator(
+    required_columns=[COL_MERGED_VALIDATION_DECISIONS, COL_VALIDATION_CANDIDATES],
+)
+def enrich_merged_validation_decisions(row: dict[str, Any]) -> dict[str, Any]:
+    """Enrich merged validation decisions with entity value and filter to known candidate IDs only."""
+    _enrich_validation_decisions(
+        row,
+        decisions_column=COL_MERGED_VALIDATION_DECISIONS,
+        candidates_column=COL_VALIDATION_CANDIDATES,
+        output_column=COL_MERGED_VALIDATED_ENTITIES,
+    )
     return row
 
 
 @custom_column_generator(
-    required_columns=[COL_TEXT, COL_MERGED_ENTITIES, COL_VALIDATED_ENTITIES],
+    required_columns=[COL_TEXT, COL_MERGED_ENTITIES, COL_MERGED_VALIDATED_ENTITIES],
     side_effect_columns=[COL_TAGGED_TEXT],
 )
 def apply_validation_and_finalize(row: dict[str, Any]) -> dict[str, Any]:
-    """Apply keep/reclass/drop decisions, expand to all occurrences, and produce final outputs."""
+    """Apply merged keep/reclass/drop decisions, expand to all occurrences, and produce final outputs."""
     text = str(row.get(COL_TEXT, ""))
     merged = _parse_entity_spans(row.get(COL_MERGED_ENTITIES, {}))
     validated = apply_validation_decisions(
         entities=merged,
-        validation_output=row.get(COL_VALIDATED_ENTITIES, {}),
+        validation_output=row.get(COL_MERGED_VALIDATED_ENTITIES, {}),
     )
     expanded = expand_entity_occurrences(text=text, entities=validated)
     row[COL_DETECTED_ENTITIES] = EntitiesSchema(entities=[entity.as_dict() for entity in expanded]).model_dump(
@@ -194,3 +196,32 @@ def _parse_entity_spans(raw_payload: object) -> list[EntitySpan]:
         )
         for e in parsed.entities
     ]
+
+
+def _enrich_validation_decisions(
+    row: dict[str, Any],
+    *,
+    decisions_column: str,
+    candidates_column: str,
+    output_column: str,
+) -> None:
+    raw_decisions = RawValidationDecisionsSchema.from_raw(row.get(decisions_column, {}))
+    candidates = ValidationCandidatesSchema.from_raw(row.get(candidates_column, {}))
+
+    candidate_lookup = {c.id: c for c in candidates.candidates}
+    valid_ids = set(candidate_lookup)
+
+    enriched = [
+        ValidatedDecisionSchema(
+            id=d.id,
+            decision=d.decision,
+            proposed_label=d.proposed_label,
+            reason=d.reason,
+            value=candidate_lookup[d.id].value,
+            label=candidate_lookup[d.id].label,
+        )
+        for d in raw_decisions.decisions
+        if d.id in valid_ids
+    ]
+
+    row[output_column] = ValidatedDecisionsSchema(decisions=enriched).model_dump(mode="json")

src/anonymizer/engine/detection/detection_workflow.py

@@ -25,6 +25,8 @@
     COL_INITIAL_TAGGED_TEXT,
     COL_LATENT_ENTITIES,
     COL_MERGED_ENTITIES,
+    COL_MERGED_VALIDATED_ENTITIES,
+    COL_MERGED_VALIDATION_DECISIONS,
     COL_RAW_DETECTED,
     COL_SEED_ENTITIES,
     COL_SEED_ENTITIES_JSON,
@@ -34,6 +36,7 @@
     COL_TAGGED_TEXT,
     COL_TEXT,
     COL_VALIDATED_ENTITIES,
+    COL_VALIDATION_CANDIDATES,
     COL_VALIDATION_DECISIONS,
     COL_VALIDATION_SKELETON,
     DEFAULT_ENTITY_LABELS,
@@ -96,7 +99,8 @@ def detect_and_validate_entities(
         data_summary: str | None = None,
         preview_num_records: int | None = None,
     ) -> EntityDetectionResult:
-        """Run the core detection pipeline: GLiNER NER, LLM validation, LLM augmentation, and finalization.
+        """Run the core detection pipeline: GLiNER NER, seed validation, LLM augmentation,
+        merged validation, and finalization.
 
         This is the primary detection workflow. It detects entities via GLiNER,
         validates/reclassifies them with an LLM (chunked across a pool of
@@ -217,6 +221,21 @@ def _build_detection_spec(
                 name=COL_MERGED_ENTITIES,
                 operation=DetectionTransformOperation.MERGE_AND_BUILD_CANDIDATES,
             ),
+            ChunkedValidationConfig(
+                name=COL_MERGED_VALIDATION_DECISIONS,
+                pool=list(validator_aliases),
+                max_entities_per_call=validation_max_entities_per_call,
+                excerpt_window_chars=validation_excerpt_window_chars,
+                single_chunk_full_text=validation_single_chunk_full_text,
+                entities_column=COL_MERGED_ENTITIES,
+                candidates_column=COL_VALIDATION_CANDIDATES,
+                prompt_template=_get_validation_prompt(data_summary=data_summary, labels=labels),
+                drop=True,
+            ),
+            DetectionTransformConfig(
+                name=COL_MERGED_VALIDATED_ENTITIES,
+                operation=DetectionTransformOperation.ENRICH_MERGED_VALIDATION_DECISIONS,
+            ),
             DetectionTransformConfig(
                 name=COL_DETECTED_ENTITIES,
                 operation=DetectionTransformOperation.APPLY_VALIDATION_AND_FINALIZE,

src/anonymizer/engine/workflow_columns/detection/config.py

@@ -14,6 +14,8 @@
     COL_INITIAL_TAGGED_TEXT,
     COL_MERGED_ENTITIES,
     COL_MERGED_TAGGED_TEXT,
+    COL_MERGED_VALIDATED_ENTITIES,
+    COL_MERGED_VALIDATION_DECISIONS,
     COL_RAW_DETECTED,
     COL_SEED_ENTITIES,
     COL_SEED_ENTITIES_JSON,
@@ -33,6 +35,7 @@ class DetectionTransformOperation(str, Enum):
     PARSE_DETECTED_ENTITIES = "parse_detected_entities"
     PREPARE_VALIDATION_INPUTS = "prepare_validation_inputs"
     ENRICH_VALIDATION_DECISIONS = "enrich_validation_decisions"
+    ENRICH_MERGED_VALIDATION_DECISIONS = "enrich_merged_validation_decisions"
     APPLY_VALIDATION_TO_SEED_ENTITIES = "apply_validation_to_seed_entities"
     MERGE_AND_BUILD_CANDIDATES = "merge_and_build_candidates"
     APPLY_VALIDATION_AND_FINALIZE = "apply_validation_and_finalize"
@@ -49,6 +52,10 @@ class DetectionTransformConfig(SingleColumnConfig):
             COL_VALIDATION_DECISIONS,
             COL_SEED_VALIDATION_CANDIDATES,
         ],
+        DetectionTransformOperation.ENRICH_MERGED_VALIDATION_DECISIONS: [
+            COL_MERGED_VALIDATION_DECISIONS,
+            COL_VALIDATION_CANDIDATES,
+        ],
         DetectionTransformOperation.APPLY_VALIDATION_TO_SEED_ENTITIES: [
             COL_TEXT,
             COL_SEED_ENTITIES,
@@ -62,13 +69,14 @@ class DetectionTransformConfig(SingleColumnConfig):
         DetectionTransformOperation.APPLY_VALIDATION_AND_FINALIZE: [
             COL_TEXT,
             COL_MERGED_ENTITIES,
-            COL_VALIDATED_ENTITIES,
+            COL_MERGED_VALIDATED_ENTITIES,
         ],
     }
     _SIDE_EFFECT_COLUMNS: ClassVar[dict[DetectionTransformOperation, list[str]]] = {
         DetectionTransformOperation.PARSE_DETECTED_ENTITIES: [COL_TAG_NOTATION],
         DetectionTransformOperation.PREPARE_VALIDATION_INPUTS: [COL_SEED_TAGGED_TEXT],
         DetectionTransformOperation.ENRICH_VALIDATION_DECISIONS: [],
+        DetectionTransformOperation.ENRICH_MERGED_VALIDATION_DECISIONS: [],
         DetectionTransformOperation.APPLY_VALIDATION_TO_SEED_ENTITIES: [
             COL_INITIAL_TAGGED_TEXT,
             COL_SEED_ENTITIES_JSON,
@@ -101,6 +109,8 @@ class ChunkedValidationConfig(SingleColumnConfig):
     excerpt_window_chars: int = Field(gt=0)
     max_parallel_chunks: int | None = Field(default=None, gt=0)
     single_chunk_full_text: bool = True
+    entities_column: str = Field(default=COL_SEED_ENTITIES, min_length=1)
+    candidates_column: str = Field(default=COL_SEED_VALIDATION_CANDIDATES, min_length=1)
     prompt_template: str = Field(repr=False)
     system_prompt: str | None = Field(default=None, repr=False)
 
@@ -112,8 +122,8 @@ def get_column_emoji() -> str:
     def required_columns(self) -> list[str]:
         return [
             COL_TEXT,
-            COL_SEED_ENTITIES,
-            COL_SEED_VALIDATION_CANDIDATES,
+            self.entities_column,
+            self.candidates_column,
             COL_TAG_NOTATION,
         ]
 

src/anonymizer/engine/workflow_columns/detection/impl.py

@@ -22,6 +22,7 @@
 from anonymizer.engine.detection.custom_columns import (
     apply_validation_and_finalize,
     apply_validation_to_seed_entities,
+    enrich_merged_validation_decisions,
     enrich_validation_decisions,
     merge_and_build_candidates,
     parse_detected_entities,
@@ -37,6 +38,7 @@
     DetectionTransformOperation.PARSE_DETECTED_ENTITIES: parse_detected_entities,
     DetectionTransformOperation.PREPARE_VALIDATION_INPUTS: prepare_validation_inputs,
     DetectionTransformOperation.ENRICH_VALIDATION_DECISIONS: enrich_validation_decisions,
+    DetectionTransformOperation.ENRICH_MERGED_VALIDATION_DECISIONS: enrich_merged_validation_decisions,
     DetectionTransformOperation.APPLY_VALIDATION_TO_SEED_ENTITIES: apply_validation_to_seed_entities,
     DetectionTransformOperation.MERGE_AND_BUILD_CANDIDATES: merge_and_build_candidates,
     DetectionTransformOperation.APPLY_VALIDATION_AND_FINALIZE: apply_validation_and_finalize,
@@ -115,6 +117,9 @@ def _params(self, models: dict[str, Any]) -> ChunkedValidationParams:
             excerpt_window_chars=self.config.excerpt_window_chars,
             max_parallel_chunks=self.config.max_parallel_chunks or _derive_max_parallel_chunks(models),
             single_chunk_full_text=self.config.single_chunk_full_text,
+            entities_column=self.config.entities_column,
+            candidates_column=self.config.candidates_column,
+            output_column=self.config.name,
             prompt_template=self.config.prompt_template,
             system_prompt=self.config.system_prompt,
         )