localm-web 0.1.0 → 0.2.0

package/dist/index.js

@@ -1,3 +1,10 @@
+const DOWNLOAD_PATTERN = /\b(fetch|download|loading from cache|cache hit|param)/i;
+const COMPILE_PATTERN = /\b(compil|shader|kernel|tensor|init|allocat|warm)/i;
+function classifyLoadPhase(text) {
+  if (DOWNLOAD_PATTERN.test(text)) return "downloading";
+  if (COMPILE_PATTERN.test(text)) return "compiling";
+  return "loading";
+}
 class LocalmWebError extends Error {
   /**
    * @param message - Human-readable description of the error.
@@ -73,10 +80,18 @@ class WebLLMEngine {
             progress: report.progress,
             text: report.text,
             loaded: 0,
-            total: 0
+            total: 0,
+            phase: classifyLoadPhase(report.text)
           });
         }
       });
+      onProgress?.({
+        progress: 1,
+        text: "Model ready.",
+        loaded: 0,
+        total: 0,
+        phase: "ready"
+      });
     } catch (err) {
       throw new ModelLoadError(`Failed to load model "${modelId}".`, err);
     }
@@ -130,6 +145,55 @@ class WebLLMEngine {
       throw new ModelLoadError("Streaming generation failed.", err);
     }
   }
+  async complete(prompt, options = {}) {
+    const engine = this.requireEngine();
+    if (options.signal?.aborted) {
+      throw new GenerationAbortedError("Generation aborted before start.");
+    }
+    const completion = await engine.completions.create({
+      ...buildSamplingParams(options),
+      prompt,
+      stream: false
+    });
+    return completion.choices[0]?.text ?? "";
+  }
+  async *streamCompletion(prompt, options = {}) {
+    const engine = this.requireEngine();
+    if (options.signal?.aborted) {
+      throw new GenerationAbortedError("Generation aborted before start.");
+    }
+    const completion = await engine.completions.create({
+      ...buildSamplingParams(options),
+      prompt,
+      stream: true
+    });
+    let index = 0;
+    let finished = false;
+    try {
+      for await (const chunk of completion) {
+        if (options.signal?.aborted) {
+          throw new GenerationAbortedError("Generation aborted by signal.");
+        }
+        const choice = chunk.choices[0];
+        const delta = choice?.text ?? "";
+        if (delta) {
+          yield { text: delta, index, done: false };
+          index += 1;
+        }
+        if (choice?.finish_reason) {
+          finished = true;
+          yield { text: "", index, done: true };
+          index += 1;
+        }
+      }
+      if (!finished) {
+        yield { text: "", index, done: true };
+      }
+    } catch (err) {
+      if (err instanceof GenerationAbortedError) throw err;
+      throw new ModelLoadError("Streaming completion failed.", err);
+    }
+  }
   async unload() {
     if (this.engine) {
       await this.engine.unload();
@@ -143,6 +207,283 @@ class WebLLMEngine {
     return this.engine;
   }
 }
+function toSerializableOptions(options = {}) {
+  const { signal: _signal, ...rest } = options;
+  return rest;
+}
+class WorkerEngine {
+  constructor(worker) {
+    this.worker = worker;
+    this.listener = (event) => this.handleMessage(event.data);
+    this.worker.addEventListener("message", this.listener);
+  }
+  nextId = 1;
+  loaded = false;
+  currentLoad = null;
+  currentLoadId = 0;
+  currentLoadProgress = void 0;
+  currentUnload = null;
+  currentUnloadId = 0;
+  pendingGenerates = /* @__PURE__ */ new Map();
+  pendingStreams = /* @__PURE__ */ new Map();
+  listener;
+  isLoaded() {
+    return this.loaded;
+  }
+  async load(modelId, onProgress) {
+    if (this.currentLoad) {
+      throw new ModelLoadError("Another load is already in progress.");
+    }
+    const id = this.allocateId();
+    this.currentLoadId = id;
+    this.currentLoadProgress = onProgress;
+    return new Promise((resolve, reject) => {
+      this.currentLoad = { resolve, reject };
+      this.send({ op: "load", id, modelId });
+    });
+  }
+  async generate(messages, options = {}) {
+    const id = this.allocateId();
+    return new Promise((resolve, reject) => {
+      this.pendingGenerates.set(id, { resolve, reject });
+      this.send({
+        op: "generate",
+        id,
+        messages,
+        options: toSerializableOptions(options)
+      });
+      options.signal?.addEventListener("abort", () => this.send({ op: "abort", id }));
+    });
+  }
+  async *stream(messages, options = {}) {
+    const id = this.allocateId();
+    const queue = [];
+    let done = false;
+    let error = null;
+    let notify = null;
+    const wakeup = () => {
+      if (notify) {
+        const fn = notify;
+        notify = null;
+        fn();
+      }
+    };
+    this.pendingStreams.set(id, {
+      push: (chunk) => {
+        queue.push(chunk);
+        wakeup();
+      },
+      end: () => {
+        done = true;
+        wakeup();
+      },
+      fail: (err) => {
+        error = err;
+        done = true;
+        wakeup();
+      }
+    });
+    this.send({
+      op: "stream",
+      id,
+      messages,
+      options: toSerializableOptions(options)
+    });
+    options.signal?.addEventListener("abort", () => this.send({ op: "abort", id }));
+    try {
+      while (true) {
+        if (queue.length > 0) {
+          const chunk = queue.shift();
+          if (chunk) yield chunk;
+          continue;
+        }
+        if (error) throw error;
+        if (done) return;
+        await new Promise((r) => {
+          notify = r;
+        });
+      }
+    } finally {
+      this.pendingStreams.delete(id);
+    }
+  }
+  async complete(prompt, options = {}) {
+    const id = this.allocateId();
+    return new Promise((resolve, reject) => {
+      this.pendingGenerates.set(id, { resolve, reject });
+      this.send({
+        op: "complete",
+        id,
+        prompt,
+        options: toSerializableOptions(options)
+      });
+      options.signal?.addEventListener("abort", () => this.send({ op: "abort", id }));
+    });
+  }
+  async *streamCompletion(prompt, options = {}) {
+    const id = this.allocateId();
+    const queue = [];
+    let done = false;
+    let error = null;
+    let notify = null;
+    const wakeup = () => {
+      if (notify) {
+        const fn = notify;
+        notify = null;
+        fn();
+      }
+    };
+    this.pendingStreams.set(id, {
+      push: (chunk) => {
+        queue.push(chunk);
+        wakeup();
+      },
+      end: () => {
+        done = true;
+        wakeup();
+      },
+      fail: (err) => {
+        error = err;
+        done = true;
+        wakeup();
+      }
+    });
+    this.send({
+      op: "stream-completion",
+      id,
+      prompt,
+      options: toSerializableOptions(options)
+    });
+    options.signal?.addEventListener("abort", () => this.send({ op: "abort", id }));
+    try {
+      while (true) {
+        if (queue.length > 0) {
+          const chunk = queue.shift();
+          if (chunk) yield chunk;
+          continue;
+        }
+        if (error) throw error;
+        if (done) return;
+        await new Promise((r) => {
+          notify = r;
+        });
+      }
+    } finally {
+      this.pendingStreams.delete(id);
+    }
+  }
+  async unload() {
+    if (!this.loaded) return;
+    if (this.currentUnload) {
+      throw new ModelLoadError("Another unload is already in progress.");
+    }
+    const id = this.allocateId();
+    this.currentUnloadId = id;
+    return new Promise((resolve, reject) => {
+      this.currentUnload = { resolve, reject };
+      this.send({ op: "unload", id });
+    });
+  }
+  /** Tear down the underlying worker. The engine is unusable after this. */
+  terminate() {
+    this.worker.removeEventListener("message", this.listener);
+    this.worker.terminate();
+    this.loaded = false;
+  }
+  allocateId() {
+    const id = this.nextId;
+    this.nextId += 1;
+    return id;
+  }
+  send(req) {
+    this.worker.postMessage(req);
+  }
+  handleMessage(msg) {
+    switch (msg.op) {
+      case "loaded":
+        if (this.currentLoad && msg.id === this.currentLoadId) {
+          this.loaded = true;
+          this.currentLoad.resolve();
+          this.currentLoad = null;
+          this.currentLoadProgress = void 0;
+        }
+        return;
+      case "progress":
+        if (msg.id === this.currentLoadId) {
+          this.currentLoadProgress?.(msg.payload);
+        }
+        return;
+      case "generated": {
+        const pending = this.pendingGenerates.get(msg.id);
+        if (pending) {
+          pending.resolve(msg.text);
+          this.pendingGenerates.delete(msg.id);
+        }
+        return;
+      }
+      case "token": {
+        const stream = this.pendingStreams.get(msg.id);
+        stream?.push(msg.chunk);
+        return;
+      }
+      case "stream-end": {
+        const stream = this.pendingStreams.get(msg.id);
+        stream?.end();
+        return;
+      }
+      case "unloaded":
+        if (this.currentUnload && msg.id === this.currentUnloadId) {
+          this.loaded = false;
+          this.currentUnload.resolve();
+          this.currentUnload = null;
+        }
+        return;
+      case "is-loaded":
+        return;
+      case "error": {
+        const err = mapError(msg.name, msg.message);
+        if (this.currentLoad && msg.id === this.currentLoadId) {
+          this.currentLoad.reject(err);
+          this.currentLoad = null;
+          this.currentLoadProgress = void 0;
+          return;
+        }
+        if (this.currentUnload && msg.id === this.currentUnloadId) {
+          this.currentUnload.reject(err);
+          this.currentUnload = null;
+          return;
+        }
+        const generate = this.pendingGenerates.get(msg.id);
+        if (generate) {
+          generate.reject(err);
+          this.pendingGenerates.delete(msg.id);
+          return;
+        }
+        const stream = this.pendingStreams.get(msg.id);
+        if (stream) {
+          stream.fail(err);
+          return;
+        }
+        return;
+      }
+    }
+  }
+}
+function mapError(name, message) {
+  switch (name) {
+    case "ModelLoadError":
+      return new ModelLoadError(message);
+    case "ModelNotLoadedError":
+      return new ModelNotLoadedError(message);
+    case "GenerationAbortedError":
+      return new GenerationAbortedError(message);
+    default: {
+      const err = new Error(message);
+      err.name = name;
+      return err;
+    }
+  }
+}
 const MODEL_PRESETS = Object.freeze({
   "phi-3.5-mini-int4": {
     id: "phi-3.5-mini-int4",
@@ -183,6 +524,15 @@ function resolveModelPreset(modelId) {
 function listSupportedModels() {
   return Object.keys(MODEL_PRESETS);
 }
+function createInferenceWorker() {
+  return new Worker(new URL(
+    /* @vite-ignore */
+    "/assets/inference.worker-CwvQtobb.js",
+    import.meta.url
+  ), {
+    type: "module"
+  });
+}
 class LMTask {
   constructor(engine, preset) {
     this.engine = engine;
@@ -198,12 +548,18 @@ class LMTask {
    */
   static async createEngine(modelId, options = {}) {
     const preset = resolveModelPreset(modelId);
-    const engine = options.engine ?? new WebLLMEngine();
+    const engine = options.engine ?? LMTask.defaultEngine(options);
     if (!engine.isLoaded()) {
       await engine.load(preset.webllmId, options.onProgress);
     }
     return { engine, preset };
   }
+  static defaultEngine(options) {
+    if (options.inWorker) {
+      return new WorkerEngine(createInferenceWorker());
+    }
+    return new WebLLMEngine();
+  }
   /** Release engine resources. Safe to call multiple times. */
   async unload() {
     await this.engine.unload();
@@ -221,6 +577,14 @@ class ChatReply {
     this.finishReason = finishReason;
   }
 }
+class CompletionResult {
+  constructor(text, prompt, tokensGenerated, finishReason) {
+    this.text = text;
+    this.prompt = prompt;
+    this.tokensGenerated = tokensGenerated;
+    this.finishReason = finishReason;
+  }
+}
 class Chat extends LMTask {
   history = [];
   systemPrompt = null;
@@ -300,6 +664,152 @@ class Chat extends LMTask {
     return messages;
   }
 }
+class Completion extends LMTask {
+  constructor(engine, preset) {
+    super(engine, preset);
+  }
+  /**
+   * Create and load a `Completion` task for the given model.
+   *
+   * @param modelId - Friendly model id from the registry (e.g. `"qwen2.5-1.5b-int4"`).
+   * @param options - Optional creation options (progress callback, engine override).
+   */
+  static async create(modelId, options = {}) {
+    const { engine, preset } = await LMTask.createEngine(modelId, options);
+    return new Completion(engine, preset);
+  }
+  /**
+   * Generate a continuation for the given prompt.
+   *
+   * @param prompt - Raw text fed to the model.
+   * @param options - Generation options.
+   * @returns A {@link CompletionResult} with the generated continuation.
+   */
+  async predict(prompt, options = {}) {
+    const text = await this.engine.complete(prompt, options);
+    return new CompletionResult(text, prompt, 0, "stop");
+  }
+  /**
+   * Stream a continuation for the given prompt as an async iterable of token
+   * chunks.
+   *
+   * @param prompt - Raw text fed to the model.
+   * @param options - Generation options including an optional `signal`.
+   */
+  async *stream(prompt, options = {}) {
+    for await (const chunk of this.engine.streamCompletion(prompt, options)) {
+      yield chunk;
+    }
+  }
+}
+let webllmCachePromise = null;
+async function loadWebLLMCacheHelpers() {
+  if (!webllmCachePromise) {
+    webllmCachePromise = import("@mlc-ai/web-llm").then((m) => ({
+      hasModelInCache: m.hasModelInCache,
+      deleteModelInCache: m.deleteModelInCache
+    }));
+  }
+  return webllmCachePromise;
+}
+async function defaultEstimate() {
+  if (typeof navigator === "undefined" || !navigator.storage?.estimate) {
+    return { usage: 0, quota: 0 };
+  }
+  const estimate = await navigator.storage.estimate();
+  return {
+    usage: estimate.usage ?? 0,
+    quota: estimate.quota ?? 0
+  };
+}
+class ModelCache {
+  hasModelHook;
+  deleteModelHook;
+  estimateHook;
+  constructor(options = {}) {
+    this.hasModelHook = options.hasModel;
+    this.deleteModelHook = options.deleteModel;
+    this.estimateHook = options.estimate ?? defaultEstimate;
+  }
+  /**
+   * Whether the model's weights are present in the browser cache.
+   *
+   * @param modelId - Friendly id from the registry.
+   * @throws UnknownModelError if `modelId` is not in the registry.
+   */
+  async has(modelId) {
+    const backendId = resolveModelPreset(modelId).webllmId;
+    const fn = this.hasModelHook ?? (await loadWebLLMCacheHelpers()).hasModelInCache;
+    return fn(backendId);
+  }
+  /**
+   * Delete a single model's weights from the browser cache. No-op when the
+   * model is not cached.
+   *
+   * @param modelId - Friendly id from the registry.
+   * @throws UnknownModelError if `modelId` is not in the registry.
+   */
+  async delete(modelId) {
+    const backendId = resolveModelPreset(modelId).webllmId;
+    const fn = this.deleteModelHook ?? (await loadWebLLMCacheHelpers()).deleteModelInCache;
+    await fn(backendId);
+  }
+  /**
+   * List the registry models that are currently cached.
+   *
+   * Iterates `MODEL_PRESETS` and probes each one. Only returns models known
+   * to the SDK — models cached by external WebLLM calls outside our registry
+   * are not included.
+   *
+   * @returns Empty list when nothing is cached.
+   */
+  async list() {
+    const fn = this.hasModelHook ?? (await loadWebLLMCacheHelpers()).hasModelInCache;
+    const probes = await Promise.all(
+      Object.values(MODEL_PRESETS).map(async (preset) => {
+        const cached = await fn(preset.webllmId);
+        if (!cached) return null;
+        const entry = {
+          id: preset.id,
+          backendId: preset.webllmId,
+          family: preset.family,
+          parameters: preset.parameters
+        };
+        return entry;
+      })
+    );
+    return probes.filter((p) => p !== null);
+  }
+  /**
+   * Delete every registry model from the cache. Useful for logout flows or
+   * "reset" buttons. Models cached outside the registry are not touched.
+   */
+  async clear() {
+    const fn = this.deleteModelHook ?? (await loadWebLLMCacheHelpers()).deleteModelInCache;
+    await Promise.all(Object.values(MODEL_PRESETS).map((p) => fn(p.webllmId)));
+  }
+  /**
+   * Aggregate storage stats from the browser. Returned numbers cover the
+   * entire origin (Cache API + IndexedDB + Service Workers + OPFS), not
+   * just our model cache — use it for "you have X of Y available" hints.
+   */
+  async estimateUsage() {
+    return this.estimateHook();
+  }
+  /**
+   * Throw a descriptive error if the given id is not in the registry.
+   * Exposed for code paths that want to validate before calling other
+   * methods (those already throw on their own).
+   *
+   * @throws UnknownModelError
+   */
+  static assertKnown(modelId) {
+    if (!(modelId in MODEL_PRESETS)) {
+      const available = Object.keys(MODEL_PRESETS).join(", ");
+      throw new UnknownModelError(`Unknown model "${modelId}". Available models: ${available}.`);
+    }
+  }
+}
 async function collectStream(stream) {
   let acc = "";
   for await (const chunk of stream) {
@@ -313,22 +823,27 @@ async function* tap(stream, onChunk) {
     yield chunk;
   }
 }
-const VERSION = "0.1.0";
+const VERSION = "0.2.0";
 export {
   BackendNotAvailableError,
   Chat,
   ChatReply,
+  Completion,
+  CompletionResult,
   GenerationAbortedError,
   LMTask,
   LocalmWebError,
   MODEL_PRESETS,
+  ModelCache,
   ModelLoadError,
   ModelNotLoadedError,
   QuotaExceededError,
   UnknownModelError,
   VERSION,
   WebGPUUnavailableError,
+  WorkerEngine,
   collectStream,
+  createInferenceWorker,
   listSupportedModels,
   resolveModelPreset,
   tap