diff --git a/packages/backend/package.json b/packages/backend/package.json index 3ed1a0072..447612f3d 100644 --- a/packages/backend/package.json +++ b/packages/backend/package.json @@ -55,6 +55,7 @@ "watch": "vite --mode development build -w" }, "dependencies": { + "@huggingface/gguf": "^0.1.7", "isomorphic-git": "^1.27.0", "mustache": "^4.2.0", "openai": "^4.52.3", diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts index 9881b90e4..6208297fa 100644 --- a/packages/backend/src/managers/modelsManager.spec.ts +++ b/packages/backend/src/managers/modelsManager.spec.ts @@ -29,6 +29,8 @@ import * as utils from '../utils/utils'; import { TaskRegistry } from '../registries/TaskRegistry'; import type { CancellationTokenRegistry } from '../registries/CancellationTokenRegistry'; import * as sha from '../utils/sha'; +import type { GGUFParseOutput } from '@huggingface/gguf'; +import { gguf } from '@huggingface/gguf'; const mocks = vi.hoisted(() => { return { @@ -45,6 +47,10 @@ const mocks = vi.hoisted(() => { }; }); +vi.mock('@huggingface/gguf', () => ({ + gguf: vi.fn(), +})); + vi.mock('../utils/podman', () => ({ getFirstRunningMachineName: mocks.getFirstRunningMachineNameMock, getPodmanCli: mocks.getPodmanCliMock, @@ -833,3 +839,99 @@ describe('downloadModel', () => { expect(mocks.onEventDownloadMock).toHaveBeenCalledTimes(2); }); }); + +describe('getModelMetadata', () => { + test('unknown model', async () => { + const manager = new ModelsManager( + 'appdir', + {} as Webview, + { + getModels: (): ModelInfo[] => [], + } as CatalogManager, + telemetryLogger, + taskRegistry, + cancellationTokenRegistryMock, + ); + + await expect(() => manager.getModelMetadata('unknown-model-id')).rejects.toThrowError( + 'model with id unknown-model-id does not exists.', + ); + }); + + test('remote model', async () => { + const manager = new ModelsManager( + 'appdir', + {} as Webview, + { + getModels: (): ModelInfo[] => [ + { + id: 'test-model-id', + url: 'dummy-url', + file: undefined, + } as unknown as ModelInfo, + ], + onCatalogUpdate: vi.fn(), + } as unknown as CatalogManager, + telemetryLogger, + taskRegistry, + cancellationTokenRegistryMock, + ); + + manager.init(); + + const fakeMetadata: Record = { + hello: 'world', + }; + + vi.mocked(gguf).mockResolvedValue({ + metadata: fakeMetadata, + } as unknown as GGUFParseOutput & { parameterCount: number }); + + const result = await manager.getModelMetadata('test-model-id'); + expect(result).toStrictEqual(fakeMetadata); + + expect(gguf).toHaveBeenCalledWith('dummy-url'); + }); + + test('local model', async () => { + const manager = new ModelsManager( + 'appdir', + { + postMessage: vi.fn(), + } as unknown as Webview, + { + getModels: (): ModelInfo[] => [ + { + id: 'test-model-id', + url: 'dummy-url', + file: { + file: 'random', + path: 'dummy-path', + }, + } as unknown as ModelInfo, + ], + onCatalogUpdate: vi.fn(), + } as unknown as CatalogManager, + telemetryLogger, + taskRegistry, + cancellationTokenRegistryMock, + ); + + manager.init(); + + const fakeMetadata: Record = { + hello: 'world', + }; + + vi.mocked(gguf).mockResolvedValue({ + metadata: fakeMetadata, + } as unknown as GGUFParseOutput & { parameterCount: number }); + + const result = await manager.getModelMetadata('test-model-id'); + expect(result).toStrictEqual(fakeMetadata); + + expect(gguf).toHaveBeenCalledWith(path.join('dummy-path', 'random'), { + allowLocalFile: true, + }); + }); +}); diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts index 56e027ee9..fc3ea5497 100644 --- a/packages/backend/src/managers/modelsManager.ts +++ b/packages/backend/src/managers/modelsManager.ts @@ -17,7 +17,7 @@ ***********************************************************************/ import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo'; -import fs from 'fs'; +import fs from 'node:fs'; import * as path from 'node:path'; import { type Webview, fs as apiFs, type Disposable, env } from '@podman-desktop/api'; import { Messages } from '@shared/Messages'; @@ -34,6 +34,8 @@ import { deleteRemoteModel, getLocalModelFile, isModelUploaded } from '../utils/ import { getFirstRunningMachineName } from '../utils/podman'; import type { CancellationTokenRegistry } from '../registries/CancellationTokenRegistry'; import { hasValidSha } from '../utils/sha'; +import type { GGUFParseOutput } from '@huggingface/gguf'; +import { gguf } from '@huggingface/gguf'; export class ModelsManager implements Disposable { #models: Map; @@ -422,4 +424,36 @@ export class ModelsManager implements Disposable { // perform download return uploader.perform(model.id); } + + async getModelMetadata(modelId: string): Promise> { + const model = this.#models.get(modelId); + if (!model) throw new Error(`model with id ${modelId} does not exists.`); + + const before = performance.now(); + const data: Record = { + 'model-id': model.url ? modelId : 'imported', // filter imported models + }; + + try { + let result: GGUFParseOutput<{ strict: false }>; + if (this.isModelOnDisk(modelId)) { + const modelPath = path.normalize(getLocalModelFile(model)); + console.debug(`[getModelMetadata] reading model ${modelPath}`); + result = await gguf(modelPath, { allowLocalFile: true }); + } else if (model.url) { + result = await gguf(model.url); + } else { + throw new Error('cannot get model metadata'); + } + console.log('result', result); + return result.metadata; + } catch (err: unknown) { + data['error'] = err; + console.error(err); + throw err; + } finally { + data['duration'] = performance.now() - before; + this.telemetry.logUsage('get-metadata'); + } + } } diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts index 76c3f0fe3..294d95fe6 100644 --- a/packages/backend/src/studio-api-impl.ts +++ b/packages/backend/src/studio-api-impl.ts @@ -221,6 +221,10 @@ export class StudioApiImpl implements StudioAPI { return this.modelsManager.getModelsInfo(); } + getModelMetadata(modelId: string): Promise> { + return this.modelsManager.getModelMetadata(modelId); + } + async getCatalog(): Promise { return this.catalogManager.getCatalog(); } diff --git a/packages/backend/vite.config.js b/packages/backend/vite.config.js index 911a37dce..559b7f5ec 100644 --- a/packages/backend/vite.config.js +++ b/packages/backend/vite.config.js @@ -49,6 +49,7 @@ const config = { rollupOptions: { external: [ '@podman-desktop/api', + '@huggingface/gguf', ...builtinModules.flatMap(p => [p, `node:${p}`]), ], output: { diff --git a/packages/shared/src/StudioAPI.ts b/packages/shared/src/StudioAPI.ts index 6aa4ac56d..3779d7760 100644 --- a/packages/shared/src/StudioAPI.ts +++ b/packages/shared/src/StudioAPI.ts @@ -58,6 +58,14 @@ export abstract class StudioAPI { * Get the information of models saved locally into the user's directory */ abstract getModelsInfo(): Promise; + + /** + * Given a modelId will return the model metadata + * @remark If the model is not available locally, a fetch request will be used to get its metadata from the header. + * @param modelId + */ + abstract getModelMetadata(modelId: string): Promise>; + /** * Delete the folder containing the model from local storage */ diff --git a/yarn.lock b/yarn.lock index 9daeba1b9..ec8a9ecd7 100644 --- a/yarn.lock +++ b/yarn.lock @@ -281,6 +281,11 @@ dependencies: "@fortawesome/fontawesome-common-types" "6.5.2" +"@huggingface/gguf@^0.1.7": + version "0.1.7" + resolved "https://registry.yarnpkg.com/@huggingface/gguf/-/gguf-0.1.7.tgz#0f5afab60f182cf27341738bff794750cacb5876" + integrity sha512-RQN1WwuusLjiBTNFuAJCUlhRejIhKt395ywnTmc+Jy8dajGwk8k7EsfgtxVqkBSTWy9D55XzCII7jUjkHEv3JA== + "@humanwhocodes/config-array@^0.11.14": version "0.11.14" resolved "https://registry.yarnpkg.com/@humanwhocodes/config-array/-/config-array-0.11.14.tgz#d78e481a039f7566ecc9660b4ea7fe6b1fec442b" @@ -4188,16 +4193,7 @@ std-env@^3.7.0: resolved "https://registry.yarnpkg.com/std-env/-/std-env-3.7.0.tgz#c9f7386ced6ecf13360b6c6c55b8aaa4ef7481d2" integrity sha512-JPbdCEQLj1w5GilpiHAx3qJvFndqybBysA3qUOnznweH4QbNYUsW/ea8QzSrnh0vNsezMMw5bcVool8lM0gwzg== -"string-width-cjs@npm:string-width@^4.2.0": - version "4.2.3" - resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz" - integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g== - dependencies: - emoji-regex "^8.0.0" - is-fullwidth-code-point "^3.0.0" - strip-ansi "^6.0.1" - -string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.3: +"string-width-cjs@npm:string-width@^4.2.0", string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.3: version "4.2.3" resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz" integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g== @@ -4249,14 +4245,7 @@ string_decoder@^1.1.1: dependencies: safe-buffer "~5.2.0" -"strip-ansi-cjs@npm:strip-ansi@^6.0.1": - version "6.0.1" - resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-6.0.1.tgz" - integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== - dependencies: - ansi-regex "^5.0.1" - -strip-ansi@^6.0.0, strip-ansi@^6.0.1: +"strip-ansi-cjs@npm:strip-ansi@^6.0.1", strip-ansi@^6.0.0, strip-ansi@^6.0.1: version "6.0.1" resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-6.0.1.tgz" integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A== @@ -4864,16 +4853,7 @@ winreg@^1.2.5: resolved "https://registry.yarnpkg.com/winreg/-/winreg-1.2.5.tgz#b650383e89278952494b5d113ba049a5a4fa96d8" integrity sha512-uf7tHf+tw0B1y+x+mKTLHkykBgK2KMs3g+KlzmyMbLvICSHQyB/xOFjTT8qZ3oeTFyU7Bbj4FzXitGG6jvKhYw== -"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0": - version "7.0.0" - resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-7.0.0.tgz" - integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q== - dependencies: - ansi-styles "^4.0.0" - string-width "^4.1.0" - strip-ansi "^6.0.0" - -wrap-ansi@^7.0.0: +"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0", wrap-ansi@^7.0.0: version "7.0.0" resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-7.0.0.tgz" integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q==