static/js/views/interactive.js

// Interactive view: step-through generation. Session id lives here; the
// tracer lives server-side. Every action re-renders from the canonical
// payload (status, text, candidates, preview, step count).

import { call, stream } from '../api.js';
import { renderPlot, clearPlot } from '../plot.js';
import { createPromptEditor, createSamplingControls } from '../prompt-editor.js';
import {
  button, checkbox, details, el, escapeHTML, field, numberInput, outputText,
  saveJSON, segmented, statusPane,
} from '../ui.js';

export function createInteractiveView() {
  let sessionId = null;
  let running = false;

  const editor = createPromptEditor({
    promptDefault: 'The future of artificial intelligence is',
  });
  const sampling = createSamplingControls();

  const logTopK = numberInput({ value: 10, min: 1 });
  const logFullProbs = checkbox('Log full probabilities', {
    hint: 'Whole-vocabulary distribution per step (large exports)',
  });
  const stopAtEos = checkbox('Stop at EOS', { checked: true });

  const initBtn = button('Initialize', { variant: 'primary' });
  const resetBtn = button('Reset');

  const status = statusPane({ lines: 2 });
  const output = outputText({ rows: 6, label: 'Output' });

  const continueTokens = numberInput({ value: 10, min: 1, max: 1000 });
  const runBtn = button('Run');
  const stopBtn = button('Stop', { variant: 'danger' });
  stopBtn.disabled = true;

  // --------------------------------------------------- candidates + preview
  const candidatesBody = el('tbody');
  const candidatesTable = el('div', { class: 'table-scroll candidates-scroll' },
    el('table', { class: 'data-table' },
      el('thead', {}, el('tr', {},
        el('th', {}, ''), el('th', {}, 'Rank'), el('th', {}, 'Token ID'),
        el('th', {}, 'Token'), el('th', {}, 'p'), el('th', {}, 'p (raw)'))),
      candidatesBody));

  const useOverride = checkbox('Token ID override', {
    hint: 'Specify an arbitrary token ID instead of selecting from the list',
  });
  const overrideId = numberInput({ value: 0, min: 0 });
  const overrideField = field('Token ID', overrideId);
  overrideField.hidden = true;
  useOverride.input.addEventListener('change', () => {
    overrideField.hidden = !useOverride.input.checked;
    candidatesTable.classList.toggle('is-disabled', useOverride.input.checked);
  });

  let selectedTokenId = null;
  function renderCandidates(candidates, previewId) {
    selectedTokenId = previewId;
    candidatesBody.replaceChildren(...candidates.map((c) => {
      const radio = el('input', { type: 'radio', name: 'jc-candidate', value: String(c.token_id) });
      radio.checked = c.token_id === previewId;
      radio.addEventListener('change', () => { selectedTokenId = c.token_id; });
      return el('tr', { class: c.token_id === previewId ? 'is-preview' : '' },
        el('td', {}, radio),
        el('td', {}, String(c.rank)),
        el('td', { class: 'font-mono' }, String(c.token_id)),
        el('td', { class: 'font-mono token-cell' }, c.text),
        el('td', { class: 'font-mono' }, c.prob.toFixed(4)),
        el('td', { class: 'font-mono' }, c.raw_prob.toFixed(4)),
      );
    }));
  }

  const undoBtn = button('Step back (undo)');
  const stepBtn = button('Next step ▸', { variant: 'primary' });

  // -------------------------------------------------------- lens panel
  const lensMode = segmented([
    { value: 'logit', label: 'Logit' },
    { value: 'jacobian', label: 'Jacobian' },
    { value: 'diff', label: 'Diff (J − logit)' },
  ], { value: 'logit', small: true });
  const lensStride = numberInput({ value: 1, min: 1 });
  const lensTopK = numberInput({ value: 50, min: 1 });
  const lensRefresh = button('Refresh lens', { size: 'sm' });
  const lensApplyIv = button('Apply Lens-view interventions to this session', { size: 'sm' });
  const lensStatus = statusPane({ lines: 1 });
  const lensPlot = el('div', { class: 'plot-host' });

  const lensPanel = details('Layer lens (current position) — experimental', el('div', { class: 'stack' },
    el('p', { class: 'field-hint' },
      'Per-layer readout of the next-token position via the logit or Jacobian '
      + 'lens. Refresh runs one extra forward pass. ⚠ Experimental — readouts '
      + 'may currently yield nonsense.'),
    el('div', { class: 'row-wrap' },
      field('Lens', lensMode), field('Layer stride', lensStride),
      field('Readouts per layer', lensTopK)),
    el('div', { class: 'btn-row' }, lensRefresh, lensApplyIv),
    lensStatus, lensPlot,
  ));

  // -------------------------------------------------- navigation + export
  const gotoStep = numberInput({ value: 0, min: 0 });
  const stepDisplay = el('span', { class: 'pill-value font-mono' }, '0');
  const gotoBtn = button('Go to step');
  const exportBtn = button('Download JSON');
  exportBtn.disabled = true;

  const root = el('div', { class: 'view-grid' },
    el('section', { class: 'card' },
      el('h2', { class: 'card-title' }, 'Prompt'),
      editor.root,
      details('Sampling', sampling.root, { open: true }),
      details('Advanced (logging & stopping)', el('div', { class: 'row-wrap' },
        field('Log top-K tokens', logTopK), logFullProbs, stopAtEos)),
      el('div', { class: 'btn-row' }, initBtn, resetBtn),
      status,
      output,
      el('div', { class: 'row-wrap row-end' },
        field('Continue for N tokens', continueTokens),
        el('div', { class: 'btn-row' }, runBtn, stopBtn)),
    ),
    el('section', { class: 'card' },
      el('h2', { class: 'card-title' }, 'Token selection'),
      el('p', { class: 'card-sub' },
        'The highlighted row is exactly the token Next step will commit — '
        + 'for sampling, the draw happens at preview time.'),
      candidatesTable,
      el('div', { class: 'row-wrap' }, useOverride, overrideField),
      el('div', { class: 'btn-row' }, undoBtn, stepBtn),
      lensPanel,
      el('h2', { class: 'card-title', style: 'margin-top:1.25rem' }, 'Navigation & export'),
      el('div', { class: 'row-wrap row-end' },
        field('Target step (0 = initial)', gotoStep),
        el('div', { class: 'pill' }, el('span', { class: 'pill-label' }, 'Current step'), stepDisplay),
        el('div', { class: 'btn-row' }, gotoBtn, exportBtn)),
    ),
  );

  // ----------------------------------------------------------------- state
  function samplingArgs() {
    return [sampling.strategy(), sampling.temperature(), sampling.topK(), sampling.topP()];
  }

  function apply(payload) {
    if (!payload.ok) {
      status.set(payload.error, 'error');
      return;
    }
    status.set(payload.status);
    output.set(payload.text);
    sessionId = payload.session_id;
    stepDisplay.textContent = String(payload.step ?? 0);
    exportBtn.disabled = !sessionId;
    if (payload.eos) {
      renderCandidates([], null);
      candidatesBody.replaceChildren(el('tr', {},
        el('td', { colspan: 6, class: 'eos-cell' }, 'EOS — generation complete')));
    } else if (payload.candidates) {
      renderCandidates(payload.candidates, payload.preview_id);
    }
  }

  initBtn.addEventListener('click', async () => {
    status.set('Initializing…', 'busy');
    apply(await call('interactive_init', [
      ...editor.values(), ...samplingArgs(), num(logTopK),
    ]).catch(errPayload));
  });

  resetBtn.addEventListener('click', async () => {
    apply(await call('interactive_reset', [sessionId]).catch(errPayload));
    candidatesBody.replaceChildren();
    clearPlot(lensPlot);
    lensStatus.set('');
  });

  stepBtn.addEventListener('click', async () => {
    status.set('Stepping…', 'busy');
    apply(await call('interactive_step', [
      sessionId, ...samplingArgs(),
      selectedTokenId, useOverride.input.checked, num(overrideId),
      num(logTopK), logFullProbs.input.checked, stopAtEos.input.checked,
    ]).catch(errPayload));
  });

  undoBtn.addEventListener('click', async () => {
    apply(await call('interactive_undo', [
      sessionId, ...samplingArgs(), num(logTopK),
    ]).catch(errPayload));
  });

  gotoBtn.addEventListener('click', async () => {
    apply(await call('interactive_goto', [
      sessionId, num(gotoStep), ...samplingArgs(), num(logTopK),
    ]).catch(errPayload));
  });

  runBtn.addEventListener('click', async () => {
    if (running) return;
    running = true;
    runBtn.disabled = true; stepBtn.disabled = true; stopBtn.disabled = false;
    try {
      const final = await stream('interactive_continue', [
        sessionId, ...samplingArgs(), num(continueTokens),
        num(logTopK), logFullProbs.input.checked, stopAtEos.input.checked,
      ], (chunk) => {
        if (chunk.type === 'progress') {
          status.set(chunk.status, 'busy');
          output.set(chunk.text);
          stepDisplay.textContent = String(chunk.step);
        }
      });
      apply(final);
    } catch (e) {
      status.set(`Error: ${escapeHTML(e.message)}`, 'error');
    } finally {
      running = false;
      runBtn.disabled = false; stepBtn.disabled = false; stopBtn.disabled = true;
    }
  });

  stopBtn.addEventListener('click', async () => {
    stopBtn.disabled = true;
    await call('interactive_stop', [sessionId]).catch(() => {});
  });

  exportBtn.addEventListener('click', async () => {
    const res = await call('interactive_export', [sessionId, ...samplingArgs()])
      .catch(errPayload);
    if (!res.ok) { status.set(res.error, 'error'); return; }
    saveJSON(res.export, 'jacobina-interactive.json');
  });

  lensRefresh.addEventListener('click', async () => {
    lensStatus.set('Computing…', 'busy');
    const res = await call('interactive_lens', [
      sessionId, lensMode.value(), num(lensStride), num(lensTopK),
    ]).catch(errPayload);
    if (!res.ok) { lensStatus.set(res.error, 'error'); clearPlot(lensPlot); return; }
    lensStatus.set(res.status);
    renderPlot(lensPlot, res.figure);
  });

  lensApplyIv.addEventListener('click', async () => {
    const res = await call('interactive_apply_interventions', [sessionId])
      .catch(errPayload);
    lensStatus.set(res.ok ? res.status : res.error, res.ok ? '' : 'error');
  });

  return { root };
}

function num(input) {
  const v = parseFloat(input.value);
  return Number.isNaN(v) ? null : v;
}

function errPayload(e) {
  return { ok: false, error: e.message || String(e) };
}