Merge pull request #569 from janhq/chore/update-start-model-api

namchuai · web-flow · commit 502ae349aa70 · 2024-05-15T18:27:29.000+07:00
chore: update start model api
diff --git a/cortex-js/constant.ts b/cortex-js/constant.ts
@@ -4,3 +4,6 @@ export const databaseFile = `${databaseName}.db`;
 
 export const defaultCortexJsHost = 'localhost';
 export const defaultCortexJsPort = 7331;
+
+export const defaultCortexCppHost = '127.0.0.1';
+export const defaultCortexCppPort = 3928;
diff --git a/cortex-js/src/domain/abstracts/engine.abstract.ts b/cortex-js/src/domain/abstracts/engine.abstract.ts
@@ -1,12 +1,15 @@
-import { Model } from '../models/model.interface';
+import { Model, ModelSettingParams } from '../models/model.interface';
 import { Extension } from './extension.abstract';
 
 export abstract class EngineExtension extends Extension {
   abstract provider: string;
 
   abstract inference(completion: any, req: any, stream: any, res?: any): void;
 
-  async loadModel(model: Model): Promise<void> {}
+  async loadModel(
+    model: Model,
+    settingParams?: ModelSettingParams,
+  ): Promise<void> {}
 
   async unloadModel(modelId: string): Promise<void> {}
 }
diff --git a/cortex-js/src/infrastructure/commanders/inference.command.ts b/cortex-js/src/infrastructure/commanders/inference.command.ts
@@ -2,35 +2,41 @@ import { ChatUsecases } from '@/usecases/chat/chat.usecases';
 import { CommandRunner, SubCommand } from 'nest-commander';
 import { CreateChatCompletionDto } from '../dtos/chat/create-chat-completion.dto';
 import { ChatCompletionRole } from '@/domain/models/message.interface';
