space_app.py

"""Hugging Face Space entrypoint -- a Gradio front end over the SAME memory core.

Why this file exists: HF Spaces no longer offers free Docker or CPU Gradio hosting. The one
free path for a personal account is a ZeroGPU Space, and ZeroGPU is Gradio-SDK-only. So the
FastAPI dashboard in ../app.py cannot run there, and this is its Gradio equivalent.

NOTHING about the memory controller changes here. This module imports `Session` from app.py
and drives the identical MemorySidecar + Agent. The write path is still local and still costs
zero LLM tokens; only the presentation layer differs.

No @spaces.GPU decorator anywhere, deliberately: the models are tiny (a 2.8M-param neural
memory, a 33M embedder, a 22M NLI cross-encoder) and run fine on CPU. GPU quota is only
consumed inside decorated functions, so this Space uses none of its daily allowance.
"""
from __future__ import annotations

# MUST BE THE FIRST REAL IMPORT, ahead of anything that touches torch.
# On ZeroGPU hardware Gradio imports the `spaces` package in a watcher thread, and that
# package hard-errors with "CUDA has been initialized before importing the `spaces`
# package" if torch got there first. We never call @spaces.GPU -- importing it early is
# purely to satisfy that ordering rule. Absent locally, hence the guard.
try:
    import spaces  # noqa: F401
except ImportError:
    pass

import html
import tempfile
import uuid
from pathlib import Path

import gradio as gr

from app import PRESETS, Session  # the exact same core the judged build used

# ---------------------------------------------------------------- palette (matches the
# original dashboard's "warm phosphor" instrument look, carried over deliberately)
FIELD, PANEL, PANEL2 = "#0c0a07", "#17120b", "#201810"
INK, INK2, MUTED = "#f8f4ec", "#c7bda8", "#8b8069"
KEPT, SKIP, SUPERSEDED = "#3ac95f", "#8b8069", "#e0a63a"
COST, SIGNAL, COOL = "#e2503a", "#f4a723", "#5a9bd4"

CSS = f"""
.gradio-container {{ background: {FIELD} !important; max-width: 1280px !important; }}
#mast {{ border-bottom: 1px solid rgba(248,244,236,.09); padding-bottom: 14px; margin-bottom: 6px; }}
.fb-mono {{ font-family: "SF Mono","Cascadia Code",ui-monospace,Menlo,Consolas,monospace; }}
footer {{ display: none !important; }}
"""


def _esc(s) -> str:
    return html.escape(str(s))


# ---------------------------------------------------------------- render helpers
def _header(st: dict) -> str:
    live = SUPERSEDED if st.get("mock") else KEPT
    label = _esc(st.get("llm_label", "linking..."))
    return f"""
<div id="mast" class="fb-mono" style="display:flex;align-items:center;gap:16px;flex-wrap:wrap;color:{INK}">
  <span style="font-weight:700;font-size:19px;letter-spacing:.02em">
    flashbulb<span style="color:{SIGNAL}">&middot;</span>memory
    <span style="color:{MUTED};font-weight:500"> // zero-token write gate</span>
  </span>
  <span style="color:{INK2};font-size:12.5px;border-left:1px solid rgba(248,244,236,.14);padding-left:16px;
               font-family:system-ui,-apple-system,sans-serif">
    drift-adaptive agent memory that decides what to keep, forget, and doubt for zero LLM tokens
  </span>
  <span style="margin-left:auto;display:flex;align-items:center;gap:9px;font-size:11.5px;color:{INK2};
               border:1px solid rgba(248,244,236,.075);border-radius:7px;padding:6px 12px;background:{PANEL}">
    <span style="width:7px;height:7px;border-radius:50%;background:{live};box-shadow:0 0 8px {live}"></span>{label}
  </span>
</div>"""


def _channels(st: dict) -> str:
    t = st.get("tokens", {})
    ours, base = t.get("surprise_write", 0), t.get("llm_extract_write", 0)
    turns, read = t.get("turns", 0), t.get("read", 0)
    saved = f"{base - ours:,} tokens saved" if base else "send a turn to accrue the baseline"

    def cell(label, body, sub, grow=True):
        return f"""<div style="flex:{'1.5' if grow else '1'};min-width:150px;background:{PANEL};padding:16px 20px">
          <div class="fb-mono" style="font-size:10px;letter-spacing:.13em;text-transform:uppercase;color:{MUTED}">{label}</div>
          <div class="fb-mono" style="font-size:31px;line-height:1.05;margin-top:6px">{body}</div>
          <div style="font-size:11.5px;color:{INK2};margin-top:5px">{sub}</div></div>"""

    ab = (f'<span style="color:{KEPT}">{ours}</span>'
          f'<span style="color:{MUTED};font-size:16px"> vs </span>'
          f'<span style="color:{COST}">{base:,}</span>')
    return f"""<div style="display:flex;flex-wrap:wrap;gap:1px;background:rgba(248,244,236,.14);
        border:1px solid rgba(248,244,236,.14);border-radius:14px;overflow:hidden;margin:14px 0">
      {cell("write tokens &mdash; ours vs llm_extract", ab, saved)}
      {cell("turns", f'<span style="color:{INK}">{turns}</span>', "same conversation, both policies", False)}
      {cell("read tokens", f'<span style="color:{COOL}">{read:,}</span>', "answering only", False)}
    </div>"""


