PII Detection and Redaction: Checksums, LLM Entity Extraction, Reversible Placeholders and a Leak Test

"""Fernleaf Energy's PII redactor. Chats go to an external summarisation service, and no personal data may go with
them. You write the functions marked TODO, one step at a time."""
import json
import re

from harness import ask_model

TYPES = ["NAME", "ADDRESS", "EMAIL", "PHONE", "CARD", "IBAN"]

EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}")
PHONE_RE = re.compile(r"(?<![\w+])(?<!\d )(?<!\d-)(?:\+353[ -]?(?:\(0\)[ -]?)?|\+44[ -]?|0)\d{1,4}(?:[ -]?\d){5,8}(?![ -]?\d)")
CARD_RE = re.compile(r"(?<!\d)\d(?:[ -]?\d){12,18}(?!\d)")
IBAN_RE = re.compile(r"\b[A-Za-z]{2}\d{2}(?: ?[A-Za-z0-9]{4}){2,7}(?: ?[A-Za-z0-9]{1,3})?\b")


def span(text, start, end, typ):
    return {"start": start, "end": end, "type": typ, "text": text[start:end]}


# ---------- Step 1: patterns, checked ----------

def luhn_ok(number):
    """True if the digits of number (spaces and dashes ignored) pass the Luhn check that every card number passes."""
    # TODO (Step 1): digits only; from the right, double every second digit (minus 9 if over 9);
    # the sum must be a multiple of 10.
    raise NotImplementedError("Step 1: write luhn_ok()")


def iban_ok(value):
    """True if value (spaces ignored, any case) is a well-formed IBAN whose check digits are right (mod 97 == 1)."""
    # TODO (Step 1): strip spaces, upper-case, check the shape, move the first 4 characters to the end,
    # letters to numbers with int(c, 36), then int(...) % 97 == 1.
    raise NotImplementedError("Step 1: write iban_ok()")


def find_pattern_pii(text):
    """EMAIL, PHONE, CARD (13-19 digits that pass Luhn) and IBAN (passes iban_ok) spans, sorted by start."""
    # TODO (Step 1): finditer with EMAIL_RE, PHONE_RE, CARD_RE and IBAN_RE; keep cards of 13-19
    # digits that pass luhn_ok() and IBANs that pass iban_ok(); span() each; sort by start.
    raise NotImplementedError("Step 1: write find_pattern_pii()")


# ---------- Step 2: names and addresses ----------

def ner_prompt(text):
    """The prompt that asks the in-house model for every person's name and every postal address in text."""
    raise NotImplementedError("ner_prompt() arrives in Step 2")


def parse_entities(reply):
    """The [{"type", "text"}] items of NAME or ADDRESS type in the model's reply; [] if it holds no JSON list."""
    raise NotImplementedError("parse_entities() arrives in Step 2")


def locate(text, entities):
    """A span for every place each entity's text occurs in text (exact match first, else ignoring case);
    entities that do not occur are dropped."""
    raise NotImplementedError("locate() arrives in Step 2")


def find_llm_pii(text):
    """NAME and ADDRESS spans found by the in-house model."""
    raise NotImplementedError("find_llm_pii() arrives in Step 2")


# ---------- Step 3: one detector, measured ----------

def merge(spans):
    """Non-overlapping spans sorted by start: overlapping spans become one span covering both, typed like the
    longer of them."""
    raise NotImplementedError("merge() arrives in Step 3")


def detect(text):
    """Every PII span in text: patterns and model, merged."""
    raise NotImplementedError("detect() arrives in Step 3")


def evaluate(rows, detector=None):
    """{type: {"gold", "found", "recall", "predicted", "precision"}} over rows with gold "pii" spans.
    A gold span is found when a predicted span of its type covers all of it; a predicted span is correct when it
    overlaps a gold span of its type. recall = found / gold, precision = correct / predicted (1.0 when 0/0)."""
    raise NotImplementedError("evaluate() arrives in Step 3")


# ---------- Step 4: redact, and put it back ----------

def normal_form(typ, value):
    """The form in which two mentions of the same thing compare equal."""
    if typ in ("PHONE", "CARD"):
        return re.sub(r"\D", "", value)[-9:]
    if typ == "IBAN":
        return value.replace(" ", "").upper()
    return " ".join(value.lower().split())


def redact(text, spans):
    """(redacted text, mapping). Each span becomes a placeholder "[TYPE_n]", numbered per type in order of first
    appearance; the same value (normal_form) always gets the same placeholder. mapping: placeholder -> the text it
    replaced the first time."""
    raise NotImplementedError("redact() arrives in Step 4")


def restore(text, mapping):
    """text with every placeholder in mapping replaced by its original."""
    raise NotImplementedError("restore() arrives in Step 4")


# ---------- Step 5: the leak test ----------

def leaked(sent, typ, value):
    """True if value, a piece of PII of type typ, appears in the text sent in any form.
    CARD, IBAN: its letters and digits, in sent's letters and digits. PHONE: its last 7 digits, in sent's digits.
    EMAIL: the part before "@" (or " at "), with " dot " read as "." in both, its letters and digits in sent's.
    NAME: any of its words of 3+ letters as a whole word of sent. ADDRESS: the part before the first comma, as a
    phrase of sent. Case never matters."""
    raise NotImplementedError("leaked() arrives in Step 5")


def leaks(sent, pii):
    """The items of pii ([{"type", "value"}]) that leaked into the text sent."""
    raise NotImplementedError("leaks() arrives in Step 5")


SPELLED_EMAIL_RE = re.compile(r"\b[a-z0-9]+(?: dot [a-z0-9]+)* at [a-z0-9-]+(?: dot [a-z0-9-]+)* dot [a-z]{2,}\b", re.I)


def find_spelled_emails(text):
    """EMAIL spans for addresses written out in words, such as "jo dot bloggs at gmail dot com"."""
    return []  # Step 5 catches addresses written out in words