Merge pull request #615 from janhq/feat/add-cortex-data-folder

namchuai · web-flow · commit 5603112f761e · 2024-05-28T15:17:13.000+07:00
feat: add cortex data folder
diff --git a/cortex-js/package.json b/cortex-js/package.json
@@ -44,6 +44,7 @@
     "class-validator": "^0.14.1",
     "cli-progress": "^3.12.0",
     "decompress": "^4.2.1",
+    "js-yaml": "^4.1.0",
     "nest-commander": "^3.13.0",
     "readline": "^1.3.0",
     "reflect-metadata": "^0.2.0",
@@ -52,7 +53,6 @@
     "sqlite3": "^5.1.7",
     "typeorm": "^0.3.20",
     "ulid": "^2.3.0",
-    "yaml": "^2.4.2",
     "uuid": "^9.0.1"
   },
   "devDependencies": {
@@ -65,6 +65,7 @@
     "@types/decompress": "^4.2.7",
     "@types/express": "^4.17.17",
     "@types/jest": "^29.5.2",
+    "@types/js-yaml": "^4.0.9",
     "@types/node": "^20.12.9",
     "@types/supertest": "^6.0.0",
     "@types/uuid": "^9.0.8",
diff --git a/cortex-js/src/app.module.ts b/cortex-js/src/app.module.ts
@@ -11,6 +11,7 @@ import { CortexModule } from './usecases/cortex/cortex.module';
 import { ConfigModule } from '@nestjs/config';
 import { env } from 'node:process';
 import { SeedService } from './usecases/seed/seed.service';
+import { FileManagerModule } from './file-manager/file-manager.module';
 
 @Module({
   imports: [
@@ -29,6 +30,7 @@ import { SeedService } from './usecases/seed/seed.service';
     AssistantsModule,
     CortexModule,
     ExtensionModule,
+    FileManagerModule,
   ],
   providers: [SeedService],
 })
diff --git a/cortex-js/src/command.module.ts b/cortex-js/src/command.module.ts
@@ -23,6 +23,7 @@ import { ModelUpdateCommand } from './infrastructure/commanders/models/model-upd
 import { AssistantsModule } from './usecases/assistants/assistants.module';
 import { CliUsecasesModule } from './infrastructure/commanders/usecases/cli.usecases.module';
 import { MessagesModule } from './usecases/messages/messages.module';
+import { FileManagerModule } from './file-manager/file-manager.module';
 
 @Module({
   imports: [
@@ -39,6 +40,7 @@ import { MessagesModule } from './usecases/messages/messages.module';
     CliUsecasesModule,
     AssistantsModule,
     MessagesModule,
+    FileManagerModule,
   ],
   providers: [
     CortexCommand,
diff --git a/cortex-js/src/domain/config/config.interface.ts b/cortex-js/src/domain/config/config.interface.ts
@@ -0,0 +1,3 @@
+export interface Config {
+  dataFolderPath: string;
+}
diff --git a/cortex-js/src/file-manager/file-manager.module.ts b/cortex-js/src/file-manager/file-manager.module.ts
@@ -0,0 +1,8 @@
+import { Module } from '@nestjs/common';
+import { FileManagerService } from './file-manager.service';
+
+@Module({
+  providers: [FileManagerService],
+  exports: [FileManagerService],
+})
+export class FileManagerModule {}
diff --git a/cortex-js/src/file-manager/file-manager.service.spec.ts b/cortex-js/src/file-manager/file-manager.service.spec.ts
@@ -0,0 +1,18 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { FileManagerService } from './file-manager.service';
+
+describe('FileManagerService', () => {
+  let service: FileManagerService;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [FileManagerService],
+    }).compile();
+
+    service = module.get<FileManagerService>(FileManagerService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+});
diff --git a/cortex-js/src/file-manager/file-manager.service.ts b/cortex-js/src/file-manager/file-manager.service.ts
@@ -0,0 +1,78 @@
+import { Config } from '@/domain/config/config.interface';
+import { Injectable } from '@nestjs/common';
+import os from 'os';
+import { join } from 'node:path';
+import { existsSync, promises } from 'node:fs';
+import yaml from 'js-yaml';
+
+@Injectable()
+export class FileManagerService {
+  private configFile = '.cortexrc';
+  private cortexDirectoryName = 'cortex';
+  private modelFolderName = 'models';
+  private cortexCppFolderName = 'cortex-cpp';
+
+  async getConfig(): Promise<Config> {
+    const homeDir = os.homedir();
+    const configPath = join(homeDir, this.configFile);
+
+    if (!existsSync(configPath)) {
+      const config = this.defaultConfig();
+      await this.createFolderIfNotExist(config.dataFolderPath);
+      await this.writeConfigFile(config);
+      return config;
+    }
+
+    try {
+      const content = await promises.readFile(configPath, 'utf8');
+      const config = yaml.load(content) as Config;
+      return config;
+    } catch (error) {
+      console.warn('Error reading config file. Using default config.');
+      console.warn(error);
+      const config = this.defaultConfig();
+      await this.createFolderIfNotExist(config.dataFolderPath);
+      await this.writeConfigFile(config);
+      return config;
+    }
+  }
+
+  private async writeConfigFile(config: Config): Promise<void> {
+    const homeDir = os.homedir();
+    const configPath = join(homeDir, this.configFile);
+
+    // write config to file as yaml
+    const configString = yaml.dump(config);
+    await promises.writeFile(configPath, configString, 'utf8');
+  }
+
+  private async createFolderIfNotExist(dataFolderPath: string): Promise<void> {
+    if (!existsSync(dataFolderPath)) {
+      await promises.mkdir(dataFolderPath, { recursive: true });
+    }
+
+    const modelFolderPath = join(dataFolderPath, this.modelFolderName);
+    const cortexCppFolderPath = join(dataFolderPath, this.cortexCppFolderName);
+    if (!existsSync(modelFolderPath)) {
+      await promises.mkdir(modelFolderPath, { recursive: true });
+    }
+    if (!existsSync(cortexCppFolderPath)) {
+      await promises.mkdir(cortexCppFolderPath, { recursive: true });
+    }
+  }
+
+  private defaultConfig(): Config {
+    // default will store at home directory
+    const homeDir = os.homedir();
+    const dataFolderPath = join(homeDir, this.cortexDirectoryName);
+
+    return {
+      dataFolderPath,
+    };
+  }
+
+  async getDataFolderPath(): Promise<string> {
+    const config = await this.getConfig();
+    return config.dataFolderPath;
+  }
+}
diff --git a/cortex-js/src/infrastructure/commanders/usecases/cli.usecases.module.ts b/cortex-js/src/infrastructure/commanders/usecases/cli.usecases.module.ts
@@ -9,6 +9,7 @@ import { CortexModule } from '@/usecases/cortex/cortex.module';
 import { ThreadsModule } from '@/usecases/threads/threads.module';
 import { AssistantsModule } from '@/usecases/assistants/assistants.module';
 import { MessagesModule } from '@/usecases/messages/messages.module';
+import { FileManagerModule } from '@/file-manager/file-manager.module';
 
 @Module({
   imports: [
@@ -19,6 +20,7 @@ import { MessagesModule } from '@/usecases/messages/messages.module';
     ThreadsModule,
     AssistantsModule,
     MessagesModule,
+    FileManagerModule,
   ],
   providers: [InitCliUsecases, ModelsCliUsecases, ChatCliUsecases],
   exports: [InitCliUsecases, ModelsCliUsecases, ChatCliUsecases],
diff --git a/cortex-js/src/infrastructure/commanders/usecases/init.cli.usecases.ts b/cortex-js/src/infrastructure/commanders/usecases/init.cli.usecases.ts
@@ -1,20 +1,24 @@
 import { createWriteStream, existsSync, rmSync } from 'fs';
-import { resolve, delimiter, join } from 'path';
+import { delimiter, join } from 'path';
 import { HttpService } from '@nestjs/axios';
 import { Presets, SingleBar } from 'cli-progress';
 import decompress from 'decompress';
 import { exit } from 'node:process';
 import { InitOptions } from '../types/init-options.interface';
 import { Injectable } from '@nestjs/common';
 import { firstValueFrom } from 'rxjs';
+import { FileManagerService } from '@/file-manager/file-manager.service';
 
 @Injectable()
 export class InitCliUsecases {
   CORTEX_RELEASES_URL = 'https://api.github.com/repos/janhq/cortex/releases';
   CUDA_DOWNLOAD_URL =
     'https://catalog.jan.ai/dist/cuda-dependencies/<version>/<platform>/cuda.tar.gz';
 
-  constructor(private readonly httpService: HttpService) {}
+  constructor(
+    private readonly httpService: HttpService,
+    private readonly fileManagerService: FileManagerService,
+  ) {}
 
   installEngine = async (
     engineFileName: string,
@@ -53,7 +57,8 @@ export class InitCliUsecases {
     }
 
     console.log(`Downloading engine file ${engineFileName}`);
-    const engineDir = resolve(this.rootDir(), 'cortex-cpp');
+    const dataFolderPath = await this.fileManagerService.getDataFolderPath();
+    const engineDir = join(dataFolderPath, 'cortex-cpp');
     if (existsSync(engineDir)) rmSync(engineDir, { recursive: true });
 
     const download = await firstValueFrom(
@@ -66,7 +71,7 @@ export class InitCliUsecases {
       process.exit(1);
     }
 
-    const destination = resolve(this.rootDir(), toDownloadAsset.name);
+    const destination = join(dataFolderPath, toDownloadAsset.name);
 
     await new Promise((resolve, reject) => {
       const writer = createWriteStream(destination);
@@ -95,10 +100,7 @@ export class InitCliUsecases {
     });
 
     try {
-      await decompress(
-        resolve(this.rootDir(), destination),
-        resolve(this.rootDir()),
-      );
+      await decompress(destination, join(dataFolderPath));
     } catch (e) {
       console.error('Error decompressing file', e);
       exit(1);
@@ -124,8 +126,6 @@ export class InitCliUsecases {
     return `${engineName}.tar.gz`;
   };
 
-  rootDir = () => resolve(__dirname, `../../../../`);
-
   cudaVersion = async () => {
     let filesCuda12: string[];
     let filesCuda11: string[];
@@ -178,11 +178,12 @@ export class InitCliUsecases {
   installCudaToolkitDependency = async (options: InitOptions) => {
     const platform = process.platform === 'win32' ? 'windows' : 'linux';
 
+    const dataFolderPath = await this.fileManagerService.getDataFolderPath();
     const url = this.CUDA_DOWNLOAD_URL.replace(
       '<version>',
       options.cudaVersion === '11' ? '11.7' : '12.0',
     ).replace('<platform>', platform);
-    const destination = resolve(this.rootDir(), 'cuda-toolkit.tar.gz');
+    const destination = join(dataFolderPath, 'cuda-toolkit.tar.gz');
 
     const download = await firstValueFrom(
       this.httpService.get(url, {
@@ -222,10 +223,7 @@ export class InitCliUsecases {
     });
 
     try {
-      await decompress(
-        resolve(this.rootDir(), destination),
-        resolve(this.rootDir(), 'cortex-cpp'),
-      );
+      await decompress(destination, join(dataFolderPath, 'cortex-cpp'));
     } catch (e) {
       console.log(e);
       exit(1);
diff --git a/cortex-js/src/infrastructure/database/database.module.ts b/cortex-js/src/infrastructure/database/database.module.ts
@@ -4,8 +4,10 @@ import { sqliteDatabaseProviders } from './sqlite-database.providers';
 import { modelProviders } from './providers/model.providers';
 import { assistantProviders } from './providers/assistant.providers';
 import { messageProviders } from './providers/message.providers';
+import { FileManagerModule } from '@/file-manager/file-manager.module';
 
 @Module({
+  imports: [FileManagerModule],
   providers: [
     ...sqliteDatabaseProviders,
     ...threadProviders,
diff --git a/cortex-js/src/infrastructure/database/sqlite-database.providers.ts b/cortex-js/src/infrastructure/database/sqlite-database.providers.ts
@@ -1,12 +1,15 @@
+import { FileManagerService } from '@/file-manager/file-manager.service';
 import { databaseFile } from 'constant';
-import { resolve } from 'path';
+import { join } from 'path';
 import { DataSource } from 'typeorm';
 
 export const sqliteDatabaseProviders = [
   {
     provide: 'DATA_SOURCE',
-    useFactory: async () => {
-      const sqlitePath = resolve(__dirname, `../../../${databaseFile}`);
+    inject: [FileManagerService],
+    useFactory: async (fileManagerService: FileManagerService) => {
+      const dataFolderPath = await fileManagerService.getDataFolderPath();
+      const sqlitePath = join(dataFolderPath, databaseFile);
       const dataSource = new DataSource({
         type: 'sqlite',
         database: sqlitePath,
diff --git a/cortex-js/src/infrastructure/providers/cortex/cortex.module.ts b/cortex-js/src/infrastructure/providers/cortex/cortex.module.ts
@@ -1,9 +1,10 @@
 import { Module } from '@nestjs/common';
 import CortexProvider from './cortex.provider';
 import { HttpModule } from '@nestjs/axios';
+import { FileManagerModule } from '@/file-manager/file-manager.module';
 
 @Module({
-  imports: [HttpModule],
+  imports: [HttpModule, FileManagerModule],
   providers: [
     {
       provide: 'CORTEX_PROVIDER',
diff --git a/cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts b/cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts
@@ -1,13 +1,14 @@
 import { Injectable } from '@nestjs/common';
 import { OAIEngineExtension } from '@/domain/abstracts/oai.abstract';
 import { PromptTemplate } from '@/domain/models/prompt-template.interface';
-import { join, resolve } from 'path';
+import { join } from 'path';
 import { Model, ModelSettingParams } from '@/domain/models/model.interface';
 import { HttpService } from '@nestjs/axios';
 import { defaultCortexCppHost, defaultCortexCppPort } from 'constant';
 import { readdirSync } from 'node:fs';
 import { normalizeModelId } from '@/infrastructure/commanders/utils/normalize-model-id';
 import { firstValueFrom } from 'rxjs';
+import { FileManagerService } from '@/file-manager/file-manager.service';
 
 /**
  * A class that implements the InferenceExtension interface from the @janhq/core package.
@@ -22,17 +23,23 @@ export default class CortexProvider extends OAIEngineExtension {
   private loadModelUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/loadmodel`;
   private unloadModelUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/unloadmodel`;
 
-  constructor(protected readonly httpService: HttpService) {
+  constructor(
+    protected readonly httpService: HttpService,
+    private readonly fileManagerService: FileManagerService,
+  ) {
     super(httpService);
   }
 
-  modelDir = () => resolve(__dirname, `../../../../models`);
+  private async getModelDirectory(): Promise<string> {
+    const dataFolderPath = await this.fileManagerService.getDataFolderPath();
+    return join(dataFolderPath, 'models');
+  }
 
   override async loadModel(
     model: Model,
     settings?: ModelSettingParams,
   ): Promise<void> {
-    const modelsContainerDir = this.modelDir();
+    const modelsContainerDir = await this.getModelDirectory();
 
     const modelFolderFullPath = join(
       modelsContainerDir,
diff --git a/cortex-js/src/main.ts b/cortex-js/src/main.ts
@@ -4,6 +4,7 @@ import { DocumentBuilder, SwaggerModule } from '@nestjs/swagger';
 import { INestApplication, ValidationPipe } from '@nestjs/common';
 import { defaultCortexJsHost, defaultCortexJsPort } from 'constant';
 import { SeedService } from './usecases/seed/seed.service';
+import { FileManagerService } from './file-manager/file-manager.service';
 
 async function bootstrap() {
   const app = await NestFactory.create(AppModule, {
@@ -13,6 +14,9 @@ async function bootstrap() {
   const seedService = app.get(SeedService);
   await seedService.seed();
 
+  const fileService = app.get(FileManagerService);
+  await fileService.getConfig();
+
   app.useGlobalPipes(
     new ValidationPipe({
       enableDebugMessages: true,
diff --git a/cortex-js/src/usecases/cortex/cortex.module.ts b/cortex-js/src/usecases/cortex/cortex.module.ts
@@ -2,9 +2,10 @@ import { Module } from '@nestjs/common';
 import { CortexUsecases } from './cortex.usecases';
 import { CortexController } from '@/infrastructure/controllers/cortex.controller';
 import { HttpModule } from '@nestjs/axios';
+import { FileManagerModule } from '@/file-manager/file-manager.module';
 
 @Module({
-  imports: [HttpModule],
+  imports: [HttpModule, FileManagerModule],
   providers: [CortexUsecases],
   controllers: [CortexController],
   exports: [CortexUsecases],
diff --git a/cortex-js/src/usecases/cortex/cortex.usecases.ts b/cortex-js/src/usecases/cortex/cortex.usecases.ts
diff --git a/cortex-js/src/usecases/models/models.module.ts b/cortex-js/src/usecases/models/models.module.ts
diff --git a/cortex-js/src/usecases/models/models.usecases.ts b/cortex-js/src/usecases/models/models.usecases.ts

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+export interface Config {`
	`2`	`+ dataFolderPath: string;`
	`3`	`+}`