| |
| """Stage E gates, run on the injection plan before anything is injected. |
| |
| lint visit_exists: every planned atom targets an existing visit and a content bucket; single_layer: one |
| synthetic layer per fact, no topic across batches; record_has_no_injection: no injected atom in the |
| evidence of a record question, by atom id or by text; atom_has_axis: every atom declares the axis that |
| will look for it -> lint_report.json |
| persona an LLM check that drops every synthetic question with an atom impossible for this patient (e.g. a |
| pregnancy test for a male patient), with its atoms and claim; the persona combines the sex in the visit |
| table with the documented conditions |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import glob |
| import json |
| import sys |
| from collections import Counter |
| from pathlib import Path |
|
|
| from pipeline.common import BASE, FIELDS, SEX_MAP, VISIT_TABLE, Progress, model, patient_id |
| from pipeline.llm import ChatClient |
| from pipeline.synthetic import load, norm, save |
|
|
| PATIENT = patient_id() |
| ENGINE = model("generation") |
|
|
|
|
| def visit_exists(reg, plan): |
| """Every planned atom targets an existing visit and a content bucket.""" |
| f = [] |
| chain = {json.loads(l)["visit_number"] for l in open(BASE / "event_chain.jsonl", encoding="utf-8")} |
| for a in plan: |
| if a.get("visit_number") not in chain: |
| f.append(f"atom {a.get('atom_id')} targets visit {a.get('visit_number')}, " |
| f"which does not exist (injection may not create visits)") |
| if a.get("field") not in FIELDS: |
| f.append(f"atom {a.get('atom_id')} field={a.get('field')!r} is not a content bucket") |
| return f |
|
|
|
|
| def single_layer(reg, plan): |
| """One synthetic layer per fact: no atom anchors on another batch's injected evidence, and no topic is |
| reused across batches.""" |
| f = [] |
| applied_atoms = {a["atom_id"]: a for a in reg["atoms"]} |
| for a in plan: |
| for dep in a.get("anchors_on") or []: |
| if dep in applied_atoms: |
| f.append(f"atom {a.get('atom_id')} anchors on injected atom {dep} " |
| f"(batch {applied_atoms[dep]['batch_id']}); only one synthetic layer per fact") |
| topic_batch = {} |
| for a in reg["atoms"] + plan: |
| t, b = a.get("topic"), a.get("batch_id") |
| if not t: |
| continue |
| if t in topic_batch and topic_batch[t] != b: |
| f.append(f"topic {t!r} used by batch {topic_batch[t]!r} and {b!r}; topics may not cross batches") |
| topic_batch[t] = b |
| return f |
|
|
|
|
| def record_has_no_injection(reg, plan): |
| """No claim with grounding=record has an injected atom in its evidence, and no record question's evidence text |
| contains the text of a planned atom (case and punctuation ignored).""" |
| f = [] |
| planned_ids = {a.get("atom_id") for a in plan} |
| for c in reg["claims"]: |
| if str(c.get("grounding")) != "record": |
| continue |
| dep = set(c.get("depends_on_atoms") or []) & (planned_ids | {a["atom_id"] for a in reg["atoms"]}) |
| if dep: |
| f.append(f"{c['qid']} is grounding=record but its evidence includes injected atoms {sorted(dep)}") |
| planned = [(a.get("atom_id"), f" {norm(a.get('text') or '')} ") for a in plan if norm(a.get("text") or "")] |
| for bank in sorted(glob.glob(str(BASE / "questions_axis*.jsonl"))): |
| for l in open(bank, encoding="utf-8"): |
| q = json.loads(l) if l.strip() else {} |
| if str(q.get("grounding")) != "record": |
| continue |
| for e in (q.get("derivation") or {}).get("evidence") or []: |
| text = f" {norm(e.get('atom') or '')} " if isinstance(e, dict) else "" |
| for aid, t in planned: |
| if text.strip() and t in text: |
| f.append(f"{q.get('qid')} is grounding=record but its evidence at v{e.get('visit')} " |
| f"contains the text of planned atom {aid}") |
| return f |
|
|
|
|
| def atom_has_axis(reg, plan): |
| """Every planned atom declares the axis whose discovery will look for it.""" |
| f = [] |
| for a in plan: |
| if not a.get("anchors_axis"): |
| f.append(f"atom {a.get('atom_id')} declares no anchors_axis, so no question can be built on it") |
| return f |
|
|
|
|
| GATES = [("visit_exists", visit_exists), ("single_layer", single_layer), |
| ("record_has_no_injection", record_has_no_injection), ("atom_has_axis", atom_has_axis)] |
|
|
|
|
| def lint() -> None: |
| """Run the gates over the registry and injection_plan.json; exit 1 on any failure.""" |
| reg = load() |
| pf = BASE / "injection_plan.json" |
| plan = json.loads(pf.read_text())["atoms"] if pf.exists() else [] |
| if not plan: |
| print(f"note: no {pf}; linting the registry alone (0 planned atoms)") |
|
|
| print(f"lint, patient {PATIENT}") |
| print(f" registry: {len(reg['batches'])} batches, {len(reg['atoms'])} atoms, {len(reg['claims'])} claims") |
| print(f" plan: {len(plan)} atoms\n") |
|
|
| fails, report = [], {} |
| for name, fn in GATES: |
| f = fn(reg, plan) |
| report[name] = f |
| fails += f |
| print(f" {'OK ' if not f else 'FAIL'} {name:<22} {len(f) or ''}") |
| for x in f[:4]: |
| print(f" {x}") |
|
|
| (BASE / "lint_report.json").write_text(json.dumps( |
| {"patient_id": PATIENT, |
| "n_planned": len(plan), "pass": not fails, "findings": report}, ensure_ascii=False, indent=1)) |
| print(f"\n{'PASS: safe to inject' if not fails else f'{len(fails)} FAILURES: do not inject'}") |
| print(f"-> {BASE/'lint_report.json'}") |
| sys.exit(1 if fails else 0) |
|
|
|
|
| def build_persona(chain): |
| """Sex from the visit table, the patient's own documented diagnoses and the most frequent tests and |
| medications.""" |
| sex = None |
| try: |
| import pandas as pd |
| df = pd.read_parquet(VISIT_TABLE) |
| sub = df[df["record_id"] == PATIENT] |
| for col in df.columns: |
| if any(k in col for k in ("sex", "gender")): |
| v = str(sub[col].dropna().iloc[0]).strip() if len(sub[col].dropna()) else None |
| sex = SEX_MAP.get(v, v) |
| break |
| except Exception as e: |
| raise RuntimeError(f"persona: cannot read the patient's sex from {VISIT_TABLE}: {e}") from e |
| if not sex: |
| raise RuntimeError(f"persona: no sex recorded for the patient in {VISIT_TABLE}") |
| dx = Counter() |
| for r in chain: |
| for d in r.get("diagnosis") or []: |
| dx[str(d)[:50]] += 1 |
| top_dx = [d for d, _ in dx.most_common(15)] |
| tm = Counter(str(a)[:50] for r in chain for a in r.get("tests_and_medication") or []) |
| return {"sex": sex, "n_visits": len(chain), |
| "date_span": f"{min(r['date'] for r in chain)}..{max(r['date'] for r in chain)}", |
| "documented_conditions": top_dx, |
| "frequent_tests_and_medications": [t for t, _ in tm.most_common(20)]} |
|
|
|
|
| def gate(cli, persona, atoms): |
| """Ask the LLM which atoms are inconsistent with the persona; returns (ok_ids, violations). A reply without a |
| violations list raises; a violation that is not an object naming a listed atom is skipped and reported.""" |
| idx = {a["atom_id"]: a for a in atoms} |
| listing = "\n".join(f' {a["atom_id"]}: {str(a.get("text"))[:120]}' for a in atoms) |
| persona_json = json.dumps(persona, ensure_ascii=False, indent=1) |
| sys_p = ( |
| "You check synthetic clinical facts for consistency with the patient's persona before they are\n" |
| "added to the patient's record. Flag only facts that are impossible or contradictory for this\n" |
| "specific person: a finding that a person of this sex cannot have, such as pregnancy, menstruation\n" |
| "or a Pap test for a man, or the prostate or PSA for a woman; or a fact that flatly contradicts the\n" |
| "documented persona, such as an allergy to a drug the patient takes. Do not flag values that are\n" |
| "merely unusual, extreme or out of range. Extreme values are fine, because this is a memory test and\n" |
| "charts have outliers. Flag only what this person cannot have. Return only JSON.") |
| usr = ( |
| f"Patient persona:\n{persona_json}\n\n" |
| f"Synthetic facts (atom id: text):\n{listing}\n\n" |
| "Return JSON in this form:\n" |
| '{"violations": [{"atom_id": "<id>", "reason": "<why it is impossible for this person>"}]}\n' |
| "Return an empty list if all facts are consistent.") |
| reply = cli.complete_json(sys_p, usr) |
| v = reply.get("violations") |
| if not isinstance(v, list): |
| raise ValueError(f"persona gate reply lacks a violations list: {str(reply)[:80]}") |
| ok = [x for x in v if isinstance(x, dict) and x.get("atom_id") in idx] |
| if len(ok) != len(v): |
| print(f" skipped {len(v) - len(ok)} violation item(s) that do not name a listed atom: " |
| f"{[x for x in v if x not in ok][:3]}") |
| bad = {x["atom_id"] for x in ok} |
| return [a for a in idx if a not in bad], ok |
|
|
|
|
| def _atom_qid(aid): |
| """Plan atom_id 'synthetic-axis7/S000_1' -> (axis '7', qid '<PID>/axis7/S000').""" |
| batch, tail = aid.split("/", 1) |
| ax = batch.split("axis")[-1] |
| suffix = tail.rsplit("_", 1)[0] |
| return ax, f"{PATIENT}/axis{ax}/{suffix}" |
|
|
|
|
| def filter_plan_and_banks(bad_atom_ids): |
| """Drop every question with a persona-impossible atom (its atoms from injection_plan.json, the question from |
| its axis bank, its claim from the registry) and return the dropped qids.""" |
| dead = {_atom_qid(a)[1] for a in bad_atom_ids} |
| pf = Path(f"{BASE}/injection_plan.json") |
| plan = json.loads(pf.read_text()) |
| plan["atoms"] = [a for a in plan["atoms"] if _atom_qid(a["atom_id"])[1] not in dead] |
| plan["n"] = len(plan["atoms"]) |
| pf.write_text(json.dumps(plan, ensure_ascii=False, indent=1)) |
| for bank in glob.glob(f"{BASE}/questions_axis*.jsonl"): |
| rows = [json.loads(l) for l in open(bank)] |
| keep = [q for q in rows if q.get("qid") not in dead] |
| if len(keep) != len(rows): |
| with open(bank, "w", encoding="utf-8") as fh: |
| for q in keep: |
| fh.write(json.dumps(q, ensure_ascii=False) + "\n") |
| reg = load() |
| reg["claims"] = [c for c in reg["claims"] if c["qid"] not in dead] |
| save(reg) |
| return dead |
|
|
|
|
| def persona() -> None: |
| """Drop synthetic atoms, and their questions, that are impossible for this patient.""" |
| cli = ChatClient(model=ENGINE, reasoning_effort="medium", max_completion_tokens=3000) |
|
|
| chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl") if l.strip()] |
| persona = build_persona(chain) |
|
|
| pf = Path(f"{BASE}/injection_plan.json") |
| atoms = json.loads(pf.read_text())["atoms"] if pf.exists() else [] |
|
|
| print(f"persona gate, patient {PATIENT} (injection_plan: {len(atoms)} atoms)") |
| print(f" persona: sex={persona['sex']}, {persona['n_visits']} visits, " |
| f"top dx: {', '.join(persona['documented_conditions'][:5])}") |
| if not atoms: |
| print(" no atoms to gate") |
| return |
|
|
| with Progress("persona gate (LLM)", total=1, out=f"{BASE}/progress_persona.json") as pg: |
| _, violations = gate(cli, persona, atoms) |
| pg.tick(ok=True) |
|
|
| (Path(f"{BASE}/persona_gate_report.json")).write_text(json.dumps( |
| {"patient_id": PATIENT, "persona": persona, "n_atoms": len(atoms), |
| "n_violations": len(violations), "violations": violations, |
| "pass": not violations}, ensure_ascii=False, indent=1)) |
| if violations: |
| print(f" {len(violations)} persona violation(s):") |
| for v in violations: |
| print(f" {v['atom_id']}: {v.get('reason')}") |
| dropped = filter_plan_and_banks({v["atom_id"] for v in violations}) |
| print(f"\n dropped {len(dropped)} question(s) with a persona-impossible atom: {sorted(dropped)}") |
| print(" continuing with the filtered plan") |
| print(f"-> {BASE}/persona_gate_report.json") |
| return |
| print(f" PASS: all {len(atoms)} atoms are consistent with the persona.") |
| print(f"-> {BASE}/persona_gate_report.json") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Stage E gates on the injection plan.") |
| sp = ap.add_subparsers(dest="cmd", required=True) |
| sp.add_parser("lint") |
| sp.add_parser("persona") |
| a = ap.parse_args() |
| lint() if a.cmd == "lint" else persona() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
