refactor(pii): type NER payload and runtime boundaries

pyproject.toml

@@ -352,13 +352,13 @@ safe-synthesizer = "nemo_safe_synthesizer.cli.cli:cli"
 # that appear unused when the package is installed. Suppress the noise.
 unused-ignore-comment = "ignore"
 
+
 [tool.ty.environment]
 extra-paths = ["typings"]
 
 [tool.ty.src]
 # Below is a list of excluded directories from ty typechecks.
 exclude = [
-    ".uv_cache", # Cache Dir in CI
-    "./docs/**/*.ipynb",
-    "./src/nemo_safe_synthesizer/pii_replacer/",
+  "./uv-cache",
+  "./docs/**/*.ipynb",
 ]

src/nemo_safe_synthesizer/data_processing/records/fragment.py

@@ -14,9 +14,9 @@
 import time
 import uuid
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any
+from typing import Any, NotRequired, TypeAlias, TypedDict
 
 from ...pii_replacer.ner.entity import Score
 from ...pii_replacer.ner.predictor import NERPrediction
@@ -32,6 +32,62 @@ class MetadataError(Exception):
 E2F = "fields_by_entity"
 
 
+class NERRawPredictionPayload(TypedDict):
+    """Raw ``NERPrediction.as_dict`` payload consumed by metadata helpers."""
+
+    text: str
+    start: int
+    end: int
+    label: str
+    source: str
+    score: float | None
+    field: NotRequired[str | None]
+    value_path: NotRequired[tuple[str | int, ...] | list[str | int] | None]
+    substring_match: NotRequired[bool | None]
+
+
+class NERFieldLabelPayload(TypedDict):
+    """Per-field label entry in the NER model metadata payload."""
+
+    start: int
+    end: int
+    label: str
+    score: float | None
+    source: str
+    text: str
+
+
+NERMetadataFieldsPayload: TypeAlias = dict[str, dict[str, dict[str, list[NERFieldLabelPayload]]]]
+
+
+class NEREntityMapPayload(TypedDict):
+    """Score-bucketed entity summary in the NER model metadata payload."""
+
+    score_high: list[str]
+    score_med: list[str]
+    score_low: list[str]
+    fields_by_entity: dict[str, list[str]]
+
+
+class NERMetadataPayload(TypedDict):
+    """API-facing NER metadata payload."""
+
+    record_id: str
+    fields: NERMetadataFieldsPayload
+    entities: NEREntityMapPayload
+    received_at: str
+
+
+NERRecordPayload: TypeAlias = dict[str, Any]
+
+
+class NERApiResponseRow(TypedDict):
+    """One API response row for dict-based NER model predictions."""
+
+    data: NERRecordPayload
+    model_metadata: NERMetadataPayload
+
+
 @dataclass
 class Metadata:
     """Merged record metadata aggregated from one or more ``MetadataFragment`` objects.
@@ -43,18 +99,23 @@ class Metadata:
 
     record_id: str
 
-    fields: dict
+    fields: NERMetadataFieldsPayload
     """Nested dict of per-field, per-fragment metadata."""
 
-    entities: dict
+    entities: NEREntityMapPayload
     """Entity map produced by ``predictions_to_dict``."""
 
     received_at: str
     """ISO-8601 timestamp of the earliest fragment."""
 
-    def as_dict(self):
+    def as_dict(self) -> NERMetadataPayload:
         """Serialize to a plain dictionary."""
-        return self.__dict__
+        return {
+            "record_id": self.record_id,
+            "fields": self.fields,
+            "entities": self.entities,
+            "received_at": self.received_at,
+        }
 
 
 @dataclass
@@ -75,6 +136,7 @@ class MetadataFragment:
     fragment_ts: str
     fragment_epoch: float
     fragment_name: str
+    fields: dict[str, dict[str, list[NERFieldLabelPayload]]] = field(init=False)
 
     def __post_init__(self):
         self.fields = defaultdict(lambda: defaultdict(list))
@@ -84,7 +146,12 @@ def fragment_datetime(self) -> datetime:
         """Fragment creation time as a ``datetime`` object."""
         return datetime.fromtimestamp(self.fragment_epoch)
 
-    def add_field_data(self, field_name: str, metadata_type: str, field_data: dict | list):
+    def add_field_data(
+        self,
+        field_name: str,
+        metadata_type: str,
+        field_data: NERFieldLabelPayload | list[NERFieldLabelPayload],
+    ) -> None:
         """Append metadata entries for a field.
 
         Args:
