typed_layer_runner.py

"""Task 37 runner — typed-layer experiment, operator-controlled (Aaron's gates).

Modes:
  sample            extract 20 preference-haystack sessions, print the facts (eyeball gate)
  ingest            extract haystack sessions -> bench_data/typed_facts.jsonl, checkpointed;
                    priority order: preference -> temporal -> multi-session -> rest
  status            progress: sessions done/total, facts minted, ETA at observed rate

Operator control (hard gates):
  - STOP: the existence of bench_data/EXPERIMENT_STOP stops the run within one session
    (~seconds). experiment-stop.bat creates it; experiment-start.bat clears it.
  - WINDOW: --until HH:MM (default 06:00; Aaron 2026-07-20: start 20:00, stop 06:00) —
    the runner self-terminates at the window end unconditionally, even mid-run.
    Crossing midnight is handled. Primary operator interface: Aaron tells Claude to
    start/stop; the .bat files are the manual backup.
  - CHECKPOINT: extraction is keyed by session-content sha1; identical sessions shared
    between questions extract once; stop/resume repeats zero work.

Zero cloud. Local qwen2.5:7b only. Live new_brain untouched — everything in bench_data/.
"""

from __future__ import annotations

import hashlib
import json
import re
import sys
import time
import urllib.request
from datetime import datetime, timedelta
from pathlib import Path

HERE = Path(__file__).resolve().parent
DATA = HERE / "bench_data" / "longmemeval_s_cleaned.json"
FACTS = HERE / "bench_data" / "typed_facts.jsonl"
STOP = HERE / "bench_data" / "EXPERIMENT_STOP"
MODEL = "qwen2.5:7b"

PRIORITY = ["single-session-preference", "temporal-reasoning", "multi-session"]

EXTRACT = """You are building a long-term memory from ONE chat session between a user and an
assistant, dated {date}. Extract the USER's persistent facts as JSON, one per line:

{{"kind": "PREFERENCE|STATE|EVENT", "statement": "<one self-contained sentence>"}}

Rules:
- PREFERENCE: what the user likes, dislikes, enjoys, wants, or habitually does.
- STATE: what the user owns, uses, or what is currently true about them (include brands,
  models, names — "the user owns a Sony A7R IV camera").
- EVENT: something the user did or that happened, WITH its date if inferable.
- Facts must be about the USER, not the assistant's advice or suggestions.
- Self-contained and concrete: a reader with no context must understand each statement.
- 0 to 8 facts per session; quality over coverage. If none, output exactly: NONE

SESSION:
{session}"""


def _qwen(prompt: str) -> str:
    payload = json.dumps({"model": MODEL, "prompt": prompt, "stream": False,
                          "options": {"temperature": 0, "num_predict": 400}}).encode()
    req = urllib.request.Request("http://localhost:11434/api/generate", data=payload,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=300) as r:
        return json.loads(r.read())["response"]


def _session_text(sess: list[dict], cap: int = 8000) -> str:
    lines = [f"{t.get('role', '?')}: {(t.get('content') or '').strip()}"
             for t in sess if (t.get("content") or "").strip()]
    return "\n".join(lines)[:cap]


def _parse_facts(raw: str) -> list[dict]:
    out = []
    for line in raw.splitlines():
        line = line.strip().strip("`")
        if not line or line.upper() == "NONE":
            continue
        m = re.search(r"\{.*\}", line)
        if not m:
            continue
        try:
            f = json.loads(m.group())
        except json.JSONDecodeError:
            continue
        if (f.get("kind") in ("PREFERENCE", "STATE", "EVENT")
                and isinstance(f.get("statement"), str) and len(f["statement"]) > 15):
            out.append({"kind": f["kind"], "statement": f["statement"][:400]})
    return out[:8]


def _work_units() -> list[dict]:
    """One unit per unique session CONTENT (sha-keyed), with all its provenance uses.
    Ordered by the priority of the neediest question that uses it."""
    questions = json.loads(DATA.read_text(encoding="utf-8"))
    units: dict = {}
    for q in questions:
        if "_abs" in q["question_id"]:
            continue
        qt = q.get("question_type", "?")
        prio = PRIORITY.index(qt) if qt in PRIORITY else len(PRIORITY)
        dates = q.get("haystack_dates") or []
        for si, (sid, sess) in enumerate(zip(q["haystack_session_ids"], q["haystack_sessions"])):
            text = _session_text(sess)
            if len(text) < 40:
                continue
            skey = hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest()
            u = units.setdefault(skey, {"skey": skey, "text": text, "prio": prio,
                                        "date": str(dates[si]) if si < len(dates) else "",
                                        "uses": []})
            u["prio"] = min(u["prio"], prio)
            u["uses"].append({"qid": q["question_id"], "sid": sid})
    return sorted(units.values(), key=lambda u: u["prio"])


