| |
| """Axis 3 (drug safety, multiple choice): LLM discovery, then question generation. |
| |
| Stage 1: the LLM reads the drug and constraint timelines of the evidence chain and returns standing drug |
| constraints (class, unsafe member, safe equivalents). Stage 2: the axis-3 rules (pipeline.prompts) turn |
| each constraint into a multiple-choice question with exactly one correct option. The gold is |
| options[].is_correct. The options are put in an order seeded by the question id and labelled A-D in that order; |
| the question text does not list them. synthesize() authors synthetic questions. |
| |
| python -m pipeline.questions.axis3 discovery -> _discover_axis3.json |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import random |
| from pathlib import Path |
|
|
| from pipeline import prompts |
| from pipeline.common import (BASE, Progress, has_text, lists_options, model, patient_context, patient_id, |
| patient_line) |
| from pipeline.llm import ChatClient, subst |
|
|
| PATIENT = patient_id() |
| ENGINE = model("generation") |
|
|
| AXIS = 3 |
| DRUG_ASPECTS = {"med_rationale_referral", "medication", "constraint_or_safety", "care_plan"} |
|
|
| DISCOVER_SYS = ( |
| "You read one patient's longitudinal record of medications and constraints and find standing\n" |
| "drug-safety constraints. A constraint is a documented reason, specific to this patient, why some\n" |
| "drugs a clinician might choose are inappropriate for this patient while a clinically reasonable\n" |
| "alternative remains fine: an adverse reaction, an allergy, an intolerance, or an interacting\n" |
| "medication that the patient takes. Return only JSON.") |
|
|
|
|
| def discover_constraints(cli, timelines): |
| """Read the drug-bearing timelines and return the patient's standing drug constraints.""" |
| rows = [] |
| for t in timelines: |
| segs = t.get("segments") or [] |
| rows.append(f"- entity={t.get('entity')!r} aspect={t.get('aspect')} dept={t.get('department')} " |
| f"status={t.get('current_status')!r} " |
| f"values={[str(s.get('value'))[:60] for s in segs[:4]]}") |
| record = "\n".join(rows) |
| usr = ( |
| "Medication and constraint record, from the evidence chain:\n" + record + "\n\n" |
| "Find every standing drug-safety constraint documented for this patient. A constraint qualifies\n" |
| "only if all of these hold:\n" |
| "- It is specific to this patient: an allergy, intolerance or adverse reaction on file, or a drug\n" |
| " the patient takes that interacts with a drug class. General guideline caution does not count.\n" |
| "- It makes at least three drugs a clinician might reach for inappropriate for this patient, while\n" |
| " at least one clinically reasonable alternative stays appropriate, so that a four-option question\n" |
| " with exactly one appropriate answer is possible.\n" |
| "- It is standing, that is, still true at later visits, and not a one-off event.\n" |
| "Exclude negated screening entries such as 'no known allergy', 'denied' or '(-)'.\n\n" |
| "For each constraint give:\n" |
| "- the drug class;\n" |
| "- at least three drugs that are unsafe for this patient;\n" |
| "- one to four safe alternatives;\n" |
| "- a realistic scenario in which the class would be prescribed to the patient;\n" |
| "- the first-line trap drugs that a clinician, or a model without memory, would reach for first.\n" |
| " These must be among the unsafe drugs, or the question is pointless.\n\n" |
| "Return JSON in this form:\n" |
| '{"constraints": [{\n' |
| ' "constraint_id": "<slug>",\n' |
| ' "implicated_class": "<class>",\n' |
| ' "constraint_text": "<the documented fact, close to its wording>",\n' |
| ' "source_entity": "<the entity exactly as listed above>",\n' |
| ' "unsafe_examples": ["<drug>", ...],\n' |
| ' "safe_examples": ["<drug>", ...],\n' |
| ' "scenario_hint": "<why the class would be prescribed>",\n' |
| ' "required_patient_info": ["<the exact memory fact needed>"],\n' |
| ' "first_line_trap": "<the drug a model without memory picks by default>"\n' |
| "}]}") |
| r = cli.complete_json(DISCOVER_SYS, usr) |
| return r.get("constraints") or [] |
|
|
|
|
| def discover() -> None: |
| """Discover the axis-3 candidates of this patient -> _discover_axis3.json.""" |
| cli = ChatClient(model=ENGINE, reasoning_effort="medium", max_completion_tokens=6000) |
|
|
| tl = [json.loads(l) for l in open(f"{BASE}/evidence_chain.jsonl") if l.strip()] |
| chain = {json.loads(l)["visit_number"]: json.loads(l) |
| for l in open(f"{BASE}/event_chain.jsonl") if l.strip()} |
| plast = max(r["date"] for r in chain.values()) |
|
|
| drug_tl = [t for t in tl if t.get("aspect") in DRUG_ASPECTS] |
|
|
| with Progress("axis3 discover (LLM)", total=2, out=f"{BASE}/progress_axis3_discover.json") as pg: |
| try: |
| seeds = discover_constraints(cli, drug_tl) |
| except Exception as e: |
| print(f" discover FAILED: {str(e)[:160]}") |
| seeds = [] |
| pg.tick(ok=bool(seeds)) |
|
|
| AX3 = subst(prompts.AXIS3_RULES) |
| COMMON = subst(prompts.AXIS3_COMMON) |
| system = subst(f"{COMMON}\n\n{AX3}") |
| candidates = [] |
| for i, s in enumerate(seeds): |
| if len(s.get("unsafe_examples") or []) < 3: |
| print(f" skip {s.get('constraint_id')}: fewer than three unsafe drugs for the traps") |
| continue |
| gt_visit, gt_date = _anchor(s, tl, chain) |
| if gt_visit is None: |
| print(f" skip {s.get('constraint_id')}: its source entity is not in the evidence chain") |
| continue |
| seed_json = json.dumps({ |
| "constraint_id": s.get("constraint_id"), "implicated_class": s.get("implicated_class"), |
| "scenario_hint": s.get("scenario_hint"), "constraint_text": s.get("constraint_text"), |
| "unsafe_examples": s.get("unsafe_examples"), "safe_examples": s.get("safe_examples"), |
| "required_patient_info": s.get("required_patient_info"), |
| "first_line_trap": s.get("first_line_trap"), |
| "gt_date": gt_date, "grounding": "record"}, ensure_ascii=False) |
| usr = ( |
| "Write the axis-3 multiple-choice drug-safety question for this seed, following every rule\n" |
| "above. Return only a JSON object of the form {\"questions\": [ ... ]} with one question, or\n" |
| "with an empty list if the rules cannot be met.\n\nSeed:\n" + seed_json) |
| try: |
| out = cli.complete_json(system, usr) |
| qs = out.get("questions") or ([out] if out.get("options") else []) |
| except Exception as e: |
| print(f" gen[{i}] FAILED: {str(e)[:120]}") |
| continue |
| for q in qs: |
| if not has_text(q): |
| print(f" reject {s.get('constraint_id')}: the reply has no question text") |
| continue |
| opts = q.get("options") or [] |
| if not (isinstance(opts, list) and all(isinstance(o, dict) for o in opts)): |
| print(f" reject {s.get('constraint_id')}: the options are not a list of objects") |
| continue |
| for o in opts: |
| if not o.get("content") and o.get("text"): |
| o["content"] = o["text"] |
| if lists_options(q.get("question"), opts): |
| print(f" reject {s.get('constraint_id')}: the question text lists the options") |
| continue |
| if sum(1 for o in opts if o.get("is_correct")) != 1: |
| print(f" reject {s.get('constraint_id')}: not exactly 1 correct option") |
| continue |
| if not all(_names(o, s.get("safe_examples") if o.get("is_correct") else s.get("unsafe_examples")) |
| for o in opts): |
| print(f" reject {s.get('constraint_id')}: an option is not one of the seed's safe or " |
| f"unsafe drugs") |
| continue |
| qid = f"{PATIENT}/axis{AXIS}/{len(candidates):03d}" |
| opts = shuffled(opts, qid) |
| safe = next(o for o in opts if o.get("is_correct")) |
| candidates.append({ |
| "qid": qid, "patient_id": PATIENT, |
| "axis": "3_drug_safety", "question": q.get("question"), "options": opts, |
| "answer": safe.get("content"), "answer_type": "multiple_choice", "grounding": "record", |
| "constraint_id": s.get("constraint_id"), |
| "evidence_visits": [gt_visit] if gt_visit else [], "gt_visit": gt_visit, "gt_date": gt_date, |
| "question_appearance_window": {"from": gt_date, "to": plast}, |
| "derivation": { |
| "evidence": [{"visit": gt_visit, "date": gt_date, "field": "evidence_chain", |
| "atom": s.get("constraint_text")}], |
| "steps": [f"Standing constraint: {s.get('constraint_text')!r}.", |
| f"Of the 4 options only {safe.get('content')!r} is appropriate given it.", |
| safe.get("explanation", "")], |
| "gold": safe.get("content"), |
| "rivals": [{"value": o.get("content"), "why_wrong": o.get("explanation")} |
| for o in opts if not o.get("is_correct")], |
| "operands": {"implicated_class": s.get("implicated_class"), |
| "constraint_text": s.get("constraint_text"), |
| "first_line_trap": s.get("first_line_trap")}, |
| "derived": False}, |
| "provenance": f"LLM-discovered constraint ({ENGINE}); appropriate/inappropriate split " |
| f"turns on recalling {s.get('constraint_id')}", |
| "trap_design": q.get("trap_design"), |
| "axis_specific": {"constraint_id": s.get("constraint_id"), |
| "required_patient_info": s.get("required_patient_info")}, |
| }) |
| pg.tick(ok=True) |
|
|
| out = {"axis": AXIS, "patient_id": PATIENT, "n_record": len(candidates), "candidates": candidates} |
| Path(f"{BASE}/_discover_axis{AXIS}.json").write_text(json.dumps(out, ensure_ascii=False, indent=1)) |
| print(f"\npatient {PATIENT} engine={ENGINE}") |
| print(f" drug/constraint timelines scanned : {len(drug_tl)}") |
| print(f" constraints discovered (LLM) : {len(seeds)}") |
| print(f" MC questions built : {len(candidates)}") |
| for c in candidates[:3]: |
| print(f" [{c['constraint_id']}] {str(c['question'])[:70]} -> {c['answer']}") |
| print(f" CANDIDATES: {len(candidates)} -> {BASE}/_discover_axis{AXIS}.json") |
|
|
|
|
| def shuffled(options, qid): |
| """The options in an order seeded by the question id, relabelled A, B, C, ... in that order.""" |
| out = [dict(o) for o in options] |
| random.Random(qid).shuffle(out) |
| for k, o in enumerate(out): |
| o["label"] = chr(ord("A") + k) |
| return out |
|
|
|
|
| def _names(option, drugs): |
| """Whether an option names one of the drugs (case-insensitive; either text may contain the other).""" |
| text = str(option.get("content") or "").strip().lower() |
| return bool(text) and any(d and (d in text or text in d) for d in (str(x).strip().lower() for x in drugs or [])) |
|
|
|
|
| def _anchor(seed, tl, chain): |
| """Anchor a constraint at the earliest visit of its source entity's timeline; (None, None) when the |
| entity is not in the evidence chain.""" |
| ent = seed.get("source_entity") |
| for t in tl: |
| if t.get("entity") == ent and t.get("segments"): |
| vf = (t["segments"][0].get("valid_from") or {}) |
| if vf.get("visit_number") in chain: |
| return vf["visit_number"], vf.get("date") |
| return None, None |
|
|
|
|
| def synthesize(cli, chain, by_v, forbidden, n, batch): |
| """Author a standing drug constraint with four options; the gold is the safe member (options[].is_correct).""" |
| out, vs = [], sorted(by_v) |
| for i in range(n): |
| host = vs[int(len(vs) * (0.15 + 0.4 * i / max(1, n)))] |
| dept = by_v[host]["department"] |
| host_date = by_v[host]["date"] |
| patient, context = patient_line(), patient_context(chain) |
| forbidden_list = "\n".join(f"- {e}" for e in forbidden[:50]) |
| r = cli.complete_json( |
| "You write one drug-safety multiple-choice question for a medical-memory benchmark. Return only JSON.", |
| f"{patient}\n{context}\n" |
| f"Visit: v{host} · {host_date} · {dept}\n\n" |
| "Write a standing drug constraint, such as an allergy, an intolerance or an interaction, that\n" |
| "makes three drugs a clinician might choose unsafe for this patient while one clinically\n" |
| "reasonable alternative stays safe. Then write four options: the three unsafe drugs and the safe\n" |
| "one. A model without memory must be drawn to an unsafe first-line choice.\n\n" |
| "Do not name, describe or hint at the constraint in the question or the options; the model must\n" |
| "recall it from memory. Never use the words \"safe\", \"unsafe\", \"risk\", \"dangerous\" or \"avoid\" in\n" |
| "the question or the option text; the options are plain drug names.\n\n" |
| "Do not use these drugs or classes, which existing answer keys already track:\n" |
| + forbidden_list + |
| "\n\nReturn JSON in this form:\n" |
| '{"atom_text": "<terse chart fact recording the constraint>",\n' |
| ' "entity": "<the constraint>",\n' |
| ' "topic": "<slug>",\n' |
| ' "question": "<the patient asks which drug to start, without listing the options>",\n' |
| ' "options": [{"label": "A", "content": "<drug>", "is_correct": false,\n' |
| ' "explanation": "<why it is or is not appropriate>"},\n' |
| " ...]}") |
| if not (has_text(r) and has_text(r, "atom_text") and has_text(r, "entity") and r.get("options")): |
| print(f" axis3[{i}] rejected: the reply lacks the question, its options, the fact or its entity") |
| continue |
| opts = r["options"] |
| if not (isinstance(opts, list) and all(isinstance(o, dict) and has_text(o, "content") for o in opts)): |
| print(f" axis3[{i}] rejected: an option is not an object with a drug name") |
| continue |
| if lists_options(r["question"], opts): |
| print(f" axis3[{i}] rejected: the question text lists the options") |
| continue |
| n_correct = sum(1 for o in opts if o.get("is_correct")) |
| if n_correct != 1: |
| print(f" axis3[{i}] rejected: {n_correct} correct options, need exactly 1") |
| continue |
| qid = f"{PATIENT}/axis3/S{len(out):03d}" |
| opts = shuffled(opts, qid) |
| safe = next(o for o in opts if o.get("is_correct")) |
| atom = {"atom_id": f"{batch}/{len(out):03d}", "batch_id": batch, "axis": "3", |
| "visit_number": host, "date": by_v[host]["date"], "department": dept, |
| "field": "doctor_extra_notes", "text": r["atom_text"], "entities": [r["entity"]], |
| "topic": r.get("topic"), "anchors_axis": "3", "anchors_on": []} |
| out.append({"atom": atom, "q": { |
| "qid": qid, "patient_id": PATIENT, "axis": "3_drug_safety", |
| "question": r["question"], "options": opts, "answer": safe["content"], |
| "answer_type": "multiple_choice", "grounding": f"synthetic:{batch}", |
| "evidence_visits": [host], "gt_visit": host, "gt_date": by_v[host]["date"], |
| "question_appearance_window": {"from": by_v[host]["date"], "to": by_v[vs[-1]]["date"]}, |
| "derivation": {"evidence": [{"visit": host, "date": by_v[host]["date"], |
| "field": "doctor_extra_notes", "atom": r["atom_text"]}], |
| "steps": [f"Standing constraint recorded v{host}: {r['atom_text']!r}.", |
| f"Of the four options only {safe['content']} is compatible with it.", |
| f"{safe.get('explanation') or ''}"], |
| "gold": safe["content"], |
| "rivals": [{"value": o.get("content"), "why_wrong": o.get("explanation")} |
| for o in opts if not o.get("is_correct")], |
| "operands": {"constraint_text": r["atom_text"], "constraint_id": r.get("topic")}, |
| "derived": False}, |
| "provenance": f"synthetic drug-safety constraint ({batch}); only the recalled constraint " |
| f"distinguishes the safe member", |
| "axis_specific": {"constraint_id": r.get("topic"), "depends_on_atoms": [atom["atom_id"]]}, |
| }}) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| discover() |
|
|
