feat: add model event

Signed-off-by: James <namnh0122@gmail.com>

Author: namchuai

cortex-js/src/app.module.ts

@@ -17,7 +17,6 @@ import { AppLoggerMiddleware } from './infrastructure/middlewares/app.logger.mid
 import { EventEmitterModule } from '@nestjs/event-emitter';
 import { DownloadManagerModule } from './download-manager/download-manager.module';
 import { EventsController } from './infrastructure/controllers/events.controller';
-import { AppController } from './infrastructure/controllers/app.controller';
 import { AssistantsController } from './infrastructure/controllers/assistants.controller';
 import { ChatController } from './infrastructure/controllers/chat.controller';
 import { EmbeddingsController } from './infrastructure/controllers/embeddings.controller';
@@ -49,7 +48,7 @@ import { ProcessController } from './infrastructure/controllers/process.controll
     DownloadManagerModule,
   ],
   controllers: [
-    AppController,
+    EventsController,
     AssistantsController,
     ChatController,
     EmbeddingsController,

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

@@ -0,0 +1,35 @@
+export type ModelId = string;
+
+const ModelLoadingEvents = [
+  'starting',
+  'stopping',
+  'started',
+  'stopped',
+  'starting-failed',
+  'stopping-failed',
+] as const;
+export type ModelLoadingEvent = (typeof ModelLoadingEvents)[number];
+
+const AllModelStates = ['starting', 'stopping', 'started'] as const;
+export type ModelState = (typeof AllModelStates)[number];
+
+export interface ModelStatus {
+  model: ModelId;
+  status: ModelState;
+  metadata: Record<string, unknown>;
+}
+
+export interface ModelEvent {
+  model: ModelId;
+  event: ModelLoadingEvent;
+  metadata: Record<string, unknown>;
+}
+
+export const EmptyModelEvent = {};
+
+export interface ModelStatusAndEvent {
+  data: {
+    status: Record<ModelId, ModelStatus>;
+    event: ModelEvent | typeof EmptyModelEvent;
+  };
+}

cortex-js/src/download-manager/download-manager.service.spec.ts

@@ -1,11 +1,14 @@
 import { Test, TestingModule } from '@nestjs/testing';
 import { DownloadManagerService } from './download-manager.service';
+import { EventEmitterModule } from '@nestjs/event-emitter';
+import { HttpModule } from '@nestjs/axios';
 
 describe('DownloadManagerService', () => {
   let service: DownloadManagerService;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
+      imports: [HttpModule, EventEmitterModule.forRoot()],
       providers: [DownloadManagerService],
     }).compile();
 

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

@@ -2,21 +2,40 @@ import {
   DownloadState,
   DownloadStateEvent,
 } from '@/domain/models/download.interface';
+import {
+  EmptyModelEvent,
+  ModelEvent,
+  ModelId,
+  ModelStatus,
+  ModelStatusAndEvent,
+} from '@/domain/models/model.event';
 import { DownloadManagerService } from '@/download-manager/download-manager.service';
+import { ModelsUsecases } from '@/usecases/models/models.usecases';
 import { Controller, Sse } from '@nestjs/common';
 import { EventEmitter2 } from '@nestjs/event-emitter';
-import { Observable, fromEvent, map, merge, of, throttleTime } from 'rxjs';
+import { ApiTags } from '@nestjs/swagger';
+import {
+  Observable,
+  combineLatest,
+  fromEvent,
+  map,
+  merge,
+  of,
+  startWith,
+  throttleTime,
+} from 'rxjs';
 
+@ApiTags('Events')
 @Controller('events')
 export class EventsController {
   constructor(
     private readonly downloadManagerService: DownloadManagerService,
+    private readonly modelsUsecases: ModelsUsecases,
     private readonly eventEmitter: EventEmitter2,
   ) {}
 
   @Sse('download')
   downloadEvent(): Observable<DownloadStateEvent> {
-    // Welcome message Observable
     const latestDownloadState$: Observable<DownloadStateEvent> = of({
       data: this.downloadManagerService.getDownloadStates(),
     });
@@ -40,4 +59,20 @@ export class EventsController {
       downloadAbortEvent$,
     ).pipe();
   }
+
+  @Sse('model')
+  modelEvent(): Observable<ModelStatusAndEvent> {
+    const latestModelStatus$: Observable<Record<ModelId, ModelStatus>> = of(
+      this.modelsUsecases.getModelStatuses(),
+    );
+
+    const modelEvent$ = fromEvent<ModelEvent>(
+      this.eventEmitter,
+      'model.event',
+    ).pipe(startWith(EmptyModelEvent));
+
+    return combineLatest([latestModelStatus$, modelEvent$]).pipe(
+      map(([status, event]) => ({ data: { status, event } })),
+    );
+  }
 }

cortex-js/src/usecases/models/models.usecases.spec.ts

@@ -6,17 +6,21 @@ import { ExtensionModule } from '@/infrastructure/repositories/extensions/extens
 import { FileManagerModule } from '@/infrastructure/services/file-manager/file-manager.module';
 import { HttpModule } from '@nestjs/axios';
 import { ModelRepositoryModule } from '@/infrastructure/repositories/models/model.module';
+import { EventEmitterModule } from '@nestjs/event-emitter';
+import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 
 describe('ModelsService', () => {
   let service: ModelsUsecases;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       imports: [
+        EventEmitterModule.forRoot(),
         DatabaseModule,
         ModelsModule,
         ExtensionModule,
         FileManagerModule,
+        DownloadManagerModule,
         HttpModule,
         ModelRepositoryModule,
       ],

cortex-js/src/usecases/models/models.usecases.ts

@@ -34,6 +34,8 @@ import {
 } from '@/utils/huggingface';
 import { DownloadType } from '@/domain/models/download.interface';
 import { DownloadManagerService } from '@/download-manager/download-manager.service';
+import { EventEmitter2 } from '@nestjs/event-emitter';
+import { ModelId, ModelStatus } from '@/domain/models/model.event';
 
 @Injectable()
 export class ModelsUsecases {
@@ -43,8 +45,11 @@ export class ModelsUsecases {
     private readonly fileManagerService: FileManagerService,
     private readonly downloadManagerService: DownloadManagerService,
     private readonly httpService: HttpService,
+    private readonly eventEmitter: EventEmitter2,
   ) {}
 
+  private activeModelStatuses: Record<ModelId, ModelStatus> = {};
+
   /**
    * Create a new model
    * @param createModelDto Model data
@@ -148,6 +153,17 @@ export class ModelsUsecases {
       };
     }
 
+    // update states and emitting event
+    this.activeModelStatuses[modelId] = {
+      model: modelId,
+      status: 'starting',
+      metadata: {},
+    };
+    this.eventEmitter.emit('model.event', {
+      id: modelId,
+      action: 'starting',
+    });
+
     const parser = new ModelParameterParser();
     const loadModelSettings: ModelSettingParams = {
       // Default settings
@@ -167,17 +183,40 @@ export class ModelsUsecases {
 
     return engine
       .loadModel(model, loadModelSettings)
+      .then(() => {
+        this.activeModelStatuses[modelId] = {
+          model: modelId,
+          status: 'started',
+          metadata: {},
+        };
+
+        this.eventEmitter.emit('model.event', {
+          id: modelId,
+          action: 'started',
+        });
+      })
       .then(() => ({
         message: 'Model loaded successfully',
         modelId,
       }))
-      .catch((e) => ({
-        message:
-          e.code === AxiosError.ERR_BAD_REQUEST
-            ? 'Model already loaded'
-            : 'Model failed to load',
-        modelId,
-      }));
+      .catch((e) => {
+        // remove the model from this.activeModelStatus.
+        delete this.activeModelStatuses[modelId];
+
+        this.eventEmitter.emit('model.event', {
+          id: modelId,
+          action: 'starting-failed',
+        });
+        console.error('Starting model failed', e.code, e.message, e.stack);
+
+        return {
+          message:
+            e.code === AxiosError.ERR_BAD_REQUEST
+              ? 'Model already loaded'
+              : 'Model failed to load',
+          modelId,
+        };
+      });
   }
 
   async stopModel(modelId: string): Promise<StartModelSuccessDto> {
@@ -194,16 +233,41 @@ export class ModelsUsecases {
       };
     }
 
+    this.activeModelStatuses[modelId] = {
+      model: modelId,
+      status: 'stopping',
+      metadata: {},
+    };
+    this.eventEmitter.emit('model.event', {
+      id: modelId,
+      action: 'stopping',
+    });
+
     return engine
       .unloadModel(modelId)
+      .then(() => {
+        delete this.activeModelStatuses[modelId];
+
+        this.eventEmitter.emit('model.event', {
+          id: modelId,
+          action: 'stopped',
+        });
+      })
       .then(() => ({
         message: 'Model is stopped',
         modelId,
       }))
-      .catch(() => ({
-        message: 'Failed to stop model',
-        modelId,
-      }));
+      .catch(() => {
+        this.eventEmitter.emit('model.event', {
+          id: modelId,
+          action: 'stopping-failed',
+        });
+
+        return {
+          message: 'Failed to stop model',
+          modelId,
+        };
+      });
   }
 
   /**
@@ -377,4 +441,8 @@ export class ModelsUsecases {
     if (modelId.includes('/')) return fetchHuggingFaceRepoData(modelId);
     else return fetchJanRepoData(modelId);
   }
+
+  getModelStatuses(): Record<ModelId, ModelStatus> {
+    return this.activeModelStatuses;
+  }
 }