diff --git a/examples/dataflow-tool/README.md b/examples/dataflow-tool/README.md new file mode 100644 index 000000000..7989bfb79 --- /dev/null +++ b/examples/dataflow-tool/README.md @@ -0,0 +1,113 @@ +# In-flow record redaction for dataflow tools + +`openmed.integrations.dataflow_tool_processor` provides a small callable for +visual dataflow tools and script processors. It accepts one record and its +flow-file attributes, redacts only configured fields through OpenMed's batch +API, and returns exactly one record plus one attribute mapping. + +The adapter keeps a model loader in a module-level cache. A long-lived worker +therefore loads a model pipeline once and reuses it across records. Cache keys +include the model and OpenMed configuration, so differently configured flows +do not share a loader. + +## Call it from an embedded Python processor + +Import the callable at script-module scope so the host keeps the module and its +pipeline cache alive between record callbacks: + +```python +from openmed.integrations.dataflow_tool_processor import script_processor + + +def transform(record, attributes): + return script_processor( + record, + attributes, + fields=("note", "patient.contact"), + policy="hipaa_safe_harbor", + method="mask", + ) +``` + +For this synthetic input: + +```python +record = { + "record_id": "synthetic-001", + "note": "Patient Jane Roe called 555-0100.", + "patient": {"contact": "jane.roe@example.org"}, + "facility": "example-clinic", +} +attributes = {"route": "clinical", "mime.type": "application/json"} + +redacted_record, output_attributes = transform(record, attributes) +``` + +the non-target fields and incoming attributes are preserved. OpenMed adds only +string-encoded counts suitable for flow-file attribute systems: + +```json +{ + "openmed.redaction.record_count": "1", + "openmed.redaction.field_count": "2", + "openmed.redaction.entity_count": "3" +} +``` + +The adapter never adds original values, replacement values, entity surfaces, +or field names to attributes. Existing attributes are passed through as-is, so +the flow itself remains responsible for preventing PHI from being placed in +incoming attributes. + +Configured fields may be top-level names or dotted mapping paths. Missing, +null, empty, and non-string targets pass through unchanged. A batch failure, +invalid item, or cardinality mismatch fails closed with a diagnostic that does +not include record content. + +## Run the persistent JSON-lines entrypoint + +Hosts that exchange data over stdin/stdout can keep one Python worker alive and +send one envelope per line. Each envelope must contain a `record` object and +may contain an `attributes` object: + +```json +{"record":{"record_id":"synthetic-001","note":"Patient Jane Roe"},"attributes":{"route":"clinical"}} +``` + +Start the worker after installing OpenMed and caching the selected PII model: + +```bash +OPENMED_OFFLINE=1 python -m openmed.integrations.dataflow_tool_processor \ + --field note \ + --field patient.contact \ + --policy hipaa_safe_harbor +``` + +The worker writes one compact JSON envelope for every non-blank input envelope, +in the same order. Keep stderr on a protected operational channel; errors +contain only a line number and a PHI-free failure category. + +## NiFi-style wiring + +For a CPython-capable scripted record processor, use the `script_processor` +callback above and map its returned record and attributes to the outgoing +FlowFile. Keep one worker/module instance per processor task to retain the +pipeline cache. + +Apache NiFi's `ExecuteScript` Python engine is Jython and cannot load CPython +packages. For NiFi, use +[`ExecuteStreamCommand`](https://nifi.apache.org/components/org.apache.nifi.processors.standard.ExecuteStreamCommand/) +or a long-lived sidecar that runs the module entrypoint. Configure the command +to receive stdin, route zero-status output to `output stream`, and route +`nonzero status` to a protected failure path. If `ExecuteStreamCommand` starts +a fresh process for each FlowFile, batch multiple record envelopes in that +FlowFile or use a resident sidecar so the model is not reloaded per record. + +Do not place raw records in command arguments, environment variables, +attributes, provenance annotations, or processor logs. Pass record content +only through stdin or the in-process callable, run with `OPENMED_OFFLINE=1` +after model provisioning, and validate the output relationship before it +reaches an external sink. + +This integration is a script adapter only. It does not provide a compiled NiFi +NAR or configure a flow controller. diff --git a/openmed/integrations/__init__.py b/openmed/integrations/__init__.py index cf18b3818..2bcdd3d23 100644 --- a/openmed/integrations/__init__.py +++ b/openmed/integrations/__init__.py @@ -6,6 +6,19 @@ redact_columnar, redact_columnar_dataset, ) +from .dataflow_tool_processor import ( + DEFAULT_DATAFLOW_TOOL_MODEL, + ENTITY_COUNT_ATTRIBUTE, + FIELD_COUNT_ATTRIBUTE, + RECORD_COUNT_ATTRIBUTE, + DataflowToolConfig, + DataflowToolProcessorError, + clear_pipeline_cache, + process_flow_file, + process_json_lines, + process_record, + script_processor, +) from .lakehouse_redact import ( LakehouseRedactionProgress, LakehouseRedactionResult, @@ -34,18 +47,28 @@ __all__ = [ "ColumnarProgress", "ColumnarRedactionResult", + "DEFAULT_DATAFLOW_TOOL_MODEL", "DEFAULT_LOG_MESSAGE_FIELDS", "DEFAULT_LOG_REDACTION_MODEL", "DEFAULT_BATCH_ID_COLUMN", "DEFAULT_SPARK_POLICY", + "ENTITY_COUNT_ATTRIBUTE", + "FIELD_COUNT_ATTRIBUTE", "LakehouseRedactionProgress", "LakehouseRedactionResult", "LogRedactorConfig", "LogRedactorError", + "RECORD_COUNT_ATTRIBUTE", "SparkDeidentifyColumn", "SparkDeidentifySink", "SparkDeidentifyStreamBuilder", + "DataflowToolConfig", + "DataflowToolProcessorError", + "clear_pipeline_cache", "deidentify_write_stream", + "process_flow_file", + "process_json_lines", + "process_record", "redact_columnar", "redact_columnar_dataset", "redact_lakehouse", @@ -53,5 +76,6 @@ "redact_log_events", "redact_ndjson_lines", "redact_ndjson_stream", + "script_processor", "write_deidentified_stream", ] diff --git a/openmed/integrations/dataflow_tool_processor.py b/openmed/integrations/dataflow_tool_processor.py new file mode 100644 index 000000000..d3c640b10 --- /dev/null +++ b/openmed/integrations/dataflow_tool_processor.py @@ -0,0 +1,525 @@ +"""Script-processor adapter for in-flow record de-identification. + +The public callable accepts one record plus its flow-file attributes and +returns a redacted copy with the attributes preserved. Only configured string +fields are submitted to OpenMed. Model loaders are cached at module scope so a +long-lived script processor can reuse its loaded pipeline across records. +""" + +from __future__ import annotations + +import argparse +import copy +import json +import sys +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from threading import RLock +from typing import Any, TextIO + +from openmed.core.pii_i18n import DEFAULT_PII_MODELS +from openmed.processing.batch import process_batch + +DEFAULT_DATAFLOW_TOOL_MODEL = DEFAULT_PII_MODELS["en"] +RECORD_COUNT_ATTRIBUTE = "openmed.redaction.record_count" +FIELD_COUNT_ATTRIBUTE = "openmed.redaction.field_count" +ENTITY_COUNT_ATTRIBUTE = "openmed.redaction.entity_count" + +ProcessBatch = Callable[..., Any] +_MISSING = object() +_RESERVED_DEIDENTIFY_KWARGS = frozenset( + { + "batch_size", + "confidence_threshold", + "continue_on_error", + "ids", + "loader", + "method", + "model_name", + "operation", + "texts", + } +) + + +class DataflowToolProcessorError(RuntimeError): + """Raised when a flow-file record cannot be redacted safely.""" + + +@dataclass(frozen=True) +class DataflowToolConfig: + """Configuration for record-level in-flow redaction. + + Args: + fields: Top-level or dotted record fields to redact. + model_name: PII model passed to batch de-identification. + method: De-identification method. + confidence_threshold: Minimum confidence for PII detection. + deidentify_kwargs: Extra keyword arguments forwarded to + :func:`openmed.processing.batch.process_batch`. + """ + + fields: tuple[str, ...] + model_name: str = DEFAULT_DATAFLOW_TOOL_MODEL + method: str = "mask" + confidence_threshold: float | None = 0.7 + deidentify_kwargs: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "fields", _normalize_fields(self.fields)) + if not isinstance(self.model_name, str) or not self.model_name.strip(): + raise ValueError("model_name must be a non-empty string") + if not isinstance(self.method, str) or not self.method.strip(): + raise ValueError("method must be a non-empty string") + + reserved = _RESERVED_DEIDENTIFY_KWARGS.intersection(self.deidentify_kwargs) + if reserved: + names = ", ".join(sorted(reserved)) + raise ValueError( + "deidentify kwargs must not override processor options: " + names + ) + + +@dataclass +class _CachedPipeline: + loader: Any + lock: RLock = field(default_factory=RLock) + + +@dataclass(frozen=True) +class _Target: + path: tuple[str, ...] + text: str + + +_PIPELINE_CACHE: dict[tuple[str, int | None], _CachedPipeline] = {} +_PIPELINE_CACHE_LOCK = RLock() + + +def process_flow_file( + record: Mapping[str, Any], + attributes: Mapping[str, str] | None = None, + *, + fields: Sequence[str] | str, + model_name: str = DEFAULT_DATAFLOW_TOOL_MODEL, + method: str = "mask", + confidence_threshold: float | None = 0.7, + config: Any | None = None, + process_batch_fn: ProcessBatch | None = None, + **deidentify_kwargs: Any, +) -> tuple[dict[str, Any], dict[str, str]]: + """Return one redacted record and its PHI-safe flow-file attributes. + + The input mappings are never mutated. Incoming attributes pass through + unchanged, while OpenMed adds only numeric record, changed-field, and + detected-entity counts. Missing, null, empty, and non-string configured + fields pass through unchanged. + + Args: + record: Record-like flow-file content. + attributes: Existing string flow-file attributes. + fields: Top-level or dotted record fields to redact. + model_name: PII model passed to batch de-identification. + method: De-identification method. + confidence_threshold: Minimum confidence for PII detection. + config: Optional OpenMed configuration used by the cached model loader. + process_batch_fn: Optional ``process_batch`` replacement, primarily for + offline tests. + **deidentify_kwargs: Extra batch de-identification options such as + ``policy``, ``lang``, or ``use_safety_sweep``. + + Returns: + A ``(record, attributes)`` tuple containing one output record for the + one input record. + + Raises: + TypeError: If the record or attributes are not mappings. + DataflowToolProcessorError: If batch de-identification fails or returns + an invalid result. + """ + + if not isinstance(record, Mapping): + raise TypeError("record must be a mapping") + + processor_config = DataflowToolConfig( + fields=_normalize_fields(fields), + model_name=model_name, + method=method, + confidence_threshold=confidence_threshold, + deidentify_kwargs=deidentify_kwargs, + ) + output = copy.deepcopy(dict(record)) + output_attributes = _copy_attributes(attributes) + targets = _collect_targets(output, processor_config.fields) + + changed_fields = 0 + entity_count = 0 + if targets: + pipeline = _get_pipeline(processor_config.model_name, config) + batch_callable = process_batch_fn or process_batch + texts = [target.text for target in targets] + kwargs = dict(processor_config.deidentify_kwargs) + kwargs.update( + { + "operation": "deidentify", + "batch_size": len(texts), + "continue_on_error": False, + "loader": pipeline.loader, + "method": processor_config.method, + "confidence_threshold": processor_config.confidence_threshold, + } + ) + + try: + with pipeline.lock: + batch_result = batch_callable( + texts, + model_name=processor_config.model_name, + config=config, + ids=[f"field_{index}" for index in range(len(texts))], + **kwargs, + ) + except Exception: + raise DataflowToolProcessorError( + "failed to redact configured record fields" + ) from None + + items = _batch_items(batch_result) + if len(items) != len(targets): + raise DataflowToolProcessorError( + "batch result count did not match configured field count" + ) + + for target, item in zip(targets, items): + redacted_text, redactions = _redacted_item(item) + _set_path(output, target.path, redacted_text) + changed_fields += int(redacted_text != target.text) + entity_count += redactions + + output_attributes[RECORD_COUNT_ATTRIBUTE] = "1" + output_attributes[FIELD_COUNT_ATTRIBUTE] = str(changed_fields) + output_attributes[ENTITY_COUNT_ATTRIBUTE] = str(entity_count) + return output, output_attributes + + +def process_record( + record: Mapping[str, Any], + attributes: Mapping[str, str] | None = None, + **kwargs: Any, +) -> tuple[dict[str, Any], dict[str, str]]: + """Alias-style record entrypoint for generic script processors.""" + + return process_flow_file(record, attributes, **kwargs) + + +def script_processor( + record: Mapping[str, Any], + attributes: Mapping[str, str] | None = None, + **kwargs: Any, +) -> tuple[dict[str, Any], dict[str, str]]: + """Process one record using a script-processor-compatible signature.""" + + return process_flow_file(record, attributes, **kwargs) + + +def process_json_lines( + input_stream: Iterable[str], + output_stream: TextIO, + *, + fields: Sequence[str] | str, + model_name: str = DEFAULT_DATAFLOW_TOOL_MODEL, + method: str = "mask", + confidence_threshold: float | None = 0.7, + config: Any | None = None, + process_batch_fn: ProcessBatch | None = None, + **deidentify_kwargs: Any, +) -> int: + """Process JSON-line flow-file envelopes through a long-lived worker. + + Each non-blank input line must contain ``record`` and may contain + ``attributes``. Exactly one output envelope is written per input envelope. + The long-lived module process reuses its cached model loader across lines. + + Returns: + Number of output envelopes written. + """ + + emitted = 0 + for line_number, raw_line in enumerate(input_stream, start=1): + if not raw_line.strip(): + continue + try: + envelope = json.loads(raw_line) + except json.JSONDecodeError: + raise DataflowToolProcessorError( + f"invalid JSON envelope at input line {line_number}" + ) from None + if not isinstance(envelope, Mapping): + raise DataflowToolProcessorError( + f"input line {line_number} must contain a JSON object" + ) + if "record" not in envelope: + raise DataflowToolProcessorError( + f"input line {line_number} is missing the record object" + ) + + try: + redacted_record, redacted_attributes = process_flow_file( + envelope["record"], + envelope.get("attributes"), + fields=fields, + model_name=model_name, + method=method, + confidence_threshold=confidence_threshold, + config=config, + process_batch_fn=process_batch_fn, + **deidentify_kwargs, + ) + except (DataflowToolProcessorError, TypeError, ValueError): + raise DataflowToolProcessorError( + f"failed to process input line {line_number}" + ) from None + + output_stream.write( + json.dumps( + { + "record": redacted_record, + "attributes": redacted_attributes, + }, + ensure_ascii=False, + separators=(",", ":"), + ) + ) + output_stream.write("\n") + emitted += 1 + + return emitted + + +def main( + argv: Sequence[str] | None = None, + *, + stdin: TextIO | None = None, + stdout: TextIO | None = None, + stderr: TextIO | None = None, +) -> int: + """Run the persistent JSON-lines script-processor entrypoint.""" + + parser = _build_parser() + args = parser.parse_args(argv) + input_stream = stdin if stdin is not None else sys.stdin + output_stream = stdout if stdout is not None else sys.stdout + error_stream = stderr if stderr is not None else sys.stderr + + try: + process_json_lines( + input_stream, + output_stream, + fields=tuple(args.field), + model_name=args.model_name, + method=args.method, + confidence_threshold=args.confidence_threshold, + lang=args.lang, + policy=args.policy, + use_safety_sweep=args.use_safety_sweep, + ) + except (DataflowToolProcessorError, TypeError, ValueError) as exc: + error_stream.write(f"dataflow record redaction failed: {exc}\n") + return 1 + return 0 + + +def clear_pipeline_cache() -> None: + """Clear cached model loaders used by long-lived script processors.""" + + with _PIPELINE_CACHE_LOCK: + _PIPELINE_CACHE.clear() + + +def _get_pipeline(model_name: str, config: Any | None) -> _CachedPipeline: + cache_key = (model_name, None if config is None else id(config)) + with _PIPELINE_CACHE_LOCK: + pipeline = _PIPELINE_CACHE.get(cache_key) + if pipeline is None: + pipeline = _CachedPipeline(loader=_create_pipeline(config)) + _PIPELINE_CACHE[cache_key] = pipeline + return pipeline + + +def _create_pipeline(config: Any | None) -> Any: + from openmed.core import ModelLoader + + return ModelLoader(config) + + +def _copy_attributes( + attributes: Mapping[str, str] | None, +) -> dict[str, str]: + if attributes is None: + return {} + if not isinstance(attributes, Mapping): + raise TypeError("attributes must be a mapping") + output = dict(attributes) + if any(not isinstance(key, str) for key in output): + raise TypeError("attribute names must be strings") + if any(not isinstance(value, str) for value in output.values()): + raise TypeError("attribute values must be strings") + return output + + +def _collect_targets(record: Mapping[str, Any], fields: Sequence[str]) -> list[_Target]: + targets: list[_Target] = [] + for field_name in fields: + path = tuple(field_name.split(".")) + value = _get_path(record, path) + if isinstance(value, str) and value: + targets.append(_Target(path=path, text=value)) + return targets + + +def _get_path(record: Mapping[str, Any], path: Sequence[str]) -> Any: + value: Any = record + for part in path: + if not isinstance(value, Mapping) or part not in value: + return _MISSING + value = value[part] + return value + + +def _set_path(record: dict[str, Any], path: Sequence[str], value: str) -> None: + parent: dict[str, Any] = record + for part in path[:-1]: + child = parent[part] + if not isinstance(child, dict): + raise DataflowToolProcessorError( + "configured record field has an invalid nested path" + ) + parent = child + parent[path[-1]] = value + + +def _batch_items(batch_result: Any) -> list[Any]: + raw_items = getattr(batch_result, "items", batch_result) + try: + return list(raw_items) + except TypeError: + raise DataflowToolProcessorError( + "batch redaction returned a non-iterable result" + ) from None + + +def _redacted_item(item: Any) -> tuple[str, int]: + if getattr(item, "success", True) is False: + raise DataflowToolProcessorError( + "one or more configured record fields could not be redacted" + ) + + result = getattr(item, "result", item) + if isinstance(result, str): + return result, 0 + if isinstance(result, Mapping): + redacted = result.get("deidentified_text") + entities = result.get("pii_entities", result.get("entities", ())) + else: + redacted = getattr(result, "deidentified_text", None) + entities = getattr( + result, + "pii_entities", + getattr(result, "entities", ()), + ) + if not isinstance(redacted, str): + raise DataflowToolProcessorError( + "batch redaction returned an invalid field result" + ) + + try: + entity_count = len(entities or ()) + except TypeError: + entity_count = 0 + return redacted, entity_count + + +def _normalize_fields(fields: Sequence[str] | str) -> tuple[str, ...]: + raw_fields = (fields,) if isinstance(fields, str) else tuple(fields) + normalized: list[str] = [] + for field_name in raw_fields: + if not isinstance(field_name, str) or not field_name.strip(): + raise ValueError("fields must contain non-empty strings") + field_name = field_name.strip() + if any(not part for part in field_name.split(".")): + raise ValueError("fields must use non-empty dotted path segments") + if field_name not in normalized: + normalized.append(field_name) + if not normalized: + raise ValueError("fields must include at least one field name") + return tuple(normalized) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=("Redact configured fields in JSON-line flow-file envelopes.") + ) + parser.add_argument( + "--field", + action="append", + required=True, + help=( + "Record field to redact. Repeat for multiple fields; dotted paths " + "such as patient.note are supported." + ), + ) + parser.add_argument( + "--model-name", + default=DEFAULT_DATAFLOW_TOOL_MODEL, + help="PII model identifier passed to OpenMed batch de-identification.", + ) + parser.add_argument( + "--method", + default="mask", + choices=("mask", "remove", "replace", "hash", "shift_dates"), + help="De-identification method.", + ) + parser.add_argument( + "--confidence-threshold", + type=float, + default=0.7, + help="Minimum PII detection confidence.", + ) + parser.add_argument( + "--lang", + default="en", + help="Language hint forwarded to de-identification.", + ) + parser.add_argument( + "--policy", + default=None, + help="Optional policy profile forwarded to de-identification.", + ) + parser.add_argument( + "--use-safety-sweep", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable or disable the regex safety-sweep fallback.", + ) + return parser + + +__all__ = [ + "DEFAULT_DATAFLOW_TOOL_MODEL", + "ENTITY_COUNT_ATTRIBUTE", + "FIELD_COUNT_ATTRIBUTE", + "RECORD_COUNT_ATTRIBUTE", + "DataflowToolConfig", + "DataflowToolProcessorError", + "ProcessBatch", + "clear_pipeline_cache", + "main", + "process_flow_file", + "process_json_lines", + "process_record", + "script_processor", +] + + +if __name__ == "__main__": # pragma: no cover - exercised via main() + raise SystemExit(main()) diff --git a/tests/unit/integrations/test_dataflow_tool_processor.py b/tests/unit/integrations/test_dataflow_tool_processor.py new file mode 100644 index 000000000..3e8a277cf --- /dev/null +++ b/tests/unit/integrations/test_dataflow_tool_processor.py @@ -0,0 +1,280 @@ +"""Offline tests for the scriptable dataflow-tool processor.""" + +from __future__ import annotations + +import io +import json +from types import SimpleNamespace +from typing import Any + +import pytest + +from openmed.integrations import dataflow_tool_processor as processor + + +def _redact(text: str) -> tuple[str, list[object]]: + replacements = { + "Jane Roe": "[NAME]", + "John Doe": "[NAME]", + "555-0100": "[PHONE]", + "jane.roe@example.org": "[EMAIL]", + } + entities: list[object] = [] + for raw_value, replacement in replacements.items(): + if raw_value in text: + text = text.replace(raw_value, replacement) + entities.append(object()) + return text, entities + + +def _batch_result(texts: list[str]) -> SimpleNamespace: + items = [] + for text in texts: + redacted, entities = _redact(text) + items.append( + SimpleNamespace( + success=True, + result=SimpleNamespace( + deidentified_text=redacted, + pii_entities=entities, + ), + ) + ) + return SimpleNamespace(items=items) + + +@pytest.fixture(autouse=True) +def _empty_pipeline_cache() -> None: + processor.clear_pipeline_cache() + yield + processor.clear_pipeline_cache() + + +def test_callable_redacts_configured_fields_and_reuses_pipeline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pipeline = object() + pipeline_loads: list[Any] = [] + calls: list[tuple[list[str], dict[str, Any]]] = [] + + def fake_create_pipeline(config: Any | None) -> object: + pipeline_loads.append(config) + return pipeline + + def fake_process_batch(texts: list[str], **kwargs: Any) -> SimpleNamespace: + calls.append((list(texts), dict(kwargs))) + return _batch_result(list(texts)) + + monkeypatch.setattr(processor, "_create_pipeline", fake_create_pipeline) + original_records = [ + { + "record_id": "synthetic-1", + "note": "Patient Jane Roe called 555-0100", + "patient": {"contact": "jane.roe@example.org"}, + "facility": "example-clinic", + }, + { + "record_id": "synthetic-2", + "note": "John Doe requested follow-up", + "patient": {"contact": None}, + "facility": "example-clinic", + }, + ] + incoming_attributes = [ + {"route": "clinical", "mime.type": "application/json"}, + {"route": "follow-up", "mime.type": "application/json"}, + ] + + outputs = [ + processor.process_flow_file( + record, + attributes, + fields=("note", "patient.contact", "missing"), + process_batch_fn=fake_process_batch, + policy="hipaa_safe_harbor", + use_safety_sweep=False, + ) + for record, attributes in zip(original_records, incoming_attributes) + ] + + assert len(outputs) == len(original_records) + first_record, first_attributes = outputs[0] + second_record, second_attributes = outputs[1] + assert first_record == { + "record_id": "synthetic-1", + "note": "Patient [NAME] called [PHONE]", + "patient": {"contact": "[EMAIL]"}, + "facility": "example-clinic", + } + assert second_record["note"] == "[NAME] requested follow-up" + assert second_record["patient"]["contact"] is None + assert original_records[0]["note"] == "Patient Jane Roe called 555-0100" + assert original_records[0]["patient"]["contact"] == "jane.roe@example.org" + + assert first_attributes == { + **incoming_attributes[0], + processor.RECORD_COUNT_ATTRIBUTE: "1", + processor.FIELD_COUNT_ATTRIBUTE: "2", + processor.ENTITY_COUNT_ATTRIBUTE: "3", + } + assert second_attributes == { + **incoming_attributes[1], + processor.RECORD_COUNT_ATTRIBUTE: "1", + processor.FIELD_COUNT_ATTRIBUTE: "1", + processor.ENTITY_COUNT_ATTRIBUTE: "1", + } + assert pipeline_loads == [None] + assert [texts for texts, _ in calls] == [ + [ + "Patient Jane Roe called 555-0100", + "jane.roe@example.org", + ], + ["John Doe requested follow-up"], + ] + assert all(kwargs["loader"] is pipeline for _, kwargs in calls) + assert all(kwargs["operation"] == "deidentify" for _, kwargs in calls) + assert all(kwargs["continue_on_error"] is False for _, kwargs in calls) + assert all(kwargs["policy"] == "hipaa_safe_harbor" for _, kwargs in calls) + assert all(kwargs["use_safety_sweep"] is False for _, kwargs in calls) + assert [kwargs["batch_size"] for _, kwargs in calls] == [2, 1] + + for attributes in (first_attributes, second_attributes): + emitted_counts = { + key: value + for key, value in attributes.items() + if key.startswith("openmed.redaction.") + } + assert emitted_counts + assert all(value.isdecimal() for value in emitted_counts.values()) + serialized = json.dumps(emitted_counts) + assert "Jane Roe" not in serialized + assert "John Doe" not in serialized + assert "555-0100" not in serialized + assert "jane.roe@example.org" not in serialized + + +def test_json_lines_entrypoint_preserves_record_count_and_attributes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pipeline_loads = 0 + + def fake_create_pipeline(config: Any | None) -> object: + nonlocal pipeline_loads + pipeline_loads += 1 + return object() + + monkeypatch.setattr(processor, "_create_pipeline", fake_create_pipeline) + envelopes = [ + { + "record": {"id": "one", "note": "Jane Roe"}, + "attributes": {"route": "a"}, + }, + { + "record": {"id": "two", "note": "Call 555-0100"}, + "attributes": {"route": "b"}, + }, + { + "record": {"id": "three", "note": "No identifiers"}, + "attributes": {"route": "c"}, + }, + ] + input_stream = io.StringIO( + "\n".join(json.dumps(envelope) for envelope in envelopes) + "\n" + ) + output_stream = io.StringIO() + + emitted = processor.process_json_lines( + input_stream, + output_stream, + fields="note", + process_batch_fn=lambda texts, **kwargs: _batch_result(list(texts)), + ) + + output_envelopes = [ + json.loads(line) for line in output_stream.getvalue().splitlines() + ] + assert emitted == len(envelopes) + assert len(output_envelopes) == len(envelopes) + assert [item["record"]["id"] for item in output_envelopes] == [ + "one", + "two", + "three", + ] + assert [item["attributes"]["route"] for item in output_envelopes] == [ + "a", + "b", + "c", + ] + assert output_envelopes[0]["record"]["note"] == "[NAME]" + assert output_envelopes[1]["record"]["note"] == "Call [PHONE]" + assert pipeline_loads == 1 + + +def test_missing_and_non_string_fields_pass_through_without_loading_pipeline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_create_pipeline(config: Any | None) -> object: + raise AssertionError("pipeline should not load without redaction targets") + + monkeypatch.setattr(processor, "_create_pipeline", fail_create_pipeline) + record = {"id": "synthetic", "note": None, "sequence": 4} + + output, attributes = processor.script_processor( + record, + {"route": "passthrough"}, + fields=("note", "sequence", "missing"), + ) + + assert output == record + assert attributes == { + "route": "passthrough", + processor.RECORD_COUNT_ATTRIBUTE: "1", + processor.FIELD_COUNT_ATTRIBUTE: "0", + processor.ENTITY_COUNT_ATTRIBUTE: "0", + } + + +def test_failures_do_not_expose_raw_field_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(processor, "_create_pipeline", lambda config: object()) + raw_value = "Patient Jane Roe called 555-0100" + + def fail_batch(texts: list[str], **kwargs: Any) -> None: + raise RuntimeError(f"model failed on {texts[0]}") + + with pytest.raises(processor.DataflowToolProcessorError) as exc_info: + processor.process_json_lines( + [ + json.dumps( + { + "record": {"note": raw_value}, + "attributes": {"route": "clinical"}, + } + ) + ], + io.StringIO(), + fields="note", + process_batch_fn=fail_batch, + ) + + message = str(exc_info.value) + assert "input line 1" in message + assert "Jane Roe" not in message + assert "555-0100" not in message + + +def test_cardinality_mismatch_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(processor, "_create_pipeline", lambda config: object()) + + with pytest.raises( + processor.DataflowToolProcessorError, + match="result count", + ): + processor.process_flow_file( + {"note": "Jane Roe"}, + fields="note", + process_batch_fn=lambda texts, **kwargs: SimpleNamespace(items=[]), + )