open-standards/wai/quantum-studio/src/assistant.js

// The Studio's assistant: a language model on this machine, with the engine
// as its tools. The model is any server with the chat-completions interface,
// tool calls included: the JouleClaw runtime, a llama.cpp server, others. By default
// it is local, and nothing leaves the machine.
//
// Tools first: the model asks the engine (lint, critique, draw, plan, run, the
// equivalence check) and may state only numbers the engine returned. Every program it proposes is checked by the
// engine before the Studio offers it: read as the engine reads it, and run on
// this device. A proposal that does not parse goes back to the model once, with
// the engine's reason, to be corrected.
//
// The reader is the author. The assistant speaks only when asked, and by default
// it asks a question before it proposes (the stance, which the reader sets). For
// Alternatives it is a source of candidates, never the judge: each one it writes
// is read, measured and compared by the engine before the Studio shows it.

const SYSTEM = `You are the assistant in Quantum Studio, an editor for OpenQASM programs.
Tools run the Studio's own engine:
- get_program: the program in the editor, the job's settings and the last result.
- lint_qasm: read a program as the engine reads it (its qubits and gates, or the first problem and its line).
- critique_qasm: the program's measures (gates, two-qubit gates, depth, T gates, whether it is Clifford, a modelled fidelity) and its smells, each with its lines.
- check_equivalent: whether two programs are the same circuit, and how the engine knows.
- classical_shortcuts: every exact classical route to simulating a program and its cost; a cheap route proves it classically easy, none being cheap proves nothing.
- draw_circuit: the circuit as text.
- plan_job: the method that would run a job, why, and the memory it needs.
- run_job: run a job on this device; counts, expectation values, amplitudes and the accuracy of each.
Use them. Never state a probability, count, amplitude, expectation value or cost that a tool did not return; quote the tool's figure.
The engine reads OpenQASM 3 and the OpenQASM 2 forms in common use: registers, the standard gates, gate definitions, ctrl and negctrl modifiers, measurement and reset anywhere, and if on measured bits. It refuses loops, inv and pow.
When you write a program, run it with run_job and also show the whole program in one \`\`\`qasm block; the Studio checks it and runs it before offering it.
For a state, a probability or an amplitude, run the program and quote the result rather than working it out yourself.
Write plain text: no Markdown headings, no LaTeX. Write states as |00⟩ and roots as √2. Qubit 0 is the rightmost bit of a basis state.
Answer briefly.`;

// How the assistant helps, which the reader chooses. Asking first is the default:
// a question the reader answers teaches more than an answer they read.
export const STANCES = {
  questions: { label: 'Ask me questions first', text: 'The user is the author and is learning. Before you propose a program or a change, ask one short question that helps them see it themselves: what they expect a line to do, or what a result shows. Answer a direct question from the tools, then ask one question back that moves them forward. Write a whole program only when they ask you to.' },
  suggest: { label: 'Suggest changes', text: 'The user is the author. Suggest changes as small edits to their program, each naming its line and the reason the tools gave; write a whole program only when they ask for one.' },
  write: { label: 'Write programs', text: 'When a change is wanted, write the whole program as described above.' },
};

