diff --git a/toolkit/components/aboutinference/content/aboutInference.js b/toolkit/components/aboutinference/content/aboutInference.js index 5f2dbc7c863f..a8f88b7f1ecd 100644 --- a/toolkit/components/aboutinference/content/aboutInference.js +++ b/toolkit/components/aboutinference/content/aboutInference.js @@ -388,7 +388,7 @@ async function updateModels() { for (const { name: model, revision } of models) { const icon = await hub.getOwnerIcon(model); - let files = await hub.listFiles({ model, revision }); + let { files } = await hub.listFiles({ model, revision }); // Create a new table for the current model let table = document.createElement("table"); diff --git a/toolkit/components/ml/content/ModelHub.sys.mjs b/toolkit/components/ml/content/ModelHub.sys.mjs index 0226fcb6338e..08a2f31a6feb 100644 --- a/toolkit/components/ml/content/ModelHub.sys.mjs +++ b/toolkit/components/ml/content/ModelHub.sys.mjs @@ -1090,69 +1090,112 @@ class IndexedDBCache { } /** - * Lists all files for a given model and revision stored in the cache. + * Lists all files for a given model and revision stored in the cache, + * and aggregates metadata from the file headers. * - * @param {object} config - * @param {?string} config.model - The model name (organization/name). + * When a `taskName` is provided, the method retrieves all model/revision + * pairs associated with that task; otherwise, it uses the provided `model` + * and `revision`. It then queries the store to retrieve file information (path + * and headers) and aggregates metadata (totalSize, lastUsed, updateDate, engineIds) + * across all files. + * + * @param {object} config - The configuration for querying the files. + * @param {?string} config.model - The model name (in "organization/name" format). * @param {?string} config.revision - The model version. - * @param {?string} config.taskName - name of the inference :wtask. - * @returns {Promise>} An array of file identifiers. + * @param {?string} config.taskName - The name of the inference task. + * @returns {Promise<{ + * files: Array<{ path: string, headers: object, engineIds: Array }>, + * metadata: { totalSize: number, lastUsed: number, updateDate: number, engineIds: Array } + * }>} An object containing: + * - files: an array of file records with their path, headers, and engine IDs. + * - metadata: aggregated metadata computed from all the files. */ async listFiles({ taskName, model, revision }) { - // When not providing taskName, we want model and revision + // When not providing taskName, both model and revision must be defined. if (!taskName && (!model || !revision)) { throw new Error("Both model and revision must be defined"); } + // Determine which model/revision pairs we want files for. let modelRevisions = [{ model, revision }]; if (taskName) { - // Get all model/revision associated to this task. - const data = await this.#getKeys({ + // Get all model/revision pairs associated with this task. + const keysData = await this.#getKeys({ storeName: this.taskStoreName, ...this.#getFileQuery({ taskName, model, revision }), }); - - modelRevisions = []; - for (const { key } of data) { - modelRevisions.push({ model: key[1], revision: key[2] }); - } + modelRevisions = keysData.map(({ key }) => ({ + model: key[1], + revision: key[2], + })); } - const filePromises = []; + // For each model/revision, query for headers data. + const fileDataPromises = modelRevisions.map(task => + this.#getData({ + storeName: this.headersStoreName, + indexName: this.#indices.modelRevisionIndex.name, + key: [task.model, task.revision], + }) + ); + const fileData = (await Promise.all(fileDataPromises)).flat(); - for (const task of modelRevisions) { - filePromises.push( - this.#getData({ - storeName: this.headersStoreName, - indexName: this.#indices.modelRevisionIndex.name, - key: [task.model, task.revision], - }) - ); - } - - const data = (await Promise.all(filePromises)).flat(); + // Initialize aggregated metadata. + let totalFileSize = 0; + let aggregatedLastUsed = 0; + let aggregatedUpdateDate = 0; + let aggregatedEngineIds = []; + // Process each file entry. const files = []; - for (const { file: path, headers } of data) { - files.push({ path, headers }); + for (const { file: path, headers } of fileData) { + // Aggregate metadata. + totalFileSize += headers.fileSize; + aggregatedLastUsed = Math.max(aggregatedLastUsed, headers.lastUsed); + aggregatedUpdateDate = Math.max( + aggregatedUpdateDate, + headers.lastUpdated + ); + if (headers.engineIds && headers.engineIds.length) { + aggregatedEngineIds = headers.engineIds; + } + files.push({ path, headers, engineIds: headers.engineIds || [] }); } - return files; + return { + files, + metadata: { + totalSize: totalFileSize, + lastUsed: aggregatedLastUsed, + updateDate: aggregatedUpdateDate, + engineIds: aggregatedEngineIds, + }, + }; } /** * Lists all models stored in the cache. * - * @returns {Promise>} An array of model identifiers. + * @returns {Promise>} + * An array of model identifiers. */ async listModels() { + // Get all keys (model/revision pairs) from the underlying store. const modelRevisions = await this.#getKeys({ storeName: this.taskStoreName, indexName: this.#indices.modelRevisionIndex.name, }); + const models = []; + // Process each key entry. for (const { key } of modelRevisions) { - models.push({ name: key[0], revision: key[1] }); + const model = key[0]; + const revision = key[1]; + + models.push({ + name: model, + revision, + }); } return models; } diff --git a/toolkit/components/ml/tests/browser/browser_ml_cache.js b/toolkit/components/ml/tests/browser/browser_ml_cache.js index 0f0d8a00e1a1..78a2315d9df9 100644 --- a/toolkit/components/ml/tests/browser/browser_ml_cache.js +++ b/toolkit/components/ml/tests/browser/browser_ml_cache.js @@ -1093,7 +1093,10 @@ add_task(async function test_listFiles() { headers, }); - const files = await cache.listFiles({ model: "org/model", revision: "v1" }); + const { files } = await cache.listFiles({ + model: "org/model", + revision: "v1", + }); const expected = [ { path: "file.txt", @@ -1104,6 +1107,7 @@ add_task(async function test_listFiles() { lastUsed: when1, lastUpdated: when1, }, + engineIds: [], }, { path: "file2.txt", @@ -1114,6 +1118,7 @@ add_task(async function test_listFiles() { lastUsed: when2, lastUpdated: when2, }, + engineIds: [], }, { path: "sub/file3.txt", @@ -1125,6 +1130,7 @@ add_task(async function test_listFiles() { lastUsed: when3, lastUpdated: when3, }, + engineIds: [], }, ]; @@ -1172,7 +1178,7 @@ add_task(async function test_listFilesUsingTaskName() { headers, }); - const files = await cache.listFiles({ taskName }); + const { files } = await cache.listFiles({ taskName }); const expected = [ { path: "file.txt", @@ -1183,6 +1189,7 @@ add_task(async function test_listFilesUsingTaskName() { lastUsed: when1, lastUpdated: when1, }, + engineIds: [], }, { path: "file2.txt", @@ -1193,6 +1200,7 @@ add_task(async function test_listFilesUsingTaskName() { lastUsed: when2, lastUpdated: when2, }, + engineIds: [], }, { path: "sub/file3.txt", @@ -1204,6 +1212,7 @@ add_task(async function test_listFilesUsingTaskName() { lastUsed: when3, lastUpdated: when3, }, + engineIds: [], }, ]; @@ -1252,7 +1261,7 @@ add_task(async function test_listFilesUsingNonExistingTaskName() { }), ]); - const files = await cache.listFiles({ taskName: "non-existing-task" }); + const { files } = await cache.listFiles({ taskName: "non-existing-task" }); Assert.deepEqual(files, []); @@ -1322,11 +1331,12 @@ add_task(async function test_initDbFromExistingEmpty() { lastUsed: when, lastUpdated: when, }, + engineIds: [], }, ]; // Ensure every table & indices is on so that we can list files - const files = await cache.listFiles({ taskName }); + const { files } = await cache.listFiles({ taskName }); Assert.deepEqual(files, expected); await deleteCache(cache); @@ -1369,7 +1379,7 @@ add_task(async function test_initDbFromExistingNoChange() { Assert.equal(cache.db.version, 2); // Ensure tables are all empty. - const files = await cache.listFiles({ taskName }); + const { files } = await cache.listFiles({ taskName }); Assert.deepEqual(files, []); @@ -1422,11 +1432,12 @@ add_task(async function test_initDbFromExistingElseWhereStoreChanges() { lastUpdated: when, lastUsed: when, }, + engineIds: [], }, ]; // Ensure every table & indices is on so that we can list files - const files = await cache2.listFiles({ taskName }); + const { files } = await cache2.listFiles({ taskName }); Assert.deepEqual(files, expected); await deleteCache(cache2); @@ -1750,25 +1761,17 @@ add_task(async function test_migrateStore_modelsDeleted() { cache = await IndexedDBCache.init({ dbName, version: 5 }); // Verify all unknown model data is deleted - let remainingFiles = await cache.listFiles({ + const { files: random } = await cache.listFiles({ model: "random/model", revision: "v1", }); - Assert.deepEqual( - remainingFiles, - [], - "All unknown model files should be deleted." - ); + Assert.deepEqual(random, [], "All unknown model files should be deleted."); - remainingFiles = await cache.listFiles({ + const { files: unknown } = await cache.listFiles({ model: "unknown/model", revision: "v2", }); - Assert.deepEqual( - remainingFiles, - [], - "All unknown model files should be deleted." - ); + Assert.deepEqual(unknown, [], "All unknown model files should be deleted."); await deleteCache(cache); }); diff --git a/toolkit/locales-preview/localModelManagement.ftl b/toolkit/locales-preview/localModelManagement.ftl index 3b26f83edbcc..ae43da05526a 100644 --- a/toolkit/locales-preview/localModelManagement.ftl +++ b/toolkit/locales-preview/localModelManagement.ftl @@ -15,4 +15,8 @@ mlmodel-description = They process your queries on your device, meaning they are more private than online AI providers, but sometimes slower depending on the performance of your device. Learn more -addon-detail-file-size-label = File size +# Label for button that when clicked removed local model +mlmodel-remove-addon-button = + .aria-label = Remove +# Label for the aggregated value of all files for a model +mlmodel-addon-detail-totalsize-label = File size diff --git a/toolkit/mozapps/extensions/content/aboutaddons.html b/toolkit/mozapps/extensions/content/aboutaddons.html index 77702576f03c..c32be38ac2ff 100644 --- a/toolkit/mozapps/extensions/content/aboutaddons.html +++ b/toolkit/mozapps/extensions/content/aboutaddons.html @@ -70,6 +70,10 @@ src="chrome://mozapps/content/extensions/components/mlmodel-list-intro.mjs" type="module" > +