feat: add detection workflow column plugins - #192
Merged
Conversation
Collaborator
Author
|
Added a local parity check for the detection-column migration. This was intentionally run outside the committed diff: it compares the old Parity harnessfrom __future__ import annotations
import copy
import json
from types import SimpleNamespace
from typing import Any
from data_designer.config.column_configs import CustomColumnConfig
from data_designer.engine.column_generators.generators.custom import CustomColumnGenerator
from anonymizer.engine.constants import (
COL_AUGMENTED_ENTITIES,
COL_DETECTED_ENTITIES,
COL_MERGED_ENTITIES,
COL_RAW_DETECTED,
COL_SEED_ENTITIES,
COL_SEED_ENTITIES_JSON,
COL_SEED_VALIDATION_CANDIDATES,
COL_TEXT,
COL_VALIDATED_ENTITIES,
COL_VALIDATION_DECISIONS,
)
from anonymizer.engine.detection.chunked_validation import ChunkedValidationParams, make_chunked_validation_generator
from anonymizer.engine.detection.custom_columns import (
apply_validation_and_finalize,
apply_validation_to_seed_entities,
enrich_validation_decisions,
merge_and_build_candidates,
parse_detected_entities,
prepare_validation_inputs,
)
from anonymizer.engine.workflow_columns.detection.config import (
ChunkedValidationConfig,
DetectionTransformConfig,
DetectionTransformOperation,
)
from anonymizer.engine.workflow_columns.detection.impl import ChunkedValidationGenerator, DetectionTransformGenerator
class FakeFacade:
def __init__(self, response: dict[str, Any]) -> None:
self.response = response
self.calls: list[dict[str, Any]] = []
def generate(self, *, prompt: str, parser, system_prompt: str | None = None, purpose: str | None = None, **kwargs):
self.calls.append(
{
"prompt": prompt,
"system_prompt": system_prompt,
"purpose": purpose,
"kwargs": kwargs,
}
)
raw = f"```json\n{json.dumps(self.response)}\n```"
return parser(raw), []
class FakeModelRegistry:
def __init__(self, facade: FakeFacade) -> None:
self.facade = facade
def get_model(self, *, model_alias: str) -> FakeFacade:
assert model_alias == "v0"
return self.facade
def resource_provider(facade: FakeFacade) -> SimpleNamespace:
return SimpleNamespace(model_registry=FakeModelRegistry(facade))
def custom_generator(name: str, fn, **kwargs: Any) -> CustomColumnGenerator:
return CustomColumnGenerator(
CustomColumnConfig(name=name, generator_function=fn, **kwargs),
resource_provider=SimpleNamespace(),
)
def compare_transform(
label: str,
*,
name: str,
fn,
operation: DetectionTransformOperation,
row: dict[str, Any],
) -> dict[str, Any]:
old = custom_generator(name, fn).generate(copy.deepcopy(row))
new = DetectionTransformGenerator(
DetectionTransformConfig(name=name, operation=operation),
resource_provider=SimpleNamespace(),
).generate(copy.deepcopy(row))
assert old == new, f"{label} mismatch:\nold={old}\nnew={new}"
print(f"PASS {label}")
return old
def compare_chunked_validation(row: dict[str, Any]) -> None:
response = {
"decisions": [
{"id": "first_name_0_5", "decision": "keep", "reason": "correct"},
{"id": "city_15_22", "decision": "keep", "reason": "correct"},
]
}
params = ChunkedValidationParams(
pool=["v0"],
max_entities_per_call=10,
excerpt_window_chars=20,
prompt_template=(
"Tagged: {{ _seed_tagged_text }}\n"
"Skeleton: {{ _validation_skeleton }}\n"
"Notation: {{ _tag_notation }}"
),
)
old_facade = FakeFacade(response)
old = CustomColumnGenerator(
CustomColumnConfig(
name=COL_VALIDATION_DECISIONS,
generator_function=make_chunked_validation_generator(["v0"]),
generator_params=params,
drop=True,
),
resource_provider=resource_provider(old_facade),
).generate(copy.deepcopy(row))
new_facade = FakeFacade(response)
new = ChunkedValidationGenerator(
ChunkedValidationConfig(
name=COL_VALIDATION_DECISIONS,
pool=["v0"],
max_entities_per_call=params.max_entities_per_call,
excerpt_window_chars=params.excerpt_window_chars,
prompt_template=params.prompt_template,
drop=True,
),
resource_provider=resource_provider(new_facade),
).generate(copy.deepcopy(row))
assert old == new, f"chunked validation mismatch:\nold={old}\nnew={new}"
assert old_facade.calls[0]["purpose"] == new_facade.calls[0]["purpose"]
print("PASS chunked validation")
def main() -> None:
text = "Alice moved to Seattle."
row: dict[str, Any] = {
COL_TEXT: text,
COL_RAW_DETECTED: json.dumps(
{
"entities": [
{"text": "Alice", "label": "first_name", "start": 0, "end": 5, "score": 0.99},
{"text": "Seattle", "label": "city", "start": 15, "end": 22, "score": 0.96},
]
}
),
}
row = compare_transform(
"parse_detected_entities",
name=COL_SEED_ENTITIES,
fn=parse_detected_entities,
operation=DetectionTransformOperation.PARSE_DETECTED_ENTITIES,
row=row,
)
row = compare_transform(
"prepare_validation_inputs",
name=COL_SEED_VALIDATION_CANDIDATES,
fn=prepare_validation_inputs,
operation=DetectionTransformOperation.PREPARE_VALIDATION_INPUTS,
row=row,
)
compare_chunked_validation(row)
row[COL_VALIDATION_DECISIONS] = {
"decisions": [
{"id": "first_name_0_5", "decision": "keep", "reason": "correct"},
{"id": "city_15_22", "decision": "reclass", "proposed_label": "location", "reason": "more precise"},
]
}
row = compare_transform(
"enrich_validation_decisions",
name=COL_VALIDATED_ENTITIES,
fn=enrich_validation_decisions,
operation=DetectionTransformOperation.ENRICH_VALIDATION_DECISIONS,
row=row,
)
row = compare_transform(
"apply_validation_to_seed_entities",
name=COL_SEED_ENTITIES_JSON,
fn=apply_validation_to_seed_entities,
operation=DetectionTransformOperation.APPLY_VALIDATION_TO_SEED_ENTITIES,
row=row,
)
row[COL_AUGMENTED_ENTITIES] = {"entities": []}
row = compare_transform(
"merge_and_build_candidates",
name=COL_MERGED_ENTITIES,
fn=merge_and_build_candidates,
operation=DetectionTransformOperation.MERGE_AND_BUILD_CANDIDATES,
row=row,
)
row = compare_transform(
"apply_validation_and_finalize",
name=COL_DETECTED_ENTITIES,
fn=apply_validation_and_finalize,
operation=DetectionTransformOperation.APPLY_VALIDATION_AND_FINALIZE,
row=row,
)
assert row[COL_DETECTED_ENTITIES]["entities"]
print("All parity checks passed.")
if __name__ == "__main__":
main()Output: |
andreatnvidia
changed the title
andreatnvidia
marked this pull request as ready for review
June 17, 2026 13:37
Collaborator
Author
|
Added a real-provider smoke/benchmark harness for this PR and ran it locally against the Brev endpoints. Endpoints (GLiNER/gpt-oss) both on brev. Interpretation:
2-record smoke: 4-record benchmark, 2 repeats: Local harnessfrom __future__ import annotations
import argparse
import json
import os
import statistics
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import pandas as pd
from data_designer.config.column_configs import CustomColumnConfig, LLMStructuredColumnConfig, LLMTextColumnConfig
from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig, ModelProvider
from data_designer.interface.data_designer import DataDesigner
from anonymizer.config.models import DetectionModelSelection
from anonymizer.engine.constants import (
COL_AUGMENTED_ENTITIES,
COL_DETECTED_ENTITIES,
COL_FINAL_ENTITIES,
COL_MERGED_ENTITIES,
COL_RAW_DETECTED,
COL_SEED_ENTITIES,
COL_SEED_ENTITIES_JSON,
COL_SEED_VALIDATION_CANDIDATES,
COL_TEXT,
COL_VALIDATED_ENTITIES,
COL_VALIDATION_DECISIONS,
_jinja,
)
from anonymizer.engine.detection.chunked_validation import ChunkedValidationParams, make_chunked_validation_generator
from anonymizer.engine.detection.custom_columns import (
apply_validation_and_finalize,
apply_validation_to_seed_entities,
enrich_validation_decisions,
merge_and_build_candidates,
parse_detected_entities,
prepare_validation_inputs,
)
from anonymizer.engine.detection.detection_workflow import (
EntityDetectionResult,
EntityDetectionWorkflow,
_get_augment_prompt,
_get_validation_prompt,
_resolve_detection_labels,
)
from anonymizer.engine.ndd.adapter import NddAdapter
from anonymizer.engine.ndd.model_loader import resolve_model_alias, resolve_model_aliases
from anonymizer.engine.schemas import AugmentedEntitiesSchema, EntitiesSchema
LABELS = [
"first_name",
"last_name",
"email",
"phone_number",
"city",
"organization",
"street_address",
"date",
"ssn",
]
ROWS = [
"Alice Johnson moved from Seattle to Denver on March 3, 2024. Call her at (555) 123-4567.",
"Bob from Acme Corp emailed bob.smith@example.com after visiting 221B Baker Street.",
"Maria Garcia lives in Miami and listed SSN 123-45-6789 on the onboarding form.",
"Chen Wei joined Contoso in Boston; his backup number is 415-555-0199.",
]
@dataclass(frozen=True)
class RunResult:
mode: str
elapsed_s: float
failed_records: int
dataframe: pd.DataFrame
class LegacyEntityDetectionWorkflow(EntityDetectionWorkflow):
def detect_and_validate_entities(
self,
dataframe: pd.DataFrame,
*,
model_configs: list[ModelConfig],
selected_models: DetectionModelSelection,
gliner_detection_threshold: float,
validation_max_entities_per_call: int,
validation_excerpt_window_chars: int,
entity_labels: list[str] | None = None,
data_summary: str | None = None,
preview_num_records: int | None = None,
) -> EntityDetectionResult:
labels = _resolve_detection_labels(entity_labels)
workflow_model_configs = self._inject_detector_params(
model_configs=model_configs,
selected_models=selected_models,
labels=labels,
gliner_detection_threshold=gliner_detection_threshold,
)
detection_alias = resolve_model_alias("entity_detector", selected_models)
validator_aliases = resolve_model_aliases("entity_validator", selected_models)
augmenter_alias = resolve_model_alias("entity_augmenter", selected_models)
validator_generator = make_chunked_validation_generator(validator_aliases)
validator_params = ChunkedValidationParams(
pool=list(validator_aliases),
max_entities_per_call=validation_max_entities_per_call,
excerpt_window_chars=validation_excerpt_window_chars,
prompt_template=_get_validation_prompt(data_summary=data_summary, labels=labels),
)
result = self._adapter.run_workflow(
dataframe,
model_configs=workflow_model_configs,
columns=[
LLMTextColumnConfig(
name=COL_RAW_DETECTED,
prompt=_jinja(COL_TEXT),
model_alias=detection_alias,
),
CustomColumnConfig(name=COL_SEED_ENTITIES, generator_function=parse_detected_entities),
CustomColumnConfig(
name=COL_SEED_VALIDATION_CANDIDATES,
generator_function=prepare_validation_inputs,
),
CustomColumnConfig(
name=COL_VALIDATION_DECISIONS,
generator_function=validator_generator,
generator_params=validator_params,
drop=True,
),
CustomColumnConfig(name=COL_VALIDATED_ENTITIES, generator_function=enrich_validation_decisions),
CustomColumnConfig(
name=COL_SEED_ENTITIES_JSON,
generator_function=apply_validation_to_seed_entities,
),
LLMStructuredColumnConfig(
name=COL_AUGMENTED_ENTITIES,
prompt=_get_augment_prompt(
data_summary=data_summary,
labels=labels,
strict_labels=entity_labels is not None,
),
model_alias=augmenter_alias,
output_format=AugmentedEntitiesSchema,
),
CustomColumnConfig(name=COL_MERGED_ENTITIES, generator_function=merge_and_build_candidates),
CustomColumnConfig(
name=COL_DETECTED_ENTITIES,
generator_function=apply_validation_and_finalize,
),
],
workflow_name="legacy-entity-detection",
preview_num_records=preview_num_records,
)
return EntityDetectionResult(dataframe=result.dataframe.copy(), failed_records=result.failed_records)
def model_providers(api_key_env: str) -> list[ModelProvider]:
if not os.environ.get(api_key_env):
raise RuntimeError(f"{api_key_env} must be set")
return [
ModelProvider(
name="gliner",
endpoint=GLINER_ENDPOINT,
provider_type="openai",
),
ModelProvider(
name="gptoss",
endpoint=GPTOSS_ENDPOINT,
provider_type="openai",
api_key=api_key_env,
),
]
def model_configs() -> list[ModelConfig]:
return [
ModelConfig(
alias="gliner-pii-detector",
model="fastino/gliner2-privacy-filter-PII-multi",
provider="gliner",
skip_health_check=True,
inference_parameters=ChatCompletionInferenceParams(
max_parallel_requests=4,
timeout=120,
),
),
ModelConfig(
alias="gpt-oss-120b",
model="gpt-oss-120b",
provider="gptoss",
skip_health_check=True,
inference_parameters=ChatCompletionInferenceParams(
max_parallel_requests=2,
max_tokens=4096,
temperature=0.0,
top_p=1.0,
timeout=300,
),
),
]
def selection() -> DetectionModelSelection:
return DetectionModelSelection(
entity_detector="gliner-pii-detector",
entity_validator=["gpt-oss-120b"],
entity_augmenter="gpt-oss-120b",
latent_detector="gpt-oss-120b",
)
def input_df(records: int) -> pd.DataFrame:
rows = [ROWS[i % len(ROWS)] for i in range(records)]
return pd.DataFrame({COL_TEXT: rows})
def run_workflow(
mode: str,
workflow_cls: type[EntityDetectionWorkflow],
providers: list[ModelProvider],
records: int,
) -> RunResult:
artifact_root = Path(tempfile.mkdtemp(prefix=f"anonymizer_{mode}_bench_"))
data_designer = DataDesigner(artifact_path=artifact_root, model_providers=providers)
workflow = workflow_cls(adapter=NddAdapter(data_designer))
start = time.perf_counter()
result = workflow.run(
input_df(records),
model_configs=model_configs(),
selected_models=selection(),
gliner_detection_threshold=0.5,
validation_max_entities_per_call=8,
validation_excerpt_window_chars=160,
entity_labels=LABELS,
data_summary="Short synthetic contact records containing names, locations, organizations, emails, phones, and IDs.",
tag_latent_entities=False,
compute_grouped_entities=True,
preview_num_records=records,
)
elapsed = time.perf_counter() - start
return RunResult(mode=mode, elapsed_s=elapsed, failed_records=len(result.failed_records), dataframe=result.dataframe)
def signatures(df: pd.DataFrame) -> list[list[tuple[str, str, int, int]]]:
result: list[list[tuple[str, str, int, int]]] = []
for raw in df[COL_FINAL_ENTITIES].tolist():
parsed = EntitiesSchema.from_raw(raw)
result.append(
sorted(
(entity.value, entity.label, entity.start_position, entity.end_position)
for entity in parsed.entities
)
)
return result
def row_counts(df: pd.DataFrame) -> list[int]:
return [len(row) for row in signatures(df)]
def emit_summary(run_results: list[RunResult]) -> None:
by_mode: dict[str, list[RunResult]] = {}
for result in run_results:
by_mode.setdefault(result.mode, []).append(result)
print("\nRuntime summary")
for mode, results in by_mode.items():
times = [result.elapsed_s for result in results]
print(
json.dumps(
{
"mode": mode,
"runs": len(times),
"elapsed_s": [round(t, 3) for t in times],
"mean_s": round(statistics.mean(times), 3),
"failed_records": [result.failed_records for result in results],
"entity_counts_last": row_counts(results[-1].dataframe),
},
sort_keys=True,
)
)
if {"legacy", "plugin"} <= by_mode.keys():
legacy_sig = signatures(by_mode["legacy"][-1].dataframe)
plugin_sig = signatures(by_mode["plugin"][-1].dataframe)
print("\nLast-run output comparison")
print(json.dumps({"exact_entity_signature_match": legacy_sig == plugin_sig}, sort_keys=True))
if legacy_sig != plugin_sig:
print(json.dumps({"legacy_last": legacy_sig, "plugin_last": plugin_sig}, default=str))
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--records", type=int, default=4)
parser.add_argument("--repeats", type=int, default=1)
parser.add_argument("--api-key-env", default="GPTOSS_API_KEY")
args = parser.parse_args()
providers = model_providers(args.api_key_env)
schedule: list[tuple[str, type[EntityDetectionWorkflow]]] = []
for _ in range(args.repeats):
schedule.extend(
[
("legacy", LegacyEntityDetectionWorkflow),
("plugin", EntityDetectionWorkflow),
]
)
results: list[RunResult] = []
for mode, cls in schedule:
print(f"\nRunning {mode} workflow on {args.records} records...")
result = run_workflow(mode, cls, providers, args.records)
results.append(result)
print(
json.dumps(
{
"mode": mode,
"elapsed_s": round(result.elapsed_s, 3),
"failed_records": result.failed_records,
"entity_counts": row_counts(result.dataframe),
},
sort_keys=True,
)
)
emit_summary(results)
if __name__ == "__main__":
main() |
Contributor
Greptile SummaryThis PR migrates the Anonymizer detection workflow off
Confidence Score: 5/5Safe to merge — the plugin migration is a clean structural refactor with no functional changes to detection behaviour, and the new async path is well-tested. The CustomColumnConfig replacement is mechanical and correct; DAG column dependencies in _REQUIRED_COLUMNS/_SIDE_EFFECT_COLUMNS match the actual function signatures; the bridge test added in this PR closes the coverage gap flagged in the previous review round. The single asyncio.gather observation is an efficiency note consistent with existing ThreadPoolExecutor behaviour, not a correctness regression. No files require blocking attention; impl.py is worth a second glance if DataDesigner's async engine becomes the default before a TaskGroup-style cancellation improvement is made. Important Files Changed
Reviews (3): Last reviewed commit: "feat: add detection workflow column plug..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
lipikaramaswamy
approved these changes
Jun 23, 2026
lipikaramaswamy
left a comment
lipikaramaswamy
left a comment
Collaborator
There was a problem hiding this comment.
Looks great! I don't think there's much overlap between files changed here and current main, but worth just doing a quick rebase.
I verified the migration preserves the existing detection workflow behavior: the old inline custom-column wiring is replaced with serializable DataDesigner plugin column configs, while the underlying detection transform and chunked validation logic remains reused.
Validation notes:
- full test suite passed
- targeted detection/chunked-validation tests passed
- format check passed
- plugin discovery / config serialization smoke passed
- real legal dataset smoke passed with real hosted detector + LLM validation/augmentation on a 3-row preview
One operational note, not a blocker here but related to other threads: chunk-level validation "parallelism" already existed before this PR. This PR preserves it through the plugin column. Local validation work is still one worker per chunk, while provider pressure is controlled indirectly by validation_max_entities_per_call, validator pool aliases, and per-alias max_parallel_requests. I wonder to what extent batching requests helps (I think this depends on the endpoint infra too?) , and also I wonder if an explicit validation_max_parallel_chunks knob should be a follow-up hardening item.
andreatnvidia
force-pushed
the
andreatgretel/feat/workflow-columns
branch
from
June 24, 2026 23:19
887c00f to
2197ca9
Compare
andreatnvidia
force-pushed
the
andreatgretel/feat/workflow-columns
branch
from
June 24, 2026 23:21
2197ca9 to
fdc397e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
data_designer.pluginsentry points.CustomColumnConfigfor parse, validation prep, chunked validation, merge, and finalize steps.plans/custom-column-plugins/, including the note that native NER transport is a separate follow-up from pluginizing custom columns.Anonymizerexport lazy so DataDesigner plugin discovery does not hit a circular import.Type of Change
Testing
make testpasses locallymake checkpasses locally (format + lint + typecheck + lock-check)Targeted validation:
.venv/bin/ruff check --fix ..venv/bin/ruff format ..venv/bin/ruff check src/anonymizer/engine/detection src/anonymizer/engine/workflow_columns tests/engine/test_detection_workflow.py tests/engine/test_chunked_validation.py.venv/bin/pytest tests/engine/test_detection_postprocess.py tests/engine/test_chunked_validation.py tests/engine/test_detection_workflow.pyanonymizer-chunked-validationandanonymizer-detection-transformload via DataDesigner registry.Documentation
make docs-buildpasses locallyNot run; this PR adds a migration plan under
plans/, not rendered docs-site content.NER Transport Note
This PR does not remove the chat-completions-compatible NER path. The detector remains an
LLMTextColumnConfig; avoiding the extra chat completions head would require a separate native detector workflow column or provider/client path.Related Issues