NPCvector/distill.py

#!/opt/homebrew/bin/python3
"""GABGOD distill — mine the forum corpus for VOICE, write ORIGINAL NPC lines.

For each (persona, topic): embed the topic (nomic on m4pro), pull the top-k
most-similar forum posts from pgvector (style reference only), then ask the
writer LLM for `lines_per_topic` original lines in the persona's voice.
Lines land in gabgod.line_bank (approved = NULL until John reviews).

Resume-safe: a (persona, topic) that already has >= lines_per_topic lines is
skipped. Heartbeat: ~/.jobs/gabgod-distill.status.

Usage:
  ./distill.py                          # everything in personas.json
  ./distill.py --persona record_keeper  # one persona
  ./distill.py --persona record_keeper --topics 2 --dry-run   # smoke test
"""
import argparse, json, re, sys, time
from datetime import datetime
from pathlib import Path

import psycopg2, requests

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "config.json").read_text())
PERSONAS = json.loads((HERE / "personas.json").read_text())


def heartbeat(done, total, note=""):
    Path.home().joinpath(".jobs").mkdir(exist_ok=True)
    Path.home().joinpath(".jobs/gabgod-distill.status").write_text(
        f"{datetime.now():%Y-%m-%d %H:%M:%S} | {done}/{total} | {note}\n")


def writer_cfg():
    w = CFG["writer"]
    return w["abliterix"] if w.get("abliterated") else w["stock"]


def embed(text):
    r = requests.post(f"{CFG['embed']['url']}/api/embeddings",
                      json={"model": CFG["embed"]["model"], "prompt": text}, timeout=60)
    r.raise_for_status()
    return r.json()["embedding"]


def retrieve(cur, vec, k, min_len, max_len):
    """Top-k similar posts across the configured corpus tables."""
    lit = "[" + ",".join(f"{x:.6f}" for x in vec) + "]"
    rows = []
    for table in CFG["retrieval"]["sources"]:
        cur.execute(
            f"""select post_id, author, body_text, '{table.split('_')[0]}' src
                from {table}
                where embedding is not null
                  and length(body_text) between %s and %s
                order by embedding <=> %s::vector
                limit %s""",
            (min_len, max_len, lit, k))
        rows += cur.fetchall()
    # keep global top-k by re-ranking is overkill; interleave and cap
    return rows[: k]


def ask_writer(persona, topic, refs):
    w = writer_cfg()
    ref_block = "\n---\n".join(r[2].strip()[:400] for r in refs[:6])  # >6x400 sends reasoning models into a death spiral
    n = CFG["lines_per_topic"]
    sys_p = (
        "You write dialogue for NPCs in a 90s Australian shopping-town video game. "
        "You are given real forum posts as a VOICE AND ATTITUDE reference. "
        "Never copy, quote, or lightly rephrase the reference text; never use usernames, "
        "real names, URLs, or identifiable stories from it. Write ORIGINAL lines only.")
    user_p = (
        f"CHARACTER: {PERSONAS[persona]['voice']}\n\n"
        f"SITUATION: {topic}\n\n"
        f"VOICE REFERENCE (do not copy from this):\n{ref_block}\n\n"
        f"Write {n} distinct things this character might SAY OUT LOUD in this situation. "
        f"One or two sentences each, spoken register, no stage directions, no numbering baggage. "
        f'Answer as a JSON array of {n} strings and nothing else.')
    last_err = None
    for attempt in range(3):
        # native ollama endpoint: think=False stops reasoning models burning the
        # whole token budget deliberating (supergemma wrote 12k chars of drafts
        # and no answer via /v1). 400 fallback for models without a think switch.
        payload = {"model": w["model"], "stream": False, "think": False,
                   "options": {"temperature": 0.85, "num_predict": 2048},
                   "messages": [{"role": "system", "content": sys_p},
                                {"role": "user", "content": user_p}]}
        r = requests.post(f"{w['url']}/api/chat", timeout=900, json=payload)
        if r.status_code == 400:
            payload.pop("think", None)
            r = requests.post(f"{w['url']}/api/chat", timeout=900, json=payload)
        r.raise_for_status()
        txt = r.json()["message"]["content"]
        try:
            start = txt.find("[")
            if start < 0:
                raise ValueError(f"no JSON array in writer output: {txt[:200]}")
            try:
                arr, _ = json.JSONDecoder().raw_decode(txt[start:])
            except json.JSONDecodeError:
                # salvage: models sometimes break an array element mid-list;
                # harvest the well-formed quoted strings instead of losing all 12
                arr = [json.loads(f'"{m}"') for m in
                       re.findall(r'"((?:[^"\\]|\\.){8,300})"', txt[start:])]
                if not arr:
                    raise
        except (ValueError, json.JSONDecodeError) as e:
            last_err = e
            continue
        lines = [s.strip() for s in arr if isinstance(s, str) and sane_line(s.strip())]
        if lines:
            return lines, w["model"]
        last_err = ValueError("all lines rejected by sanity filter")
    raise last_err


def sane_line(s):
    """Reject model garbage: wrong language, fragments, essays."""
    non_ascii = sum(1 for ch in s if ord(ch) > 127)
    return 8 <= len(s) <= 300 and non_ascii <= 3


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--persona")
    ap.add_argument("--topics", type=int, help="only first N topics (smoke test)")
    ap.add_argument("--dry-run", action="store_true", help="print, don't insert")
    args = ap.parse_args()

    names = [args.persona] if args.persona else list(PERSONAS)
    jobs = [(p, t) for p in names for t in
            (PERSONAS[p]["topics"][: args.topics] if args.topics else PERSONAS[p]["topics"])]

    conn = psycopg2.connect(CFG["db"])
    conn.autocommit = True
    cur = conn.cursor()
    R = CFG["retrieval"]
    target = CFG["lines_per_topic"]

    done = 0
    for persona, topic in jobs:
        heartbeat(done, len(jobs), f"{persona} :: {topic[:50]}")
        cur.execute("""select count(*) from gabgod.line_bank
                       where persona=%s and situation=%s and writer_model=%s""",
                    (persona, topic, writer_cfg()["model"]))
        if cur.fetchone()[0] >= target and not args.dry_run:
            done += 1
            continue
        try:
            refs = retrieve(cur, embed(topic), R["top_k"], R["min_len"], R["max_len"])
            lines, model = ask_writer(persona, topic, refs)
        except Exception as e:
            print(f"FAIL {persona}/{topic[:40]}: {e}", file=sys.stderr)
            heartbeat(done, len(jobs), f"FAIL {persona}: {e}")
            time.sleep(5)
            continue
        ids = [r[0] for r in refs]
        for ln in lines:
            if args.dry_run:
                print(f"  [{persona}] {ln}")
            else:
                cur.execute(
                    """insert into gabgod.line_bank
                       (persona, situation, line, writer_model, source_post_ids)
                       values (%s,%s,%s,%s,%s) on conflict do nothing""",
                    (persona, topic, ln, model, ids))
        done += 1
    heartbeat(done, len(jobs), "DONE")
    print(f"DONE {done}/{len(jobs)} persona-topics")


if __name__ == "__main__":
    main()