open-standards/aac/aac-rs/third_party/litert-lm/worker.js

// Sentence helper worker for the AAC reference app: one on-device language
// model (Gemma 4 E2B, Google's web build) on Google's LiteRT-LM web runtime,
// text in, text out.
//
// Nothing leaves the device. The worker never fetches: the model's parts come
// from the model store (Cache API "aac-models-v1", each part SHA-256 checked
// when it was stored), streamed one after another into GPU memory, and the
// runtime's WebAssembly is handed over by src/bindings/ai_helper.js, which
// read it from the model store too. So the helper keeps working offline.
//
// This file is bundled with @litert-lm/core and @litertjs/wasm-utils into one
// classic worker script (ai-worker.js) by third_party/litert-lm/vendor.sh: the
// runtime loads its Emscripten glue with importScripts, which module workers
// do not have.
//
// Messages in (each with an `id`, answered once):
//   {type: "init", loader, wasm}   loader: Blob URL of the runtime's glue
//                                  script; wasm: its ArrayBuffer
//   {type: "load", key, parts, maxTokens}
//                                  parts: the model's part URLs, in order
//   {type: "generate", preface, user, maxOutputTokens}  -> {text}
//                                  preface: [{role, content}...]
//   {type: "cancel"}               stops the generation running now
//   {type: "unload"}
// Replies: {id, ok: true, ...} or {id, ok: false, error}. "cancel" is
// answered at once; the others in the order they came.
//
// Vendored into vendor/litert-lm-<version>-<revision>/ by vendor.sh. The
// bundle's bytes are pinned in crates/aac-web/src/ai_helper.rs
// (RUNTIME_FILES): change this file only together with a new
// LITERT_LM_VENDOR_REV and its new size and hash there.

import { Backend, Engine, SamplerType, getGlobalLiteRtLmPromise, loadLiteRtLm } from "@litert-lm/core";

const CACHE_NAME = "aac-models-v1";

let engine = null; // the loaded Engine
let key = ""; // which model `engine` runs
let running = null; // the Conversation generating now, if any

function describe(err) {
    return String(err && err.message ? err.message : err);
}

async function init(msg) {
    if (getGlobalLiteRtLmPromise()) {
        await getGlobalLiteRtLmPromise();
        return {};
    }
    // The glue would fetch its WebAssembly by name; Emscripten's
    // instantiateWasm hook hands it the bytes instead, so it never fetches.
    // (This build reads instantiateWasm, not wasmBinary.)
    let failed;
    const instantiateFailed = new Promise((_, reject) => { failed = reject; });
    self.Module = {
        instantiateWasm(imports, receive) {
            WebAssembly.instantiate(msg.wasm, imports).then((r) => receive(r.instance, r.module), failed);
            return {};
        },
    };
    // The runtime takes a URL ending in ".js" as the glue itself (anything
    // else as a directory to look in); a fragment does not change what a
    // Blob URL points at.
    await Promise.race([loadLiteRtLm(`${msg.loader}#litertlm.js`), instantiateFailed]);
    return {};
}

// The model's parts, read from the model store one after another, as one
// stream: the runtime streams it into GPU memory, so no copy of the whole
// model is ever held here.
function partsStream(cache, urls) {
    let next = 0;
    let reader = null;
    return new ReadableStream({
        async pull(controller) {
            for (;;) {
                if (!reader) {
                    if (next >= urls.length) {
                        controller.close();
                        return;
                    }
                    const res = await cache.match(urls[next]);
                    next += 1;
                    if (!res || !res.body) {
                        throw new Error(`part ${next} of ${urls.length} of the model is not on this device`);
                    }
                    reader = res.body.getReader();
                }
                const step = await reader.read();
                if (step.done) {
                    reader = null;
                    continue;
                }
                controller.enqueue(step.value);
                return;
            }
        },
        cancel() {
            if (reader) reader.cancel().catch(() => {});
        },
    });
}

async function load(msg) {
    if (engine && key === msg.key) return {};
    await unload();
    if (typeof caches === "undefined") throw new Error("this browser keeps no model store in workers");
    const cache = await caches.open(CACHE_NAME);
    engine = await Engine.create({
        model: partsStream(cache, msg.parts),
        backend: Backend.GPU_ARTISAN,
        mainExecutorSettings: { maxNumTokens: msg.maxTokens },
    });
    key = msg.key;
    return {};
}

async function generate(msg) {
    if (!engine) throw new Error("the model is not loaded");
    const conversation = await engine.createConversation({
        preface: { messages: msg.preface },
        sessionConfig: {
            // The most likely words, every time: the same request gives the
            // same suggestions.
            samplerParams: { type: SamplerType.GREEDY },
            maxOutputTokens: msg.maxOutputTokens,
        },
    });
    running = conversation;
    try {
        let text = "";
        // getReader(), not for-await: Safari 26 has no async iterator on
        // ReadableStream.
        const reader = conversation.sendMessageStreaming(msg.user).getReader();
        for (;;) {
            const { done, value } = await reader.read();
            if (done) break;
            for (const item of (value && value.content) || []) {
                if (item.type === "text") text += item.text;
            }
        }
        return { text };
    } finally {
        running = null;
        await conversation.delete().catch(() => {});
    }
}

async function unload() {
    const old = engine;
    engine = null;
    key = "";
    if (old) await old.delete().catch(() => {});
    return {};
}

const HANDLERS = { init, load, generate, unload };

// One at a time, in order: the runtime runs one thing at a time anyway.
let queue = Promise.resolve();

self.onmessage = (event) => {
    const msg = event.data || {};
    if (msg.type === "cancel") {
        try { if (running) running.cancel(); } catch (_) { /* already finished */ }
        self.postMessage({ id: msg.id, ok: true });
        return;
    }
    const handler = HANDLERS[msg.type];
    queue = queue.then(async () => {
        if (!handler) {
            self.postMessage({ id: msg.id, ok: false, error: `unknown message ${msg.type}` });
            return;
        }
        try {
            const reply = await handler(msg);
            self.postMessage({ id: msg.id, ok: true, ...reply });
        } catch (e) {
            self.postMessage({ id: msg.id, ok: false, error: describe(e) });
        }
    });
};