chromium/ash/webui/recorder_app_ui/resources/core/on_device_model/types.ts

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

import {ReadonlySignal} from '../reactive/signal.js';

/**
 * Model installation state.
 *
 * This is also used by SODA.
 */
// prettier-ignore
export type ModelState = {
  kind: 'error'|'installed'|'notInstalled'|'unavailable',
}|{
  kind: 'installing',

  /**
   * A number between 0 to 100 indicating the progress of the download / install
   * of SODA or on-device model.
   */
  progress: number,
};

/**
 * Possible error types from model responses.
 */
export enum ModelResponseError {
  // General error.
  GENERAL = 'GENERAL',

  // Filtered by T&S on the request or response string.
  UNSAFE = 'UNSAFE',
}

// prettier-ignore
export type ModelResponse<T> = {
  kind: 'error',
  error: ModelResponseError,
}|{
  kind: 'success',
  result: T,
};

export abstract class ModelLoader<T> {
  /**
   * The state of the model.
   */
  abstract state: ReadonlySignal<ModelState>;

  /**
   * Loads the model.
   */
  abstract load(): Promise<Model<T>>;

  /**
   * Requests download of the given model.
   */
  download(): void {
    // TODO(pihsun): There's currently no way of requesting download of the
    // model but not load it, so we load the model (which downloads the model)
    // and then immediately unloads it. Check the performance overhead and
    // consider adding another API for only downloading the model if the
    // overhead is large.
    void this.load().then((model) => {
      model.close();
    });
  }
}

export interface Model<T> {
  /**
   * Returns the model response based on content.
   */
  execute(content: string): Promise<ModelResponse<T>>;

  /**
   * Closes the model connection.
   *
   * This should release resources used by the model, and no further call of
   * other calls should happen after this.
   */
  close(): void;
}