alwaysai/src/core/model/model-package-cache.ts

The alwaysAI command-line interface (CLI)

import { join, dirname } from 'path';
import { rename, createWriteStream, existsSync, createReadStream } from 'fs';
import { promisify } from 'util';
import mkdirp = require('mkdirp');
import pump = require('pump');
import rimraf from 'rimraf';

import { ModelId } from './model-id';
import { MODEL_PACKAGE_CACHE_DIR } from '../../paths';
import { RandomString } from '../../util';
import { getSystemId } from '../../infrastructure';
import { Readable } from 'node:stream';

export const modelPackageCache = {
  has(id: string, version: number) {
    const modelPackagePath = ModelPackagePath(id, version);
    return existsSync(modelPackagePath);
  },

  read(id: string, version: number) {
    const modelPackagePath = ModelPackagePath(id, version);
    return createReadStream(modelPackagePath);
  },

  async clear() {
    await rimraf(ModelPackageDirectory());
  },

  async write(id: string, version: number, readable: Readable) {
    const modelPackagePath = ModelPackagePath(id, version);
    mkdirp.sync(dirname(modelPackagePath));
    const tmpFilePath = `${modelPackagePath}.${RandomString()}.tmp`;
    try {
      await new Promise<void>((resolve, reject) => {
        const writeable = createWriteStream(tmpFilePath);
        pump(readable, writeable, (err) => {
          if (err) {
            reject(err);
          } else {
            resolve();
          }
        });
      });
      await promisify(rename)(tmpFilePath, modelPackagePath);
    } catch (exception) {
      try {
        // Attempt to delete the new files since they could be corrupted
        await rimraf(tmpFilePath);
        await rimraf(modelPackagePath);
      } finally {
        // Do nothing
      }
      throw exception;
    }
  },

  async remove(id: string, version: number) {
    const modelPackagePath = ModelPackagePath(id, version);
    await rimraf(modelPackagePath);
  }
};

export function ModelPackageDirectory() {
  return join(MODEL_PACKAGE_CACHE_DIR, getSystemId());
}

export function ModelExists(modelId: string) {
  return existsSync(join(ModelPackageDirectory(), modelId));
}

function ModelPackagePath(id: string, version: number) {
  const { publisher, name } = ModelId.parse(id);
  return join(ModelPackageDirectory(), publisher, name, `${version}.tar.gz`);
}