code_release/code/construction/pipeline/llm.py

#!/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()