def _signal_svg(st: dict) -> str:
    """The surprise trace against the adaptive threshold -- the gate, made visible."""
    log = st.get("log", [])[-28:]
    w, h, pad = 620, 132, 16
    if not log:
        body = f'<text x="{w/2}" y="{h/2}" fill="{MUTED}" font-size="12" text-anchor="middle">no observations yet</text>'
    else:
        vals = [e["surprise"] for e in log] + [e["threshold"] for e in log]
        lo, hi = min(vals), max(vals)
        span = (hi - lo) or 1.0
        step = (w - 2 * pad) / max(len(log), 1)

        def y(v):
            return h - pad - ((v - lo) / span) * (h - 2 * pad)

        bars, thr = [], []
        for i, e in enumerate(log):
            x = pad + i * step + step * 0.5
            col = {"novel": KEPT, "update": SUPERSEDED}.get(e["reason"], SKIP)
            bars.append(f'<rect x="{x-step*0.3:.1f}" y="{y(e["surprise"]):.1f}" width="{max(step*0.6,2):.1f}" '
                        f'height="{max(h-pad-y(e["surprise"]),1):.1f}" fill="{col}" opacity=".85" rx="1"/>')
            thr.append(f"{x:.1f},{y(e['threshold']):.1f}")
        body = "".join(bars) + f'<polyline points="{" ".join(thr)}" fill="none" stroke="{SIGNAL}" stroke-width="1.6" stroke-dasharray="4 3"/>'
    return f"""<div style="border:1px solid rgba(248,244,236,.14);border-radius:14px;background:{PANEL};overflow:hidden">
      <div style="padding:13px 18px 0">
        <span class="fb-mono" style="font-size:10.5px;letter-spacing:.16em;text-transform:uppercase;color:{SIGNAL}">surprise signal</span>
        <span style="font-size:12px;color:{INK2};margin-left:10px">bar = surprise &middot; dashed = adaptive threshold &middot;
        <span style="color:{KEPT}">novel</span> / <span style="color:{SUPERSEDED}">update</span> / <span style="color:{MUTED}">skipped</span></span>
      </div>
      <svg viewBox="0 0 {w} {h}" style="width:100%;display:block">{body}</svg></div>"""


def _memory_html(st: dict) -> str:
    recs = st.get("records", [])
    if not recs:
        return f'<div style="color:{MUTED};text-align:center;padding:26px 0;font-size:13px">memory is empty</div>'
    rows = []
    for r in recs:
        ret = r["retention"]
        dead = r["superseded_by"] is not None or ret < 0.15
        col = SUPERSEDED if dead else KEPT
        tag = "forgotten" if dead else f"conf {ret:.2f}"
        rows.append(f"""<div style="display:flex;align-items:center;gap:10px;padding:7px 0;
              border-bottom:1px solid rgba(248,244,236,.075)">
            <div style="flex:1;min-width:0;color:{INK2 if dead else INK};font-size:13px;
                 {'text-decoration:line-through;opacity:.6' if dead else ''}">{_esc(r["text"])}</div>
            <div style="width:82px;height:6px;background:rgba(248,244,236,.075);border-radius:99px;overflow:hidden">
              <div style="width:{max(ret,0)*100:.0f}%;height:100%;background:{col}"></div></div>
            <div class="fb-mono" style="font-size:10.5px;color:{col};width:74px;text-align:right">{tag}</div>
          </div>""")
    return "".join(rows)


def _render(sess: Session):
    st = sess.state()
    return _header(st), _channels(st), _signal_svg(st), _memory_html(st)


# ---------------------------------------------------------------- event handlers
def _new_session():
    return Session()


