chore: auto load model on /chat/completions request (#900)

Author: louis-jan

cortex-js/src/domain/abstracts/engine.abstract.ts

@@ -3,6 +3,10 @@ import stream from 'stream';
 import { Model, ModelSettingParams } from '../../domain/models/model.interface';
 import { Extension } from './extension.abstract';
 
+/**
+ * This class should be extended by any class that represents an engine extension.
+ * It provides methods for loading and unloading models, and for making inference requests.
+ */
 export abstract class EngineExtension extends Extension {
   abstract onLoad(): void;
 
@@ -12,16 +16,42 @@ export abstract class EngineExtension extends Extension {
 
   initalized: boolean = false;
 
+  /**
+   * Makes an inference request to the engine.
+   * @param dto
+   * @param headers
+   */
   abstract inference(
     dto: any,
     headers: Record<string, string>,
   ): Promise<stream.Readable | any>;
 
+  /**
+   * Checks if a model is running by the engine
+   * This method should check run-time status of the model
+   * Since the model can be corrupted during the run-time
+   * This method should return false if the model is not running
+   * @param modelId
+   */
+  async isModelRunning(modelId: string): Promise<boolean> {
+    return true;
+  }
+
+  /**
+   * Loads a model into the engine.
+   * There are model settings such as `ngl` and `ctx_len` that can be passed to the engine.
+   * Applicable for local engines only
+   * @param model
+   * @param settingParams
+   */
   async loadModel(
     model: Model,
     settingParams?: ModelSettingParams,
   ): Promise<void> {}
 
+  /**
+   * Unloads a model from the engine.
+   * @param modelId
+   */
   async unloadModel(modelId: string): Promise<void> {}
-  
 }

cortex-js/src/domain/models/model.interface.ts

@@ -40,6 +40,10 @@ export interface ModelSettingParams {
    * The number of layers to load onto the GPU for acceleration.
    */
   ngl?: number;
+
+  /**
+   * Support embedding or not (legacy)
+   */
   embedding?: boolean;
 
   /**
@@ -117,6 +121,11 @@ export interface ModelSettingParams {
    * To enable mmap, default is true
    */
   use_mmap?: boolean;
+
+  /**
+   * Model type we want to use: llm or embedding, default value is llm (latest llama.cpp update)
+   */
+  model_type?: string;
 }
 
 /**

cortex-js/src/infrastructure/constants/cortex.ts

@@ -7,6 +7,8 @@ export const defaultCortexJsPort = 1337;
 
 export const defaultCortexCppHost = '127.0.0.1';
 export const defaultCortexCppPort = 3929;
+
+export const defaultEmbeddingModel = 'nomic-embed-text-v1';
 // CORTEX CPP
 export const CORTEX_CPP_EMBEDDINGS_URL = (
   host: string = defaultCortexCppHost,
@@ -50,4 +52,4 @@ export const CUDA_DOWNLOAD_URL =
 
 export const telemetryServerUrl = 'https://telemetry.jan.ai';
 
-export const MIN_CUDA_VERSION = '12.3';
+export const MIN_CUDA_VERSION = '12.3';

cortex-js/src/infrastructure/controllers/chat.controller.spec.ts

@@ -9,6 +9,7 @@ import { DownloadManagerModule } from '@/infrastructure/services/download-manage
 import { EventEmitterModule } from '@nestjs/event-emitter';
 import { TelemetryModule } from '@/usecases/telemetry/telemetry.module';
 import { FileManagerModule } from '../services/file-manager/file-manager.module';
+import { ModelsModule } from '@/usecases/models/models.module';
 
 describe('ChatController', () => {
   let controller: ChatController;
@@ -25,6 +26,7 @@ describe('ChatController', () => {
         EventEmitterModule.forRoot(),
         TelemetryModule,
         FileManagerModule,
+        ModelsModule,
       ],
       controllers: [ChatController],
       providers: [ChatUsecases],

cortex-js/src/infrastructure/controllers/chat.controller.ts

@@ -35,27 +35,21 @@ export class ChatController {
   ) {
     const { stream } = createChatDto;
 
-    if (stream) {
-      this.chatService
-        .inference(createChatDto, extractCommonHeaders(headers))
-        .then((stream) => {
+    this.chatService
+      .inference(createChatDto, extractCommonHeaders(headers))
+      .then((response) => {
+        if (stream) {
           res.header('Content-Type', 'text/event-stream');
-          stream.pipe(res);
-        })
-        .catch((error) =>
-          res.status(error.statusCode ?? 400).send(error.message),
-        );
-    } else {
-      res.header('Content-Type', 'application/json');
-      this.chatService
-        .inference(createChatDto, extractCommonHeaders(headers))
-        .then((response) => {
+          response.pipe(res);
+        } else {
+          res.header('Content-Type', 'application/json');
           res.json(response);
-        })
-        .catch((error) =>
-          res.status(error.statusCode ?? 400).send(error.message),
-        );
-    }
+        }
+      })
+      .catch((error) =>
+        res.status(error.statusCode ?? 400).send(error.message),
+      );
+
     this.telemetryUsecases.addEventToQueue({
       name: EventName.CHAT,
       modelId: createChatDto.model,

cortex-js/src/infrastructure/controllers/embeddings.controller.spec.ts

@@ -9,6 +9,7 @@ import { DownloadManagerModule } from '@/infrastructure/services/download-manage
 import { EventEmitterModule } from '@nestjs/event-emitter';
 import { TelemetryModule } from '@/usecases/telemetry/telemetry.module';
 import { FileManagerModule } from '../services/file-manager/file-manager.module';
+import { ModelsModule } from '@/usecases/models/models.module';
 
 describe('EmbeddingsController', () => {
   let controller: EmbeddingsController;
@@ -25,6 +26,7 @@ describe('EmbeddingsController', () => {
         EventEmitterModule.forRoot(),
         TelemetryModule,
         FileManagerModule,
+        ModelsModule,
       ],
       controllers: [EmbeddingsController],
       providers: [ChatUsecases],

cortex-js/src/infrastructure/controllers/embeddings.controller.ts

@@ -1,4 +1,4 @@
-import { Body, Controller, Post, HttpCode } from '@nestjs/common';
+import { Body, Controller, Post, HttpCode, Res } from '@nestjs/common';
 import { ChatUsecases } from '@/usecases/chat/chat.usecases';
 import { ApiOperation, ApiTags, ApiResponse } from '@nestjs/swagger';
 import { CreateEmbeddingsDto } from '../dtos/embeddings/embeddings-request.dto';

cortex-js/src/infrastructure/dtos/messages/create-message.dto.ts

@@ -21,6 +21,7 @@ export class CreateMessageDto implements Partial<Message> {
     example: 'user',
     description: 'The sources of the messages.',
   })
+  @IsString()
   role: 'user' | 'assistant';
 
   @ApiProperty({

cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts

@@ -1,10 +1,11 @@
-import { Injectable } from '@nestjs/common';
+import { HttpStatus, Injectable } from '@nestjs/common';
 import { OAIEngineExtension } from '@/domain/abstracts/oai.abstract';
 import { PromptTemplate } from '@/domain/models/prompt-template.interface';
 import { join } from 'path';
 import { Model, ModelSettingParams } from '@/domain/models/model.interface';
 import { HttpService } from '@nestjs/axios';
 import {
+  CORTEX_CPP_MODELS_URL,
   defaultCortexCppHost,
   defaultCortexCppPort,
 } from '@/infrastructure/constants/cortex';
@@ -13,6 +14,11 @@ import { normalizeModelId } from '@/utils/normalize-model-id';
 import { firstValueFrom } from 'rxjs';
 import { FileManagerService } from '@/infrastructure/services/file-manager/file-manager.service';
 
+export interface ModelStatResponse {
+  object: string;
+  data: any;
+}
+
 @Injectable()
 export default class CortexProvider extends OAIEngineExtension {
   apiUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/chat_completion`;
@@ -28,11 +34,12 @@ export default class CortexProvider extends OAIEngineExtension {
 
   constructor(
     protected readonly httpService: HttpService,
-    private readonly fileManagerService: FileManagerService,
+    protected readonly fileManagerService: FileManagerService,
   ) {
     super(httpService);
   }
 
+  // Override the inference method to make an inference request to the engine
   override async loadModel(
     model: Model,
     settings?: ModelSettingParams,
@@ -92,6 +99,30 @@ export default class CortexProvider extends OAIEngineExtension {
     ).then(); // pipe error or void instead of throwing
   }
 
+  // Override the isModelRunning method to check if the model is running
+  override async isModelRunning(modelId: string): Promise<boolean> {
+    const configs = await this.fileManagerService.getConfig();
+
+    return firstValueFrom(
+      this.httpService.get(
+        CORTEX_CPP_MODELS_URL(configs.cortexCppHost, configs.cortexCppPort),
+      ),
+    )
+      .then((res) => {
+        const data = res.data as ModelStatResponse;
+        if (
+          res.status === HttpStatus.OK &&
+          data &&
+          Array.isArray(data.data) &&
+          data.data.length > 0
+        ) {
+          return data.data.find((e) => e.id === modelId);
+        }
+        return false;
+      })
+      .catch(() => false);
+  }
+
   private readonly promptTemplateConverter = (
     promptTemplate: string,
   ): PromptTemplate => {

cortex-js/src/usecases/chat/chat.module.ts

@@ -6,6 +6,7 @@ import { ModelRepositoryModule } from '@/infrastructure/repositories/models/mode
 import { HttpModule } from '@nestjs/axios';
 import { TelemetryModule } from '../telemetry/telemetry.module';
 import { FileManagerModule } from '@/infrastructure/services/file-manager/file-manager.module';
+import { ModelsModule } from '../models/models.module';
 
 @Module({
   imports: [
@@ -15,6 +16,7 @@ import { FileManagerModule } from '@/infrastructure/services/file-manager/file-m
     HttpModule,
     TelemetryModule,
     FileManagerModule,
+    ModelsModule,
   ],
   controllers: [],
   providers: [ChatUsecases],