esh-pfi-infrastructure/scripts/r49-corpus/generate_incumbent_arm.py

"""R49 H02 — the INCUMBENT arm: the live gen seat, style-prompted.

H02 is explicit that this arm is not optional: what a trained voice adapter
displaces is not the unadapted base model, it is a large instruct model asked
nicely to write like the author, which is free and already deployed. Comparing
only against the base flatters the adapter.

⚠ Two things recorded rather than glossed:

1. **The backing model, not the alias.** `gen` is a gateway alias and has pointed
   at different concrete models over time -- counting by an alias once inflated an
   exposure figure 4.7x on this fleet. The concrete model is resolved at run START
   and again at run END, and both go in the artefact.

2. **The harness differs from the other arms, unavoidably.** The base and adapted
   arms are local transformers on gx10; the incumbent is a served NVFP4 27B on
   ana-ml2 reached over the gateway, and it is an INSTRUCT model receiving a style
   instruction where the others are base models receiving none. That asymmetry IS
   the comparison H02 wants -- prompted imitation against trained voice -- but it
   means this arm is not harness-matched to the others and must not be reported as
   if it were.
"""
from __future__ import annotations
import argparse, json, os, sys, time, urllib.request
from pathlib import Path

GATEWAY = "http://10.250.50.70:4000"

STYLE_SYSTEM = (
    "You are continuing a passage from a Victorian novel by Charlotte Brontë. "
    "Write in her voice: first-person retrospective narration, long periodic "
    "sentences with subordinate clauses, concrete physical detail, moral "
    "self-examination, and direct address of feeling without modern idiom. "
    "Continue the passage exactly where it stops. Do not summarise, do not "
    "comment, do not use headings or lists — write only the continuation prose."
)


def post(path: str, payload: dict, key: str) -> dict:
    req = urllib.request.Request(
        GATEWAY + path, data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
    with urllib.request.urlopen(req, timeout=180) as r:
        return json.loads(r.read())


def resolve(alias: str, key: str) -> str | None:
    req = urllib.request.Request(GATEWAY + "/v1/model/info",
                                 headers={"Authorization": f"Bearer {key}"})
    with urllib.request.urlopen(req, timeout=30) as r:
        for m in json.loads(r.read()).get("data", []):
            if m.get("model_name") == alias:
                return (m.get("litellm_params") or {}).get("model")
    return None


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompts", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--alias", default="gen")
    ap.add_argument("--max-new-tokens", type=int, default=400)
    ap.add_argument("--temperature", type=float, default=0.9)
    ap.add_argument("--top-p", type=float, default=0.95)
    a = ap.parse_args()

    key = os.environ.get("LITELLM_KEY") or Path(
        os.path.expanduser("~/.config/litellm/infra-ops-key")).read_text().strip()

    resolved_start = resolve(a.alias, key)
    print(f"[arm] alias {a.alias!r} resolved at START -> {resolved_start}", flush=True)
    if not resolved_start:
        raise SystemExit(f"REFUSING: alias {a.alias!r} does not resolve; refusing to record an alias as provenance")

    prompts = [json.loads(l) for l in Path(a.prompts).read_text(encoding="utf-8").splitlines()]
    out = Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    with out.open("w", encoding="utf-8") as fh:
        for i, p in enumerate(prompts):
            r = post("/v1/chat/completions", {
                "model": a.alias,
                "messages": [{"role": "system", "content": STYLE_SYSTEM},
                             {"role": "user", "content": p["prompt"]}],
                "max_tokens": a.max_new_tokens, "temperature": a.temperature,
                "top_p": a.top_p}, key)
            cont = r["choices"][0]["message"]["content"]
            fh.write(json.dumps({
                "arm": "incumbent-style-prompted",
                "work": p["work"], "copy": p["copy"], "chapter": p["chapter"],
                "prompt": p["prompt"], "prompt_tokens": p["prompt_tokens"],
                "continuation": cont,
                "completion_tokens": (r.get("usage") or {}).get("completion_tokens"),
                # ⚠ The gateway echoes the ALIAS here, not the concrete model. Keep
                # it labelled as the alias and stamp the resolved model beside it,
                # so a row read on its own cannot record an alias as provenance.
                "alias_echoed_by_gateway": r.get("model"),
                "backing_model_resolved": resolved_start,
                "backing_model_resolved_date": time.strftime("%Y-%m-%d"),
                "sampler": {"temperature": a.temperature, "top_p": a.top_p,
                            "max_new_tokens": a.max_new_tokens},
                "style_system_prompt": STYLE_SYSTEM,
            }, ensure_ascii=False) + "\n")
            if (i + 1) % 6 == 0:
                print(f"[arm] {i+1}/{len(prompts)}  {time.time()-t0:.0f}s", flush=True)

    resolved_end = resolve(a.alias, key)
    meta = {"alias": a.alias, "resolved_at_start": resolved_start, "resolved_at_end": resolved_end,
            "stable_across_run": resolved_start == resolved_end,
            "resolved_date": time.strftime("%Y-%m-%d"),
            "harness": {"path": "LiteLLM gateway -> vLLM seat ana-ml2:8015",
                        "note": "NOT harness-matched to the gx10 local-transformers arms; "
                                "instruct model receiving a style instruction vs base models receiving none"},
            "records": len(prompts)}
    Path(str(out) + ".meta.json").write_text(json.dumps(meta, indent=2))
    print(f"[arm] resolved at END -> {resolved_end}  stable={resolved_start == resolved_end}", flush=True)
    print(f"[arm] -> {out} in {time.time()-t0:.0f}s", flush=True)
    if resolved_start != resolved_end:
        print("[arm] ⚠ THE ALIAS MOVED MID-RUN -- this arm's provenance is split", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())