export const TOOLS = [
  { type: 'function', function: { name: 'get_program', description: 'The program in the editor, the job settings, and a summary of the last result.', parameters: { type: 'object', properties: {} } } },
  { type: 'function', function: { name: 'lint_qasm', description: 'Read an OpenQASM program as the engine reads it: its qubits, gates and measurements, or the first problem with its line.', parameters: { type: 'object', required: ['qasm'], properties: { qasm: { type: 'string' } } } } },
  { type: 'function', function: { name: 'critique_qasm', description: 'A program\'s measures (gates by kind, two-qubit gates, depth, two-qubit depth, T count, whether it is a Clifford circuit, a fidelity under reference error rates, which is a model and not a run) and its smells: gates that cancel, rotations that merge or do nothing, idle or unmeasured qubits, each with its lines.', parameters: { type: 'object', required: ['qasm'], properties: { qasm: { type: 'string' } } } } },
  { type: 'function', function: { name: 'check_equivalent', description: 'Whether two programs are the same circuit up to a global phase, and how the engine knows: equal tableaus (Clifford, any size), every basis state (to 8 qubits), random inputs (to 22 qubits, evidence rather than proof), or unchecked. b defaults to the program in the editor.', parameters: { type: 'object', required: ['a'], properties: { a: { type: 'string' }, b: { type: 'string' } } } } },
  { type: 'function', function: { name: 'classical_shortcuts', description: 'Every exact classical route the engine knows for simulating a program (stabilizer tableau, cut-rank sum, exact matrix product state at a rigorous bond bound, light cone, tensor network when thorough, state vector), each with its memory and the log2 of its operations, and the cheapest for each question. A cheap route proves the program classically easy; none being cheap proves nothing about hardness.', parameters: { type: 'object', properties: { qasm: { type: 'string' }, thorough: { type: 'boolean' } } } } },
  { type: 'function', function: { name: 'draw_circuit', description: 'A program\'s circuit drawn as text.', parameters: { type: 'object', required: ['qasm'], properties: { qasm: { type: 'string' } } } } },
  {
    type: 'function',
    function: {
      name: 'plan_job',
      description: 'The method a job would run on and why, with the memory it needs, without running it.',
      parameters: { type: 'object', required: ['qasm'], properties: { qasm: { type: 'string' }, shots: { type: 'integer' }, observables: { type: 'array', items: { type: 'string' } }, method: { type: 'string' } } },
    },
  },
  {
    type: 'function',
    function: {
      name: 'run_job',
      description: 'Run a job on this device: measurement counts, Pauli expectation values (as "Z0 Z1") and/or amplitudes (bit strings, highest qubit first), with the method used and its accuracy.',
      parameters: {
        type: 'object',
        required: ['qasm'],
        properties: { qasm: { type: 'string' }, shots: { type: 'integer' }, observables: { type: 'array', items: { type: 'string' } }, amplitudes: { type: 'array', items: { type: 'string' } }, method: { type: 'string' } },
      },
    },
  },
];

const clip = (s, n = 6000) => (s.length > n ? s.slice(0, n) + `… (${s.length - n} more characters)` : s);

// A result as a model reads it: what was computed, not the receipt.
export function summary(job) {
  const r = job?.result || {};
  const counts = r.counts ? Object.entries(r.counts).sort((a, b) => b[1] - a[1]).slice(0, 16) : undefined;
  return {
    method: job?.method, why: job?.why, accuracy: r.accuracy, qubits: r.qubits, gates: r.gates,
    counts: counts && Object.fromEntries(counts), distinct_outcomes: r.distinct_outcomes,
    expectations: r.expectations, amplitudes: r.amplitudes?.slice(0, 16), dynamic: r.dynamic, warning: r.warning,
  };
}

// The ```qasm (or ```openqasm) blocks of an answer.
export function proposals(text) {
  const out = [];
  const re = /```(?:qasm|openqasm|openqasm3)?\s*\n([\s\S]*?)```/gi;
  for (let m; (m = re.exec(text)); ) if (/\b(qubit|qreg|OPENQASM)\b/.test(m[1])) out.push(m[1].trim() + '\n');
  return out;
}

// A model behind a chat-completions endpoint, called from this page.
export class HttpModel {
  constructor(endpoint, name) { this.endpoint = endpoint.replace(/\/+$/, ''); this.name = name; }
  async chat(body) {
    const res = await fetch(`${this.endpoint}/chat/completions`, { method: 'POST', headers: { 'content-type': 'application/json' }, body: JSON.stringify(body) });
    const v = await res.json().catch(() => ({}));
    if (!res.ok) throw new Error(v.error?.message || v.error || `the model answered ${res.status}`);
    return v;
  }
  async models() {
    const res = await fetch(`${this.endpoint}/models`);
    if (!res.ok) throw new Error(`the model server answered ${res.status}`);
    return (await res.json()).data?.map((m) => m.id) ?? [];
  }
}