def _done_keys() -> set:
    if not FACTS.exists():
        return set()
    return {json.loads(l)["skey"] for l in FACTS.read_text(encoding="utf-8").splitlines() if l.strip()}


def _deadline(until: str) -> datetime:
    now = datetime.now()
    h, m = (int(x) for x in until.split(":"))
    dl = now.replace(hour=h, minute=m, second=0, microsecond=0)
    return dl if dl > now else dl + timedelta(days=1)


def cmd_sample() -> int:
    units = [u for u in _work_units() if u["prio"] == 0][:20]
    print(f"extraction-quality sample: {len(units)} preference-haystack sessions\n")
    for i, u in enumerate(units):
        facts = _parse_facts(_qwen(EXTRACT.format(date=u["date"] or "unknown", session=u["text"])))
        print(f"--- session {i + 1} (dated {u['date'] or '?'}):")
        if not facts:
            print("    (no facts extracted)")
        for f in facts:
            print(f"    [{f['kind']:10}] {f['statement']}")
    print("\nEYEBALL GATE: facts must be concrete, user-owned, self-contained. "
          "If they read like vague summaries or assistant advice, iterate the prompt "
          "BEFORE any night burn.")
    return 0


def cmd_ingest(until: str) -> int:
    dl = _deadline(until)
    units = _work_units()
    done = _done_keys()
    todo = [u for u in units if u["skey"] not in done]
    print(f"[{datetime.now():%H:%M:%S}] ingest: {len(done)}/{len(units)} done, "
          f"{len(todo)} to go, window ends {dl:%H:%M}", flush=True)
    t0, n0 = time.time(), 0
    with FACTS.open("a", encoding="utf-8") as out:
        for u in todo:
            if STOP.exists():
                print(f"[{datetime.now():%H:%M:%S}] STOP marker — clean exit "
                      f"({n0} sessions this run)", flush=True)
                return 0
            if datetime.now() >= dl:
                print(f"[{datetime.now():%H:%M:%S}] window end — clean exit "
                      f"({n0} sessions this run)", flush=True)
                return 0
            try:
                facts = _parse_facts(_qwen(EXTRACT.format(date=u["date"] or "unknown",
                                                          session=u["text"])))
            except Exception as e:
                print(f"  ! extraction error ({e}) — session skipped, will retry next run",
                      flush=True)
                time.sleep(5)
                continue
            out.write(json.dumps({"skey": u["skey"], "date": u["date"], "prio": u["prio"],
                                  "uses": u["uses"], "facts": facts}) + "\n")
            out.flush()
            n0 += 1
            if n0 % 50 == 0:
                rate = n0 / max(time.time() - t0, 1)
                print(f"  {len(done) + n0}/{len(units)} ({rate * 3600:.0f}/hr)", flush=True)
    print(f"[{datetime.now():%H:%M:%S}] ALL WORK COMPLETE ({n0} sessions this run)", flush=True)
    return 0


def cmd_status() -> int:
    units = _work_units()
    done = _done_keys()
    n_facts = sum(1 for l in FACTS.read_text(encoding="utf-8").splitlines()
                  for _ in json.loads(l).get("facts", [])) if FACTS.exists() else 0
    by_prio: dict = {}
    for u in units:
        p = u["prio"]
        by_prio.setdefault(p, [0, 0])
        by_prio[p][1] += 1
        if u["skey"] in done:
            by_prio[p][0] += 1
    names = {0: "preference", 1: "temporal", 2: "multi-session", 3: "rest"}
    print(f"typed-layer ingest: {len(done)}/{len(units)} sessions, {n_facts} facts minted")
    for p in sorted(by_prio):
        d, t = by_prio[p]
        print(f"  {names.get(p, p):13} {d}/{t}")
    print(f"stop marker: {'PRESENT (stopped)' if STOP.exists() else 'absent (clear to run)'}")
    return 0


def main() -> int:
    mode = sys.argv[1] if len(sys.argv) > 1 else "status"
    until = sys.argv[sys.argv.index("--until") + 1] if "--until" in sys.argv else "06:00"
    if mode == "sample":
        return cmd_sample()
    if mode == "ingest":
        return cmd_ingest(until)
    return cmd_status()


if __name__ == "__main__":
    raise SystemExit(main())