diff --git a/cortex-js/src/infrastructure/commanders/models/model-pull.command.ts b/cortex-js/src/infrastructure/commanders/models/model-pull.command.ts index 03139a48f..defdaf670 100644 --- a/cortex-js/src/infrastructure/commanders/models/model-pull.command.ts +++ b/cortex-js/src/infrastructure/commanders/models/model-pull.command.ts @@ -16,6 +16,7 @@ import { Engines } from '../types/engine.interface'; import { CortexUsecases } from '@/usecases/cortex/cortex.usecases'; import { BaseCommand } from '../base.command'; import { Presets, SingleBar } from 'cli-progress'; +import { downloadModelProgress } from '@/utils/pull-model'; @SubCommand({ name: 'pull', @@ -56,42 +57,7 @@ export class ModelPullCommand extends BaseCommand { exit(1); }); - const response = await this.cortex.events.downloadEvent(); - - const rl = require('readline').createInterface({ - input: stdin, - output: stdout, - }); - - rl.on('SIGINT', () => { - console.log('\nStopping download...'); - process.emit('SIGINT'); - }); - process.on('SIGINT', async () => { - await this.cortex.models.abortDownload(modelId); - exit(1); - }); - - const progressBar = new SingleBar({}, Presets.shades_classic); - progressBar.start(100, 0); - - for await (const stream of response) { - if (stream.length) { - const data = stream[0] as any; - - if (data.status === 'downloaded') break; - - let totalBytes = 0; - let totalTransferred = 0; - data.children.forEach((child: any) => { - totalBytes += child.size.total; - totalTransferred += child.size.transferred; - }); - progressBar.update(Math.floor((totalTransferred / totalBytes) * 100)); - } - } - progressBar.stop(); - rl.close(); + await downloadModelProgress(this.cortex, modelId); const existingModel = await this.cortex.models.retrieve(modelId); const engine = existingModel?.engine || Engines.llamaCPP; diff --git a/cortex-js/src/infrastructure/commanders/run.command.ts b/cortex-js/src/infrastructure/commanders/run.command.ts index c014cde69..3c98c5a57 100644 --- a/cortex-js/src/infrastructure/commanders/run.command.ts +++ b/cortex-js/src/infrastructure/commanders/run.command.ts @@ -10,6 +10,7 @@ import { checkModelCompatibility } from '@/utils/model-check'; import { BaseCommand } from './base.command'; import { isRemoteEngine } from '@/utils/normalize-model-id'; import { ChatClient } from './services/chat-client'; +import { downloadModelProgress } from '@/utils/pull-model'; type RunOptions = { threadId?: string; @@ -62,6 +63,7 @@ export class RunCommand extends BaseCommand { checkingSpinner.fail(e.message ?? e); exit(1); }); + await downloadModelProgress(this.cortex, modelId); } // Second check if model is available diff --git a/cortex-js/src/utils/pull-model.ts b/cortex-js/src/utils/pull-model.ts new file mode 100644 index 000000000..17f29bcbd --- /dev/null +++ b/cortex-js/src/utils/pull-model.ts @@ -0,0 +1,42 @@ +import { Presets, SingleBar } from "cli-progress"; +import { Cortex } from "cortexso-node"; +import { exit, stdin, stdout } from 'node:process'; + +export const downloadModelProgress = async (cortex: Cortex, modelId: string) => { + const response = await cortex.events.downloadEvent(); + + const rl = require('readline').createInterface({ + input: stdin, + output: stdout, + }); + + rl.on('SIGINT', () => { + console.log('\nStopping download...'); + process.emit('SIGINT'); + }); + process.on('SIGINT', async () => { + await cortex.models.abortDownload(modelId); + exit(1); + }); + + const progressBar = new SingleBar({}, Presets.shades_classic); + progressBar.start(100, 0); + + for await (const stream of response) { + if (stream.length) { + const data = stream[0] as any; + + if (data.status === 'downloaded') break; + + let totalBytes = 0; + let totalTransferred = 0; + data.children.forEach((child: any) => { + totalBytes += child.size.total; + totalTransferred += child.size.transferred; + }); + progressBar.update(Math.floor((totalTransferred / (totalBytes || 1)) * 100)); + } + } + progressBar.stop(); + rl.close(); +}; \ No newline at end of file