update

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

Author: namchuai

cortex-js/src/app.module.ts

@@ -16,6 +16,7 @@ import { FileManagerModule } from './file-manager/file-manager.module';
 import { AppLoggerMiddleware } from './infrastructure/middlewares/app.logger.middleware';
 import { EventEmitterModule } from '@nestjs/event-emitter';
 import { AppController } from './infrastructure/controllers/app.controller';
+import { DownloadManagerModule } from './download-manager/download-manager.module';
 
 @Module({
   imports: [
@@ -37,6 +38,7 @@ import { AppController } from './infrastructure/controllers/app.controller';
     ExtensionModule,
     FileManagerModule,
     ModelRepositoryModule,
+    DownloadManagerModule,
   ],
   controllers: [AppController],
   providers: [SeedService],

cortex-js/src/command.module.ts

@@ -29,6 +29,7 @@ import { KillCommand } from './infrastructure/commanders/kill.command';
 import { PresetCommand } from './infrastructure/commanders/presets.command';
 import { EmbeddingCommand } from './infrastructure/commanders/embeddings.command';
 import { EventEmitterModule } from '@nestjs/event-emitter';
+import { DownloadManagerModule } from './download-manager/download-manager.module';
 
 @Module({
   imports: [
@@ -47,6 +48,7 @@ import { EventEmitterModule } from '@nestjs/event-emitter';
     AssistantsModule,
     MessagesModule,
     FileManagerModule,
+    DownloadManagerModule,
   ],
   providers: [
     CortexCommand,

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

@@ -1,6 +1,6 @@
 export class DownloadState {
   /**
-   * The id of a particular download. This will be an uuid.
+   * The id of a particular download. Being used to prevent duplication of downloads.
    */
   id: string;
 
@@ -12,7 +12,7 @@ export class DownloadState {
   /**
    * The type of download.
    */
-  type: 'model' | 'miscelanous';
+  type: DownloadType;
 
   /**
    * The status of the download.
@@ -65,3 +65,8 @@ export class DownloadItem {
 export interface DownloadStateEvent {
   data: DownloadState[];
 }
+
+export enum DownloadType {
+  Model = 'model',
+  Miscelanous = 'miscelanous',
+}

cortex-js/src/download-manager/download-manager.module.ts

@@ -0,0 +1,10 @@
+import { Module } from '@nestjs/common';
+import { DownloadManagerService } from './download-manager.service';
+import { HttpModule } from '@nestjs/axios';
+
+@Module({
+  imports: [HttpModule],
+  providers: [DownloadManagerService],
+  exports: [DownloadManagerService],
+})
+export class DownloadManagerModule {}

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

@@ -0,0 +1,18 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { DownloadManagerService } from './download-manager.service';
+
+describe('DownloadManagerService', () => {
+  let service: DownloadManagerService;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [DownloadManagerService],
+    }).compile();
+
+    service = module.get<DownloadManagerService>(DownloadManagerService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+});

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

@@ -0,0 +1,186 @@
+import {
+  DownloadItem,
+  DownloadState,
+  DownloadStatus,
+  DownloadType,
+} from '@/domain/models/download.interface';
+import { HttpService } from '@nestjs/axios';
+import { Injectable } from '@nestjs/common';
+import { EventEmitter2 } from '@nestjs/event-emitter';
+import { createWriteStream } from 'node:fs';
+import { firstValueFrom } from 'rxjs';
+
+@Injectable()
+export class DownloadManagerService {
+  private allDownloadStates: DownloadState[] = [];
+
+  constructor(
+    private readonly httpService: HttpService,
+    private readonly eventEmitter: EventEmitter2,
+  ) {
+    // start emitting download state each 500ms
+    setInterval(() => {
+      this.eventEmitter.emit('download.event', this.allDownloadStates);
+    }, 500);
+  }
+
+  async submitDownloadRequest(
+    downloadId: string,
+    title: string,
+    downloadType: DownloadType,
+    urlToDestination: Record<string, string>,
+  ) {
+    if (
+      this.allDownloadStates.find(
+        (downloadState) => downloadState.id === downloadId,
+      )
+    ) {
+      return;
+    }
+
+    const downloadItems: DownloadItem[] = Object.keys(urlToDestination).map(
+      (url) => {
+        const destination = urlToDestination[url];
+        const downloadItem: DownloadItem = {
+          id: destination,
+          time: {
+            elapsed: 0,
+            remaining: 0,
+          },
+          size: {
+            total: 0,
+            transferred: 0,
+          },
+          status: DownloadStatus.Downloading,
+        };
+
+        return downloadItem;
+      },
+    );
+
+    const downloadState: DownloadState = {
+      id: downloadId,
+      title: title,
+      type: downloadType,
+      status: DownloadStatus.Downloading,
+      children: downloadItems,
+    };
+
+    this.allDownloadStates.push(downloadState);
+
+    Object.keys(urlToDestination).forEach((url) => {
+      const destination = urlToDestination[url];
+      this.downloadFile(downloadId, url, destination);
+    });
+  }
+
+  private async downloadFile(
+    downloadId: string,
+    url: string,
+    destination: string,
+  ) {
+    const response = await firstValueFrom(
+      this.httpService.get(url, {
+        responseType: 'stream',
+      }),
+    );
+    // check if response is success
+    if (!response) {
+      throw new Error('Failed to download model');
+    }
+
+    const writer = createWriteStream(destination);
+    const totalBytes = response.headers['content-length'];
+
+    // update download state
+    // TODO: separte this duplicate code to a function
+    const currentDownloadState = this.allDownloadStates.find(
+      (downloadState) => downloadState.id === downloadId,
+    );
+    if (!currentDownloadState) {
+      return;
+    }
+    const downloadItem = currentDownloadState?.children.find(
+      (downloadItem) => downloadItem.id === destination,
+    );
+    if (downloadItem) {
+      downloadItem.size.total = totalBytes;
+    }
+
+    let transferredBytes = 0;
+
+    writer.on('finish', () => {
+      const currentDownloadState = this.allDownloadStates.find(
+        (downloadState) => downloadState.id === downloadId,
+      );
+      if (!currentDownloadState) {
+        return;
+      }
+
+      // update current child status to downloaded, find by destination as id
+      const downloadItem = currentDownloadState?.children.find(
+        (downloadItem) => downloadItem.id === destination,
+      );
+      if (downloadItem) {
+        downloadItem.status = DownloadStatus.Downloaded;
+      }
+
+      const allChildrenDownloaded = currentDownloadState?.children.every(
+        (downloadItem) => downloadItem.status === DownloadStatus.Downloaded,
+      );
+
+      if (allChildrenDownloaded) {
+        currentDownloadState.status = DownloadStatus.Downloaded;
+        // TODO: notify if download success so that client can auto refresh
+        // remove download state if all children is downloaded
+        this.allDownloadStates = this.allDownloadStates.filter(
+          (downloadState) => downloadState.id !== downloadId,
+        );
+      }
+    });
+
+    writer.on('error', (error) => {
+      const currentDownloadState = this.allDownloadStates.find(
+        (downloadState) => downloadState.id === downloadId,
+      );
+      if (!currentDownloadState) {
+        return;
+      }
+
+      const downloadItem = currentDownloadState?.children.find(
+        (downloadItem) => downloadItem.id === destination,
+      );
+      if (downloadItem) {
+        downloadItem.status = DownloadStatus.Error;
+        downloadItem.error = error.message;
+      }
+
+      currentDownloadState.status = DownloadStatus.Error;
+      currentDownloadState.error = error.message;
+
+      // TODO: notify if download error
+      // remove download state if all children is downloaded
+      this.allDownloadStates = this.allDownloadStates.filter(
+        (downloadState) => downloadState.id !== downloadId,
+      );
+    });
+
+    response.data.on('data', (chunk: any) => {
+      transferredBytes += chunk.length;
+
+      const currentDownloadState = this.allDownloadStates.find(
+        (downloadState) => downloadState.id === downloadId,
+      );
+      if (!currentDownloadState) return;
+
+      const downloadItem = currentDownloadState?.children.find(
+        (downloadItem) => downloadItem.id === destination,
+      );
+      if (downloadItem) {
+        downloadItem.size.transferred = transferredBytes;
+      }
+    });
+
+    response.data.pipe(writer);
+  }
+}

cortex-js/src/infrastructure/repositories/model/model.module.ts

@@ -4,9 +4,15 @@ import { HttpModule } from '@nestjs/axios';
 import { ModelRepository } from '@/domain/repositories/model.interface';
 import { ModelRepositoryImpl } from './model.repository';
 import { FileManagerModule } from '@/file-manager/file-manager.module';
+import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 
 @Module({
-  imports: [CortexProviderModule, HttpModule, FileManagerModule],
+  imports: [
+    CortexProviderModule,
+    HttpModule,
+    FileManagerModule,
+    DownloadManagerModule,
+  ],
   providers: [
     {
       provide: ModelRepository,

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

@@ -7,6 +7,7 @@ import { ExtensionModule } from '@/infrastructure/repositories/extensions/extens
 import { HttpModule } from '@nestjs/axios';
 import { FileManagerModule } from '@/file-manager/file-manager.module';
 import { ModelRepositoryModule } from '@/infrastructure/repositories/model/model.module';
+import { DownloadManagerModule } from '@/download-manager/download-manager.module';
 
 @Module({
   imports: [
@@ -16,6 +17,7 @@ import { ModelRepositoryModule } from '@/infrastructure/repositories/model/model
     HttpModule,
     FileManagerModule,
     ModelRepositoryModule,
+    DownloadManagerModule,
   ],
   controllers: [ModelsController],
   providers: [ModelsUsecases],