Merge pull request #697 from janhq/feat/add-model-event

feat: add model event

Author: namchuai

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/infrastructure/commanders/test/helpers.command.spec.ts

@@ -3,7 +3,11 @@ import { spy, Stub, stubMethod } from 'hanbi';
 import { CommandTestFactory } from 'nest-commander-testing';
 import { CommandModule } from '@/command.module';
 import { LogService } from '@/infrastructure/commanders/test/log.service';
+import { FileManagerService } from '@/infrastructure/services/file-manager/file-manager.service';
+
 import axios from 'axios';
+import { join } from 'path';
+import { rmSync } from 'fs';
 
 let commandInstance: TestingModule,
   exitSpy: Stub<typeof process.exit>,
@@ -24,9 +28,29 @@ beforeEach(
         .overrideProvider(LogService)
         .useValue({ log: spy().handler })
         .compile();
-      res();
       stdoutSpy.reset();
       stderrSpy.reset();
+
+      const fileService =
+        await commandInstance.resolve<FileManagerService>(FileManagerService);
+
+      // Attempt to create test folder
+      await fileService.writeConfigFile({
+        dataFolderPath: join(__dirname, 'test_data'),
+      });
+      res();
+    }),
+);
+
+afterEach(
+  () =>
+    new Promise<void>(async (res) => {
+      // Attempt to clean test folder
+      rmSync(join(__dirname, 'test_data'), {
+        recursive: true,
+        force: true,
+      });
+      res();
     }),
 );
 
@@ -44,19 +68,28 @@ describe('Helper commands', () => {
     timeout,
   );
 
-  test('Chat with option -m', async () => {
-    const logMock = stubMethod(console, 'log');
+  test(
+    'Chat with option -m',
+    async () => {
+      const logMock = stubMethod(console, 'log');
 
-    await CommandTestFactory.run(commandInstance, [
-      'chat',
-      // '-m',
-      // 'hello',
-      // '>output.txt',
-    ]);
-    expect(logMock.firstCall?.args[0]).toBe("Inorder to exit, type 'exit()'.");
-    // expect(exitSpy.callCount).toBe(1);
-    // expect(exitSpy.firstCall?.args[0]).toBe(1);
-  });
+      await CommandTestFactory.run(commandInstance, [
+        'run',
+        'tinyllama',
+        // '-m',
+        // 'hello',
+        // '>output.txt',
+      ]);
+      expect(
+        logMock.firstCall?.args[0] === "Inorder to exit, type 'exit()'." ||
+          logMock.firstCall?.args[0] ===
+            'Model tinyllama not found. Try pulling model...',
+      ).toBeTruthy();
+      // expect(exitSpy.callCount).toBe(1);
+      // expect(exitSpy.firstCall?.args[0]).toBe(1);
+    },
+    timeout,
+  );
 
   test('Show / kill running models', async () => {
     const tableMock = stubMethod(console, 'table');

cortex-js/src/infrastructure/commanders/test/models.command.spec.ts

@@ -5,6 +5,7 @@ import { CommandModule } from '@/command.module';
 import { join } from 'path';
 import { rmSync } from 'fs';
 import { timeout } from '@/infrastructure/commanders/test/helpers.command.spec';
+import { FileManagerService } from '@/infrastructure/services/file-manager/file-manager.service';
 
 let commandInstance: TestingModule;
 
@@ -17,6 +18,14 @@ beforeEach(
         // .overrideProvider(LogService)
         // .useValue({})
         .compile();
+      const fileService =
+        await commandInstance.resolve<FileManagerService>(FileManagerService);
+
+      // Attempt to create test folder
+      await fileService.writeConfigFile({
+        dataFolderPath: join(__dirname, 'test_data'),
+      });
+
       res();
     }),
 );

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

@@ -7,19 +7,22 @@ import { ModelRepositoryModule } from '../repositories/models/model.module';
 import { HttpModule } from '@nestjs/axios';
 import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 import { EventEmitterModule } from '@nestjs/event-emitter';
+import { TelemetryModule } from '@/usecases/telemetry/telemetry.module';
 
 describe('ChatController', () => {
   let controller: ChatController;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       imports: [
+        EventEmitterModule.forRoot(),
         DatabaseModule,
         ExtensionModule,
         ModelRepositoryModule,
         HttpModule,
         DownloadManagerModule,
         EventEmitterModule.forRoot(),
+        TelemetryModule,
       ],
       controllers: [ChatController],
       providers: [ChatUsecases],

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

@@ -7,19 +7,22 @@ import { ExtensionModule } from '../repositories/extensions/extension.module';
 import { HttpModule } from '@nestjs/axios';
 import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 import { EventEmitterModule } from '@nestjs/event-emitter';
+import { TelemetryModule } from '@/usecases/telemetry/telemetry.module';
 
 describe('EmbeddingsController', () => {
   let controller: EmbeddingsController;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       imports: [
+        EventEmitterModule.forRoot(),
         DatabaseModule,
         ModelRepositoryModule,
         ExtensionModule,
         HttpModule,
         DownloadManagerModule,
         EventEmitterModule.forRoot(),
+        TelemetryModule,
       ],
       controllers: [EmbeddingsController],
       providers: [ChatUsecases],

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/infrastructure/controllers/models.controller.spec.ts

@@ -9,20 +9,26 @@ import { CortexUsecases } from '@/usecases/cortex/cortex.usecases';
 import { ModelRepositoryModule } from '../repositories/models/model.module';
 import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 import { EventEmitterModule } from '@nestjs/event-emitter';
+import { TelemetryModule } from '@/usecases/telemetry/telemetry.module';
+import { UtilModule } from '@/util/util.module';
 
 describe('ModelsController', () => {
   let controller: ModelsController;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       imports: [
+        EventEmitterModule.forRoot(),
         DatabaseModule,
         ExtensionModule,
         FileManagerModule,
         HttpModule,
+        DownloadManagerModule,
         ModelRepositoryModule,
         DownloadManagerModule,
         EventEmitterModule.forRoot(),
+        TelemetryModule,
+        UtilModule,
       ],
       controllers: [ModelsController],
       providers: [ModelsUsecases, CortexUsecases],

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

@@ -14,6 +14,7 @@ describe('ChatService', () => {
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       imports: [
+        EventEmitterModule.forRoot(),
         DatabaseModule,
         ExtensionModule,
         ModelRepositoryModule,

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

@@ -8,21 +8,28 @@ import { HttpModule } from '@nestjs/axios';
 import { ModelRepositoryModule } from '@/infrastructure/repositories/models/model.module';
 import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 import { EventEmitterModule } from '@nestjs/event-emitter';
+import { TelemetryModule } from '../telemetry/telemetry.module';
+import { UtilModule } from '@/util/util.module';
 
 describe('ModelsService', () => {
   let service: ModelsUsecases;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       imports: [
+        EventEmitterModule.forRoot(),
         DatabaseModule,
         ModelsModule,
         ExtensionModule,
         FileManagerModule,
+        DownloadManagerModule,
         HttpModule,
         ModelRepositoryModule,
         DownloadManagerModule,
         EventEmitterModule.forRoot(),
+        TelemetryModule,
+        TelemetryModule,
+        UtilModule,
       ],
       providers: [ModelsUsecases],
       exports: [ModelsUsecases],