// The same, called by the desktop app (no browser origin rules in the way).
export class DesktopModel {
  constructor(endpoint, name) { this.endpoint = endpoint.replace(/\/+$/, ''); this.name = name; this.invoke = window.__TAURI__.core.invoke; }
  async chat(body) {
    const v = JSON.parse(await this.invoke('chat', { endpoint: this.endpoint, body: JSON.stringify(body) }));
    if (v.error && !v.choices) throw new Error(v.error.message || v.error);
    return v;
  }
  async models() {
    const v = JSON.parse(await this.invoke('models', { endpoint: this.endpoint }));
    if (v.error && !v.data) throw new Error(v.error.message || v.error);
    return v.data?.map((m) => m.id) ?? [];
  }
}

export class Assistant {
  // engine: the Studio's local engine; context(): {qasm, job, last} as the editor holds them;
  // stance(): a key of STANCES.
  constructor({ engine, model, context, stance = () => 'questions' }) {
    this.engine = engine;
    this.model = model;
    this.context = context;
    this.stance = stance;
    this.history = [];
  }

  // Candidates for Alternatives, from one call without tools and outside the
  // conversation: [{idea, qasm}]. The engine checks each; nothing here does.
  async muse(prompt, temperature) {
    const res = await this.model.chat({ model: this.model.name || undefined, temperature, messages: [{ role: 'system', content: MUSE }, { role: 'user', content: prompt }] });
    return ideas(unlatex((res.choices?.[0]?.message?.content || '').replace(/<think>[\s\S]*?<\/think>\s*/g, '')));
  }
  improve(qasm, k = 3) {
    return this.muse(`Write ${k} programs that are the same circuit as this one (the same unitary up to a global phase, and the same measurements) but cheaper: fewer two-qubit gates, less depth or fewer T gates. Before each, one line that starts "Idea:" and says what it changes. Each program whole, in its own \`\`\`qasm block.\n\n\`\`\`qasm\n${qasm}\n\`\`\``, 0.4);
  }
  explore(qasm, aim, k = 4) {
    return this.muse(`Imagine ${k} different programs in the spirit of this one${aim ? `, toward this aim: ${aim}` : ''}. Make them differ from each other in kind, not only in their angles: another structure, another entangler, another symmetry, another use of the qubits. Before each, one line that starts "Idea:" and says what it tries. Each program whole, in its own \`\`\`qasm block.\n\n\`\`\`qasm\n${qasm}\n\`\`\``, 0.9);
  }

  async tool(name, args) {
    const ctx = this.context();
    const job = (a) => {
      const j = { qasm: a.qasm };
      for (const k of ['shots', 'observables', 'amplitudes', 'method']) if (a[k] != null && !(Array.isArray(a[k]) && !a[k].length)) j[k] = a[k];
      if (!j.shots && !j.observables && !j.amplitudes) j.shots = 1000;
      return j;
    };
    switch (name) {
      case 'get_program': return { qasm: ctx.qasm, job: ctx.job, last_result: ctx.last ? summary(ctx.last) : null };
      case 'lint_qasm': return this.engine.lint(args.qasm ?? ctx.qasm);
      case 'critique_qasm': return this.engine.critique(args.qasm ?? ctx.qasm);
      case 'check_equivalent': return this.engine.equivalent(args.a ?? ctx.qasm, args.b ?? ctx.qasm);
      case 'classical_shortcuts': return this.engine.shortcuts(args.qasm ?? ctx.qasm, !!args.thorough);
      case 'draw_circuit': return { text: (await this.engine.draw(args.qasm ?? ctx.qasm, 80)).text };
      case 'plan_job': return this.engine.plan(job({ ...args, qasm: args.qasm ?? ctx.qasm }));
      case 'run_job': return summary(await this.engine.run(job({ ...args, qasm: args.qasm ?? ctx.qasm })));
      default: throw new Error(`no tool ${name}`);
    }
  }

