fix: handle multi download model, uninstall script (#932)

marknguyen1302 · web-flow · commit 095415234fd8 · 2024-07-29T14:23:41.000+07:00
diff --git a/cortex-js/package.json b/cortex-js/package.json
@@ -34,7 +34,8 @@
     "build:dev": "yarn build && run-script-os && npm link",
     "build:dev:windows": "echo 'Windows build complete'",
     "build:dev:macos": "chmod +x ./dist/src/command.js",
-    "build:dev:linux": "chmod +x ./dist/src/command.js"
+    "build:dev:linux": "chmod +x ./dist/src/command.js",
+    "preuninstall": "node ./uninstall.js"
   },
   "dependencies": {
     "@cortexso/cortex.js": "^0.1.3",
diff --git a/cortex-js/src/infrastructure/commanders/chat.command.ts b/cortex-js/src/infrastructure/commanders/chat.command.ts
@@ -16,8 +16,9 @@ import { FileManagerService } from '../services/file-manager/file-manager.servic
 import { isRemoteEngine } from '@/utils/normalize-model-id';
 import { Cortex } from '@cortexso/cortex.js';
 import { ChatClient } from './services/chat-client';
-import { downloadModelProgress } from '@/utils/pull-model';
+import { downloadProgress } from '@/utils/download-progress';
 import { CortexClient } from './services/cortex.client';
+import { DownloadType } from '@/domain/models/download.interface';
 
 type ChatOptions = {
   threadId?: string;
@@ -92,7 +93,7 @@ export class ChatCommand extends BaseCommand {
     ) {
       console.log('Downloading engine...');
       await this.cortex.engines.init(engine);
-      await downloadModelProgress(this.cortex);
+      await downloadProgress(this.cortex, undefined, DownloadType.Engine)
     }
 
     if (!message) options.attach = true;
@@ -107,7 +108,7 @@ export class ChatCommand extends BaseCommand {
     );
 
     const preset = await this.fileService.getPreset(options.preset);
-
+    
     return this.cortex.models.start(modelId, preset).then(() =>
       this.chatClient.chat(
         modelId,
diff --git a/cortex-js/src/infrastructure/commanders/engines/engines-init.command.ts b/cortex-js/src/infrastructure/commanders/engines/engines-init.command.ts
@@ -9,6 +9,7 @@ import { BaseCommand } from '../base.command';
 import { defaultInstallationOptions } from '@/utils/init';
 import { Presets, SingleBar } from 'cli-progress';
 import { CortexClient } from '../services/cortex.client';
+import ora from 'ora';
 
 @SubCommand({
   name: '<name> init',
@@ -44,9 +45,11 @@ export class EnginesInitCommand extends BaseCommand {
     const host = configs.cortexCppHost;
     const port = configs.cortexCppPort;
     // Should stop cortex before installing engine
+    const stopCortexSpinner = ora('Stopping cortex...').start();
     if (await this.cortexUsecases.healthCheck(host, port)) {
       await this.cortexUsecases.stopCortex();
     }
+    stopCortexSpinner.succeed('Cortex stopped');
     console.log(`Installing engine ${engine}...`);
     await this.cortex.engines.init(engine, params);
     const response = await this.cortex.events.downloadEvent();
@@ -68,6 +71,9 @@ export class EnginesInitCommand extends BaseCommand {
       }
     }
     progressBar.stop();
+    const startCortexSpinner = ora('Starting cortex...').start();
+    await this.cortexUsecases.startCortex();
+    startCortexSpinner.succeed('Cortex started');
     console.log('Engine installed successfully');
     process.exit(0);
   }
diff --git a/cortex-js/src/infrastructure/commanders/models/model-pull.command.ts b/cortex-js/src/infrastructure/commanders/models/model-pull.command.ts
@@ -15,8 +15,9 @@ import { checkModelCompatibility } from '@/utils/model-check';
 import { Engines } from '../types/engine.interface';
 import { CortexUsecases } from '@/usecases/cortex/cortex.usecases';
 import { BaseCommand } from '../base.command';
-import { downloadModelProgress } from '@/utils/pull-model';
+import { downloadProgress } from '@/utils/download-progress';
 import { CortexClient } from '../services/cortex.client';
+import { DownloadType } from '@/domain/models/download.interface';
 
 @SubCommand({
   name: 'pull',
@@ -58,7 +59,7 @@ export class ModelPullCommand extends BaseCommand {
       exit(1);
     });
 
-    await downloadModelProgress(this.cortex, modelId);
+    await downloadProgress(this.cortex, modelId);
 
     const existingModel = await this.cortex.models.retrieve(modelId);
     const engine = existingModel?.engine || Engines.llamaCPP;
@@ -70,7 +71,7 @@ export class ModelPullCommand extends BaseCommand {
       console.log('\n');
       console.log('Downloading engine...');
       await this.cortex.engines.init(engine);
-      await downloadModelProgress(this.cortex);
+      await downloadProgress(this.cortex, undefined, DownloadType.Engine);
     }
     this.telemetryUsecases.sendEvent(
       [
diff --git a/cortex-js/src/infrastructure/commanders/models/model-start.command.ts b/cortex-js/src/infrastructure/commanders/models/model-start.command.ts
@@ -11,8 +11,9 @@ import { Engines } from '../types/engine.interface';
 import { checkModelCompatibility } from '@/utils/model-check';
 import { BaseCommand } from '../base.command';
 import { isRemoteEngine } from '@/utils/normalize-model-id';
-import { downloadModelProgress } from '@/utils/pull-model';
+import { downloadProgress } from '@/utils/download-progress';
 import { CortexClient } from '../services/cortex.client';
+import { DownloadType } from '@/domain/models/download.interface';
 
 type ModelStartOptions = {
   attach: boolean;
@@ -73,7 +74,7 @@ export class ModelStartCommand extends BaseCommand {
     ) {
       console.log('Downloading engine...');
       await this.cortex.engines.init(engine);
-      await downloadModelProgress(this.cortex);
+      await downloadProgress(this.cortex, undefined, DownloadType.Engine);
     }
 
     // Attached - stdout logs
diff --git a/cortex-js/src/infrastructure/commanders/run.command.ts b/cortex-js/src/infrastructure/commanders/run.command.ts
@@ -10,8 +10,9 @@ import { checkModelCompatibility } from '@/utils/model-check';
 import { BaseCommand } from './base.command';
 import { isRemoteEngine } from '@/utils/normalize-model-id';
 import { ChatClient } from './services/chat-client';
-import { downloadModelProgress } from '@/utils/pull-model';
+import { downloadProgress } from '@/utils/download-progress';
 import { CortexClient } from './services/cortex.client';
+import { DownloadType } from '@/domain/models/download.interface';
 
 type RunOptions = {
   threadId?: string;
@@ -65,7 +66,7 @@ export class RunCommand extends BaseCommand {
         checkingSpinner.fail(e.message ?? e);
         exit(1);
       });
-      await downloadModelProgress(this.cortex, modelId);
+      await downloadProgress(this.cortex, modelId);
     }
 
     // Second check if model is available
@@ -84,7 +85,7 @@ export class RunCommand extends BaseCommand {
     ) {
       console.log('Downloading engine...');
       await this.cortex.engines.init(engine);
-      await downloadModelProgress(this.cortex);
+      await downloadProgress(this.cortex, undefined, DownloadType.Engine);
     }
 
     const startingSpinner = ora('Loading model...').start();
diff --git a/cortex-js/src/infrastructure/controllers/engines.controller.ts b/cortex-js/src/infrastructure/controllers/engines.controller.ts
@@ -7,8 +7,10 @@ import {
   Post,
   Body,
   Patch,
+  Res,
 } from '@nestjs/common';
 import { ApiOperation, ApiParam, ApiTags, ApiResponse } from '@nestjs/swagger';
+import { Response } from 'express';
 import { TransformInterceptor } from '../interceptors/transform.interceptor';
 import { EnginesUsecases } from '@/usecases/engines/engines.usecase';
 import { EngineDto, InitEngineDto } from '../dtos/engines/engines.dto';
@@ -78,11 +80,15 @@ export class EnginesController {
     description: 'The unique identifier of the engine.',
   })
   @Post(':name(*)/init')
-  initialize(@Param('name') name: string, @Body() body: InitEngineDto | undefined) {
+  initialize(@Param('name') name: string, @Body() body: InitEngineDto | undefined, @Res() res: Response) {
+    try{
     this.initUsescases.installEngine(body, 'latest', name, true);
-    return {
-      message: 'Engine initialization started successfully.',
-    };
+    res.json({
+        message: 'Engine initialization started successfully.',
+    })
+  } catch (error) {
+    res.status(400).send(error.message);
+  }
   }
 
   @HttpCode(200)
diff --git a/cortex-js/src/usecases/engines/engines.usecase.ts b/cortex-js/src/usecases/engines/engines.usecase.ts
@@ -79,9 +79,11 @@ export class EnginesUsecases {
     force: boolean = false,
   ): Promise<any> => {
     // Use default option if not defined
+
     if (!options && engine === Engines.llamaCPP) {
       options = await defaultInstallationOptions();
     }
+    const installPackages = [];
     // Ship Llama.cpp engine by default
     if (
       !existsSync(
@@ -93,7 +95,7 @@ export class EnginesUsecases {
         engine === Engines.llamaCPP &&
         (options?.vulkan ||
           (options?.runMode === 'GPU' && options?.gpuType !== 'Nvidia'));
-      await this.installAcceleratedEngine(version, engine, [
+          installPackages.push(this.installAcceleratedEngine(version, engine, [
         process.platform === 'win32'
           ? '-windows'
           : process.platform === 'darwin'
@@ -116,7 +118,7 @@ export class EnginesUsecases {
             ? '-arm64'
             : '-amd64'
           : '',
-      ]);
+      ]));
     }
 
     if (
@@ -126,12 +128,12 @@ export class EnginesUsecases {
         options?.gpuType === 'Nvidia' &&
         !options?.vulkan)
     )
-      await this.installCudaToolkitDependency(
+      installPackages.push(this.installCudaToolkitDependency(
         engine === Engines.tensorrtLLM
           ? MIN_CUDA_VERSION
           : options?.cudaVersion,
-      );
-
+      ));
+    await Promise.all(installPackages);
     // Update states
     await this.extensionRepository.findOne(engine).then((e) => {
       if (e) e.status = EngineStatus.READY;
@@ -254,7 +256,6 @@ export class EnginesUsecases {
       console.log(
         `Could not find engine file for platform ${process.platform}`,
       );
-      exit(1);
     }
 
     const engineDir = await this.fileManagerService.getCortexCppEnginePath();
diff --git a/cortex-js/src/usecases/models/models.usecases.ts b/cortex-js/src/usecases/models/models.usecases.ts
@@ -119,7 +119,14 @@ export class ModelsUsecases {
     }
 
     const modelFolder = join(modelsContainerDir, normalizeModelId(id));
+    const model = await this.getModelOrThrow(id);
+    const engine = (await this.extensionRepository.findOne(
+      model!.engine ?? Engines.llamaCPP,
+    )) as EngineExtension | undefined;
 
+    if (engine) {
+      await engine.unloadModel(id, model.engine || Engines.llamaCPP).catch(() => {}); // Silent fail
+    }
     return this.modelRepository
       .remove(id)
       .then(
diff --git a/cortex-js/src/utils/download-progress.ts b/cortex-js/src/utils/download-progress.ts
@@ -1,8 +1,9 @@
 import { Presets, SingleBar } from "cli-progress";
 import { Cortex } from "@cortexso/cortex.js";
 import { exit, stdin, stdout } from 'node:process';
+import { DownloadType } from "@/domain/models/download.interface";
 
-export const downloadModelProgress = async (cortex: Cortex, downloadId?: string) => {
+export const downloadProgress = async (cortex: Cortex, downloadId?: string, downloadType?: DownloadType) => {
     const response = await cortex.events.downloadEvent();
 
     const rl = require('readline').createInterface({
@@ -27,6 +28,7 @@ export const downloadModelProgress = async (cortex: Cortex, downloadId?: string)
     for await (const stream of response) {
       if (stream.length) {
         const data = stream[0] as any;
+        if (downloadId && data.id !== downloadId || downloadType && data.type !== downloadType) continue;
 
         if (data.status === 'downloaded') break;
 
diff --git a/cortex-js/uninstall.js b/cortex-js/uninstall.js
@@ -0,0 +1,32 @@
+const { promises } = require('node:fs');
+const os = require('os');
+const fs = require('fs');
+const path = require('path');
+const yaml = require('js-yaml');
+
+const uninstall = async () => {
+    const cortexConfigPath = path.join(os.homedir(), '.cortexrc');
+    if (fs.existsSync(cortexConfigPath)) {
+        const content = await promises.readFile(cortexConfigPath, 'utf8');
+        const config = yaml.load(content);
+        if(!config) {
+            return;
+        }   
+        const { dataFolderPath } = config;
+        const modelsFolderPath = path.join(dataFolderPath, 'models');
+        // remove all data in data folder path except models
+        const files = fs.readdirSync(dataFolderPath);
+        for (const file of files) {
+            const fileStat = fs.statSync(path.join(dataFolderPath, file));
+            if (file !== 'models') {
+                if (fileStat.isDirectory()) {
+                    fs.rmSync(path.join(dataFolderPath, file), { recursive: true });
+                } else {
+                    fs.unlinkSync(path.join(dataFolderPath, file));
+                }
+            }
+        }
+    }
+};
+
+uninstall();
