@agorapete/wllama 3.5.1-q2.0

package/src/model-manager.ts

@@ -0,0 +1,359 @@
+import CacheManager, {
+  type CacheEntryMetadata,
+  type CacheEntry,
+  type DownloadOptions,
+} from './cache-manager';
+import { isString, isValidGgufFile, sumArr } from './utils';
+import { WllamaError, type WllamaLogger } from './wllama';
+
+const DEFAULT_PARALLEL_DOWNLOADS = 3;
+
+export interface ModelSource {
+  /**
+   * URL to the GGUF file. If the model is splitted, pass the URL to the first shard. Mmproj file should also be provided via this array if you want to use multimodal.
+   */
+  url: string;
+  mmprojUrl?: string;
+}
+
+/**
+ * Callback function to track download progress
+ */
+export type DownloadProgressCallback = (opts: {
+  /**
+   * Number of bytes loaded (sum of all shards)
+   */
+  loaded: number;
+  /**
+   * Total number of bytes (sum of all shards)
+   */
+  total: number;
+}) => any;
+
+/**
+ * Status of the model validation
+ */
+export enum ModelValidationStatus {
+  VALID = 'valid',
+  INVALID = 'invalid',
+  DELETED = 'deleted',
+}
+
+/**
+ * Parameters for ModelManager constructor
+ */
+export interface ModelManagerParams {
+  cacheManager?: CacheManager;
+  logger?: WllamaLogger;
+  /**
+   * Number of parallel downloads
+   *
+   * Default: 3
+   */
+  parallelDownloads?: number | undefined;
+  /**
+   * Allow offline mode
+   *
+   * Default: false
+   */
+  allowOffline?: boolean | undefined;
+}
+
+/**
+ * Model class
+ *
+ * One model can have multiple shards, each shard is a GGUF file.
+ */
+export class Model {
+  private modelManager: ModelManager;
+  constructor(
+    modelManager: ModelManager,
+    url: string,
+    mmprojUrl?: string,
+    savedFiles?: CacheEntry[]
+  ) {
+    this.modelManager = modelManager;
+    this.url = url;
+    this.mmprojUrl = mmprojUrl;
+    if (savedFiles) {
+      // this file is already in cache
+      this.files = this.getAllFiles(savedFiles);
+      this.size = sumArr(this.files.map((f) => f.metadata.originalSize));
+    } else {
+      // this file is not in cache, we are about to download it
+      this.files = [];
+      this.size = 0;
+    }
+  }
+  /**
+   * URL to the GGUF file (in case it contains multiple shards, the URL should point to the first shard)
+   *
+   * This URL will be used to identify the model in the cache. There can't be 2 models with the same URL.
+   */
+  url: string;
+  /**
+   * URL to mmproj file, if exists
+   */
+  mmprojUrl?: string | undefined;
+  /**
+   * Size in bytes (total size of all shards).
+   *
+   * A value of -1 means the model is deleted from the cache. You must call `ModelManager.downloadModel` to re-download the model.
+   */
+  size: number;
+  /**
+   * List of all shards in the cache, sorted by original URL (ascending order)
+   */
+  files: CacheEntry[];
+  /**
+   * Open and get a list of all shards as Blobs
+   */
+  async open(): Promise<Blob[]> {
+    if (this.size === -1) {
+      throw new WllamaError(
+        `Model is deleted from the cache; Call ModelManager.downloadModel to re-download the model`,
+        'load_error'
+      );
+    }
+    const blobs: Blob[] = [];
+    for (const file of this.files) {
+      const blob = await this.modelManager.cacheManager.open(file.name);
+      if (!blob) {
+        throw new Error(
+          `Failed to open file ${file.name}; Hint: the model may be invalid, please refresh it`
+        );
+      }
+      blobs.push(blob);
+    }
+    return blobs;
+  }
+  /**
+   * Validate the model files.
+   *
+   * If the model is invalid, the model manager will not be able to use it. You must call `refresh` to re-download the model.
+   *
+   * Cases that model is invalid:
+   * - The model is deleted from the cache
+   * - The model files are missing (or the download is interrupted)
+   */
+  validate(): ModelValidationStatus {
+    let nbShards = ModelManager.parseModelUrl(this.url).length;
+    if (this.mmprojUrl) {
+      nbShards += 1;
+    }
+    if (this.size === -1) {
+      return ModelValidationStatus.DELETED;
+    }
+    if (this.size < 16 || this.files.length !== nbShards) {
+      return ModelValidationStatus.INVALID;
+    }
+    for (const file of this.files) {
+      if (!file.metadata || file.metadata.originalSize !== file.size) {
+        return ModelValidationStatus.INVALID;
+      }
+    }
+    return ModelValidationStatus.VALID;
+  }
+  /**
+   * In case the model is invalid, call this function to re-download the model
+   */
+  async refresh(options: DownloadOptions = {}): Promise<void> {
+    const urls = ModelManager.parseModelUrl(this.url);
+    if (this.mmprojUrl) {
+      urls.push(this.mmprojUrl);
+    }
+    const works = urls.map((url, index) => ({
+      url,
+      index,
+    }));
+    this.modelManager.logger.debug('Downloading model files:', urls);
+    const nParallel =
+      this.modelManager.params.parallelDownloads ?? DEFAULT_PARALLEL_DOWNLOADS;
+    const totalSize = await this.getTotalDownloadSize(urls);
+    const loadedSize: number[] = [];
+    const worker = async () => {
+      while (works.length > 0) {
+        const w = works.shift();
+        if (!w) break;
+        await this.modelManager.cacheManager.download(w.url, {
+          ...options,
+          metadataAdditional: {
+            originalURL: w.url,
+            mmprojURL: this.mmprojUrl,
+          } satisfies Partial<CacheEntryMetadata>,
+          progressCallback: ({ loaded }) => {
+            loadedSize[w.index] = loaded;
+            options.progressCallback?.({
+              loaded: sumArr(loadedSize),
+              total: totalSize,
+            });
+          },
+        });
+      }
+    };
+    const promises: Promise<void>[] = [];
+    for (let i = 0; i < nParallel; i++) {
+      promises.push(worker());
+      loadedSize.push(0);
+    }
+    await Promise.all(promises);
+    this.files = this.getAllFiles(await this.modelManager.cacheManager.list());
+    this.size = this.files.reduce((acc, f) => acc + f.metadata.originalSize, 0);
+  }
+  /**
+   * Remove the model from the cache
+   */
+  async remove(): Promise<void> {
+    this.files = this.getAllFiles(await this.modelManager.cacheManager.list());
+    await this.modelManager.cacheManager.deleteMany(
+      (f) => !!this.files.find((file) => file.name === f.name)
+    );
+    this.size = -1;
+  }
+
+  private getAllFiles(savedFiles: CacheEntry[]): CacheEntry[] {
+    const allUrls = new Set(ModelManager.parseModelUrl(this.url));
+    if (this.mmprojUrl) {
+      allUrls.add(this.mmprojUrl);
+    }
+    // console.log({allUrls, savedFiles});
+    const allFiles: CacheEntry[] = [];
+    for (const url of allUrls) {
+      const file = savedFiles.find((f) => f.metadata.originalURL === url);
+      if (!file) {
+        throw new Error(`Model file not found: ${url}`);
+      }
+      allFiles.push(file);
+    }
+    allFiles.sort((a, b) =>
+      a.metadata.originalURL.localeCompare(b.metadata.originalURL)
+    );
+    return allFiles;
+  }
+
+  private async getTotalDownloadSize(urls: string[]): Promise<number> {
+    const responses = await Promise.all(
+      urls.map((url) => fetch(url, { method: 'HEAD' }))
+    );
+    const sizes = responses.map((res) =>
+      Number(res.headers.get('content-length') || '0')
+    );
+    return sumArr(sizes);
+  }
+}
+
+export class ModelManager {
+  // The CacheManager singleton, can be accessed by user
+  public cacheManager: CacheManager;
+
+  public params: ModelManagerParams;
+  public logger: WllamaLogger;
+
+  constructor(params: ModelManagerParams = {}) {
+    this.cacheManager = params.cacheManager || new CacheManager();
+    this.params = params;
+    this.logger = params.logger || console;
+  }
+
+  /**
+   * Parses a model URL and returns an array of URLs based on the following patterns:
+   * - If the input URL is an array, it returns the array itself.
+   * - If the input URL is a string in the `gguf-split` format, it returns an array containing the URL of each shard in ascending order.
+   * - Otherwise, it returns an array containing the input URL as a single element array.
+   * @param modelUrl URL or list of URLs
+   */
+  static parseModelUrl(modelUrl: string | string[]): string[] {
+    if (Array.isArray(modelUrl)) {
+      return modelUrl;
+    }
+    const urlPartsRegex = /-(\d{5})-of-(\d{5})\.gguf(?:\?.*)?$/;
+    const queryMatch = modelUrl.match(/\.gguf(\?.*)?$/);
+    const queryParams = queryMatch?.[1] ?? '';
+    const matches = modelUrl.match(urlPartsRegex);
+    if (!matches) {
+      return [modelUrl];
+    }
+    const baseURL = modelUrl.replace(urlPartsRegex, '');
+    const total = matches[2];
+    const paddedShardIds = Array.from({ length: Number(total) }, (_, index) =>
+      (index + 1).toString().padStart(5, '0')
+    );
+    return paddedShardIds.map(
+      (current) => `${baseURL}-${current}-of-${total}.gguf${queryParams}`
+    );
+  }
+
+  /**
+   * Get all models in the cache
+   */
+  async getModels(opts: { includeInvalid?: boolean } = {}): Promise<Model[]> {
+    const cachedFiles = await this.cacheManager.list();
+    let models: Model[] = [];
+    for (const file of cachedFiles) {
+      const shards = ModelManager.parseModelUrl(file.metadata.originalURL);
+      const mmprojUrl = file.metadata.mmprojURL;
+      const isFirstShard =
+        shards.length === 1 || shards[0] === file.metadata.originalURL;
+      if (isFirstShard) {
+        models.push(
+          new Model(this, file.metadata.originalURL, mmprojUrl, cachedFiles)
+        );
+      }
+    }
+    if (!opts.includeInvalid) {
+      models = models.filter(
+        (m) => m.validate() === ModelValidationStatus.VALID
+      );
+    }
+    return models;
+  }
+
+  /**
+   * Download a model from the given URL.
+   *
+   * The URL must end with `.gguf`
+   */
+  async downloadModel(
+    sourceOrURL: ModelSource | string,
+    options: DownloadOptions = {}
+  ): Promise<Model> {
+    const source: ModelSource = isString(sourceOrURL)
+      ? { url: sourceOrURL as string }
+      : (sourceOrURL as ModelSource);
+    if (!isValidGgufFile(source.url)) {
+      throw new WllamaError(
+        `Invalid model URL: ${source.url}; URL must ends with ".gguf"`,
+        'download_error'
+      );
+    }
+    const model = new Model(this, source.url, source.mmprojUrl);
+    const validity = model.validate();
+    if (validity !== ModelValidationStatus.VALID) {
+      await model.refresh(options);
+    }
+    return model;
+  }
+
+  /**
+   * Get a model from the cache or download it if it's not available.
+   */
+  async getModelOrDownload(
+    source: ModelSource,
+    options: DownloadOptions = {}
+  ): Promise<Model> {
+    const models = await this.getModels();
+    const model = models.find((m) => m.url === source.url);
+    if (model) {
+      options.progressCallback?.({ loaded: model.size, total: model.size });
+      return model;
+    }
+    return this.downloadModel(source, options);
+  }
+
+  /**
+   * Remove all models from the cache
+   */
+  async clear(): Promise<void> {
+    await this.cacheManager.clear();
+  }
+}

