diff --git a/src/anonymizer/engine/constants.py b/src/anonymizer/engine/constants.py index 3cb0876a..4775b8e5 100644 --- a/src/anonymizer/engine/constants.py +++ b/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" diff --git a/src/anonymizer/engine/detection/chunked_validation.py b/src/anonymizer/engine/detection/chunked_validation.py index f74e54bb..331baa37 100644 --- a/src/anonymizer/engine/detection/chunked_validation.py +++ b/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,15 +629,21 @@ 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.") @@ -638,8 +651,8 @@ def make_chunked_validation_generator(pool: list[str]) -> Any: @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), diff --git a/src/anonymizer/engine/detection/custom_columns.py b/src/anonymizer/engine/detection/custom_columns.py index 059d82ac..3e826991 100644 --- a/src/anonymizer/engine/detection/custom_columns.py +++ b/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") diff --git a/src/anonymizer/engine/detection/detection_workflow.py b/src/anonymizer/engine/detection/detection_workflow.py index c01ff781..24ef1183 100644 --- a/src/anonymizer/engine/detection/detection_workflow.py +++ b/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, diff --git a/src/anonymizer/engine/workflow_columns/detection/config.py b/src/anonymizer/engine/workflow_columns/detection/config.py index a9661818..5d3e2cb0 100644 --- a/src/anonymizer/engine/workflow_columns/detection/config.py +++ b/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, ] diff --git a/src/anonymizer/engine/workflow_columns/detection/impl.py b/src/anonymizer/engine/workflow_columns/detection/impl.py index 7b0dae38..3345b3d7 100644 --- a/src/anonymizer/engine/workflow_columns/detection/impl.py +++ b/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, ) diff --git a/tests/engine/test_chunked_validation.py b/tests/engine/test_chunked_validation.py index 4a4207a0..c84e4e4c 100644 --- a/tests/engine/test_chunked_validation.py +++ b/tests/engine/test_chunked_validation.py @@ -20,7 +20,9 @@ from data_designer.engine.models.clients.errors import SyncClientUnavailableError from anonymizer.engine.constants import ( + COL_MERGED_ENTITIES, COL_MERGED_TAGGED_TEXT, + COL_MERGED_VALIDATION_DECISIONS, COL_SEED_ENTITIES, COL_SEED_VALIDATION_CANDIDATES, COL_TAG_NOTATION, @@ -202,7 +204,7 @@ def test_orders_by_start_then_end_then_id(self) -> None: def test_missing_seed_raises_with_triage_hint(self) -> None: candidates = _candidates_schema(("missing", "x", "y")) - with pytest.raises(ValueError, match="merge_and_build_candidates or prepare_validation_inputs"): + with pytest.raises(ValueError, match="upstream candidate builder"): order_candidates_by_position(candidates, []) @@ -459,6 +461,69 @@ def test_single_chunk_single_alias_dispatches_once_and_merges(self) -> None: decisions = out[COL_VALIDATION_DECISIONS]["decisions"] assert {d["id"]: d["decision"] for d in decisions} == {"a": "keep", "b": "drop"} + def test_supports_non_default_entity_and_candidate_columns(self) -> None: + text = "Alice met on 12 February 1980." + merged_entities = EntitiesSchema( + entities=[ + { + "id": "first_name_0_5", + "value": "Alice", + "label": "first_name", + "start_position": 0, + "end_position": 5, + "score": 1.0, + "source": "detector", + }, + { + "id": "api_key_13_29", + "value": "12 February 1980", + "label": "api_key", + "start_position": 13, + "end_position": 29, + "score": 1.0, + "source": "augmenter", + }, + ] + ).model_dump(mode="json") + candidates = ValidationCandidatesSchema( + candidates=[ + ValidationCandidateSchema(id="first_name_0_5", value="Alice", label="first_name"), + ValidationCandidateSchema(id="api_key_13_29", value="12 February 1980", label="api_key"), + ] + ).model_dump(mode="json") + row = { + COL_TEXT: text, + COL_MERGED_ENTITIES: merged_entities, + COL_VALIDATION_CANDIDATES: candidates, + COL_TAG_NOTATION: TagNotation.xml.value, + } + + facade = FakeFacade( + "v0", + response={ + "decisions": [ + {"id": "first_name_0_5", "decision": "keep"}, + {"id": "api_key_13_29", "decision": "drop"}, + ] + }, + ) + params = ChunkedValidationParams( + pool=["v0"], + max_entities_per_call=10, + excerpt_window_chars=100, + entities_column=COL_MERGED_ENTITIES, + candidates_column=COL_VALIDATION_CANDIDATES, + output_column=COL_MERGED_VALIDATION_DECISIONS, + prompt_template=_MINIMAL_TEMPLATE, + ) + + out = chunked_validate_row(row, params, {"v0": facade}) + + decisions = { + d["id"]: d["decision"] for d in out[COL_MERGED_VALIDATION_DECISIONS]["decisions"] + } + assert decisions == {"first_name_0_5": "keep", "api_key_13_29": "drop"} + def test_single_chunk_sends_single_chunk_tagged_text_not_windowed_excerpt(self) -> None: """Single-chunk rows must receive the fully tagged document, not a windowed excerpt. diff --git a/tests/engine/test_detection_custom_columns.py b/tests/engine/test_detection_custom_columns.py index 1a42f226..ac7fc4fe 100644 --- a/tests/engine/test_detection_custom_columns.py +++ b/tests/engine/test_detection_custom_columns.py @@ -17,6 +17,8 @@ COL_AUGMENTED_ENTITIES, COL_DETECTED_ENTITIES, COL_MERGED_ENTITIES, + COL_MERGED_VALIDATED_ENTITIES, + COL_MERGED_VALIDATION_DECISIONS, COL_RAW_DETECTED, COL_SEED_ENTITIES, COL_SEED_VALIDATION_CANDIDATES, @@ -30,6 +32,7 @@ from anonymizer.engine.detection.custom_columns import ( _parse_entity_spans, apply_validation_and_finalize, + enrich_merged_validation_decisions, enrich_validation_decisions, merge_and_build_candidates, parse_detected_entities, @@ -146,8 +149,89 @@ def test_apply_validation_and_finalize_handles_malformed_merged_entities() -> No row: dict[str, Any] = { COL_TEXT: "Alice works at Acme.", COL_MERGED_ENTITIES: ["bad-shape"], - COL_VALIDATED_ENTITIES: {"decisions": []}, + COL_MERGED_VALIDATED_ENTITIES: {"decisions": []}, } result = apply_validation_and_finalize(row) assert result[COL_DETECTED_ENTITIES] == {"entities": []} + + +def test_enrich_merged_validation_decisions_uses_merged_candidates() -> None: + row = { + COL_MERGED_VALIDATION_DECISIONS: { + "decisions": [ + {"id": "id2", "decision": "drop", "proposed_label": "", "reason": "hallucinated"}, + ] + }, + COL_VALIDATION_CANDIDATES: { + "candidates": [ + {"id": "id1", "value": "Alice", "label": "first_name", "context_before": "", "context_after": ""}, + {"id": "id2", "value": "12 February 1980", "label": "api_key", "context_before": "", "context_after": ""}, + ] + }, + } + result = enrich_merged_validation_decisions(row) + decisions = result[COL_MERGED_VALIDATED_ENTITIES]["decisions"] + assert decisions == [ + { + "id": "id2", + "value": "12 February 1980", + "label": "api_key", + "decision": "drop", + "proposed_label": "", + "reason": "hallucinated", + } + ] + + +def test_apply_validation_and_finalize_drops_augmented_entity_when_merged_validation_drops_it() -> None: + row: dict[str, Any] = { + COL_TEXT: "Alice met on 12 February 1980.", + COL_MERGED_ENTITIES: { + "entities": [ + { + "id": "first_name_0_5", + "value": "Alice", + "label": "first_name", + "start_position": 0, + "end_position": 5, + "score": 0.95, + "source": "detector", + }, + { + "id": "api_key_13_29", + "value": "12 February 1980", + "label": "api_key", + "start_position": 13, + "end_position": 29, + "score": 1.0, + "source": "augmenter", + }, + ] + }, + COL_MERGED_VALIDATED_ENTITIES: { + "decisions": [ + { + "id": "first_name_0_5", + "value": "Alice", + "label": "first_name", + "decision": "keep", + "proposed_label": "", + "reason": "real first name", + }, + { + "id": "api_key_13_29", + "value": "12 February 1980", + "label": "api_key", + "decision": "drop", + "proposed_label": "", + "reason": "hallucinated augmentation", + }, + ] + }, + } + + result = apply_validation_and_finalize(row) + entities = result[COL_DETECTED_ENTITIES]["entities"] + assert len(entities) == 1 + assert entities[0]["value"] == "Alice" diff --git a/tests/engine/test_detection_workflow.py b/tests/engine/test_detection_workflow.py index 84857314..3394239c 100644 --- a/tests/engine/test_detection_workflow.py +++ b/tests/engine/test_detection_workflow.py @@ -21,6 +21,8 @@ COL_FINAL_ENTITIES, COL_LATENT_ENTITIES, COL_MERGED_ENTITIES, + COL_MERGED_VALIDATED_ENTITIES, + COL_MERGED_VALIDATION_DECISIONS, COL_SEED_ENTITIES, COL_SEED_ENTITIES_JSON, COL_SEED_VALIDATION_CANDIDATES, @@ -28,6 +30,7 @@ COL_TAGGED_TEXT, COL_TEXT, COL_VALIDATED_ENTITIES, + COL_VALIDATION_CANDIDATES, COL_VALIDATION_DECISIONS, DEFAULT_ENTITY_LABELS, ) @@ -516,6 +519,7 @@ def test_detection_workflow_uses_plugin_transform_columns( COL_VALIDATED_ENTITIES: DetectionTransformOperation.ENRICH_VALIDATION_DECISIONS, COL_SEED_ENTITIES_JSON: DetectionTransformOperation.APPLY_VALIDATION_TO_SEED_ENTITIES, COL_MERGED_ENTITIES: DetectionTransformOperation.MERGE_AND_BUILD_CANDIDATES, + COL_MERGED_VALIDATED_ENTITIES: DetectionTransformOperation.ENRICH_MERGED_VALIDATION_DECISIONS, COL_DETECTED_ENTITIES: DetectionTransformOperation.APPLY_VALIDATION_AND_FINALIZE, } for name, operation in expected_operations.items(): @@ -592,6 +596,17 @@ def test_validation_column_is_chunked_validation_plugin( COL_TAG_NOTATION, } + merged_validation_col = _find_column(columns, COL_MERGED_VALIDATION_DECISIONS) + assert isinstance(merged_validation_col, ChunkedValidationConfig) + assert merged_validation_col.entities_column == COL_MERGED_ENTITIES + assert merged_validation_col.candidates_column == COL_VALIDATION_CANDIDATES + assert set(merged_validation_col.required_columns) == { + COL_TEXT, + COL_MERGED_ENTITIES, + COL_VALIDATION_CANDIDATES, + COL_TAG_NOTATION, + } + def test_validator_pool_kwargs_thread_through_to_plugin_config( stub_detector_model_configs: list[ModelConfig],