  // Ask; `on(event)` hears each step: {type: 'tool', name, args, result} and
  // finally {type: 'answer', text, proposals: [{qasm, lint, job}]}.
  async ask(text, on = () => {}) {
    const ctx = this.context();
    const used = []; // programs the model passed to the tools, other than the editor's
    const ranWith = new Map(); // program → the settings the model ran it with
    const heard = [ctx.qasm, JSON.stringify(ctx.job ?? {}), JSON.stringify(ctx.last ? summary(ctx.last) : {})]; // what the engine and the editor told it
    const messages = [
      { role: 'system', content: `${SYSTEM}\n${(STANCES[this.stance()] || STANCES.questions).text}` },
      ...this.history,
      { role: 'user', content: `${text}\n\n(The program in the editor:)\n\`\`\`qasm\n${ctx.qasm}\n\`\`\`` },
    ];
    let retried = false;
    for (let round = 0; round < 8; round++) {
      // The first turn must ask the engine something, so every answer starts from it.
      const res = await this.model.chat({ model: this.model.name || undefined, messages, tools: TOOLS, tool_choice: round === 0 ? 'required' : 'auto', temperature: 0.2 });
      const msg = res.choices?.[0]?.message;
      if (!msg) throw new Error('the model gave no answer');
      messages.push({ role: 'assistant', content: msg.content ?? '', ...(msg.tool_calls?.length ? { tool_calls: msg.tool_calls } : {}) });
      if (msg.tool_calls?.length) {
        for (const call of msg.tool_calls) {
          let args = {};
          try { args = JSON.parse(call.function?.arguments || '{}'); } catch (e) {}
          let result;
          try { result = await this.tool(call.function?.name, args); } catch (e) { result = { error: e.message }; }
          if (typeof args.qasm === 'string' && args.qasm.trim() && args.qasm.trim() !== ctx.qasm.trim()) used.push(args.qasm.trim() + '\n');
          if (call.function?.name === 'run_job') {
            const eff = pick(args, ['shots', 'observables', 'amplitudes', 'method']);
            if (!eff.shots && !eff.observables && !eff.amplitudes) eff.shots = 1000; // as the tool ran it
            ranWith.set((args.qasm ?? ctx.qasm).trim(), eff);
          }
          on({ type: 'tool', name: call.function?.name, args, result });
          messages.push({ role: 'tool', tool_call_id: call.id, content: clip(JSON.stringify(result)) });
          heard.push(JSON.stringify(args), JSON.stringify(result)); // what it asked, and what it was told
        }
        continue;
      }
      const answer = unlatex((msg.content || '').replace(/<think>[\s\S]*?<\/think>\s*/g, '').trim());
      // Check each proposal as the engine reads it; a broken one goes back once.
      // The answer's programs, and the last one the model ran or checked through a tool
      // when the answer shows none: either way, read by the engine here.
      const written = proposals(answer);
      const shown = written.length ? written : used.slice(-1);
      const checked = [];
      for (const qasm of shown) checked.push({ qasm, lint: await this.engine.lint(qasm) });
      const broken = checked.find((p) => !p.lint.ok);
      if (broken && !retried) {
        retried = true;
        on({ type: 'retry', error: broken.lint.error });
        messages.push({ role: 'user', content: `The engine could not read the program you proposed: ${broken.lint.error}. Give the corrected whole program in one \`\`\`qasm block.` });
        continue;
      }
      for (const p of checked) {
        if (!p.lint.ok) continue;
        // Run as the model ran it, so the counts it quoted are the counts shown; else
        // with the editor's settings.
        const theirs = ranWith.get(p.qasm.trim());
        const settings = theirs ?? { ...(ctx.job?.shots ? {} : { shots: 1000 }), ...pick(ctx.job, ['shots', 'observables']) };
        try { p.job = await this.engine.run({ qasm: p.qasm, ...settings }); } catch (e) { p.error = e.message; }
      }
      this.history.push({ role: 'user', content: text }, { role: 'assistant', content: answer });
      this.history = this.history.slice(-12);
      on({ type: 'answer', text: answer, proposals: checked, usage: res.usage, unsourced: unsourced(answer, heard.join(' ')) });
      return;
    }
    throw new Error('the model kept calling tools without answering');
  }
}

