feat(models): adding backend logic to collect model metadata

Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com>

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",

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<string, string> = {
+      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<string, string> = {
+      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,
+    });
+  });
+});

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<string, ModelInfo>;
@@ -422,4 +424,36 @@ export class ModelsManager implements Disposable {
     // perform download
     return uploader.perform(model.id);
   }
+
+  async getModelMetadata(modelId: string): Promise<Record<string, unknown>> {
+    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<string, unknown> = {
+      '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');
+    }
+  }
 }

packages/backend/src/studio-api-impl.ts

@@ -221,6 +221,10 @@ export class StudioApiImpl implements StudioAPI {
     return this.modelsManager.getModelsInfo();
   }
 
+  getModelMetadata(modelId: string): Promise<Record<string, unknown>> {
+    return this.modelsManager.getModelMetadata(modelId);
+  }
+
   async getCatalog(): Promise<ApplicationCatalog> {
     return this.catalogManager.getCatalog();
   }

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: {

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<ModelInfo[]>;
+
+  /**
+   * 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<Record<string, unknown>>;
+
   /**
    * Delete the folder containing the model from local storage
    */

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==