container topic inference

2025-07-20 00:37:51 +02:00
parent 2a3e0b6b6d
commit dc14a7d718
29 changed files with 1117 additions and 12 deletions
@@ -0,0 +1,151 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+"use strict";
+
+const { createEngine } = ChromeUtils.importESModule("chrome://global/content/ml/EngineProcess.sys.mjs");
+
+const ML_TASK_FEATURE_EXTRACTION = "feature-extraction";
+const ML_TASK_TEXT2TEXT = "text2text-generation";
+
+const SMART_TAB_GROUPING_CONFIG = {
+    embedding: {
+        dtype: "q8",
+        timeoutMS: 2 * 60 * 1000, // 2 minutes
+        taskName: ML_TASK_FEATURE_EXTRACTION,
+        featureId: "smart-tab-embedding",
+        backend: "onnx",
+    },
+    topicGeneration: {
+        dtype: "q8",
+        timeoutMS: 2 * 60 * 1000, // 2 minutes
+        taskName: ML_TASK_TEXT2TEXT,
+        featureId: "smart-tab-topic",
+        backend: "onnx",
+    },
+    // dataConfig: {
+    //     titleKey: "label",
+    //     descriptionKey: "description",
+    // },
+    // clustering: {
+    //     dimReductionMethod: null, // Not completed.
+    //     clusterImplementation: CLUSTER_METHODS.KMEANS,
+    //     clusteringTriesPerK: 3,
+    //     anchorMethod: ANCHOR_METHODS.FIXED,
+    //     pregroupedHandlingMethod: PREGROUPED_HANDLING_METHODS.EXCLUDE,
+    //     pregroupedSilhouetteBoost: 2, // Relative weight of the cluster's score and all other cluster's combined
+    //     suggestOtherTabsMethod: SUGGEST_OTHER_TABS_METHODS.NEAREST_NEIGHBOR,
+    // },
+};
+
+/**
+ * Generate model input from keywords and documents
+ * @param {string []} keywords
+ * @param {string []} documents
+ */
+function createModelInput(keywords, documents) {
+    if (!keywords || keywords.length === 0) {
+        return `Topic from keywords: titles: \n${documents.join(" \n")}`;
+    }
+    return `Topic from keywords: ${keywords.join(", ")}. titles: \n${documents.join(" \n")}`;
+}
+
+/**
+ * One artifact of the LLM output is that sometimes words are duplicated
+ * This function cuts the phrase when it sees the first duplicate word.
+ * Handles simple singluar / plural duplicates (-s only).
+ * @param {string} phrase Input phrase
+ * @returns {string} phrase cut before any duplicate word
+ */
+function cutAtDuplicateWords(phrase) {
+    if (!phrase.length) {
+        return phrase;
+    }
+    const wordsSet = new Set();
+    const wordList = phrase.split(" ");
+    for (let i = 0; i < wordList.length; i++) {
+        let baseWord = wordList[i].toLowerCase();
+        if (baseWord.length > 3) {
+            if (baseWord.slice(-1) === "s") {
+                baseWord = baseWord.slice(0, -1);
+            }
+        }
+        if (wordsSet.has(baseWord)) {
+            // We are seeing a baseWord word. Exit with just the words so far and don't
+            // add any new words
+            return wordList.slice(0, i).join(" ");
+        }
+        wordsSet.add(baseWord);
+    }
+    return phrase; // return original phrase
+}
+
+
+/**
+   *
+   * @param {MLEngine} engine the engine to check
+   * @return {boolean} true if the engine has not been initialized or closed
+   */
+function isEngineClosed(engine) {
+    return !engine || engine?.engineStatus === "closed";
+}
+
+this.ml = class extends ExtensionAPI {
+    getAPI(context) {
+        return {
+            experiments: {
+                ml: {
+                    async containerTopic(keywords, documents) {
+                        if (isEngineClosed(this.topicEngine)) {
+                            const {
+                                featureId,
+                                engineId,
+                                dtype,
+                                taskName,
+                                timeoutMS,
+                                modelId,
+                                modelRevision,
+                                backend,
+                            } = SMART_TAB_GROUPING_CONFIG.topicGeneration;
+
+                            let initData = {
+                                featureId,
+                                engineId,
+                                dtype,
+                                taskName,
+                                timeoutMS,
+                                modelId,
+                                modelRevision,
+                                backend,
+                            };
+
+                            this.topicEngine = await createEngine(initData);
+                        }
+
+                        const inputArgs = createModelInput(
+                            keywords,
+                            documents
+                        );
+                        const requestInfo = {
+                            inputArgs,
+                            runOptions: {
+                                max_length: 6,
+                            },
+                        };
+                        const request = {
+                            args: [requestInfo.inputArgs],
+                            options: requestInfo.runOptions,
+                        };
+
+                        const res = await this.topicEngine.run(request);
+
+                        const generated = cutAtDuplicateWords((res[0]["generated_text"] || "").trim());
+
+                        return generated;
+                    }
+                }
+            }
+        };
+    }
+};
@@ -0,0 +1,29 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+"use strict";
+
+const { KeywordExtractor } = ChromeUtils.importESModule(
+    "chrome://global/content/ml/NLPUtils.sys.mjs"
+);
+
+this.nlp = class extends ExtensionAPI {
+    getAPI(context) {
+        return {
+            experiments: {
+                nlp: {
+                    async extractKeywords(corpus, maxKeywords = 3) {
+                        try {
+                            const keywordExtractor = new KeywordExtractor();
+                            const keywords = keywordExtractor.fitTransform(corpus, maxKeywords);
+                            return keywords;
+                        } catch (error) {
+                            throw new ExtensionError(`Keyword extraction failed: ${error.message}`);
+                        }
+                    },
+                }
+            }
+        };
+    }
+};
@@ -0,0 +1,38 @@
+'use strict';
+
+const port = browser.runtime.connectNative("mlEngine");
+
+function sendJsonResultForRequest(id) {
+    return function (result) {
+        port.postMessage({
+            "id": id,
+            "status": "success",
+            "result": result
+        })
+    }
+}
+
+function sendErrorForRequest(id) {
+    return function (error) {
+        console.error(error);
+        port.postMessage({
+            "id": id,
+            "status": "error",
+            "error": error
+        });
+    }
+}
+
+port.onMessage.addListener(async (message) => {
+    let requestId = message["id"]
+    switch (message["action"]) {
+        case "getContainerTopic":
+            const documents = message["args"];
+            const keywords = await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")]);
+
+            browser.experiments.ml.containerTopic(keywords[0], documents)
+                .then(sendJsonResultForRequest(requestId))
+                .catch(sendErrorForRequest(requestId))
+            break
+    }
+});
@@ -0,0 +1,61 @@
+{
+  "manifest_version": 2,
+  "name": "ml-engine",
+  "version": "1.0",
+  "description": "WebLibre ML Engine",
+  "browser_specific_settings": {
+    "gecko": {
+      "id": "ml-engine@weblibre.eu"
+    }
+  },
+  "experiment_apis": {
+    "nlp": {
+      "schema": "schema.json",
+      "parent": {
+        "scopes": [
+          "addon_parent"
+        ],
+        "script": "api/nlp.js",
+        "paths": [
+          [
+            "experiments",
+            "nlp"
+          ]
+        ]
+      }
+    },
+    "ml": {
+      "schema": "schema.json",
+      "parent": {
+        "scopes": [
+          "addon_parent"
+        ],
+        "script": "api/ml.js",
+        "paths": [
+          [
+            "experiments",
+            "ml"
+          ]
+        ]
+      }
+    }
+  },
+  "background": {
+    "scripts": [
+      "background.js"
+    ]
+  },
+  "optional_permissions": [
+    "trialML"
+  ],
+  "permissions": [
+    "nativeMessaging",
+    "nativeMessagingFromContent",
+    "geckoViewAddons",
+    "cookies",
+    "menus",
+    "scripting",
+    "storage",
+    "<all_urls>"
+  ]
+}
@@ -0,0 +1,61 @@
+[
+    {
+        "namespace": "experiments.nlp",
+        "description": "Natural Language Processing utilities",
+        "functions": [
+            {
+                "name": "extractKeywords",
+                "type": "function",
+                "description": "Extract keywords from a corpus of text documents",
+                "async": true,
+                "parameters": [
+                    {
+                        "name": "corpus",
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "description": "Array of text documents to extract keywords from"
+                    },
+                    {
+                        "name": "maxKeywords",
+                        "type": "integer",
+                        "optional": true,
+                        "default": 3,
+                        "description": "Maximum number of keywords to extract per document"
+                    }
+                ]
+            }
+        ]
+    },
+    {
+        "namespace": "experiments.ml",
+        "description": "Machine Learning utilities",
+        "functions": [
+            {
+                "name": "containerTopic",
+                "type": "function",
+                "description": "Generate topic from keywords and documents using ML engine",
+                "async": true,
+                "parameters": [
+                    {
+                        "name": "keywords",
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "description": "Array of keywords to generate topic from"
+                    },
+                    {
+                        "name": "documents",
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "description": "Array of document titles/content"
+                    }
+                ]
+            }
+        ]
+    }
+]
@@ -9,6 +9,7 @@ import eu.weblibre.flutter_mozilla_components.feature.ContainerProxyFeature
 import eu.weblibre.flutter_mozilla_components.feature.CookieManagerFeature
 import eu.weblibre.flutter_mozilla_components.feature.PrefManagerFeature
 import eu.weblibre.flutter_mozilla_components.feature.BrowserExtensionFeature
