| |
| """Axis 7 (cross-department synthesis): LLM discovery, then templates. |
| |
| The LLM reads the chain in overlapping windows of chart facts, each grouped by department, and returns pairs |
| of facts in different departments that must be joined to answer, where neither department's record alone suffices (e.g. a referral whose |
| outcome is recorded in another department). Questions and golds use Q_CLOSURE / A_CLOSURE. |
| synthesize() authors synthetic questions. |
| |
| python -m pipeline.questions.axis7 discovery -> _discover_axis7.json |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import re |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| from pipeline.common import BASE, FIELDS, Progress, has_text, model, patient_context, patient_id, patient_line |
| from pipeline.llm import ChatClient |
|
|
| PATIENT = patient_id() |
| ENGINE = model("generation") |
|
|
| AXIS = 7 |
| SCAN_FIELDS = ["doctor_plans", "doctor_extra_notes", "doctor_notes", "health_advice", "tests_and_medication", |
| "diagnosis"] |
| WINDOW, OVERLAP = 400, 100 |
|
|
| Q_CLOSURE = ("On {ref_date} my {ref_dept} doctor said {ref_atom_verbatim}. Did that ever happen, and " |
| "what actually happened at that follow-up — for example, anything found or done?") |
| A_CLOSURE = ('Yes. {ref_dept} recorded "{ref_atom_verbatim}" on {ref_date}, and {closer_dept} recorded ' |
| '"{closer_atom_verbatim}" on {closer_date}: {outcome_summary}') |
|
|
| def _sentence(text): |
| """The text with a closing full stop.""" |
| text = str(text or "").strip() |
| return text if not text or text[-1] in ".!?" else text + "." |
|
|
|
|
| DISCOVER_SYS = ( |
| "You read one patient's chart, grouped by department, and find cross-department joins. A join is a\n" |
| "question that can be answered only by combining a fact from department A with a fact from\n" |
| "department B: neither department's record alone is enough, and B's record does not cite A's. The\n" |
| "typical case: department A records a referral or a plan, and its outcome is recorded later in\n" |
| "department B. Return only JSON.") |
|
|
|
|
| def quoted_field(text, rec): |
| """The field of the visit's atom that contains the quoted text (case and spacing ignored), or None.""" |
| def norm(t): |
| return re.sub(r"\s+", " ", str(t or "")).strip().lower().rstrip(".;, ") |
| q = norm(text) |
| return next((f for f in FIELDS for a in rec.get(f) or [] if q and q in norm(a)), None) |
|
|
|
|
| def discover_joins(cli, chain, by_v): |
| """Find cross-department joins over windows of WINDOW chart facts in visit order; consecutive windows share |
| OVERLAP facts, and each window is shown grouped by department.""" |
| rows = [(r["visit_number"], r["date"], r["department"], f, str(a)[:120]) |
| for r in sorted(chain, key=lambda x: x["visit_number"]) for f in SCAN_FIELDS for a in r.get(f) or []] |
| found, seen = [], set() |
| for start in range(0, max(len(rows) - OVERLAP, 1), WINDOW - OVERLAP): |
| by_dept = defaultdict(list) |
| for v, d, dept, f, t in rows[start:start + WINDOW]: |
| by_dept[dept].append(f" v{v} {d} {f}: {t}") |
| chart_by_department = "\n\n".join(f"### {dept}\n" + "\n".join(lines) for dept, lines in by_dept.items()) |
| joins = _discover_window(cli, chart_by_department) |
| if not isinstance(joins, list): |
| joins = [joins] |
| n_bad = sum(1 for j in joins if not isinstance(j, dict)) |
| if n_bad: |
| print(f" skipped {n_bad} join(s) that are not objects") |
| for j in (j for j in joins if isinstance(j, dict)): |
| k = (j.get("ref_visit"), j.get("closer_visit"), str(j.get("ref_atom_verbatim")).lower()) |
| if k not in seen: |
| seen.add(k) |
| found.append(j) |
| return found |
|
|
|
|
| def _discover_window(cli, chart_by_department): |
| """One discovery call over a window of chart facts.""" |
| usr = ( |
| "Chart by department:\n" + chart_by_department + "\n\n" |
| "Find cross-department joins. Each join must:\n" |
| "- involve two different departments, A and B;\n" |
| "- have a fact in A, such as a referral, a plan or a request, whose outcome or resolution is a fact\n" |
| " in B;\n" |
| "- be impossible to answer from either department alone: B's note does not cite A's referral or\n" |
| " plan.\n" |
| "Do not invent links; report only joins that the chart facts actually support. Copy both facts\n" |
| "exactly as they appear above, without the visit, date and field.\n\n" |
| "Return JSON in this form:\n" |
| '{"joins": [{\n' |
| ' "ref_visit": <int>,\n' |
| ' "ref_dept": "<A>",\n' |
| ' "ref_atom_verbatim": "<A\'s fact>",\n' |
| ' "closer_visit": <int>,\n' |
| ' "closer_dept": "<B>",\n' |
| ' "closer_atom_verbatim": "<B\'s fact>",\n' |
| ' "outcome_summary": "<one line: what came of it>",\n' |
| ' "topic": "<slug>"\n' |
| "}]}") |
| try: |
| return cli.complete_json(DISCOVER_SYS, usr).get("joins") or [] |
| except Exception as e: |
| print(f" discover window FAILED: {str(e)[:120]}") |
| return [] |
|
|
|
|
| def discover() -> None: |
| """Discover the axis-7 candidates of this patient -> _discover_axis7.json.""" |
| cli = ChatClient(model=ENGINE, reasoning_effort="medium", max_completion_tokens=7000) |
|
|
| chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl") if l.strip()] |
| by_v = {r["visit_number"]: r for r in chain} |
| plast = max(r["date"] for r in chain) |
|
|
| with Progress("axis7 discover (LLM)", total=1, out=f"{BASE}/progress_axis7_discover.json") as pg: |
| joins = discover_joins(cli, chain, by_v) |
| pg.tick(ok=bool(joins)) |
|
|
| candidates = [] |
| for j in joins: |
| rv, cv = j.get("ref_visit"), j.get("closer_visit") |
| if (rv not in by_v or cv not in by_v or cv <= rv |
| or by_v[rv]["department"] == by_v[cv]["department"]): |
| continue |
| ref_field = quoted_field(j.get("ref_atom_verbatim"), by_v[rv]) |
| closer_field = quoted_field(j.get("closer_atom_verbatim"), by_v[cv]) |
| if not (ref_field and closer_field): |
| print(f" skip v{rv}->v{cv}: a quoted fact is not in its visit") |
| continue |
| ops = {"ref_date": by_v[rv]["date"], "ref_dept": by_v[rv]["department"].split(" (")[0], |
| "ref_atom_verbatim": j.get("ref_atom_verbatim"), |
| "closer_date": by_v[cv]["date"], "closer_dept": by_v[cv]["department"].split(" (")[0], |
| "closer_atom_verbatim": j.get("closer_atom_verbatim"), |
| "outcome_summary": _sentence(j.get("outcome_summary", ""))} |
| candidates.append({ |
| "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT, |
| "axis": "7_cross_dept_synthesis", |
| "question": Q_CLOSURE.format(**ops), "answer": A_CLOSURE.format(**ops), |
| "answer_type": "open_exact", "grounding": "record", |
| "evidence_visits": [rv, cv], "gt_visit": cv, "gt_date": by_v[cv]["date"], |
| "question_appearance_window": {"from": by_v[cv]["date"], "to": plast}, |
| "derivation": { |
| "evidence": [{"visit": rv, "date": by_v[rv]["date"], "field": ref_field, |
| "atom": j.get("ref_atom_verbatim")}, |
| {"visit": cv, "date": by_v[cv]["date"], "field": closer_field, |
| "atom": j.get("closer_atom_verbatim")}], |
| "steps": [f"{ops['ref_dept']} v{rv}: {j.get('ref_atom_verbatim')!r} — a request, " |
| f"no outcome here.", |
| f"{ops['closer_dept']} v{cv}: {j.get('closer_atom_verbatim')!r} — the outcome, " |
| f"another dept.", |
| "The outcome note does not cite the referral; answering requires joining both."], |
| "gold": A_CLOSURE.format(**ops), "rivals": [], "operands": ops, "derived": True}, |
| "provenance": f"LLM-discovered cross-dept join ({ENGINE}); A={ops['ref_dept']} B={ops['closer_dept']}", |
| "axis_specific": {"A_dept": ops["ref_dept"], "B_dept": ops["closer_dept"], "topic": j.get("topic")}, |
| }) |
|
|
| 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" cross-dept joins discovered (LLM): {len(joins)}") |
| print(f" CANDIDATES (record pairs) : {len(candidates)}") |
| for c in candidates[:3]: |
| print(f" {c['axis_specific']['A_dept']} -> {c['axis_specific']['B_dept']}: {c['question'][:55]}") |
| print(f" -> {BASE}/_discover_axis{AXIS}.json") |
|
|
|
|
| def synthesize(cli, chain, by_v, forbidden, n, batch): |
| """Author department A's referral and department B's outcome as two atoms; the join is the question.""" |
| out, vs = [], sorted(by_v) |
| for i in range(n): |
| a_host = vs[int(len(vs) * (0.15 + 0.5 * i / max(1, n)))] |
| a_dept = by_v[a_host]["department"] |
| later = [v for v in vs if v > a_host and by_v[v]["department"] != a_dept] |
| if not later: |
| continue |
| b_host = later[min(len(later) - 1, 12)] |
| b_dept = by_v[b_host]["department"] |
| a_date, b_date = by_v[a_host]["date"], by_v[b_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 two terse chart facts for one referral story in a health record. Return only JSON.", |
| f"{patient}\n{context}\n" |
| f"A: v{a_host} · {a_date} · {a_dept} — records the referral\n" |
| f"B: v{b_host} · {b_date} · {b_dept} — records the outcome\n\n" |
| "A refers the patient to B about a topic, and B later records what came of it. Neither fact may\n" |
| "mention the other department's note: the point is that the join itself is not recorded. Write\n" |
| "in terse chart style, such as 'refer GI for dyspepsia', not in prose.\n\n" |
| "Do not use these, which existing answer keys already track:\n" + forbidden_list + |
| "\n\nReturn JSON in this form:\n" |
| '{"ref_text": "<A\'s referral fact>",\n' |
| ' "outcome_text": "<B\'s outcome fact>",\n' |
| ' "outcome_summary": "<one line: what came of it>",\n' |
| ' "entity": "<topic>",\n' |
| ' "topic": "<slug>"}') |
| missing = [k for k in ("ref_text", "outcome_text", "outcome_summary") if not has_text(r, k)] |
| if missing: |
| print(f" axis7[{i}] rejected: the reply lacks {', '.join(missing)}") |
| continue |
| a1 = {"atom_id": f"{batch}/{len(out)*2:03d}", "batch_id": batch, "axis": "7", |
| "visit_number": a_host, "date": by_v[a_host]["date"], "department": a_dept, |
| "field": "doctor_plans", "text": r["ref_text"], "entities": [r.get("entity", "")], |
| "topic": r.get("topic"), "anchors_axis": "7", "anchors_on": []} |
| a2 = {**a1, "atom_id": f"{batch}/{len(out)*2+1:03d}", "visit_number": b_host, |
| "date": by_v[b_host]["date"], "department": b_dept, "field": "doctor_notes", |
| "text": r["outcome_text"]} |
| gold = (f"Yes. {a_dept.split(' (')[0]} recorded \"{r['ref_text']}\" on {by_v[a_host]['date']}, and " |
| f"{b_dept.split(' (')[0]} recorded \"{r['outcome_text']}\" on {by_v[b_host]['date']}: " |
| f"{_sentence(r['outcome_summary'])}") |
| out.append({"atom": a1, "atom2": a2, "q": { |
| "qid": f"{PATIENT}/axis7/S{len(out):03d}", "patient_id": PATIENT, "axis": "7_cross_dept_synthesis", |
| "question": (f"On {by_v[a_host]['date']} my {a_dept.split(' (')[0]} doctor said " |
| f"{r['ref_text']}. Did that ever happen, and what actually happened at that follow-up " |
| f"— for example, anything found or done?"), |
| "answer": gold, "answer_type": "open_exact", "grounding": f"synthetic:{batch}", |
| "evidence_visits": [a_host, b_host], "gt_visit": b_host, "gt_date": by_v[b_host]["date"], |
| "question_appearance_window": {"from": by_v[b_host]["date"], "to": by_v[vs[-1]]["date"]}, |
| "derivation": {"evidence": [ |
| {"visit": a_host, "date": by_v[a_host]["date"], "field": "doctor_plans", "atom": r["ref_text"]}, |
| {"visit": b_host, "date": by_v[b_host]["date"], "field": "doctor_notes", "atom": r["outcome_text"]}], |
| "steps": [f"{a_dept} v{a_host}: {r['ref_text']!r} — a referral with no outcome in this dept.", |
| f"{b_dept} v{b_host}: {r['outcome_text']!r} — the outcome, in a DIFFERENT dept.", |
| "The outcome note does not cite the referral; answering requires joining both."], |
| "gold": gold, "rivals": [], |
| "operands": {"ref_visit": a_host, "ref_date": by_v[a_host]["date"], "ref_dept": a_dept, |
| "ref_atom_verbatim": r["ref_text"], "target_dept": b_dept, |
| "closer_date": by_v[b_host]["date"], "closer_dept": b_dept, |
| "closer_plans_verbatim": r["outcome_text"]}, |
| "derived": True}, |
| "provenance": f"synthetic cross-dept join ({batch}); A={a_dept} B={b_dept}, join unrecorded", |
| "axis_specific": {"A_dept": a_dept, "B_dept": b_dept, |
| "depends_on_atoms": [a1["atom_id"], a2["atom_id"]]}, |
| }}) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| discover() |
|
|
