Prompting for clinical tasks

SYSTEM = """You extract information from clinic notes for a diabetes registry.
Return only JSON with exactly these keys:
  "diabetes_type": one of "type 1", "type 2", "gestational", "other", "not stated"
  "latest_hba1c_percent": a number, or null if not stated
  "on_insulin": true, false, or null if not stated
Use only information written in the note. Do not guess."""

EXAMPLE_NOTE = "54F, T2DM for 8 yrs on metformin. HbA1c today 8.1%. No insulin."
EXAMPLE_ANSWER = '{"diabetes_type": "type 2", "latest_hba1c_percent": 8.1, "on_insulin": false}'

def build_messages(note):
    return [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": f"Note:\n{EXAMPLE_NOTE}"},     # one worked example
        {"role": "assistant", "content": EXAMPLE_ANSWER},
        {"role": "user", "content": f"Note:\n{note}"},
    ]

messages = build_messages("Synthetic note: 31M, type 1 since age 12, basal-bolus insulin. A1c 7.4% (Mar).")
for m in messages:
    print(f"[{m['role']}] " + m["content"].replace("\n", " ")[:70])

# [system] You extract information from clinic notes for a diabetes registry. Ret
# [user] Note: 54F, T2DM for 8 yrs on metformin. HbA1c today 8.1%. No insulin.
# [assistant] {"diabetes_type": "type 2", "latest_hba1c_percent": 8.1, "on_insulin":
# [user] Note: Synthetic note: 31M, type 1 since age 12, basal-bolus insulin. A

import json

ALLOWED_TYPES = {"type 1", "type 2", "gestational", "other", "not stated"}

def validate(raw):
    """Return (record, problems). Never trust model output without checking it."""
    try:
        record = json.loads(raw)
    except json.JSONDecodeError:
        return None, ["not valid JSON"]
    problems = []
    if set(record) != {"diabetes_type", "latest_hba1c_percent", "on_insulin"}:
        problems.append(f"unexpected keys: {sorted(record)}")
    if record.get("diabetes_type") not in ALLOWED_TYPES:
        problems.append(f"diabetes_type not allowed: {record.get('diabetes_type')!r}")
    a1c = record.get("latest_hba1c_percent")
    if a1c is not None and not (isinstance(a1c, (int, float)) and 3 <= a1c <= 20):
        problems.append(f"implausible HbA1c: {a1c!r}")
    if record.get("on_insulin") not in (True, False, None):
        problems.append("on_insulin must be true, false or null")
    return record, problems

outputs = [
    '{"diabetes_type": "type 1", "latest_hba1c_percent": 7.4, "on_insulin": true}',
    'Sure! Here is the JSON: {"diabetes_type": "type 2"}',
    '{"diabetes_type": "Type II", "latest_hba1c_percent": 74, "on_insulin": "yes"}',
]
for raw in outputs:
    record, problems = validate(raw)
    print("OK" if not problems else "REJECT", problems)

# OK []
# REJECT ['not valid JSON']
# REJECT ["diabetes_type not allowed: 'Type II'", 'implausible HbA1c: 74', 'on_insulin must be true, false or null']

gold = [
    {"diabetes_type": "type 1", "latest_hba1c_percent": 7.4, "on_insulin": True},
    {"diabetes_type": "type 2", "latest_hba1c_percent": None, "on_insulin": False},
    {"diabetes_type": "gestational", "latest_hba1c_percent": 5.6, "on_insulin": None},
    {"diabetes_type": "type 2", "latest_hba1c_percent": 9.2, "on_insulin": True},
]
prompt_v1 = [   # validated outputs from the first prompt version
    {"diabetes_type": "type 1", "latest_hba1c_percent": 7.4, "on_insulin": True},
    {"diabetes_type": "type 2", "latest_hba1c_percent": 7.0, "on_insulin": False},  # invented value
    {"diabetes_type": "type 2", "latest_hba1c_percent": 5.6, "on_insulin": False},  # wrong type, guessed
    {"diabetes_type": "type 2", "latest_hba1c_percent": 9.2, "on_insulin": True},
]
prompt_v2 = [   # after adding "Do not guess" and the null rules
    {"diabetes_type": "type 1", "latest_hba1c_percent": 7.4, "on_insulin": True},
    {"diabetes_type": "type 2", "latest_hba1c_percent": None, "on_insulin": False},
    {"diabetes_type": "gestational", "latest_hba1c_percent": 5.6, "on_insulin": None},
    {"diabetes_type": "type 2", "latest_hba1c_percent": 9.2, "on_insulin": False},  # missed insulin
]

def field_accuracy(predictions):
    fields = gold[0].keys()
    return {f: sum(p[f] == g[f] for p, g in zip(predictions, gold)) / len(gold) for f in fields}

for name, preds in [("prompt v1", prompt_v1), ("prompt v2", prompt_v2)]:
    print(name, field_accuracy(preds))

# prompt v1 {'diabetes_type': 0.75, 'latest_hba1c_percent': 0.75, 'on_insulin': 0.75}
# prompt v2 {'diabetes_type': 1.0, 'latest_hba1c_percent': 1.0, 'on_insulin': 0.75}