def _turn(msg: str, chat: list, sess: Session):
    msg = (msg or "").strip()
    if not msg:
        return chat, "", *_render(sess)
    chat = chat + [{"role": "user", "content": msg}]
    r = sess.chat(msg)
    for tc in (r.tool_calls or []):
        name = _esc(tc.get("name", "recall_memory"))
        q = _esc(tc.get("query", ""))
        res = _esc(tc.get("result", "") or "(nothing relevant)")
        chat.append({"role": "assistant",
                     "content": f"`{name}(\"{q}\")` → {res}",
                     "metadata": {"title": "Agent - function call"}})
    chat.append({"role": "assistant", "content": r.answer})
    return chat, "", *_render(sess)


def _preset(name: str, chat: list, sess: Session):
    for line in PRESETS[name]:
        chat, _, *_ = _turn(line, chat, sess)
    return chat, "", *_render(sess)


def _reset(sess: Session):
    sess.reset()
    return [], "", *_render(sess)


_SNAPS: dict[str, str] = {}


def _snapshot(sess: Session):
    sid = getattr(sess, "sid", None) or uuid.uuid4().hex
    sess.sid = sid
    path = str(Path(tempfile.gettempdir()) / f"flashbulb_{sid}.pt")
    sess.sidecar.save(path)
    _SNAPS[sid] = path
    n = len(sess.sidecar._active_records())
    return gr.update(value=f"snapshot saved - {n} records"), *_render(sess)


def _restore(sess: Session):
    sid = getattr(sess, "sid", None)
    path = _SNAPS.get(sid or "")
    if not path or not Path(path).exists():
        return gr.update(value="no snapshot yet - click Snapshot first"), *_render(sess)
    sess.restore(path)
    n = len(sess.sidecar._active_records())
    return gr.update(value=f"restored from snapshot - {n} records"), *_render(sess)


# ---------------------------------------------------------------- UI
with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="flashbulb - drift-adaptive agent memory") as demo:
    sess = gr.State(_new_session)

    header = gr.HTML()
    channels = gr.HTML()

    with gr.Row():
        with gr.Column(scale=5):
            chat = gr.Chatbot(type="messages", height=330, show_label=False)
            with gr.Row():
                box = gr.Textbox(placeholder="e.g. I live in Ames, Iowa.", show_label=False,
                                 scale=5, container=False)
                send = gr.Button("Send", variant="primary", scale=1)
            with gr.Row():
                b_dup = gr.Button("Duplicate / poison", size="sm")
                b_upd = gr.Button("Knowledge update", size="sm")
                b_rst = gr.Button("Reset", size="sm")
        with gr.Column(scale=4):
            signal = gr.HTML()
            gr.Markdown("### Memory &mdash; retention, and what got forgotten")
            memory = gr.HTML()
            with gr.Row():
                b_snap = gr.Button("Snapshot", size="sm")
                b_rest = gr.Button("Restore", size="sm")
            boundary = gr.Markdown("*cross-session boundary: Snapshot &rarr; Reset &rarr; Restore*")

    gr.Markdown(
        "The memory decision (what to keep, forget, and doubt) is made **locally for zero LLM tokens** "
        "from two signals: online-TTT surprise and a local NLI contradiction. The model is only used to "
        "answer and to drive recall via function-calling. "
        "[Source and benchmarks](https://github.com/somtri/flashbulb)"
    )

    OUT = [header, channels, signal, memory]
    demo.load(lambda s: _render(s), inputs=sess, outputs=OUT)
    send.click(_turn, [box, chat, sess], [chat, box, *OUT])
    box.submit(_turn, [box, chat, sess], [chat, box, *OUT])
    b_dup.click(lambda c, s: _preset("duplicate", c, s), [chat, sess], [chat, box, *OUT])
    b_upd.click(lambda c, s: _preset("knowledge_update", c, s), [chat, sess], [chat, box, *OUT])
    b_rst.click(_reset, [sess], [chat, box, *OUT])
    b_snap.click(_snapshot, [sess], [boundary, *OUT])
    b_rest.click(_restore, [sess], [boundary, *OUT])

# Pinned to Gradio 5 (see hf/requirements.space.txt). Gradio 6 was tried first and failed on
# Spaces three ways: its SSR Node proxy died on startup, the `spaces` package rejected being
# imported after torch, and launch() then stopped blocking so the process exited immediately.
# Gradio 5 is what nearly every Space runs; css/theme belong on Blocks here, not on launch().
if __name__ == "__main__":
    # ssr_mode=False is REQUIRED on Spaces, on Gradio 5 as well as 6. Both default SSR on,
    # which fronts Python with a Node proxy, and on this Space that proxy dies immediately
    # ("Stopping Node.js server..."). A local run cannot reproduce it: with no Node installed
    # Gradio silently skips SSR, so it looks fine here and dies there. Do not remove this.
    demo.launch(ssr_mode=False)