package/src/storage/cos.test.ts

@@ -0,0 +1,83 @@
+import { test, expect, beforeEach } from 'vitest';
+import { COSBackend, mockCOS } from './cos';
+
+async function randomBufAndHash(): Promise<{
+  buf: Uint8Array;
+  sha256: string;
+}> {
+  const buf = crypto.getRandomValues(new Uint8Array(256));
+  const hashBuf = await crypto.subtle.digest('SHA-256', buf);
+  const sha256 = Array.from(new Uint8Array(hashBuf))
+    .map((b) => b.toString(16).padStart(2, '0'))
+    .join('');
+  return { buf, sha256 };
+}
+
+function bufStream(buf: Uint8Array): ReadableStream<Uint8Array> {
+  return new ReadableStream({
+    start(controller) {
+      controller.enqueue(buf.slice());
+      controller.close();
+    },
+  });
+}
+
+test.sequential('write then read without hint falls back to OPFS', async () => {
+  const backend = new COSBackend();
+  const { buf } = await randomBufAndHash();
+  const key = 'test-no-hint';
+
+  await backend.write(key, bufStream(buf));
+  const blob = await backend.read(key);
+  expect(blob).not.toBeNull();
+  expect(new Uint8Array(await blob!.arrayBuffer())).toEqual(buf);
+
+  await backend.delete(key);
+});
+
+beforeEach(() => {
+  mockCOS();
+});
+
+test.sequential('write with hint stores in COS only', async () => {
+  const backend = new COSBackend();
+  expect(backend.isSupported()).toBe(true);
+
+  const { buf, sha256 } = await randomBufAndHash();
+  const hint = { sha256 };
+  const key = 'test-with-hint';
+
+  await backend.write(key, bufStream(buf), hint);
+
+  // read back via hint → should hit COS
+  const blob = await backend.read(key, hint);
+  expect(blob).not.toBeNull();
+  expect(new Uint8Array(await blob!.arrayBuffer())).toEqual(buf);
+
+  // without hint → OPFS fallback → not found (was written to COS only)
+  const blobOpfs = await backend.read(key);
+  expect(blobOpfs).toBeNull();
+
+  await backend.delete(key);
+});
+
+test.sequential('getSize with hint reflects COS size', async () => {
+  const backend = new COSBackend();
+  const { buf, sha256 } = await randomBufAndHash();
+  const hint = { sha256 };
+  const key = 'test-size-hint';
+
+  await backend.write(key, bufStream(buf), hint);
+
+  const size = await backend.getSize(key, hint);
+  expect(size).toBe(buf.byteLength);
+
+  await backend.delete(key);
+});
+
+test.sequential('read missing key returns null', async () => {
+  const backend = new COSBackend();
+  const { sha256 } = await randomBufAndHash();
+  const blob = await backend.read('non-existent-key', { sha256 });
+  expect(blob).toBeNull();
+});