@@ -97,12 +164,10 @@ def add_field_data(self, field_name: str, metadata_type: str, field_data: dict |
         """
         if isinstance(field_data, list):
             self.fields[field_name][metadata_type].extend(field_data)
-        elif isinstance(field_data, dict):
-            self.fields[field_name][metadata_type].append(field_data)
         else:
-            raise TypeError("field_data must be a dict or list, got ", type(field_data))
+            self.fields[field_name][metadata_type].append(field_data)
 
-    def as_dict(self):
+    def as_dict(self) -> dict[str, Any]:
         """Serialize to a plain dictionary."""
         return self.__dict__
 
@@ -126,15 +191,29 @@ def merge_fragments(*fragments, ts: str | None = None) -> Metadata:
     else:
         record_id = fragments[0].record_id
 
-    # todo(dn): there might be a better way to build up this object
-    merged_fragment = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
+    merged_fragment: dict[str, dict[str, dict[str, list[NERFieldLabelPayload]]]] = defaultdict(
+        lambda: defaultdict(lambda: defaultdict(list))
+    )
     ts = ts or min([f.fragment_datetime for f in fragments]).isoformat() + "Z"
     fragment: MetadataFragment
     for fragment in fragments:
         for field_name, field_data in fragment.fields.items():
             for meta_type, meta_data in field_data.items():
                 merged_fragment[field_name][fragment.fragment_name][meta_type].extend(meta_data)
-    return Metadata(record_id=record_id, fields=merged_fragment, received_at=ts, entities={})
+    fields: NERMetadataFieldsPayload = {
+        field_name: {
+            fragment_name: {meta_type: list(meta_data) for meta_type, meta_data in fragment_data.items()}
+            for fragment_name, fragment_data in field_data.items()
+        }
+        for field_name, field_data in merged_fragment.items()
+    }
+    empty_entities: NEREntityMapPayload = {
+        SCORE_HIGH: [],
+        SCORE_MED: [],
+        SCORE_LOW: [],
+        E2F: {},
+    }
+    return Metadata(record_id=record_id, fields=fields, received_at=ts, entities=empty_entities)
 
 
 def fragment_for_record(record_id: str, fragment_name: str) -> MetadataFragment:
@@ -154,7 +233,7 @@ def predictions_to_dict(
     *,
     high_score: float = Score.HIGH,
     med_score: float = Score.MED,
-) -> tuple[dict, dict]:
+) -> tuple[dict[str, list[NERFieldLabelPayload]], NEREntityMapPayload]:
     """Aggregate NER predictions into per-field results and an entity map.
 
     Groups predictions by field and builds a score-bucketed entity map::
@@ -174,13 +253,11 @@ def predictions_to_dict(
     Returns:
         A tuple of (predictions_by_field, entity_map).
     """
-    entity_map: dict[str, Any] = {
-        SCORE_HIGH: set(),
-        SCORE_MED: set(),
-        SCORE_LOW: set(),
-        E2F: defaultdict(set),
-    }
-    predictions_by_key = defaultdict(list)
+    high_entities: set[str] = set()
+    medium_entities: set[str] = set()
+    low_entities: set[str] = set()
+    fields_by_entity: dict[str, set[str]] = defaultdict(set)
+    predictions_by_key: dict[str, list[NERFieldLabelPayload]] = defaultdict(list)
     for prediction in predictions:
         if prediction.field is None:
             continue
@@ -198,28 +275,30 @@ def predictions_to_dict(
         # no score is emitted. Predictions here could be
         # hit or miss so we throw it into medium
         if prediction.score is None:
-            entity_map[SCORE_MED].add(prediction.label)
+            medium_entities.add(prediction.label)
         elif prediction.score >= high_score:
-            entity_map[SCORE_HIGH].add(prediction.label)
+            high_entities.add(prediction.label)
         elif prediction.score >= med_score:
-            entity_map[SCORE_MED].add(prediction.label)
+            medium_entities.add(prediction.label)
         else:
-            entity_map[SCORE_LOW].add(prediction.label)
-        entity_map[E2F][prediction.label].add(prediction.field)
+            low_entities.add(prediction.label)
+        fields_by_entity[prediction.label].add(prediction.field)
     for _, preds in predictions_by_key.items():
         preds.sort(key=lambda p: p["start"])
-    for level in (SCORE_HIGH, SCORE_MED, SCORE_LOW):
-        entity_map[level] = list(entity_map[level])
-    for entity, _set in entity_map[E2F].items():
-        entity_map[E2F][entity] = list(_set)
+    entity_map: NEREntityMapPayload = {
+        SCORE_HIGH: list(high_entities),
+        SCORE_MED: list(medium_entities),
+        SCORE_LOW: list(low_entities),
+        E2F: {entity: list(fields) for entity, fields in fields_by_entity.items()},
+    }
     return predictions_by_key, entity_map
 
 
 def fragment_from_ner_predictions(
     fragment_name: str,
     predictions: list[NERPrediction],
     record_id: str,
-) -> tuple[MetadataFragment, dict]:
+) -> tuple[MetadataFragment, NEREntityMapPayload]:
     """Build a ``MetadataFragment`` and entity map from NER predictions.
 
     Args:
@@ -244,9 +323,9 @@ def fragment_from_ner_predictions(
     return fragment, ent_map
 
 
-def build_ner_metadata(preds: list[dict]) -> Metadata:
-    """Construct a ``Metadata`` object from raw prediction dicts."""
-    ner_preds = [NERPrediction.from_dict(p) for p in preds]
+def build_ner_metadata(preds: list[NERRawPredictionPayload]) -> NERMetadataPayload:
+    """Construct an API-facing metadata payload from raw prediction dicts."""
+    ner_preds = [NERPrediction.from_dict(dict(p)) for p in preds]
     fragment, ent_map = fragment_from_ner_predictions(
         "ner",
         ner_preds,
@@ -257,7 +336,11 @@ def build_ner_metadata(preds: list[dict]) -> Metadata:
     return meta.as_dict()
 
 
-def create_ner_api_response(records: list[dict], predictions: list[list[dict]], pure_dict: bool = False) -> list[dict]:
+def create_ner_api_response(
+    records: list[NERRecordPayload],
+    predictions: list[list[NERRawPredictionPayload]],
+    pure_dict: bool = False,
+) -> list[NERApiResponseRow]:
     """Build an API-compatible list of ``{data, model_metadata}`` dicts.
 
     Args:
@@ -268,10 +351,20 @@ def create_ner_api_response(records: list[dict], predictions: list[list[dict]],
     Returns:
         List of dicts, each containing ``data`` and ``model_metadata`` keys.
     """
-    out = [
+    out: list[NERApiResponseRow] = [
         {"data": record, "model_metadata": build_ner_metadata(prediction)}
         for record, prediction in zip(records, predictions)
     ]
     if pure_dict:
-        return json.loads(json.dumps(out))
+        data_rows = json.loads(json.dumps([row["data"] for row in out]))
+        if not isinstance(data_rows, list):
+            raise TypeError("expected JSON round-trip to preserve response row list")
+        rows: list[NERApiResponseRow] = []
+        for data, row in zip(data_rows, out):
+            if not isinstance(data, dict):
+                raise TypeError("expected JSON round-trip to preserve response data dictionaries")
+            record: NERRecordPayload = {str(key): value for key, value in data.items()}
+            response_row: NERApiResponseRow = {"data": record, "model_metadata": row["model_metadata"]}
+            rows.append(response_row)
+        return rows
     return out