-import { stdout } from 'process';
+import { exit, stdin, stdout } from 'node:process';
 import * as readline from 'node:readline/promises';
 import { ChatStreamEvent } from '@/domain/abstracts/oai.abstract';
 import { ChatCompletionMessage } from '../dtos/chat/chat-completion-message.dto';
 
 @SubCommand({ name: 'chat' })
 export class InferenceCommand extends CommandRunner {
-  exitClause = 'exit()';
-  userIndicator = '>> ';
-  exitMessage = 'Bye!';
+  private exitClause = 'exit()';
+  private userIndicator = '>> ';
+  private exitMessage = 'Bye!';
 
   constructor(private readonly chatUsecases: ChatUsecases) {
     super();
   }
 
-  async run(): Promise<void> {
+  async run(input: string[]): Promise<void> {
+    if (input.length == 0) {
+      console.error('Please provide a model id.');
+      exit(1);
+    }
+
+    const modelId = input[0];
     console.log(`Inorder to exit, type '${this.exitClause}'.`);
     const messages: ChatCompletionMessage[] = [];
 
     const rl = readline.createInterface({
-      input: process.stdin,
-      output: process.stdout,
+      input: stdin,
+      output: stdout,
       prompt: this.userIndicator,
     });
     rl.prompt();
 
     rl.on('close', () => {
       console.log(this.exitMessage);
-      process.exit(0);
+      exit(0);
     });
 
     rl.on('line', (userInput: string) => {
@@ -46,7 +52,7 @@ export class InferenceCommand extends CommandRunner {
 
       const chatDto: CreateChatCompletionDto = {
         messages,
-        model: 'TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF',
+        model: modelId,
         stream: true,
         max_tokens: 2048,
         stop: [],
diff --git a/cortex-js/src/infrastructure/commanders/init.command.ts b/cortex-js/src/infrastructure/commanders/init.command.ts
@@ -4,6 +4,7 @@ import { resolve } from 'path';
 import { HttpService } from '@nestjs/axios';
 import { Presets, SingleBar } from 'cli-progress';
 import decompress from 'decompress';
+import { exit } from 'node:process';
 
 @SubCommand({ name: 'init', aliases: ['setup'] })
 export class InitCommand extends CommandRunner {
@@ -41,7 +42,7 @@ export class InitCommand extends CommandRunner {
 
     if (!res?.data) {
       console.log('Failed to fetch releases');
-      process.exit(1);
+      exit(1);
     }
 
     let release = res?.data;
@@ -56,7 +57,7 @@ export class InitCommand extends CommandRunner {
 
     if (!toDownloadAsset) {
       console.log(`Could not find engine file ${engineFileName}`);
-      process.exit(1);
+      exit(1);
     }
 
     console.log(`Downloading engine file ${engineFileName}`);
@@ -107,9 +108,9 @@ export class InitCommand extends CommandRunner {
       );
     } catch (e) {
       console.log(e);
-      process.exit(1);
+      exit(1);
     }
-    process.exit(0);
+    exit(0);
   };
 
   parseEngineFileName = (options: {
diff --git a/cortex-js/src/infrastructure/commanders/pull.command.ts b/cortex-js/src/infrastructure/commanders/pull.command.ts
@@ -3,6 +3,7 @@ import { CommandRunner, SubCommand } from 'nest-commander';
 import { CreateModelDto } from '../dtos/models/create-model.dto';
 import { ModelFormat } from '@/domain/models/model.interface';
 import { Presets, SingleBar } from 'cli-progress';
+import { exit } from 'node:process';
 
 const AllQuantizations = [
   'Q3_K_S',
@@ -49,7 +50,7 @@ export class PullCommand extends CommandRunner {
       bar.update(progress);
     });
     console.log('\nDownload complete!');
-    process.exit(0);
+    exit(0);
   }
 
   async pullHuggingFaceModel(modelId: string) {
diff --git a/cortex-js/src/infrastructure/commanders/start.command.ts b/cortex-js/src/infrastructure/commanders/start.command.ts
@@ -1,9 +1,10 @@
 import { CortexUsecases } from '@/usecases/cortex/cortex.usecases';
 import { ModelsUsecases } from '@/usecases/models/models.usecases';
 import { CommandRunner, SubCommand } from 'nest-commander';
-import { LoadModelDto } from '../dtos/models/load-model.dto';
 import { resolve } from 'path';
 import { existsSync } from 'fs';
+import { Model, ModelSettingParams } from '@/domain/models/model.interface';
+import { exit } from 'node:process';
 
 @SubCommand({ name: 'start', aliases: ['run'] })
 export class StartCommand extends CommandRunner {
@@ -15,29 +16,32 @@ export class StartCommand extends CommandRunner {
   }
 
   async run(input: string[]): Promise<void> {
+    if (input.length === 0) {
+      console.error('Model ID is required');
+      exit(1);
+    }
     const modelId = input[0];
+    const model = await this.getModelOrStop(modelId);
 
-    if (!modelId) {
-      console.log('Model ID is required');
-      return;
-    }
     return this.startCortex()
-      .then(() => this.startModel(modelId))
+      .then(() => this.startModel(model.id))
       .then(console.log)
       .catch(console.error);
   }
 
   private async startCortex() {
     if (!existsSync(resolve(this.rootDir(), 'cortex-cpp'))) {
       console.log('Please init the cortex by running cortex init command!');
-      process.exit(0);
+      exit(0);
     }
     const host = '127.0.0.1';
-    const port = '3928';
+    const port = 3928;
     return this.cortexUsecases.startCortex(host, port);
   }
+
   private async startModel(modelId: string) {
-    const settings = {
+    // TODO: NamH remove these hardcoded value
+    const settings: ModelSettingParams = {
       cpu_threads: 10,
       ctx_len: 2048,
       embedding: false,
@@ -48,8 +52,16 @@ export class StartCommand extends CommandRunner {
       ai_prompt: '\n### Response:',
       ngl: 100,
     };
-    const loadModelDto: LoadModelDto = { modelId, settings };
-    return this.modelsUsecases.startModel(loadModelDto);
+    return this.modelsUsecases.startModel(modelId, settings);
+  }
+
+  private async getModelOrStop(modelId: string): Promise<Model> {
+    const model = await this.modelsUsecases.findOne(modelId);
+    if (!model) {
+      console.debug('Model not found');
+      exit(1);
+    }
+    return model;
   }
 
   rootDir = () => resolve(__dirname, `../../../`);
diff --git a/cortex-js/src/infrastructure/controllers/models.controller.ts b/cortex-js/src/infrastructure/controllers/models.controller.ts
@@ -12,9 +12,9 @@ import { ModelsUsecases } from '@/usecases/models/models.usecases';
 import { CreateModelDto } from '@/infrastructure/dtos/models/create-model.dto';
 import { UpdateModelDto } from '@/infrastructure/dtos/models/update-model.dto';
 import { ApiResponse, ApiTags } from '@nestjs/swagger';
-import { LoadModelSuccessDto } from '@/infrastructure/dtos/models/load-model-success.dto';
-import { LoadModelDto } from '@/infrastructure/dtos/models/load-model.dto';
+import { StartModelSuccessDto } from '@/infrastructure/dtos/models/start-model-success.dto';
 import { DownloadModelDto } from '@/infrastructure/dtos/models/download-model.dto';
+import { ModelSettingParamsDto } from '../dtos/models/model-setting-params.dto';
 
 @ApiTags('Models')
 @Controller('models')
@@ -30,11 +30,14 @@ export class ModelsController {
   @ApiResponse({
     status: 200,
     description: 'The model has been loaded successfully.',
-    type: LoadModelSuccessDto,
+    type: StartModelSuccessDto,
   })
-  @Post('load')
-  load(@Body() loadModelDto: LoadModelDto) {
-    return this.modelsService.startModel(loadModelDto);
+  @Post(':modelId/start')
+  startModel(
+    @Param('modelId') modelId: string,
+    @Body() settings: ModelSettingParamsDto,
+  ) {
+    return this.modelsService.startModel(modelId, settings);
   }
 
   @Post('download')
diff --git a/cortex-js/src/infrastructure/dtos/cortex/start-cortex.dto.ts b/cortex-js/src/infrastructure/dtos/cortex/start-cortex.dto.ts
@@ -1,10 +1,12 @@
-import { IsIP, IsString } from 'class-validator';
+import { IsIP, IsNumber, IsString, Max, Min } from 'class-validator';
 
 export class StartCortexDto {
   @IsString()
   @IsIP()
   host: string;
 
-  @IsString()
-  port: string;
+  @IsNumber()
+  @Min(0)
+  @Max(65535)
+  port: number;
 }
diff --git a/cortex-js/src/infrastructure/dtos/models/load-model.dto.ts b/cortex-js/src/infrastructure/dtos/models/load-model.dto.ts
diff --git a/cortex-js/src/infrastructure/dtos/models/start-model-success.dto.ts b/cortex-js/src/infrastructure/dtos/models/start-model-success.dto.ts
@@ -1,6 +1,6 @@
 import { IsString } from 'class-validator';
 
-export class LoadModelSuccessDto {
+export class StartModelSuccessDto {
   @IsString()
   message: string;
 
diff --git a/cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts b/cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts
@@ -2,25 +2,23 @@ import { Injectable } from '@nestjs/common';
 import { OAIEngineExtension } from '@/domain/abstracts/oai.abstract';
 import { PromptTemplate } from '@/domain/models/prompt-template.interface';
 import { basename, join, resolve } from 'path';
-import { Model } from '@/domain/models/model.interface';
+import { Model, ModelSettingParams } from '@/domain/models/model.interface';
 import { ConfigService } from '@nestjs/config';
 import { HttpService } from '@nestjs/axios';
+import { defaultCortexCppHost, defaultCortexCppPort } from 'constant';
 
 /**
  * A class that implements the InferenceExtension interface from the @janhq/core package.
  * The class provides methods for initializing and stopping a model, and for making inference requests.
  * It also subscribes to events emitted by the @janhq/core package and handles new message requests.
  */
-const LOCAL_HOST = '127.0.0.1';
-const NITRO_DEFAULT_PORT = 3928;
-const NITRO_HTTP_SERVER_URL = `http://${LOCAL_HOST}:${NITRO_DEFAULT_PORT}`;
-const LOAD_MODEL_URL = `${NITRO_HTTP_SERVER_URL}/inferences/server/loadmodel`;
-const UNLOAD_MODEL_URL = `${NITRO_HTTP_SERVER_URL}/inferences/server/unloadmodel`;
-
 @Injectable()
 export default class CortexProvider extends OAIEngineExtension {
   provider: string = 'cortex';
-  apiUrl = 'http://127.0.0.1:3928/inferences/server/chat_completion';
+  apiUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/chat_completion`;
+
+  private loadModelUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/loadmodel`;
+  private unloadModelUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/unloadmodel`;
 
   constructor(
     private readonly configService: ConfigService,
@@ -29,7 +27,10 @@ export default class CortexProvider extends OAIEngineExtension {
     super(httpService);
   }
 
-  override async loadModel(model: Model): Promise<void> {
+  override async loadModel(
+    model: Model,
+    settings?: ModelSettingParams,
+  ): Promise<void> {
     const modelsContainerDir =
       this.configService.get<string>('CORTEX_MODELS_DIR') ??
       resolve('./models');
@@ -49,6 +50,7 @@ export default class CortexProvider extends OAIEngineExtension {
       model: model.id,
       cpu_threads: cpuThreadCount,
       ...model.settings,
+      ...settings,
       llama_model_path: modelBinaryLocalPath,
       ...(model.settings.mmproj && {
         mmproj: join(modelFolderFullPath, model.settings.mmproj),
@@ -67,12 +69,12 @@ export default class CortexProvider extends OAIEngineExtension {
       modelSettings.ai_prompt = prompt.ai_prompt;
     }
 
-    await this.httpService.post(LOAD_MODEL_URL, modelSettings).toPromise();
+    await this.httpService.post(this.loadModelUrl, modelSettings).toPromise();
   }
 
   override async unloadModel(modelId: string): Promise<void> {
     await this.httpService
-      .post(UNLOAD_MODEL_URL, { model: modelId })
+      .post(this.unloadModelUrl, { model: modelId })
       .toPromise();
   }
 
diff --git a/cortex-js/src/usecases/cortex/cortex.usecases.ts b/cortex-js/src/usecases/cortex/cortex.usecases.ts
@@ -12,7 +12,7 @@ export class CortexUsecases {
 
   async startCortex(
     host: string,
-    port: string,
+    port: number,
   ): Promise<CortexOperationSuccessfullyDto> {
     if (this.cortexProcess) {
       return {
@@ -21,9 +21,9 @@ export class CortexUsecases {
       };
     }
 
-    const args: string[] = ['1', host, port];
+    const args: string[] = ['1', host, `${port}`];
+
     // go up one level to get the binary folder, have to also work on windows
-    // const binaryFolder = join(binaryPath, '..');
     this.cortexProcess = spawn(
       join(__dirname, '../../../cortex-cpp/cortex-cpp'),
       args,
diff --git a/cortex-js/src/usecases/models/models.usecases.ts b/cortex-js/src/usecases/models/models.usecases.ts
