nyaaorick's picture
feat: publish everything-webgpu package, engine source and documentation
1944112 verified
Raw
History Blame Contribute Delete
54.1 kB
/**
* 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<string, EnginePool>}
*/
#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<string, EnginePool>}
*/
#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<object>} [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<Array<{modelId: string, source: string, model: string,
* contextWindow?: number, vramRequiredMB?: number, sizeBytes?: number}>>}
*/
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<import("./device.js").DeviceProbe>}
*/
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<object>, warnings: Array<object>}>}
*/
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<object>} 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<CompletionResult>}
*/
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<CompletionResult & {toolCalls?: Array<object>}>}
*/
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<object>} input
* @param {object} [opts] anything `complete()` takes, plus `onDelta` to stream
* @returns {Promise<string>}
*/
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<number[][]>} 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<Partial<CompletionRequest>>}} payload
* @param {(item: BatchItem) => void} [onItem] called as each item lands
* @returns {Promise<Array<BatchItem>>}
*/
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`, `<input webkitdirectory>` 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;
}