feat: watch models and engines update for proper data retrieval (#960)

Co-authored-by: marknguyen1302 <nguyenvu1302.work@gmail.com>
janhq · Aug 1, 2024 · 9d011e3 · 9d011e3
1 parent 6645678
commit 9d011e3
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 4 deletions.
diff --git a/cortex-js/src/extensions/anthropic.engine.ts b/cortex-js/src/extensions/anthropic.engine.ts
@@ -62,7 +62,7 @@ export default class AnthropicEngineExtension extends OAIEngineExtension {
     const system = data.messages.find((m: any) => m.role === 'system');
     const messages = data.messages.filter((m: any) => m.role !== 'system');
     return {
-      system: system.content,
+      system: system?.content ?? '',
       messages,
       ...pick(data, ['model', 'stream', 'max_tokens']),
     };

diff --git a/cortex-js/src/infrastructure/repositories/extensions/extension.repository.ts b/cortex-js/src/infrastructure/repositories/extensions/extension.repository.ts
@@ -4,7 +4,7 @@ import { Extension } from '@/domain/abstracts/extension.abstract';
 import { readdir, lstat } from 'fs/promises';
 import { join } from 'path';
 import { FileManagerService } from '@/infrastructure/services/file-manager/file-manager.service';
-import { existsSync } from 'fs';
+import { existsSync, mkdirSync, watch } from 'fs';
 import { Engines } from '@/infrastructure/commanders/types/engine.interface';
 import { OAIEngineExtension } from '@/domain/abstracts/oai.abstract';
 import { HttpService } from '@nestjs/axios';
@@ -27,6 +27,16 @@ export class ExtensionRepositoryImpl implements ExtensionRepository {
   ) {
     this.loadCoreExtensions();
     this.loadExternalExtensions();
+
+    // Watch engine folder only for now
+    fileService.getCortexCppEnginePath().then((path) => {
+      if (!existsSync(path)) mkdirSync(path);
+      watch(path, (eventType, filename) => {
+        this.extensions.clear();
+        this.loadCoreExtensions();
+        this.loadExternalExtensions();
+      });
+    });
   }
   /**
    * Persist extension to the extensions map

diff --git a/cortex-js/src/infrastructure/repositories/models/model.repository.ts b/cortex-js/src/infrastructure/repositories/models/model.repository.ts
@@ -10,6 +10,7 @@ import {
   readdirSync,
   rmSync,
   writeFileSync,
+  watch,
 } from 'fs';
 import { load, dump } from 'js-yaml';
 import { isLocalModel, normalizeModelId } from '@/utils/normalize-model-id';
@@ -25,6 +26,12 @@ export class ModelRepositoryImpl implements ModelRepository {
 
   constructor(private readonly fileService: FileManagerService) {
     this.loadModels();
+    fileService.getModelsPath().then((path) => {
+      if (!existsSync(path)) mkdirSync(path);
+      watch(path, (eventType, filename) => {
+        this.loadModels(true);
+      });
+    });
   }
 
   /**
@@ -121,11 +128,14 @@ export class ModelRepositoryImpl implements ModelRepository {
    * This would load all the models from the models folder
    * @returns the list of models
    */
-  private async loadModels(): Promise<Model[]> {
-    if (this.loaded) return Array.from(this.models.values());
+  private async loadModels(forceReload: boolean = false): Promise<Model[]> {
+    if (this.loaded && !forceReload) return Array.from(this.models.values());
     const modelsPath =
       process.env.EXTENSIONS_PATH ?? (await this.fileService.getModelsPath());
 
+    this.models.clear();
+    this.fileModel.clear();
+
     if (!existsSync(modelsPath)) return [];
 
     const modelFiles = readdirSync(modelsPath)