/** * The engine, with no transport attached. * * This is what a developer embedding a local model in their own page holds: * one object, plain async methods, no message vocabulary. `browser.runtime` is * not referenced anywhere in this file or anything it imports — the WebExtension * message/port router that used to live here now sits behind * `src/adapters/webext.js` and calls these same methods. * * Everything runs through one EnginePool, which owns priority, cancellation and * fan-out. Nothing here decides what runs when. * * Three things are injected, and they are exactly the three places the host * environment leaked into the engine: * * - `store` a ModelStore over a StorageAdapter, because * `browser.storage.local` does not exist in a page. * - `workerUrl` defaults to `new URL("./engine-worker.js", import.meta.url)`, * which Vite, webpack 5 and esbuild all understand and which * also resolves correctly on `moz-extension://`. It replaces * `browser.runtime.getURL` rather than sitting beside it. * - `loadWebLLM` defaults to a static relative `import()`, so bundlers can see * it, and a host with its own patched build can override it. * It stays *dynamic* on purpose: the ~6 MB WebLLM bundle is * fetched when a model is loaded, not when the page paints. * * Weights arrive by any of three routes — prebuilt (WebLLM's HuggingFace list), * remote (`registerModel` with any base URL you host), or injected * (`ingestModelFolder`, no network at all). `load()` resolves across all three; * see `model-store.js`. Pass `prebuilt: false` for a build that must never * reach the network for a model. */ import { ENGINE_STATE, JOB_KIND, PRIORITY, UNLOAD_LEVEL, WORKER_CONFIGURE } from "./constants.js"; import { chatFacade } from "./chat.js"; import { environmentFacade } from "./environment.js"; import { canRun, probeDevice, projectSpeed, rankModels } from "./device.js"; import { ERROR, EngineError, asEngineError } from "./errors.js"; import { filesFromDataTransfer, filesFromInput, ingestModelFolder } from "./ingest.js"; import { ModelStore, SOURCE, groupKeysByScope, isInjected, toAppConfig } from "./model-store.js"; import { clampSteps } from "./multistep.js"; import { EnginePool } from "./pool.js"; import { prefetchModel } from "./prefetch.js"; import { ask, conversation, ghostText } from "./recipes.js"; import { SOURCE_KIND, classifySource, isDataTransfer, isFileList, nearMatches } from "./sources.js"; /** * @typedef {object} CompletionRequest * The OpenAI generation fields WebLLM already speaks, plus the scheduling * fields that are what this engine adds over calling WebLLM directly. * @property {Array<{role: string, content: string}>} messages * @property {string} [modelId] load this model first if it is not the live one * @property {string} [id] job id; also what `cancel(id)` takes * @property {number} [temperature] * @property {number} [max_tokens] * @property {object} [response_format] * @property {object} [extra_body] * @property {string} [task] the unit that owns an engine; a whole batch shares one * @property {string} [session] a later job with this key supersedes the earlier one * @property {"interactive"|"normal"|"background"} [priority] * @property {boolean} [preemptible] may be interrupted by an `interactive` job */ /** * @typedef {object} CompletionResult * @property {string} text * @property {object} [usage] * @property {"stop"|"length"|"abort"} [finishReason] WebLLM's own values * @property {true} [cancelled] superseded or explicitly cancelled * @property {true} [preempted] an `interactive` job took the slot; `text` is partial */ /** * @typedef {CompletionRequest & {index: number, engineIndex: number, * startedAt: number, finishedAt: number, error?: string}} BatchItem */ const DEFAULT_WORKER_URL = () => new URL("./engine-worker.js", import.meta.url); const DEFAULT_LOAD_WEBLLM = () => import("../../vendor/web-llm.js"); /** * Turn "the package is not wired into your build" into a sentence that says so. * * `vendor/web-llm.js` is a **build product**, not a checked-in file, so the two * ways to arrive here are both install-shaped rather than runtime-shaped: a git * dependency whose `prepare` never ran, or a source checkout where `npm run * build` was never run. Left alone this surfaced as a bare * `Cannot find module '.../vendor/web-llm.js'` under the code * `GENERATION_FAILED` — wrong twice over, since nothing had begun generating * and the path named is ours, not the caller's. */ async function loadBundle(loadWebLLM) { try { return await loadWebLLM(); } catch (err) { const message = String(err?.message ?? err); // Only a resolution failure means "not built". A bundle that throws while // *evaluating* is a real crash and must keep its own stack. if (!/Cannot find module|Failed to (fetch|resolve)|ERR_MODULE_NOT_FOUND|dynamically imported module/i.test(message)) { throw err; } throw new EngineError( ERROR.PACKAGE_INCOMPLETE, "everything-webgpu is installed but its WebLLM bundle (vendor/web-llm.js) is missing. " + "That file is generated, not checked in — run `npm run build` in the package, " + "or reinstall so its `prepare` script runs.", { cause: "vendor-bundle-missing", underlying: message }, ); } } export class ScheduledEngine { #store; #workerUrl; #loadWebLLM; #prebuilt; #chat = null; #environment = null; #probe = null; /** This machine's achieved decode bandwidth, learned from the first generation. */ #decodeBytesPerSecond = 0; /** modelId -> weight bytes, for projections. */ #modelBytes = new Map(); /** * Resident models: modelId -> EnginePool. More than one may be up at once — * a text model beside a vision model, say — which is why this is a map and * not a field. Each entry holds a full copy of its weights, so residency is * budget-gated in `load()`. * @type {Map} */ #pools = new Map(); /** Which resident model an unaddressed request goes to. */ #current = null; /** In-flight loads, per model, so two models can come up concurrently. */ #loading = new Map(); /** * Pools that are still loading. `load()` holds the pool in a local until it * is ready, which left an in-flight download unreachable — so aborting one * had nothing to tear down. See `load({ signal })`. * @type {Map} */ #loadingPools = new Map(); #listeners = new Set(); #state = { status: ENGINE_STATE.IDLE, modelId: null, progress: null, error: null, pool: { size: 0, busy: 0, queued: 0 }, /** Model ids with a live pool. `modelId` is whichever of them is current. */ resident: [], /** Latest decode probe from an engine worker; see multistep.js. */ decode: null, }; /** * @param {object} opts * @param {ModelStore | import("./model-store.js").StorageAdapter} opts.store * a ModelStore, or a bare StorageAdapter to wrap in one * @param {string | URL} [opts.workerUrl] * @param {() => Promise} [opts.loadWebLLM] * @param {boolean} [opts.prebuilt] expose WebLLM's 163 HuggingFace-hosted * models, downloaded on first load. Default true. Set false for an * offline-only build: `load()` then resolves registered models and nothing * else, and an unknown id fails before the WebLLM bundle is even fetched. */ constructor({ store, workerUrl, loadWebLLM, prebuilt = true } = {}) { if (!store) { throw new EngineError(ERROR.BAD_REQUEST, "ScheduledEngine needs a `store` (ModelStore or StorageAdapter)."); } this.#store = store instanceof ModelStore ? store : new ModelStore(store); this.#workerUrl = workerUrl ?? DEFAULT_WORKER_URL(); // Wrapped once here rather than at each of the seven `#loadWebLLM()` call // sites: a missing bundle is the same failure whichever verb reached it // first, and a site added later gets the good error for free. const load = loadWebLLM ?? DEFAULT_LOAD_WEBLLM; this.#loadWebLLM = () => loadBundle(load); this.#prebuilt = prebuilt; } /** The ModelStore, so a host can drive the registry without a second handle. */ get store() { return this.#store; } /** * `chat.completions.create()`, the WebLLM/OpenAI shape. See `chat.js`. * * Built once and cached: callers hold on to `engine.chat.completions` the way * they did with WebLLM, and a fresh object each access would break that. */ get chat() { this.#chat ??= chatFacade(this); return this.#chat; } /** * `environment()` — the read-only report, with `environment.measure()` on it. * * Cached like `chat` so a caller can hold on to it. Writes are `configure()`; * see `environment.js` for why those are separate verbs. */ get environment() { this.#environment ??= environmentFacade(this); return this.#environment; } get state() { return { ...this.#state }; } get hasWebGPU() { return Boolean(globalThis.navigator?.gpu); } /** * @param {(state: object) => void} listener called immediately, then on change * @returns {() => void} unsubscribe */ subscribe(listener) { this.#listeners.add(listener); listener(this.state); return () => this.#listeners.delete(listener); } /** The current model's pool, or null. */ get #pool() { return this.#current ? (this.#pools.get(this.#current) ?? null) : null; } /** Model ids with a live pool right now. */ get resident() { return [...this.#pools.keys()]; } /** * Choose which resident model unaddressed requests go to. * * Distinct from `load()` on purpose: this is free and instant, because the * weights are already up. `load()` is what costs. */ use(modelId) { if (!this.#pools.has(modelId)) { throw new EngineError( ERROR.UNKNOWN_MODEL, `"${modelId}" is not resident. Resident: ${this.resident.join(", ") || "none"}. Call load() first.`, { modelId, resident: this.resident }, ); } this.#current = modelId; this.#syncState(); return this.state; } /** Registered models only — cheap, no bundle load. */ listModels() { return this.#store.list(); } /** * Everything `load()` would accept, normalised: registered models first, then * WebLLM's prebuilt list. * * Costs a WebLLM bundle fetch when `prebuilt` is on, because the list lives * inside it. `listModels()` is the cheap call if you only care about what this * app registered. * * @returns {Promise>} */ async listAvailableModels() { const registered = await this.#store.list(); const own = registered.map((r) => ({ modelId: r.model_id, source: r.source ?? SOURCE.REMOTE, model: r.model, contextWindow: r.overrides?.context_window_size, vramRequiredMB: r.vram_required_MB, sizeBytes: r.sizeBytes, })); if (!this.#prebuilt) return own; const { prebuiltAppConfig, functionCallingModelIds } = await this.#loadWebLLM(); const toolCalling = new Set(functionCallingModelIds ?? []); const owned = new Set(own.map((m) => m.modelId)); const rest = prebuiltAppConfig.model_list .filter((e) => !owned.has(e.model_id)) .map((e) => ({ modelId: e.model_id, source: SOURCE.PREBUILT, model: e.model, contextWindow: e.overrides?.context_window_size, vramRequiredMB: e.vram_required_MB, // WebLLM ships the list; it is not derivable from the id. toolCalling: toolCalling.has(e.model_id), })); return [...own, ...rest]; } /** * What this machine will admit to: WebGPU, adapter, `shader-f16`, the five * limits that matter, storage quota. Cached — hardware does not change * mid-session, and `requestAdapter()` is not free. * @returns {Promise} */ probe() { this.#probe ??= probeDevice(); return this.#probe; } /** * Whether a model will run here, before anything is downloaded. * @param {string} modelId * @returns {Promise<{ok: boolean, blockers: Array, warnings: Array}>} */ async canRun(modelId) { const [probe, available] = await Promise.all([this.probe(), this.listAvailableModels()]); const found = available.find((m) => m.modelId === modelId); if (!found) { throw new EngineError(ERROR.UNKNOWN_MODEL, `Model "${modelId}" is not one this engine can load.`, { modelId, }); } return canRun( { model_id: found.modelId, vram_required_MB: found.vramRequiredMB, sizeBytes: found.sizeBytes }, probe, ); } /** * Which models this device should actually be asked to run, best first. * * The prebuilt list spans 239 MB to 31 GB; this is the answer to the first * question a developer has and the one they have least basis to answer. * * @param {{maxVramMB?: number, needsVision?: boolean, needsToolCalling?: boolean, * prefer?: "quality" | "speed"}} [opts] */ async recommendModels({ needsToolCalling = false, ...opts } = {}) { const [probe, appConfig] = await Promise.all([this.probe(), this.#appConfig()]); let list = appConfig.model_list; if (needsToolCalling) { const { functionCallingModelIds } = await this.#loadWebLLM(); const ids = new Set(functionCallingModelIds ?? []); list = list.filter((m) => ids.has(m.model_id)); } return rankModels(list, { probe, ...opts }); } /** * Is this model's data on disk, so a load would need no network? * * Routes by who knows the keys. We wrote an injected model's artifacts and * hold the manifest, so `verify()` answers exactly — including a `"partial"` * verdict WebLLM cannot give. Everything else was fetched by WebLLM, which * derives the keys as its loader did, so `hasModelInCache` is the answer. * * @returns {Promise<"cached" | "partial" | "absent">} */ async cacheState(modelId) { const record = await this.#store.get(modelId); if (isInjected(record)) { const { ok, missing } = await this.#store.verify(record); if (ok) return "cached"; const total = Object.values(groupKeysByScope(record)).flat().length; return missing.length >= total ? "absent" : "partial"; } const { hasModelInCache } = await this.#loadWebLLM(); return (await hasModelInCache(modelId, await this.#appConfig())) ? "cached" : "absent"; } /** * Download a model into the cache **without building an engine**. * * For warming during onboarding: the bytes land while the user is still * reading, and the later `load()` is a cache read. WebLLM cannot express this * — `reload()` instantiates the wasm and needs a GPU before it fetches a * single shard — so this is ours. See `prefetch.js` for the URL-derivation * risk and the oracle that closes it. * * Needs no WebGPU at all, which is the other half of the point: an app can * warm the cache on a machine it has not yet decided can run the model. * * @param {string} modelId * @param {{signal?: AbortSignal, onProgress?: Function}} [opts] */ async prefetch(modelId, { signal, onProgress } = {}) { // Before anything, including the already-cached shortcut: a caller who // aborted wants an abort, not a success they have to inspect to distrust. if (signal?.aborted) { throw new EngineError(ERROR.ABORTED, `Prefetch of "${modelId}" was aborted before it began.`, { modelId, }); } const record = await this.#store.get(modelId); // An injected model's bytes were written before it was ever registered; // there is no URL to fetch from and nothing to do. if (isInjected(record)) { const { ok } = await this.#store.verify(record); if (ok) return { modelId, files: 0, bytes: 0, alreadyCached: true }; throw new EngineError( ERROR.CACHE_INCOMPLETE, `"${modelId}" was injected from a folder, so it cannot be re-fetched. Re-register the folder.`, { modelId }, ); } const appConfig = await this.#appConfig(); const entry = appConfig.model_list.find((m) => m.model_id === modelId); if (!entry) { const near = nearMatches(modelId, appConfig.model_list.map((m) => m.model_id)); throw new EngineError( ERROR.UNKNOWN_MODEL, `Cannot prefetch "${modelId}": it is neither registered nor prebuilt. ` + (near.length ? `Did you mean ${near.map((id) => `"${id}"`).join(", ")}? ` : ""), { modelId, ...(near.length ? { near } : {}) }, ); } const { hasModelInCache } = await this.#loadWebLLM(); if (await hasModelInCache(modelId, appConfig).catch(() => false)) { return { modelId, files: 0, bytes: 0, alreadyCached: true }; } const result = await prefetchModel({ modelId, record: entry, signal, onProgress }); // The oracle. `hasModelInCache` derives its keys through the very function // `prefetch.js` mirrors, so this is the one check that can tell a correct // prefetch from one that wrote a cache the loader will never read. Without // it, a wrong key looks exactly like success and costs the user a second // download of the whole model. if (!(await hasModelInCache(modelId, appConfig).catch(() => false))) { throw new EngineError( ERROR.GENERATION_FAILED, `Prefetch of "${modelId}" wrote ${result.files} artifacts, but WebLLM still reports the ` + "model as uncached — the derived cache keys do not match the ones its loader looks for. " + "Treat the cache as cold; load() will re-download. This is what a change to WebLLM's " + "URL scheme looks like from here.", { modelId, files: result.files }, ); } return result; } /** * Free a model's bytes and **keep the registry entry**, so it stays a model * this engine knows how to get again — the distinction from * `store.remove()`, which forgets the URL a remote model would need. * * Delegates for remote and prebuilt models: `deleteModelAllInfoInCache` is * WebLLM's, covers tensors + wasm + config, and is maintained upstream. */ async evict(modelId) { if (this.#pools.has(modelId)) await this.unload(modelId); return this.#evictBytes(modelId); } /** * The byte-freeing half of `evict()`, with no pool handling. * * Split out so `unload(id, "cache")` can reach it without going back through * `evict()` → `unload()`, which would re-enter this class for a pool that has * just been torn down. */ async #evictBytes(modelId) { const record = await this.#store.get(modelId); if (isInjected(record)) return this.#store.evictInjected(modelId); const { deleteModelAllInfoInCache } = await this.#loadWebLLM(); await deleteModelAllInfoInCache(modelId, await this.#appConfig()); return { freedKeys: null }; } /** The merged model list WebLLM's cache helpers key off. */ async #appConfig() { const registered = await this.#store.list(); const prebuilt = this.#prebuilt ? (await this.#loadWebLLM()).prebuiltAppConfig : null; return toAppConfig(registered, prebuilt); } /** * Forget a model entirely: free its bytes **and** drop the registry entry. * * `evict()` first, because that is what knows how to reach the bytes for each * source — and it has to happen before the record is deleted, since for a * remote model the record holds the only URL those bytes can be derived from. * Deleting the entry first would strand them in Cache Storage permanently. */ async remove(modelId) { const freed = await this.evict(modelId); await this.#store.remove(modelId); return freed; } /** * Projected decode throughput for a model, in tokens per second. * * `basis: "measured"` once anything has actually decoded on this machine — * the engine then knows its own achieved bandwidth and every projection is * device-specific. Before that, `basis: "extrapolated"` from a reference * machine, which is a starting point and says so. * * Decode is memory-bandwidth-bound, so this is close to the whole story: * time per token scales with weight bytes and little else. * * @param {string} [modelId] defaults to the current model */ async estimateSpeed(modelId = this.#current) { let bytes = modelId ? this.#modelBytes.get(modelId) : 0; if (!bytes && modelId) { const found = (await this.listAvailableModels()).find((m) => m.modelId === modelId); bytes = found?.sizeBytes ?? (found?.vramRequiredMB ?? 0) * 1024 * 1024; } return { modelId, ...projectSpeed(bytes, this.#decodeBytesPerSecond) }; } /** * What is actually switched on right now, as opposed to what the device could * support. * * The distinction matters for KV reuse in particular: `probe().kvReuse` is a * device capability, but the decision is taken inside the engine worker, * which is the authority. A caller debugging "why is my second turn slow" * needs the decision, not the capability. */ async features() { const probe = await this.probe(); const settings = await this.#store.getSettings(); const pool = this.#pool?.status(); // The configured K, which is not the same question as whether bursts run. // `decodeSteps` stays the knob's value even when the fast path is off: it is // what `configure()` would be changing, and reporting 1 here would make // `environment()` advise `configure({ decodeSteps: 15 })` for a fault that // setting cannot fix. const decodeSteps = this.#state.decode?.steps ?? settings.decodeSteps; // Non-empty means the live pipeline failed the contract and decoding was // routed to stock single-step. The worker keeps posting the *configured* K // either way, so `steps > 1` cannot answer this — asking it alone reported // a disabled fast path as healthy, which is the one case this call exists for. const multiStepOff = this.#state.decode?.multiStepOff; return { kvReuse: Boolean(probe.kvReuse), shaderF16: Boolean(probe.features?.shaderF16), decodeSteps, multiStepDecoding: decodeSteps > 1 && !multiStepOff?.length, /** The missing tvmjs internals, or null when the fast path is intact. */ multiStepOff: multiStepOff?.length ? multiStepOff : null, engines: pool?.size ?? 0, maxEngines: pool?.maxSize ?? settings.engineCount, resident: this.resident, // dispatches per flush: >1 means consecutive kernel launches are sharing a // compute pass, i.e. the build-time batching patch is in effect. Observed, // not asserted — a build with NO_PASS_MERGE=1 reports ~1. computePassBatching: this.#state.decode?.flushes ? this.#state.decode.dispatches / this.#state.decode.flushes : null, decode: this.#state.decode ?? null, }; } /** * Register a model. Two shapes, one call, and the difference is only where * the bytes come from: * * ```js * // fetched from a base URL you host — an HF repo, a CDN, your own origin * await engine.registerModel({ * modelId: "my-model", * model: "/models/my-model/", * modelLib: "/models/my-model/my-model-webgpu.wasm", * }); * * // read off disk. No network connection at any point, ever. * await engine.registerModel({ modelId: "my-model", files: entries }); * ``` * * Both end up as one `model_list` entry that WebLLM's own loader resolves the * same way — the local one only differs in that its base URL is minted on * `.invalid` and its cache is populated before the loader ever looks. * * That origin is the *mechanism* of the offline guarantee, not a marker of * it: `.invalid` is reserved by RFC 6761 and can never resolve, so there is * no code path — no bug, no eviction, no future refactor — by which a local * model reaches the network. It fails with a DNS error instead. * * `files` is `{ path, file }[]`; `filesFromDataTransfer` and * `filesFromInput` build it from a drop event or a directory picker. */ async registerModel(spec) { if (spec?.files) { if (spec.model || spec.modelLib) { throw new EngineError( ERROR.BAD_REQUEST, "registerModel takes either `files` (local, never fetched) or `model`/`modelLib` (a base URL to fetch), not both.", ); } return ingestModelFolder(spec.files, { store: this.#store, modelId: spec.modelId, modelType: spec.modelType, onProgress: spec.onProgress, }); } return this.#store.registerModel(spec); } // ---------------------------------------------------------------- engine --- #setState(patch) { Object.assign(this.#state, patch); const snapshot = this.state; for (const listener of [...this.#listeners]) { try { listener(snapshot); } catch { /* a subscriber that throws must not stall the engine */ } } } #assertWebGPU() { if (!this.hasWebGPU) { throw new EngineError( ERROR.NO_WEBGPU, "WebGPU is unavailable in this context. On macOS Firefox, set dom.webgpu.enabled=true " + "(and gfx.webgpu.ignore-blocklist=true if your GPU is blocklisted) in about:config, then restart Firefox.", ); } } /** * Bring a model up, whatever form you have it in. * * One entry point for all three routes, because from a caller's side "load a * model" is one intention and having to know which of `load`, * `registerModel` and `ingestModelFolder` to reach for is a decision the * library can make for them: * * ```js * load("Llama-3.2-1B-Instruct-q4f16_1-MLC") // prebuilt or registered id * load("https://huggingface.co/mlc-ai/Foo", { modelLib }) // a URL you host * load({ model, modelLib }) // the same, explicit * load({ files }) | load(fileList) | load(dataTransfer) // a folder, no network * ``` * * `registerModel` and `ingestModelFolder` remain, unchanged, as the low-level * primitives — this composes them rather than replacing them. * * **A URL always needs `modelLib`.** It is not guessed; see `sources.js` for * the measurement behind that. **`defer: true`** registers the source and * stops there, returning the record instead of the state — the manager's * drop-now-load-later flow. * * Additive residency: a model already resident stays resident, so switching * back to it costs nothing. That is only safe while the weights fit, so * `keepResident: false` (the default) unloads whatever else is up first — * the old single-model behaviour, and the safe one on a 16 GB machine. * Pass `keepResident: true` to hold both, having checked the budget yourself * with `canRun()`. * * @param {string | object} src an id, a URL, `{model, modelLib}`, or a folder * @param {{keepResident?: boolean, signal?: AbortSignal, defer?: boolean, * id?: string, modelLib?: string, modelType?: string, contextWindow?: number, * vramRequiredMB?: number, onProgress?: Function}} [opts] * @returns {Promise} the engine state, or the registry record when `defer` */ async load(src, opts = {}) { const source = classifySource(src, opts); if (source.kind === SOURCE_KIND.ID) { if (opts.defer) { throw new EngineError( ERROR.BAD_REQUEST, `\`defer\` registers a source without loading it, but "${source.modelId}" is an id — ` + "there is nothing to register. Drop `defer`, or pass a URL or a folder.", { modelId: source.modelId }, ); } return this.#loadById(source.modelId, opts); } const record = await this.#register(source, opts); if (opts.defer) return record; return this.#loadById(record.model_id, opts); } /** Turns a classified non-id source into a registry record. */ async #register(source, opts) { if (source.kind === SOURCE_KIND.FILES) { return ingestModelFolder(await toEntries(source.files), { store: this.#store, modelId: source.modelId, modelType: opts.modelType, onProgress: opts.onProgress, }); } return this.#store.registerModel({ modelId: source.modelId, model: source.model, modelLib: source.modelLib, modelType: opts.modelType, contextWindow: opts.contextWindow, vramRequiredMB: opts.vramRequiredMB, }); } /** * Bring a registered or prebuilt id up and make it current. * * Cancellation is WebLLM's, not ours: `unload()` aborts the `reloadController` * whose signal it threads through every artifact fetch. And resume is free — * WebLLM caches each artifact as it arrives and skips what is already present, * so a re-`load()` picks up where the abort left off. Call `evict()` to * discard a partial download instead. * * @param {string} modelId * @param {{keepResident?: boolean, signal?: AbortSignal}} [opts] */ async #loadById(modelId, { keepResident = false, signal } = {}) { if (signal?.aborted) { throw new EngineError(ERROR.ABORTED, `Load of "${modelId}" was aborted before it began.`, { modelId, }); } if (this.#pools.has(modelId)) return this.use(modelId); const inFlight = this.#loading.get(modelId); if (inFlight) { await inFlight.catch(() => {}); return this.#pools.has(modelId) ? this.use(modelId) : this.state; } const loading = (async () => { this.#assertWebGPU(); const models = await this.#store.list(); const registered = models.find((m) => m.model_id === modelId); // Only an injected model can be *unrecoverably* broken by eviction, and // checking costs nothing, so it gates the load — before the WebLLM bundle // is fetched, so the common offline failure stays fast. if (isInjected(registered)) { const { ok, missing } = await this.#store.verify(registered); if (!ok) { throw new EngineError( ERROR.CACHE_INCOMPLETE, `Cache for "${modelId}" is incomplete (${missing.length} artifact(s) evicted, e.g. ${missing[0].split("/").pop()}). Re-register the model folder.`, { modelId, missing }, ); } } if (!registered && !this.#prebuilt) { const near = nearMatches(modelId, models.map((m) => m.model_id)); throw new EngineError( ERROR.UNKNOWN_MODEL, `Model "${modelId}" is not registered, and prebuilt models are disabled. ` + (near.length ? `Did you mean ${near.map((id) => `"${id}"`).join(", ")}? ` : "") + `Call load(url, { modelLib }) or load({ files }) first.`, { modelId, prebuilt: false, ...(near.length ? { near } : {}) }, ); } const { engineCount, decodeSteps } = await this.#store.getSettings(); this.#setState({ status: ENGINE_STATE.LOADING, modelId, error: null, progress: { text: "Starting", progress: 0 }, }); // Each resident model is a full copy of its weights. Nothing reports free // VRAM to a page (AI.md, "The pool grows, it is not sized"), so the // default is to make room rather than to gamble on it fitting. if (!keepResident) { for (const id of [...this.#pools.keys()]) await this.#unloadOne(id); } const { CreateWebWorkerMLCEngine, prebuiltAppConfig } = await this.#loadWebLLM(); const appConfig = toAppConfig(models, this.#prebuilt ? prebuiltAppConfig : null); if (!appConfig.model_list.some((m) => m.model_id === modelId)) { // A typo'd id is the single most likely way to arrive here, and the // fix is almost always visible in the list we are already holding. const near = nearMatches(modelId, appConfig.model_list.map((m) => m.model_id)); throw new EngineError( ERROR.UNKNOWN_MODEL, `Model "${modelId}" is neither registered nor in WebLLM's prebuilt list. ` + (near.length ? `Did you mean ${near.map((id) => `"${id}"`).join(", ")}? ` : "") + `Use listAvailableModels() to see what this engine can load.`, { modelId, ...(near.length ? { near } : {}) }, ); } const pool = new EnginePool({ size: engineCount, createEngine: async (_index, onProgress) => { const worker = new Worker(this.#workerUrl, { type: "module" }); // Listener, not `onmessage`: WebLLM claims `onmessage` for its own RPC. worker.addEventListener("message", (event) => { if (event.data?.ewgpuStats) this.#setState({ decode: event.data.ewgpuStats }); }); // A worker whose script 404s does not throw from `new Worker()` — it // fires one `error` event and is then simply silent, so WebLLM's // handshake below never resolves and the load hangs until the caller // gives up. That is the exact shape of the Vite dep-optimizer bug this // names: esbuild copies `new URL("./engine-worker.js", // import.meta.url)` into `.vite/deps/` verbatim, where the sibling // file does not exist. Racing the handshake against this turns a hang // into a sentence. const workerFailed = new Promise((_, reject) => { worker.addEventListener("error", (event) => { reject( new EngineError( ERROR.PACKAGE_INCOMPLETE, `The decode worker failed to load from ${this.#workerUrl}. ` + "If you are on Vite, its dependency pre-bundler rewrote the worker URL to a " + "path that does not exist — add `optimizeDeps: { exclude: [\"everything-webgpu\"] }` " + "to vite.config.js, or pass `workerUrl` yourself.", { cause: "worker-unreachable", workerUrl: String(this.#workerUrl), underlying: event.message }, ), ); }); }); // Sent before WebLLM's own handshake so the first token already decodes // multi-step; worker message order guarantees it arrives first. worker.postMessage({ kind: WORKER_CONFIGURE, decodeSteps }); const engine = await Promise.race([ CreateWebWorkerMLCEngine(worker, modelId, { appConfig, initProgressCallback: onProgress, }), workerFailed, ]); // The worker owns the decode loop, so runtime knobs go straight to it // rather than through WebLLM's request path. engine.configure = (patch) => worker.postMessage({ kind: WORKER_CONFIGURE, ...patch }); // Tear the realm down with the engine, not just the model. const unloadEngine = engine.unload.bind(engine); engine.unload = async () => { await unloadEngine().catch(() => {}); worker.terminate(); }; return engine; }, onStateChange: () => { if (this.#current === modelId) this.#syncState(); }, }); // Reachable from here on, so an abort has something to tear down. The // listener bumps the pool's generation, which is what stops an engine // that finishes building *after* the abort from installing itself. this.#loadingPools.set(modelId, pool); if (signal) signal.addEventListener("abort", () => void this.#unloadOne(modelId), { once: true }); const abortedNow = async () => { await this.#unloadOne(modelId); throw new EngineError(ERROR.ABORTED, `Load of "${modelId}" was aborted.`, { modelId }); }; // Checked on both sides of the load. Before: an abort that landed while // the registry was being read must not start a download at all. After: an // abort during the download has already torn the pool down, and this is // what turns that into a rejection rather than a silent no-op. if (signal?.aborted) return abortedNow(); await pool.load((progress) => this.#setState({ progress })); if (signal?.aborted) return abortedNow(); const entry = appConfig.model_list.find((m) => m.model_id === modelId); this.#modelBytes.set( modelId, registered?.sizeBytes ?? (entry?.vram_required_MB ?? 0) * 1024 * 1024, ); this.#pools.set(modelId, pool); this.#current = modelId; this.#setState({ status: ENGINE_STATE.READY, progress: null, error: null }); this.#syncState(); return this.state; })(); this.#loading.set(modelId, loading); try { return await loading; } catch (err) { const engineError = asEngineError(err); this.#setState({ status: this.#pools.size ? ENGINE_STATE.READY : ENGINE_STATE.ERROR, progress: null, error: engineError.message, errorCode: engineError.code, }); // A failed load must not leave `modelId` pointing at the model that did // not come up — it was set optimistically when LOADING began. this.#syncState(); throw engineError; } finally { this.#loading.delete(modelId); this.#loadingPools.delete(modelId); } } /** * Let a model go, at one of two depths. * * ```js * unload() // the current model's VRAM; cached bytes stay * unload(id) // that model's VRAM * unload(id, "cache") // and delete its cached bytes, keeping the registry entry * ``` * * At `"vram"` the bytes stay on disk, so loading it again costs no network — * that is what makes switching back cheap, and the difference between this * and `remove()`. * * **A bare `unload()` frees only the current model**, not every resident one. * `unloadAll()` is the explicit form for that: freeing everything is the more * destructive of the two readings and should have to be asked for by name. * * @param {string} [modelId] defaults to the current model. Omit both this and * any resident model to no-op. * @param {"vram"|"cache"} [level] */ async unload(modelId = this.#current, level = UNLOAD_LEVEL.VRAM) { if (!Object.values(UNLOAD_LEVEL).includes(level)) { throw new EngineError( ERROR.BAD_REQUEST, `unload() level must be ${Object.values(UNLOAD_LEVEL).map((l) => `"${l}"`).join(" or ")}, ` + `not "${level}". To forget the model entirely, use remove().`, { level }, ); } if (modelId) { await this.#unloadOne(modelId); if (level === UNLOAD_LEVEL.CACHE) await this.#evictBytes(modelId); } this.#setState({ status: this.#pools.size ? ENGINE_STATE.READY : ENGINE_STATE.IDLE, progress: null, error: null, }); this.#syncState(); return this.state; } /** Unload every resident model. */ async unloadAll() { for (const id of [...this.#pools.keys()]) await this.#unloadOne(id); this.#setState({ status: ENGINE_STATE.IDLE, progress: null, error: null }); this.#syncState(); return this.state; } async #unloadOne(modelId) { const pool = this.#pools.get(modelId) ?? this.#loadingPools.get(modelId); if (!pool) return; this.#pools.delete(modelId); this.#loadingPools.delete(modelId); if (this.#current === modelId) this.#current = this.#pools.keys().next().value ?? null; await pool.unload(); } /** * Learn this machine's achieved decode bandwidth from a finished generation. * * WebLLM already measures throughput and puts it on every response as * `usage.extra.decode_tokens_per_s` — the pool asks for usage and was * discarding it. Tokens per second times weight bytes is bytes per second, * which is the one number a projection for *any other* model needs. * * The worker's own decode probe is not this and does not replace it: it * splits a burst into CPU-encode and GPU-sync, which is what the multi-step * and compute-pass work is measured against. Nothing upstream reports that. */ #calibrate(usage, modelId) { const tps = usage?.extra?.decode_tokens_per_s; const bytes = this.#modelBytes.get(modelId); if (!tps || !bytes) return; this.#decodeBytesPerSecond = tps * bytes; } /** Recomputes the parts of `state` that are views onto the current pool. */ #syncState() { this.#setState({ modelId: this.#current, resident: this.resident, pool: this.#pool?.status() ?? { size: 0, busy: 0, queued: 0 }, }); } /** Loads on demand so callers can just ask for a completion. */ async #ensurePool(modelId) { // A request naming a resident model goes straight to it, without disturbing // which model is current — routing is not switching. if (modelId && this.#pools.has(modelId)) return this.#pools.get(modelId); if (modelId) await this.load(modelId); if (!this.#pool) { const fallback = modelId ?? this.#current ?? (await this.#store.list())[0]?.model_id; if (!fallback) { throw new EngineError( ERROR.NO_MODEL, "No model is registered yet. Call registerModel(), or load() a prebuilt id.", ); } await this.load(fallback); } return modelId ? (this.#pools.get(modelId) ?? this.#pool) : this.#pool; } async #buildParams(payload) { const settings = await this.#store.getSettings(); const messages = Array.isArray(payload.messages) ? [...payload.messages] : []; if (messages.length === 0) { throw new EngineError(ERROR.BAD_REQUEST, "`messages` must be a non-empty array."); } if (settings.systemPrompt && !messages.some((m) => m.role === "system")) { messages.unshift({ role: "system", content: settings.systemPrompt }); } return { messages, temperature: payload.temperature ?? settings.temperature, max_tokens: payload.max_tokens ?? settings.maxTokens, ...(payload.response_format ? { response_format: payload.response_format } : {}), ...(payload.extra_body ? { extra_body: payload.extra_body } : {}), }; } /** * One completion. * * Named `complete` rather than `chat` so `engine.chat.completions.create()` * — the WebLLM-shaped facade, Phase 2 — can take that name without a rename. * * @param {CompletionRequest} payload * @param {(delta: string) => void} [onChunk] called per streamed text delta * @returns {Promise} */ async complete(payload, onChunk) { // Adapts the raw chunk stream to the documented string callback. The guard // matters: raw chunks include a role-only first chunk and a finish-only // last chunk, which callers of this signature have never seen. return this.completeRaw( payload, onChunk && ((chunk) => { const delta = chunk?.choices?.[0]?.delta?.content; if (delta) onChunk(delta); }), ); } /** * `complete()`, but the callback receives WebLLM's chunk verbatim. * * Exists so the `chat.completions.create()` facade can pass chunks straight * through instead of rebuilding an envelope — which is what dropped * `tool_calls`, flattened `logprobs` and restamped `created`. * * @param {CompletionRequest} payload * @param {(chunk: object) => void} [onRawChunk] * @returns {Promise}>} */ async completeRaw(payload, onRawChunk) { const pool = await this.#ensurePool(payload.modelId); const result = unwrap( await pool.submit({ ...scheduling(payload), id: payload.id, params: await this.#buildParams(payload), onChunk: onRawChunk, }), ); this.#calibrate(result.usage, this.#current); return { text: result.text, usage: result.usage, finishReason: result.finishReason, ...(result.toolCalls ? { toolCalls: result.toolCalls } : {}), ...(result.cancelled ? { cancelled: true } : {}), ...(result.preempted ? { preempted: true } : {}), }; } // ------------------------------------------- the three shapes, as verbs --- // // `complete()` expresses all three. These exist because the scheduling is the // part that is easy to get wrong and invisible when you do — see recipes.js. /** * One question, one answer, nothing kept. * * ```js * const answer = await engine.ask("Summarise this in one line:\n" + doc); * ``` * * @param {string | Array} input * @param {object} [opts] anything `complete()` takes, plus `onDelta` to stream * @returns {Promise} */ ask(input, opts) { return ask(this, input, opts); } /** * A multi-turn conversation that keeps its own history. * * ```js * const chat = engine.conversation({ system: "You are terse." }); * await chat.say("hello"); * await chat.say("and again?"); // remembers * ``` * * @param {object} [opts] `system`, `keep`, plus `complete()` defaults */ conversation(opts) { return conversation(this, opts); } /** * Ghost text, with the debounce/supersede/drop-if-stale discipline built in * and the prompt left to you. * * ```js * const ghost = engine.ghostText({ prompt: (before) => `Continue:\n${before}` }); * editor.on("input", async () => { * const hint = await ghost.suggest(editor.textBefore()); * if (hint !== null) render(hint); // null means a newer keystroke won * }); * editor.on("blur", () => ghost.cancel()); * ``` * * @param {object} opts must include `prompt` */ ghostText(opts) { return ghostText(this, opts); } /** * Embed text into vectors, through the same scheduler as everything else. * * ```js * const [vector] = await engine.embed("a sentence", { modelId: EMBED_MODEL }); * const vectors = await engine.embed(["one", "two"], { modelId: EMBED_MODEL }); * ``` * * **Needs an embedding model**, not a chat model — `snowflake-arctic-embed-*` * in WebLLM's prebuilt list, from 239 MB. They are separate models, so this * usually names `modelId` explicitly and holds it resident alongside a chat * model with `load(id, { keepResident: true })`. * * Returns bare vectors because that is what a caller does arithmetic on; the * OpenAI envelope is available as `embedRaw()` for anyone porting code that * expects `data[].embedding`. * * **A running embedding cannot be interrupted.** Cancellation and preemption * work by making a decode loop break out; one forward pass has no loop, so a * `cancel()` that lands after the job starts marks it cancelled but does not * stop it. Queued embeddings supersede and cancel normally. This is tolerable * because an embedding is milliseconds where a completion is seconds — but it * is a weaker guarantee than `complete()` gives, so it is stated rather than * discovered. * * @param {string | string[]} input * @param {{modelId?: string, task?: string, session?: string, * priority?: string, preemptible?: boolean, id?: string}} [opts] * @returns {Promise} one vector per input, in order */ async embed(input, opts = {}) { const { data } = await this.embedRaw(input, opts); return data.map((d) => d.embedding); } /** `embed()`, returning WebLLM's OpenAI-shaped envelope untouched. */ async embedRaw(input, opts = {}) { const texts = Array.isArray(input) ? input : [input]; if (texts.length === 0 || texts.some((t) => typeof t !== "string")) { throw new EngineError( ERROR.BAD_REQUEST, "embed() takes a string or a non-empty array of strings.", { received: Array.isArray(input) ? `array of ${input.length}` : typeof input }, ); } const pool = await this.#ensurePool(opts.modelId); const result = unwrap( await pool.submit({ ...scheduling(opts), id: opts.id, kind: JOB_KIND.EMBEDDING, params: { input: texts }, }), ); return { data: result.embeddings ?? [], usage: result.usage }; } /** * Independent prompts, fanned across the pool. This is the only way to beat * the ~10 tok/s single-stream ceiling, so anything embarrassingly parallel * (translating a page, labelling a list) should arrive here rather than as a * loop of `complete` calls. * * @param {CompletionRequest & {requests: Array>}} payload * @param {(item: BatchItem) => void} [onItem] called as each item lands * @returns {Promise>} */ async batch(payload, onItem = () => {}) { const requests = payload.requests; if (!Array.isArray(requests) || requests.length === 0) { throw new EngineError(ERROR.BAD_REQUEST, "`requests` must be a non-empty array."); } const pool = await this.#ensurePool(payload.modelId); const sched = scheduling(payload); // One batch is one task, however many requests it is: "translate this page" // should hold one engine, not every engine. The pool reserves its last free // slot for a different task, so ghost-text never queues behind the page. const task = payload.task ?? `batch-${payload.id ?? crypto.randomUUID()}`; return Promise.all( requests.map(async (req, index) => { const merged = { ...payload, ...req, requests: undefined }; const result = await pool.submit({ ...sched, ...scheduling(merged), session: req.session, // a batch shares no session unless an item names one // An item that names its own session is its own task again. task: req.task ?? req.session ?? task, params: await this.#buildParams(merged), }); const item = { index, engineIndex: result.engineIndex, startedAt: result.startedAt, finishedAt: result.finishedAt, ...(result.error ? { error: result.error } : { text: result.text, usage: result.usage, finishReason: result.finishReason, ...(result.toolCalls ? { toolCalls: result.toolCalls } : {}), ...(result.cancelled ? { cancelled: true } : {}), }), }; onItem(item); return item; }), ); } /** * Cancels by job id or by session key. * @param {string} idOrSession * @returns {number} how many jobs it stopped */ cancel(idOrSession) { let stopped = 0; for (const pool of this.#pools.values()) stopped += pool.cancel(idOrSession); return stopped; } /** * Applies a runtime knob to the running pool and persists it as the default. * * `decodeSteps` is the multi-step decode width (AI.md, "Multi-step decoding"). * It takes effect on the next burst — no reload — which is what makes sweeping * it to find this machine's tick boundary cheap. */ async configure(patch) { const applied = {}; if (patch.decodeSteps !== undefined) applied.decodeSteps = clampSteps(patch.decodeSteps); if (patch.engineCount !== undefined) { const n = Math.round(Number(patch.engineCount)); if (!Number.isFinite(n) || n < 1) { throw new EngineError( ERROR.BAD_REQUEST, `engineCount must be a positive integer, not ${JSON.stringify(patch.engineCount)}.`, { engineCount: patch.engineCount }, ); } applied.engineCount = n; } if (Object.keys(applied).length === 0) { // Naming the knobs matters: this is the error a caller hits after // `environment()` told them something was operable, so it has to agree // with that report about what the operable things are. throw new EngineError( ERROR.BAD_REQUEST, "`configure` needs at least one setting. Operable: `decodeSteps`, `engineCount`.", { operable: ["decodeSteps", "engineCount"] }, ); } await this.#store.setSettings(applied); // Only `decodeSteps` is hot. `engineCount` is persisted and read when a pool // is built, so a live pool keeps the size it came up with — `environment()` // reports that gap rather than pretending the change took effect. let engines = 0; if (applied.decodeSteps !== undefined) { for (const pool of this.#pools.values()) engines += pool.configure({ decodeSteps: applied.decodeSteps }); } return { settings: applied, engines }; } } /** * Normalise every folder shape a caller might hold into `{path, file}[]`. * * A drop event gives a `DataTransfer`, `` gives a * `FileList`, and a caller who has already unpacked one gives the entries. All * three mean "this folder", so `load()` accepts all three rather than making * the caller find the right converter first. */ async function toEntries(files) { if (isDataTransfer(files)) return filesFromDataTransfer(files); if (isFileList(files)) return filesFromInput(files); return files; } /** Scheduling metadata is per-request; the pool, not the caller, acts on it. */ function scheduling(payload) { return { task: payload.task, session: payload.session, priority: payload.priority ?? PRIORITY.NORMAL, preemptible: payload.preemptible, }; } function unwrap(result) { // The pool resolves rather than rejects, so a failed generation arrives as a // string on the result. It has no code of its own by then. if (result.error) throw asEngineError(result.error, ERROR.GENERATION_FAILED); return result; }