+import eu.weblibre.flutter_mozilla_components.feature.MLEngineFeature
 import eu.weblibre.flutter_mozilla_components.pigeons.BrowserExtensionEvents
 import mozilla.components.browser.engine.gecko.GeckoEngine
 import mozilla.components.browser.engine.gecko.fetch.GeckoViewFetchClient
@@ -39,7 +40,7 @@ object EngineProvider {
            builder.extensionsWebAPIEnabled(true)

            // Disable output for now to improve performance
-            //builder.consoleOutput(true)
+            builder.consoleOutput(true)

            runtime = GeckoRuntime.create(context, builder.build())
        }
@@ -57,6 +58,7 @@ object EngineProvider {
            PrefManagerFeature.install(it)
            ContainerProxyFeature.install(it)
            BrowserExtensionFeature.install(it, extensionEvents)
+            MLEngineFeature.install(it)
        }
    }

@@ -27,6 +27,7 @@ import eu.weblibre.flutter_mozilla_components.pigeons.GeckoDownloadsApi
 import eu.weblibre.flutter_mozilla_components.pigeons.GeckoEngineSettingsApi
 import eu.weblibre.flutter_mozilla_components.pigeons.GeckoFindApi
 import eu.weblibre.flutter_mozilla_components.pigeons.GeckoIconsApi
