diff --git a/cortex-js/src/infrastructure/controllers/models.controller.ts b/cortex-js/src/infrastructure/controllers/models.controller.ts index c9b1e4ca5..851a9f91a 100644 --- a/cortex-js/src/infrastructure/controllers/models.controller.ts +++ b/cortex-js/src/infrastructure/controllers/models.controller.ts @@ -120,7 +120,7 @@ export class ModelsController { }) @Get('download/:modelId(*)') - downloadModel(@Param('modelId') modelId: string, @Query('fileName') fileName: string) { + downloadModel(@Param('modelId') modelId: string, @Query('fileName') fileName: string, @Query('persistedModelId') persistedModelId?: string) { this.modelsUsecases.pullModel(modelId, false, (files) => { return new Promise(async (resolve, reject) => { const file = files @@ -130,7 +130,7 @@ export class ModelsController { } return resolve(file); }); - }).then(() => this.telemetryUsecases.addEventToQueue({ + }, persistedModelId).then(() => this.telemetryUsecases.addEventToQueue({ name: EventName.DOWNLOAD_MODEL, modelId, }) @@ -174,7 +174,7 @@ export class ModelsController { description: 'The unique identifier of the model.', }) @Get('pull/:modelId(*)') - pullModel(@Param('modelId') modelId: string, @Query('fileName') fileName: string) { + pullModel(@Param('modelId') modelId: string, @Query('fileName') fileName: string, @Query('persistedModelId') persistedModelId?: string) { this.modelsUsecases.pullModel(modelId, false, (files) => { return new Promise(async (resolve, reject) => { const file = files @@ -184,7 +184,7 @@ export class ModelsController { } return resolve(file); }); - }).then(() => this.telemetryUsecases.addEventToQueue({ + }, persistedModelId).then(() => this.telemetryUsecases.addEventToQueue({ name: EventName.DOWNLOAD_MODEL, modelId, }) diff --git a/cortex-js/src/usecases/models/models.usecases.ts b/cortex-js/src/usecases/models/models.usecases.ts index ab900cc02..9c3824a17 100644 --- a/cortex-js/src/usecases/models/models.usecases.ts +++ b/cortex-js/src/usecases/models/models.usecases.ts @@ -321,12 +321,14 @@ export class ModelsUsecases { * @param modelId */ async pullModel( - modelId: string, + originModelId: string, inSequence: boolean = true, selection?: ( siblings: HuggingFaceRepoSibling[], ) => Promise, + persistedModelId?: string, ) { + const modelId = persistedModelId ?? originModelId; const existingModel = await this.findOne(modelId); if (isLocalModel(existingModel?.files)) { throw new BadRequestException('Model already exists'); @@ -341,10 +343,10 @@ export class ModelsUsecases { const modelFolder = join(modelsContainerDir, normalizeModelId(modelId)); await promises.mkdir(modelFolder, { recursive: true }).catch(() => {}); - let files = (await fetchJanRepoData(modelId)).siblings; + let files = (await fetchJanRepoData(originModelId)).siblings; // HuggingFace GGUF Repo - Only one file is downloaded - if (modelId.includes('/') && selection && files.length) { + if (originModelId.includes('/') && selection && files.length) { try { files = [await selection(files)]; } catch (e) { @@ -409,8 +411,6 @@ export class ModelsUsecases { model.model = modelId; if (!(await this.findOne(modelId))) await this.create(model); } else { - // Fallback if model.yml is not found & is a GGUF file - const data = await this.fetchModelMetadata(modelId); await this.populateHuggingFaceModel(modelId, files[0]); const model = await this.findOne(modelId); if (model) {