AI as a collaborator

/** Chrome's built-in model (the Prompt API). Not in TypeScript's DOM types yet. */
interface LanguageModelSession {
  prompt(input: string, options?: { signal?: AbortSignal }): Promise<string>;
  measureInputUsage?(input: string, options?: { signal?: AbortSignal }): Promise<number>;
  inputQuota?: number;
  inputUsage?: number;
  destroy(): void;
}
type Expectation = { type: 'text'; languages: string[] }[];
interface LanguageModelOptions {
  expectedInputs?: Expectation;
  expectedOutputs?: Expectation;
}
interface LanguageModelStatic {
  availability(
    options?: LanguageModelOptions,
  ): Promise<'unavailable' | 'downloadable' | 'downloading' | 'available'>;
  create(
    options?: LanguageModelOptions & {
      initialPrompts?: { role: 'system'; content: string }[];
      monitor?: (m: EventTarget) => void;
      signal?: AbortSignal;
    },
  ): Promise<LanguageModelSession>;
}
declare const LanguageModel: LanguageModelStatic | undefined;

// Availability is per configuration: ask with the same options you'll create with.
const ENGLISH = {
  expectedInputs: [{ type: 'text' as const, languages: ['en'] }],
  expectedOutputs: [{ type: 'text' as const, languages: ['en'] }],
};

/** Local: asks the browser whether the model is there, and sends nothing. */
export async function modelStatus() {
  if (typeof LanguageModel === 'undefined') return 'unavailable';
  return LanguageModel.availability(ENGLISH);
}

/** Only from a click (Ask): creating a 'downloadable' model starts its download. */
export async function ask(system: string, prompt: string, signal: AbortSignal) {
  const model = await LanguageModel!.create({
    ...ENGLISH,
    initialPrompts: [{ role: 'system', content: system }],
    signal,
    monitor: (m) =>
      m.addEventListener('downloadprogress', (e) => {
        console.log(`Downloading: ${Math.round((e as ProgressEvent).loaded * 100)}%`);
      }),
  });
  try {
    // Does the passage fit the context window? (Where the browser can tell.)
    if (model.measureInputUsage !== undefined && model.inputQuota !== undefined) {
      const usage = await model.measureInputUsage(prompt, { signal });
      if (usage > model.inputQuota - (model.inputUsage ?? 0)) return null;
    }
    return await model.prompt(prompt, { signal });
  } finally {
    model.destroy();
  }
}

/**
 * A WebLLM model in a module worker, so the page never loads WebLLM (about
 * 2 MB compressed) and typing stays smooth while the model writes. Start it
 * only from a click: `new Worker(new URL('./webllm-worker.ts', import.meta.url), { type: 'module' })`.
 */
import { hasModelInCache, MLCEngine } from '@mlc-ai/web-llm';

const MODEL = 'Qwen2.5-0.5B-Instruct-q4f16_1-MLC'; // about 300 MB, from Hugging Face
const scope = self as unknown as { postMessage(message: unknown): void };

const engine = new MLCEngine({
  initProgressCallback: ({ progress, text }) =>
    scope.postMessage({ type: 'progress', progress, text }),
});

type Request = { type: 'has' } | { type: 'load' } | { type: 'ask'; system: string; prompt: string };

addEventListener('message', async (event: MessageEvent<Request>) => {
  const request = event.data;
  if (request.type === 'has') {
    // Ask the person before a download: is the model already in this browser?
    scope.postMessage({ type: 'has', value: await hasModelInCache(MODEL) });
  } else if (request.type === 'load') {
    await engine.reload(MODEL); // downloads what the cache lacks, then compiles for WebGPU
    scope.postMessage({ type: 'loaded' });
  } else {
    const reply = await engine.chat.completions.create({
      messages: [
        { role: 'system', content: request.system },
        { role: 'user', content: request.prompt },
      ],
      max_tokens: 1024,
    });
    const choice = reply.choices[0];
    // `length`: the reply ran out of room; don't write half an answer.
    scope.postMessage({
      type: 'reply',
      text: choice?.message.content,
      finish: choice?.finish_reason,
    });
  }
});

/**
 * The reply as Markdown to write: trimmed, without the passage markers echoed
 * back, and capped. A fence is unwrapped only when it is unambiguously a
 * wrapper: tagged `markdown` or `md`, around the whole reply, with no other
 * fence inside. An untagged fence is a code block, and stays one.
 */
export function cleanReply(reply: string): string {
  let text = reply.replace(/\r\n?/g, '\n').trim();
  const fenced = /^```(?:markdown|md)[ \t]*\n([\s\S]*)\n```$/i.exec(text);
  if (fenced !== null && !/^[ \t]*```/m.test(fenced[1]!)) {
    text = fenced[1]!.trim();
  }
  text = text
    .replace(/^<<<\n?/, '')
    .replace(/\n?>>>$/, '')
    .trim();
  return text.slice(0, MAX_REPLY);
}

/**
 * A lossless picture of the blocks (every node's JSON: text, formats, styles,
 * alignment, indents), to tell whether anything changed since. Markdown can't:
 * it drops underline, alignment and runs of spaces.
 */
export function $snapshotBlocks(blocks: readonly ElementNode[]): string {
  const walk = (node: LexicalNode): unknown => {
    const json = node.getLatest().exportJSON() as Record<string, unknown>;
    if ($isElementNode(node)) {
      json.children = node.getChildren().map(walk);
    }
    return json;
  };
  return JSON.stringify(blocks.map(walk));
}

import { buildEditorFromExtensions } from '@lexical/extension';
import { RichTextExtension } from '@lexical/rich-text';
import { createYDoc, YCollabExtension } from '@lexical/y';
import * as Y from '@y/y';
import { configExtension, defineExtension } from 'lexical';

/**
 * A second author inside the tab: its own Y.Doc (a fresh client id, so blame
 * can name it) and a headless editor, synced with the note both ways.
 */
export function createPeer(note: Y.Doc) {
  const doc = createYDoc();
  const fromNote = Symbol('from-note');
  const fromPeer = Symbol('from-peer');
  Y.applyUpdate(doc, Y.encodeStateAsUpdate(note), fromNote); // start from the note as it is
  const toPeer = (update: Uint8Array, origin: unknown) => {
    if (origin !== fromPeer) Y.applyUpdate(doc, update, fromNote);
  };
  const toNote = (update: Uint8Array, origin: unknown) => {
    if (origin !== fromNote) Y.applyUpdate(note, update, fromPeer); // never echo back
  };
  note.on('update', toPeer);
  doc.on('update', toNote);

  const editor = buildEditorFromExtensions(
    defineExtension({
      name: 'my-app/peer',
      // Already synced (a copy of the note), and no undo: a peer never undoes.
      dependencies: [
        RichTextExtension,
        configExtension(YCollabExtension, {
          doc,
          schemaVersion: 1, // the same as the note's other clients
          synced: true,
          captureTimeout: null,
        }),
      ],
    }),
  );
  return {
    doc,
    /** Edits as the peer: its writes reach the note before this returns. */
    edit(fn: () => void) {
      editor.update(fn, { discrete: true });
    },
    /** Leaves; what it wrote stays. */
    dispose() {
      note.off('update', toPeer);
      doc.off('update', toNote);
      editor.dispose();
      doc.destroy();
    },
  };
}