phi.py

"""
Patient identifiers — YashTech Ltd. AI Labs.

Three rules, enforced here rather than trusted to the UI:

  1. **Identifiers never reach the model.** They label the report; they are not
     clinical context. MedGemma runs on shared multi-tenant hardware, so a name
     in a prompt is a name on someone else's GPU. `scrub()` strips them from
     free text on the way in.
  2. **Pseudonymised by default.** The report header shows a stable token such
     as `PT-K4X9M2QF`, derived by HMAC from the identifier and a secret salt.
     Same patient, same token, across reports and sessions — which is what makes
     a pseudonym useful — without ever writing the identifier down.
  3. **No mapping table.** Re-identification is by recomputing the token from an
     identifier you already hold, not by looking it up. Nothing to leak, nothing
     to subject-access-request, nothing to delete.

Showing a real identifier on a report is a deliberate act by someone holding the
`show_identifiers` permission, and the report says so in its header.

`PSEUDONYM_SALT` must be set as a secret. Without it a random salt is generated
at boot: tokens stay unguessable but stop being stable across restarts, and the
admin console says so rather than quietly producing a different token for the
same patient tomorrow.
"""

from __future__ import annotations

import base64
import dataclasses
import hashlib
import hmac
import os
import re
import secrets

PSEUDONYMISE = "pseudonymise"
OMIT = "omit"
FULL = "full"

MODES = (
    (PSEUDONYMISE, "Pseudonymised token (recommended)"),
    (OMIT, "No identifiers at all"),
    (FULL, "Show as entered (requires permission)"),
)
MODE_CHOICES = [(label, value) for value, label in MODES]
DEFAULT_MODE = PSEUDONYMISE

_EPHEMERAL_SALT = secrets.token_hex(32)


def salt() -> tuple[bytes, bool]:
    """(salt, is_stable). An unset salt is usable but not stable across restarts."""
    configured = os.getenv("PSEUDONYM_SALT", "").strip()
    if configured:
        return configured.encode("utf-8"), True
    return _EPHEMERAL_SALT.encode("utf-8"), False


def salt_status() -> str:
    _key, stable = salt()
    return (
        "configured — tokens are stable across restarts"
        if stable
        else "NOT SET — tokens change on every restart; set PSEUDONYM_SALT as a secret"
    )


def _normalise(value: str) -> str:
    """So `NHS 123 456 7890`, `nhs1234567890` and `123-456-7890` agree."""
    return re.sub(r"[^a-z0-9]", "", (value or "").strip().lower())


def pseudonym(identifier: str, prefix: str = "PT") -> str:
    key, _stable = salt()
    norm = _normalise(identifier)
    if not norm:
        return ""
    digest = hmac.new(key, norm.encode("utf-8"), hashlib.sha256).digest()
    token = base64.b32encode(digest).decode("ascii").rstrip("=")[:8]
    return f"{prefix}-{token}"


@dataclasses.dataclass(frozen=True)
class Patient:
    patient_id: str = ""
    name: str = ""
    mode: str = DEFAULT_MODE

    @property
    def supplied(self) -> bool:
        return bool(self.patient_id.strip() or self.name.strip())

    @property
    def token(self) -> str:
        """Stable token from the ID if there is one, otherwise from the name."""
        return pseudonym(self.patient_id) or pseudonym(self.name, prefix="PN")

    def terms(self) -> list[str]:
        """Strings that must never appear in a prompt."""
        out = []
        for value in (self.patient_id, self.name):
            value = (value or "").strip()
            if len(value) >= 3:
                out.append(value)
                out.extend(p for p in re.split(r"[\s,]+", value) if len(p) >= 3)
        return sorted(set(out), key=len, reverse=True)


def resolve(patient: Patient, may_show_identifiers: bool) -> tuple[Patient, str | None]:
    """Downgrade `full` to pseudonymised when the user lacks the permission."""
    if patient.mode == FULL and not may_show_identifiers:
        return (
            dataclasses.replace(patient, mode=PSEUDONYMISE),
            "Your role cannot show identifiers on a report — exported pseudonymised instead.",
        )
    if patient.mode not in {m for m, _ in MODES}:
        return dataclasses.replace(patient, mode=DEFAULT_MODE), None
    return patient, None


def header_lines(patient: Patient) -> list[str]:
    """Markdown bullets for the report header. Never raises on bad input."""
    if not patient.supplied or patient.mode == OMIT:
        return ["- **Patient:** not recorded (identifiers omitted)"]

    if patient.mode == FULL:
        lines = ["- **Patient identifiers: SHOWN IN FULL — handle as clinical record**"]
        if patient.patient_id.strip():
            lines.append(f"- **Patient ID:** {patient.patient_id.strip()}")
        if patient.name.strip():
            lines.append(f"- **Patient name:** {patient.name.strip()}")
        lines.append(f"- **Pseudonym:** {patient.token}")
        return lines

    lines = [f"- **Patient:** {patient.token} (pseudonymised)"]
    _key, stable = salt()
    if not stable:
        lines.append("- **Note:** pseudonym salt not configured — this token is not reproducible")
    return lines


def scrub(text: str, patient: Patient) -> str:
    """Remove identifiers from text heading for the model.

    Deliberately blunt. This is a backstop for a name typed into the clinical
    indication box, not a de-identification engine: it cannot catch an
    identifier nobody entered into the patient fields.
    """
    if not text or not patient.supplied:
        return text
    cleaned = text
    for term in patient.terms():
        cleaned = re.sub(re.escape(term), "[redacted]", cleaned, flags=re.IGNORECASE)
    return cleaned


def prompt_context(patient: Patient) -> str:
    """What the model is allowed to know: nothing identifying."""
    return ""