fix: remove model also remove yaml file

janhq · Jun 5, 2024 · 2adee26 · 2adee26
1 parent f7284c5
commit 2adee26
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/cortex-js/src/infrastructure/commanders/usecases/models.cli.usecases.ts b/cortex-js/src/infrastructure/commanders/usecases/models.cli.usecases.ts
@@ -180,6 +180,11 @@ export class ModelsCliUsecases {
    * @param modelId
    */
   async pullModel(modelId: string) {
+    if (await this.modelsUsecases.findOne(modelId)) {
+      console.error('Model already exists');
+      process.exit(1);
+    }
+
     if (modelId.includes('/') || modelId.includes(':')) {
       await this.pullHuggingFaceModel(modelId);
     }

diff --git a/cortex-js/src/infrastructure/commanders/utils/normalize-model-id.ts b/cortex-js/src/infrastructure/commanders/utils/normalize-model-id.ts
@@ -1,3 +1,3 @@
 export const normalizeModelId = (modelId: string): string => {
-  return modelId.replace(':', '-');
+  return modelId.replace(':', '-').replace('/', '-');
 };
diff --git a/cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts b/cortex-js/src/infrastructure/providers/cortex/cortex.provider.ts
@@ -12,7 +12,7 @@ import { FileManagerService } from '@/file-manager/file-manager.service';
 
 @Injectable()
 export default class CortexProvider extends OAIEngineExtension {
-  provider: string = 'cortex';
+  provider: string = 'cortex.llamacpp';
   apiUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/chat_completion`;
 
   private loadModelUrl = `http://${defaultCortexCppHost}:${defaultCortexCppPort}/inferences/server/loadmodel`;

diff --git a/cortex-js/src/infrastructure/repositories/model/model.repository.ts b/cortex-js/src/infrastructure/repositories/model/model.repository.ts
@@ -8,6 +8,7 @@ import {
   mkdirSync,
   readFileSync,
   readdirSync,
+  rmSync,
   writeFileSync,
 } from 'fs';
 import { load, dump } from 'js-yaml';
@@ -103,8 +104,13 @@ export class ModelRepositoryImpl implements ModelRepository {
    * This would remove the model yaml file from the models folder
    * @param id model id
    */
-  remove(id: string): Promise<void> {
+  async remove(id: string): Promise<void> {
     this.models.delete(id);
+    const yamlFilePath = join(
+      await this.fileService.getModelsPath(),
+      this.fileModel.get(id) ?? id,
+    );
+    if (existsSync(yamlFilePath)) rmSync(yamlFilePath);
     return Promise.resolve();
   }