| |
| """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() |
|
|