package/src/storage/cos.ts

@@ -0,0 +1,171 @@
+import type { StorageBackend, StorageFileHint } from './index';
+import { OPFSBackend } from './opfs';
+
+interface CrossOriginStorageRequestFileHandleHash {
+  value: string;
+  algorithm: string;
+}
+
+interface CrossOriginStorageManager {
+  requestFileHandle(
+    hash: CrossOriginStorageRequestFileHandleHash,
+    options?: { create?: boolean; origins?: string[] | string }
+  ): Promise<FileSystemFileHandle>;
+}
+
+declare global {
+  interface Navigator {
+    readonly crossOriginStorage?: CrossOriginStorageManager;
+  }
+}
+
+function makeHash(key: string): CrossOriginStorageRequestFileHandleHash {
+  return { algorithm: 'SHA-256', value: key };
+}
+
+// internal, non-standard implementation
+class COSInternalBackend implements StorageBackend {
+  isSupported(): boolean {
+    return (
+      typeof navigator !== 'undefined' && 'crossOriginStorage' in navigator
+    );
+  }
+
+  // IMPORTANT: key must be SHA-256 hash of the data
+  async read(key: string): Promise<Blob | null> {
+    try {
+      const handle = await navigator.crossOriginStorage!.requestFileHandle(
+        makeHash(key)
+      );
+      return handle.getFile();
+    } catch {
+      return null;
+    }
+  }
+
+  // IMPORTANT: key must be SHA-256 hash of the data
+  async write(key: string, stream: ReadableStream): Promise<void> {
+    const handle = await navigator.crossOriginStorage!.requestFileHandle(
+      makeHash(key),
+      { create: true }
+    );
+    const writable = await (handle as any).createWritable();
+    const reader = stream.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        await writable.write(value);
+      }
+    } finally {
+      await writable.close();
+    }
+  }
+
+  // IMPORTANT: key must be SHA-256 hash of the data
+  async getSize(key: string): Promise<number> {
+    try {
+      const handle = await navigator.crossOriginStorage!.requestFileHandle(
+        makeHash(key)
+      );
+      const file = await handle.getFile();
+      return file.size;
+    } catch {
+      return -1;
+    }
+  }
+
+  async list(): Promise<Array<{ key: string; size: number }>> {
+    throw new Error('not implemented');
+  }
+
+  async delete(_key: string): Promise<void> {
+    throw new Error('not implemented');
+  }
+}
+
+/**
+ * Storage backend that uses the Cross-Origin Storage API
+ * Metadata is stored in OPFS, while the actual data is stored in COS
+ * If hint.sha256 is provided, it will be used as the key for COS, otherwise fallback to OPFS
+ */
+export class COSBackend implements StorageBackend {
+  private cos = new COSInternalBackend();
+  private priv = new OPFSBackend();
+
+  isSupported(): boolean {
+    return this.priv.isSupported();
+  }
+
+  async read(key: string, hint?: StorageFileHint): Promise<Blob | null> {
+    if (hint?.sha256 && this.cos.isSupported()) {
+      const blob = await this.cos.read(hint.sha256);
+      if (blob) return blob;
+    }
+    return this.priv.read(key);
+  }
+
+  async write(
+    key: string,
+    stream: ReadableStream,
+    hint?: StorageFileHint
+  ): Promise<void> {
+    if (hint?.sha256 && this.cos.isSupported()) {
+      await this.cos.write(hint.sha256, stream);
+    } else {
+      await this.priv.write(key, stream);
+    }
+  }
+
+  async getSize(key: string, hint?: StorageFileHint): Promise<number> {
+    if (hint?.sha256 && this.cos.isSupported()) {
+      const size = await this.cos.getSize(hint.sha256);
+      if (size !== -1) return size;
+    }
+    return this.priv.getSize(key);
+  }
+
+  async list(): Promise<Array<{ key: string; size: number }>> {
+    return this.priv.list();
+  }
+
+  async delete(key: string): Promise<void> {
+    return this.priv.delete(key);
+  }
+}
+
+// used for testing only
+export function mockCOS(): void {
+  const store = new Map<string, Blob>();
+
+  (navigator as any).crossOriginStorage = {
+    async requestFileHandle(
+      { value }: CrossOriginStorageRequestFileHandleHash,
+      options?: { create?: boolean }
+    ): Promise<FileSystemFileHandle> {
+      if (!options?.create && !store.has(value)) {
+        throw new DOMException('File not found', 'NotFoundError');
+      }
+      return {
+        getFile() {
+          const blob = store.get(value);
+          if (!blob) throw new DOMException('File not found', 'NotFoundError');
+          return Promise.resolve(new File([blob], value));
+        },
+        createWritable() {
+          const chunks: BlobPart[] = [];
+          return Promise.resolve({
+            write(chunk: BlobPart) {
+              chunks.push(chunk);
+              return Promise.resolve();
+            },
+            close() {
+              store.set(value, new Blob(chunks));
+              return Promise.resolve();
+            },
+          });
+        },
+      } as unknown as FileSystemFileHandle;
+    },
+  } satisfies CrossOriginStorageManager;
+}

package/src/storage/index.ts

@@ -0,0 +1,40 @@
+export interface StorageFileHint {
+  sha256: string;
+}
+
+export interface StorageBackend {
+  isSupported(): boolean;
+
+  /**
+   * Read a file from storage by key.
+   * @returns Blob, or null if the key does not exist
+   */
+  read(key: string, hint?: StorageFileHint): Promise<Blob | null>;
+
+  /**
+   * Write a ReadableStream to storage under the given key.
+   * Overwrites any existing content for that key.
+   */
+  write(
+    key: string,
+    stream: ReadableStream,
+    hint?: StorageFileHint
+  ): Promise<void>;
+
+  /**
+   * Get the stored size of a file in bytes.
+   * @returns number of bytes, or -1 if the key does not exist
+   */
+  getSize(key: string, hint?: StorageFileHint): Promise<number>;
+
+  /**
+   * List all keys currently in storage.
+   * Includes metadata keys - callers are responsible for filtering.
+   */
+  list(): Promise<Array<{ key: string; size: number }>>;
+
+  /**
+   * Delete a single entry by key. No-op if the key does not exist.
+   */
+  delete(key: string): Promise<void>;
+}