#!/usr/bin/env python3 """LLM access for the construction stages: the JSON client and the template-driven LLM steps. extract --manifest M --out O --model X per-visit event extraction, one call per visit step --step S (--manifest M | --input I) --out O --model X evidence (evidence-chain timelines), tags (change tags), axis1/axis2 (question generation) The system message is the step's templates (pipeline.prompts) filled with the patient's facts; the user message is the step's task, the reply envelope and the output schema. The step's input is inlined at its marker, in the task or (axis 1) in the system message. """ from __future__ import annotations import argparse import json import os import re import sys import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from pipeline import prompts from pipeline.common import AXIS1_FORM_MIX, AXIS1_TARGET, AXIS2_TARGET, Progress def extract_json_object(text: str) -> dict: """Parse the first balanced top-level JSON object out of `text`.""" text = text.strip() if text.startswith("```"): text = re.sub(r"^```[a-zA-Z]*\n?|\n?```$", "", text).strip() try: return json.loads(text) except json.JSONDecodeError: pass start = text.find("{") while start != -1: depth, instr, esc = 0, False, False for i in range(start, len(text)): c = text[i] if instr: if esc: esc = False elif c == "\\": esc = True elif c == '"': instr = False else: if c == '"': instr = True elif c == "{": depth += 1 elif c == "}": depth -= 1 if depth == 0: try: return json.loads(text[start:i + 1]) except json.JSONDecodeError: break start = text.find("{", start + 1) raise ValueError("no parseable JSON object in model output") _REPAIR = ("\n\nYour previous reply was not valid JSON. Reply with ONLY a single valid JSON object, no prose, " "no code fences.") RATE_LIMIT_HITS = {"chat": 0} def _note_rate_limit(who: str) -> None: RATE_LIMIT_HITS[who] += 1 n = RATE_LIMIT_HITS[who] if n <= 5 or n % 25 == 0: print(f"[rate-limit] {who} 429/throttle #{n}", file=sys.stderr, flush=True) _usage_lock = threading.Lock() def _log_usage(model, resp): try: u = getattr(resp, "usage", None) if u is None: return cd = getattr(u, "completion_tokens_details", None) pdt = getattr(u, "prompt_tokens_details", None) rec = {"model": model, "prompt": int(getattr(u, "prompt_tokens", 0) or 0), "completion": int(getattr(u, "completion_tokens", 0) or 0), "reasoning": int(getattr(cd, "reasoning_tokens", 0) or 0) if cd else 0, "cached": int(getattr(pdt, "cached_tokens", 0) or 0) if pdt else 0, "total": int(getattr(u, "total_tokens", 0) or 0), "phase": " ".join([os.path.basename(sys.argv[0])] + sys.argv[1:2]) if sys.argv else "?", "patient": os.environ.get("MEMEBENCH_PATIENT", "?")} path = os.path.join(os.environ.get("MEMEBENCH_WORK", "."), "usage.jsonl") with _usage_lock: with open(path, "a", encoding="utf-8") as fh: fh.write(json.dumps(rec, ensure_ascii=False) + "\n") except Exception: pass class ChatClient: def __init__(self, model: str, reasoning_effort: str | None = "medium", max_completion_tokens: int = 8000, max_retries: int = 4): import openai if not os.environ.get("OPENAI_API_KEY"): raise RuntimeError("OPENAI_API_KEY not set") self._client = openai.OpenAI() self.model = model self.reasoning_effort = reasoning_effort self.max_completion_tokens = max_completion_tokens self.max_retries = max_retries def complete_json(self, system: str, user: str) -> dict: messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] effort_ok = fmt_ok = True last_err = None for attempt in range(self.max_retries): try: kw = dict(model=self.model, messages=messages, max_completion_tokens=self.max_completion_tokens) if fmt_ok: kw["response_format"] = {"type": "json_object"} if self.reasoning_effort and effort_ok: kw["reasoning_effort"] = self.reasoning_effort resp = self._client.chat.completions.create(**kw) _log_usage(self.model, resp) return extract_json_object(resp.choices[0].message.content or "") except Exception as e: last_err = e msg = str(e).lower() if "429" in msg or "rate limit" in msg or "rate_limit" in msg or "overloaded" in msg: _note_rate_limit("chat") if "reasoning_effort" in msg: effort_ok = False if "response_format" in msg: fmt_ok = False if attempt == self.max_retries - 2: messages = messages + [{"role": "user", "content": _REPAIR}] time.sleep(1.5 * (attempt + 1)) raise RuntimeError(f"ChatClient failed after {self.max_retries} tries: {last_err}") PLAST = os.environ.get("MEMEBENCH_PLAST") PFIRST = os.environ.get("MEMEBENCH_PFIRST") NVISITS = os.environ.get("MEMEBENCH_NVISITS") PDESC = os.environ.get("MEMEBENCH_PDESC") STEPS = { "extract": {"system": [prompts.EXTRACT_RUBRIC], "task": prompts.EXTRACT_TASK, "schema": prompts.EXTRACT_SCHEMA, "mode": "batch", "top_key": "records", "stamp": ["visit_number"], "envelope": "Return only the JSON event record of this visit, with visit_number and the field " "arrays, nothing else."}, "evidence": {"system": [prompts.EVIDENCE_RUBRIC], "task": prompts.EVIDENCE_TASK, "schema": prompts.EVIDENCE_SCHEMA, "mode": "batch", "top_key": "timelines", "stamp": ["department", "batch_id"], "envelope": 'Return only a JSON object of the form {"timelines": [ ... ]}, nothing else.\n' 'Each timeline must include:\n' '- "entity": a string\n' '- "aspect": one of the values listed above\n' '- "source_field": the source bucket the entity came from: diagnosis,\n' ' patient_self_description, tests_and_medication, doctor_notes, doctor_plans,\n' ' health_advice or doctor_extra_notes\n' '- "current_status": a short string saying where the entity stands as of the last ' 'visit\n' '- "segments": each with "value", "valid_from" {date, visit_number}, "valid_until"\n' ' ({date, visit_number} or "current"), "status" (current, resolved, superseded or\n' ' historical), "supporting_visits" [ints] and "superseded_by"\n' ' ({value, date, visit_number} or null)\n' '"notes" is optional.'}, "tags": {"system": [prompts.TAGS_RUBRIC], "task": prompts.TAGS_TASK, "schema": None, "mode": "batch", "top_key": "tags", "envelope": 'Return ONLY a JSON object of the form {"tags": [ ... ]}, nothing else. Each tag must ' 'contain: visit_number, tag ("change" or "no_change"), noise_score (0--3), and note.'}, } for _n, _p in ((1, "AXIS1"), (2, "AXIS2")): STEPS[f"axis{_n}"] = {"system": [getattr(prompts, f"{_p}_COMMON"), getattr(prompts, f"{_p}_RULES")], "task": getattr(prompts, f"{_p}_TASK"), "schema": getattr(prompts, f"{_p}_SCHEMA"), "mode": "single", "top_key": "questions", "envelope": 'Return only a JSON object of the form {"questions": [ ... ]}, nothing else.'} def _require_facts() -> None: """Exit unless the patient facts of patient.env were in the environment.""" missing = [k for k, v in (("MEMEBENCH_PLAST", PLAST), ("MEMEBENCH_PFIRST", PFIRST), ("MEMEBENCH_NVISITS", NVISITS), ("MEMEBENCH_PDESC", PDESC)) if v is None] if missing: sys.exit(f"{', '.join(missing)} not set; run the step through build_patient.sh " f"(it sources the patient's patient.env)") def subst(t): """Resolve the prompt placeholders: the patient's facts and the question counts of steps axis1 and axis2.""" _require_facts() return (t.replace("${PLAST}", PLAST).replace("${PFIRST}", PFIRST) .replace("${NVISITS}", NVISITS).replace("${PDESC}", PDESC) .replace("${AXIS1_TARGET}", AXIS1_TARGET).replace("${AXIS1_FORM_MIX}", AXIS1_FORM_MIX) .replace("${AXIS2_TARGET}", AXIS2_TARGET)) def inline(text, content, unit): """Inline the input at its path marker and fill per-batch fields. The marker may live in the system message (axis 1 puts ${SEEDS} inside its common block) or in the task (axis 2 and the chain steps). """ for pv in ("${SEEDS}", "${b.path}", "${m.path}"): text = text.replace(pv, "\n" + content + "\n") if unit: for k, v in unit.items(): text = text.replace("${b." + k + "}", str(v)).replace("${m." + k + "}", str(v)) return text def build_prompts(step, content, unit=None): """System and user message of one call of a step.""" system = "\n\n".join(subst(t) for t in step["system"]) tail = subst(step["task"]).strip("\n ").strip() system = inline(system, content, unit) tail = inline(tail, content, unit) user = tail + "\n\n" + step["envelope"] if step["schema"]: user += ("\n\nYour JSON must follow this schema exactly. Use these field names as written, and include\n" "every field listed in each `required`:\n" + subst(step["schema"])) return system, user def extract(a): """Per-visit event extraction with the extraction rubric, in parallel, with progress.""" step = STEPS["extract"] units = json.load(open(a.manifest)) if a.limit: units = units[: a.limit] def one(u): t0 = time.time() try: content = open(u["path"]).read() sys_p, usr_p = build_prompts(step, content, u) client = ChatClient(model=a.model, reasoning_effort="medium", max_completion_tokens=8000) rec = client.complete_json(sys_p, usr_p) rec["visit_number"] = u["visit_number"] return {"ok": True, "visit": u["visit_number"], "rec": rec, "secs": time.time() - t0} except Exception as e: return {"ok": False, "visit": u["visit_number"], "err": str(e)[:300], "secs": time.time() - t0} print(f"extract: {len(units)} visits, engine={a.model}, {a.workers} workers", flush=True) prog_path = a.progress or str(Path(a.out).with_name("progress_extract.json")) res = [] with Progress(f"extract [{a.model}]", total=len(units), out=prog_path) as prog: with ThreadPoolExecutor(max_workers=a.workers) as ex: futs = [ex.submit(one, u) for u in units] for f in as_completed(futs): r = f.result() res.append(r) prog.tick(ok=r["ok"]) if len(res) % 50 == 0: _dump_records(a.out, res) ok = [r for r in res if r["ok"]] bad = [r for r in res if not r["ok"]] _dump_records(a.out, res) if bad: print(f" {len(bad)} FAILED:", flush=True) for r in bad[:10]: print(f" v{r['visit']}: {r['err'][:140]}", flush=True) Path(a.out).with_name("failed_extract.json").write_text( json.dumps([{"visit": r["visit"], "err": r["err"]} for r in bad], indent=1)) print(f" -> {a.out} ({len(ok)} records)", flush=True) def _dump_records(out, res): recs = sorted((r["rec"] for r in res if r["ok"]), key=lambda r: r["visit_number"]) Path(out).write_text(json.dumps({"records": recs, "n": len(recs)}, ensure_ascii=False)) def step(a): """Run a manifest-driven step (evidence, tags) or a single-call step (axis1, axis2).""" step = STEPS[a.step] if step["mode"] == "single": if not a.input: sys.exit(f"--input is required for single-mode step {a.step} " f"(the seed file, _seed_{a.step}.json)") print(f"{a.step}: single call, engine={a.model}", flush=True) print(f" patient facts: {PDESC}, {NVISITS} visits " f"{PFIRST}..{PLAST}", flush=True) content = open(a.input, encoding="utf-8").read() sys_p, usr_p = build_prompts(step, content) t0 = time.time() client = ChatClient(model=a.model, reasoning_effort="medium", max_completion_tokens=16000) d = client.complete_json(sys_p, usr_p) items = d.get(step["top_key"], d if isinstance(d, list) else []) Path(a.out).write_text(json.dumps({step["top_key"]: items, "n": len(items)}, ensure_ascii=False)) print(f" -> {a.out} ({len(items)} {step['top_key']}, {time.time()-t0:.0f}s)", flush=True) return if not a.manifest: sys.exit(f"--manifest is required for batch-mode step {a.step}") units = json.load(open(a.manifest)) if a.limit: units = units[: a.limit] def one(i, u): t0 = time.time() try: content = open(u["path"], encoding="utf-8").read() sys_p, usr_p = build_prompts(step, content, u) client = ChatClient(model=a.model, reasoning_effort="medium", max_completion_tokens=16000) d = client.complete_json(sys_p, usr_p) got = d.get(step["top_key"], d if isinstance(d, list) else []) if isinstance(got, dict): got = [got] for r in got: for f in step.get("stamp", []): if f in u: r[f] = u[f] return {"ok": True, "i": i, "got": got, "secs": time.time() - t0} except Exception as e: return {"ok": False, "i": i, "got": [], "err": str(e)[:300], "secs": time.time() - t0} print(f"{a.step}: {len(units)} units, engine={a.model}, {a.workers} workers", flush=True) prog_path = a.progress or str(Path(a.out).with_name(f"progress_{a.step}.json")) res = [] with Progress(f"{a.step} [{a.model}]", total=len(units), out=prog_path) as prog: with ThreadPoolExecutor(max_workers=a.workers) as ex: futs = [ex.submit(one, i, u) for i, u in enumerate(units)] for f in as_completed(futs): r = f.result() res.append(r) prog.tick(ok=r["ok"]) if len(res) % 5 == 0: _dump_step(a.out, res, step) _dump_step(a.out, res, step) bad = [r for r in res if not r["ok"]] n = sum(len(r["got"]) for r in res if r["ok"]) failp = Path(a.out).with_name(f"failed_{a.step}.json") if not bad and failp.exists(): failp.unlink() if bad: print(f" {len(bad)} FAILED:", flush=True) for r in bad[:10]: print(f" unit {r['i']} ({units[r['i']].get('department','?')}): {r['err'][:140]}", flush=True) failp.write_text(json.dumps([{"i": r["i"], "unit": units[r["i"]], "err": r["err"]} for r in bad], indent=1, ensure_ascii=False)) print(f" -> {a.out} ({n} {step['top_key']})", flush=True) def _dump_step(out, res, step): merged = [] for r in sorted(res, key=lambda x: x["i"]): if r["ok"]: merged.extend(r["got"]) Path(out).write_text(json.dumps({step["top_key"]: merged, "n": len(merged)}, ensure_ascii=False)) def main(): ap = argparse.ArgumentParser(description="The LLM steps of the construction pipeline.") sp = ap.add_subparsers(dest="cmd", required=True) e = sp.add_parser("extract") e.add_argument("--manifest", required=True) e.add_argument("--out", required=True) e.add_argument("--model", required=True) e.add_argument("--workers", type=int, default=8) e.add_argument("--limit", type=int) e.add_argument("--progress", help="progress JSON path (default: alongside --out)") s = sp.add_parser("step") s.add_argument("--step", required=True, choices=[k for k in STEPS if k != "extract"]) s.add_argument("--manifest", help="batch steps (evidence, tags): the manifest of units") s.add_argument("--input", help="single steps (axis1, axis2): the seed file to inline") s.add_argument("--out", required=True) s.add_argument("--model", required=True) s.add_argument("--workers", type=int, default=6) s.add_argument("--limit", type=int) s.add_argument("--progress") a = ap.parse_args() extract(a) if a.cmd == "extract" else step(a) if __name__ == "__main__": main()