+import eu.weblibre.flutter_mozilla_components.pigeons.GeckoMlApi
 import eu.weblibre.flutter_mozilla_components.pigeons.GeckoPrefApi
 import eu.weblibre.flutter_mozilla_components.pigeons.GeckoSelectionActionController
 import eu.weblibre.flutter_mozilla_components.pigeons.GeckoSelectionActionEvents
@@ -152,6 +153,7 @@ class GeckoBrowserApiImpl : GeckoBrowserApi {
        GeckoTabsApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoTabsApiImpl())
        GeckoIconsApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoIconsApiImpl())
        GeckoCookieApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoCookieApiImpl())
+        GeckoMlApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoMlApiImpl())
        GeckoPrefApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoPrefApiImpl())
        GeckoContainerProxyApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoContainerProxyApiImpl())
        GeckoFindApi.setUp(_flutterPluginBinding.binaryMessenger, GeckoFindApiImpl())
@@ -0,0 +1,27 @@
+package eu.weblibre.flutter_mozilla_components.api
+
+import eu.weblibre.flutter_mozilla_components.feature.MLEngineFeature
+import eu.weblibre.flutter_mozilla_components.feature.ResultConsumer
+import eu.weblibre.flutter_mozilla_components.pigeons.GeckoMlApi
+import org.json.JSONArray
+import org.json.JSONObject
+
+class GeckoMlApiImpl : GeckoMlApi {
+    private fun List<String>?.toJson(): JSONArray {
+        return JSONArray().apply {
+            this@toJson?.forEach { put(it) }
+        }
+    }
+
+    override fun getContainerTopic(titles: List<String>, callback: (Result<String>) -> Unit) {
+        MLEngineFeature.scheduleRequest("getContainerTopic", titles.toJson(), object : ResultConsumer<JSONObject> {
+            override fun success(result: JSONObject) {
+                callback(Result.success(result.getString("result")))
+            }
+
+            override fun error(errorCode: String, errorMessage: String?, errorDetails: Any?) {
+                callback(Result.failure(Exception("$errorCode $errorMessage $errorDetails")))
+            }
+        })
+    }
+}
@@ -0,0 +1,126 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/.
+ */
+
+package eu.weblibre.flutter_mozilla_components.feature
+
+import androidx.annotation.VisibleForTesting
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import kotlinx.coroutines.withContext
+import mozilla.components.concept.engine.webextension.MessageHandler
+import mozilla.components.concept.engine.webextension.Port
+import mozilla.components.concept.engine.webextension.WebExtension
+import mozilla.components.concept.engine.webextension.WebExtensionRuntime
+import mozilla.components.support.base.log.logger.Logger
+import mozilla.components.support.webextensions.BuiltInWebExtensionController
+import org.json.JSONObject
+
+object MLEngineFeature {
+    private val logger = Logger("ml-engine")
+
+    private const val ML_ENGINE_REPORTER_EXTENSION_ID = "ml-engine@weblibre.eu"
+    private const val ML_ENGINE_REPORTER_EXTENSION_URL = "resource://android/assets/extensions/ml_engine/"
+    private const val ML_ENGINE_REPORTER_MESSAGING_ID = "mlEngine"
+
+    private var nextRequestId: Int = 0
+    private val requestHandlers = HashMap<Int, ResultConsumer<JSONObject>>()
+    private val mutex = Mutex()
+
+    @VisibleForTesting
+    // This is an internal var to make it mutable for unit testing purposes only
+    internal var extensionController = BuiltInWebExtensionController(
+        ML_ENGINE_REPORTER_EXTENSION_ID,
+        ML_ENGINE_REPORTER_EXTENSION_URL,
+        ML_ENGINE_REPORTER_MESSAGING_ID,
+    )
+
+    fun scheduleRequest(command: String, args: Any, callback: ResultConsumer<JSONObject>) {
+        val message = JSONObject()
+        message.put("action", command);
+        message.put("args", args)
+
+        runBlocking {
+            withContext(Dispatchers.Default) {
+                mutex.withLock {
+                    message.put("id", nextRequestId)
+
+                    requestHandlers[nextRequestId] = callback
+
+                    nextRequestId += 1
+
+                    extensionController.sendBackgroundMessage(message)
+                }
+            }
+        }
+    }
+
+    private class PrefManagerReporterBackgroundMessageHandler() : MessageHandler {
+        override fun onPortMessage(message: Any, port: Port) {
+            runBlocking {
+                withContext(Dispatchers.Default) {
+                    mutex.withLock {
+                        val messageJSON = message as JSONObject;
+
+                        val requestId = messageJSON.getInt("id")
+                        val status = messageJSON.getString("status")
+                        if (status == "success") {
+                            requestHandlers[requestId]?.success(message)
+                        } else {
+                            requestHandlers[requestId]?.error(
+                                "ML Engine",
+                                "Failed to perform operation",
+                                message.getString("error")
+                            )
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    /**
+     * Installs the web extension in the runtime through the WebExtensionRuntime install method
+     *
+     * @param runtime a WebExtensionRuntime.
+     * @param productName a custom product name used to automatically label reports. Defaults to
+     * "android-components".
+     */
+    fun install(runtime: WebExtensionRuntime) {
+        extensionController.registerBackgroundMessageHandler(
+            PrefManagerReporterBackgroundMessageHandler(),
+        )
+        extensionController.install(
+            runtime,
+            onSuccess = {
+                logger.debug("Installed ml-engine webextension: ${it.id}")
+
+                grantPermissions(runtime, it)
+            },
+            onError = { throwable ->
+                logger.error("Failed to install ml-engine webextension: ", throwable)
+            },
+        )
+    }
+
+    private fun grantPermissions(runtime: WebExtensionRuntime, extension: WebExtension) {
+        val permissions = listOf("trialML")
+        val origins = emptyList<String>() // Add any host permissions if needed
+
+        runtime.addOptionalPermissions(
+            ML_ENGINE_REPORTER_EXTENSION_ID,
+            permissions,
+            origins,
+            onSuccess = { grantedExtension ->
+                logger.debug("Successfully granted permissions to extension: ${grantedExtension.id}")
+            },
+            onError = { throwable ->
+                logger.error("Failed to grant permissions to extension: ", throwable)
+            }
+        )
+    }
+}
@@ -1,4 +1,4 @@
-// Autogenerated from Pigeon (v25.3.2), do not edit directly.
+// Autogenerated from Pigeon (v25.5.0), do not edit directly.
 // See also: https://pub.dev/packages/pigeon
@file:Suppress("UNCHECKED_CAST", "ArrayInDataClass")

@@ -3424,6 +3424,42 @@ interface GeckoPrefApi {
  }
 }
 /** Generated interface from Pigeon that represents a handler of messages from Flutter. */
+interface GeckoMlApi {
+  fun getContainerTopic(titles: List<String>, callback: (Result<String>) -> Unit)
+
+  companion object {
+    /** The codec used by GeckoMlApi. */
+    val codec: MessageCodec<Any?> by lazy {
+      GeckoPigeonCodec()
+    }
+    /** Sets up an instance of `GeckoMlApi` to handle messages through the `binaryMessenger`. */
+    @JvmOverloads
+    fun setUp(binaryMessenger: BinaryMessenger, api: GeckoMlApi?, messageChannelSuffix: String = "") {
+      val separatedMessageChannelSuffix = if (messageChannelSuffix.isNotEmpty()) ".$messageChannelSuffix" else ""
+      run {
+        val channel = BasicMessageChannel<Any?>(binaryMessenger, "dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.getContainerTopic$separatedMessageChannelSuffix", codec)
+        if (api != null) {
+          channel.setMessageHandler { message, reply ->
+            val args = message as List<Any?>
+            val titlesArg = args[0] as List<String>
+            api.getContainerTopic(titlesArg) { result: Result<String> ->
+              val error = result.exceptionOrNull()
+              if (error != null) {
+                reply.reply(GeckoPigeonUtils.wrapError(error))
+              } else {
+                val data = result.getOrNull()
+                reply.reply(GeckoPigeonUtils.wrapResult(data))
+              }
+            }
+          }
+        } else {
+          channel.setMessageHandler(null)
+        }
+      }
+    }
+  }
+}
+/** Generated interface from Pigeon that represents a handler of messages from Flutter. */
 interface GeckoBrowserExtensionApi {
  fun getMarkdown(htmlList: List<String>, callback: (Result<List<Any>>) -> Unit)