// The muse's brief: candidates, not answers, and no figures.
const MUSE = `You propose OpenQASM 3 programs for Quantum Studio's Alternatives. The Studio's engine reads, measures and compares every program you write; you do not judge them. Write each program whole, in its own \`\`\`qasm block, each after one line that starts "Idea:". State no counts, probabilities or costs. Use only what the engine reads: registers, the standard gates, gate definitions, ctrl and negctrl modifiers, measurement and reset, and if on measured bits; no loops, inv or pow.`;

// The ideas in a muse's answer: each \`\`\`qasm block with the "Idea:" line
// before it.
export function ideas(text) {
  const out = [];
  const re = /\`\`\`(?:qasm|openqasm|openqasm3)?\s*\n([\s\S]*?)\`\`\`/gi;
  let from = 0;
  for (let m; (m = re.exec(text)); ) {
    const before = text.slice(from, m.index);
    from = m.index + m[0].length;
    if (!/\b(qubit|qreg|OPENQASM)\b/.test(m[1])) continue;
    const said = [...before.matchAll(/^[\s>*#-]*(?:\d+[.)]\s*)?\**idea\**\s*[:.]\**\s*(.+)$/gim)].pop();
    out.push({ idea: said ? said[1].replace(/\*+/g, '').trim().replace(/\.$/, '') : '', qasm: m[1].trim() + '\n' });
  }
  return out;
}

// The figures in an answer's prose that no tool returned and the program does not
// hold: a decimal, a percentage, or a whole number of two digits or more.
export function unsourced(answer, heard) {
  const prose = answer.replace(/```[\s\S]*?```/g, ' ');
  const figures = prose.match(/(?<![\w.])\d+(?:\.\d+)?%?|(?<![\w.])\.\d+%?/g) ?? [];
  const out = [];
  for (const f of figures) {
    const bare = f.replace(/%$/, '');
    if (!/[.%]/.test(f) && bare.length < 2) continue; // a qubit's index, a step's number
    if (heard.includes(bare) || out.includes(f)) continue;
    out.push(f);
  }
  return out;
}

// LaTeX a model writes despite being asked not to, as plain Unicode.
export function unlatex(t) {
  return t
    .replace(/\$+/g, '')
    .replace(/\\frac\{([^{}]*)\}\{\\sqrt\{([^{}]*)\}\}/g, '$1/√$2')
    .replace(/\\frac\{([^{}]*)\}\{([^{}]*)\}/g, '($1)/($2)')
    .replace(/\\sqrt\{([^{}]*)\}/g, '√$1')
    .replace(/\\rangle/g, '⟩').replace(/\\langle/g, '⟨')
    .replace(/\\(Phi|Psi|psi|phi|theta|pi|alpha|beta|gamma|otimes|cdot|pm)\b/g, (_, w) => ({ Phi: 'Φ', Psi: 'Ψ', psi: 'ψ', phi: 'φ', theta: 'θ', pi: 'π', alpha: 'α', beta: 'β', gamma: 'γ', otimes: '⊗', cdot: '·', pm: '±' })[w])
    .replace(/\^\{?\+\}?/g, '⁺').replace(/\^\{?-\}?/g, '⁻')
    .replace(/\\(left|right)/g, '');
}

function pick(o, keys) {
  const out = {};
  for (const k of keys) if (o && o[k] != null && !(Array.isArray(o[k]) && !o[k].length)) out[k] = o[k];
  return out;
}