feat: tab suggestions

2025-07-24 23:47:34 +02:00
parent f0d17e78d7
commit 47151a2ba7
30 changed files with 1293 additions and 180 deletions
@@ -46,9 +46,9 @@ const SMART_TAB_GROUPING_CONFIG = {
 */
 function createModelInput(keywords, documents) {
    if (!keywords || keywords.length === 0) {
-        return `Topic from keywords: titles: \n${documents.join(" \n")}`;
+        return `Topic from keywords: titles: \n${documents.slice(0, 3).join(" \n")}`;
    }
-    return `Topic from keywords: ${keywords.join(", ")}. titles: \n${documents.join(" \n")}`;
+    return `Topic from keywords: ${keywords.join(", ")}. titles: \n${documents.slice(0, 3).join(" \n")}`;
 }

 /**
@@ -81,7 +81,6 @@ function cutAtDuplicateWords(phrase) {
    return phrase; // return original phrase
 }

-
 /**
   *
   * @param {MLEngine} engine the engine to check
@@ -96,7 +95,29 @@ this.ml = class extends ExtensionAPI {
        return {
            experiments: {
                ml: {
-                    async containerTopic(keywords, documents) {
+                    async generateEmbeddings(textToEmbedList) {
+                        const inputData = {
+                            inputArgs: textToEmbedList,
+                            runOptions: {
+                                pooling: "mean",
+                                normalize: true,
+                            },
+                        };
+
+                        if (isEngineClosed(this.embeddingEngine)) {
+                            this.embeddingEngine = await createEngine(SMART_TAB_GROUPING_CONFIG.embedding);
+                        }
+
+                        const request = {
+                            args: [inputData.inputArgs],
+                            options: inputData.runOptions,
+                        };
+
+                        const generated = await this.embeddingEngine.run(request);
+
+                        return JSON.stringify(generated);
+                    },
+                    async predictTopic(keywords, documents) {
                        if (isEngineClosed(this.topicEngine)) {
                            const {
                                featureId,
@@ -26,15 +26,25 @@ function sendErrorForRequest(id) {
 port.onMessage.addListener(async (message) => {
    let requestId = message["id"]
    switch (message["action"]) {
-        case "getContainerTopic":
+        case "predictDocumentTopic": {
            const documents = message["args"];
            const keywords = (documents.length > 1)
                ? await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")])
                : [[]];

-            browser.experiments.ml.containerTopic(keywords[0], documents)
+            browser.experiments.ml.predictTopic(keywords[0], documents)
                .then(sendJsonResultForRequest(requestId))
-                .catch(sendErrorForRequest(requestId))
-            break
+                .catch(sendErrorForRequest(requestId));
+
+            break;
+        }
+        case "generateDocumentEmbeddings": {
+            const documents = message["args"];
+            await browser.experiments.ml.generateEmbeddings(documents)
+                .then(sendJsonResultForRequest(requestId))
+                .catch(sendErrorForRequest(requestId));
+
+            break;
+        }
    }
 });
@@ -33,7 +33,23 @@
        "description": "Machine Learning utilities",
        "functions": [
            {
-                "name": "containerTopic",
+                "name": "generateEmbeddings",
+                "type": "function",
+                "description": "Generate embeddings for a list of text strings using ML engine",
+                "async": true,
+                "parameters": [
+                    {
+                        "name": "textToEmbedList",
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "description": "Array of text strings to generate embeddings for"
+                    }
+                ]
+            },
+            {
+                "name": "predictTopic",
                "type": "function",
                "description": "Generate topic from keywords and documents using ML engine",
                "async": true,
@@ -13,8 +13,8 @@ class GeckoMlApiImpl : GeckoMlApi {
        }
    }

-    override fun getContainerTopic(titles: List<String>, callback: (Result<String>) -> Unit) {
-        MLEngineFeature.scheduleRequest("getContainerTopic", titles.toJson(), object : ResultConsumer<JSONObject> {
+    override fun predictDocumentTopic(documents: List<String>, callback: (Result<String>) -> Unit) {
+        MLEngineFeature.scheduleRequest("predictDocumentTopic", documents.toJson(), object : ResultConsumer<JSONObject> {
            override fun success(result: JSONObject) {
                callback(Result.success(result.getString("result")))
            }
@@ -24,4 +24,38 @@ class GeckoMlApiImpl : GeckoMlApi {
            }
        })
    }
+
+    override fun generateDocumentEmbeddings(
+        documents: List<String>,
+        callback: (Result<List<Any?>>) -> Unit
+    ) {
+        MLEngineFeature.scheduleRequest("generateDocumentEmbeddings", documents.toJson(), object : ResultConsumer<JSONObject> {
+            override fun success(result: JSONObject) {
+                try {
+                    val encodedResult = result.getString("result")
+                    val decodedJsonArray = JSONArray(encodedResult)
+                    val embeddings = mutableListOf<List<Double>>()
+
+                    for (i in 0 until decodedJsonArray.length()) {
+                        val embeddingArray = decodedJsonArray.getJSONArray(i)
+                        val embedding = mutableListOf<Double>()
+
+                        for (j in 0 until embeddingArray.length()) {
+                            embedding.add(embeddingArray.getDouble(j))
+                        }
+                        embeddings.add(embedding)
+                    }
+
+                    callback(Result.success(embeddings))
+                } catch (e: Exception) {
+                    callback(Result.failure(e))
+                }
+            }
+
+            override fun error(errorCode: String, errorMessage: String?, errorDetails: Any?) {
+                callback(Result.failure(Exception("$errorCode $errorMessage $errorDetails")))
+            }
+        })
+    }
+
 }
@@ -3425,7 +3425,8 @@ 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)
+  fun predictDocumentTopic(documents: List<String>, callback: (Result<String>) -> Unit)
+  fun generateDocumentEmbeddings(documents: List<String>, callback: (Result<List<Any?>>) -> Unit)

  companion object {
    /** The codec used by GeckoMlApi. */
@@ -3437,12 +3438,32 @@ interface GeckoMlApi {
    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)
+        val channel = BasicMessageChannel<Any?>(binaryMessenger, "dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.predictDocumentTopic$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 documentsArg = args[0] as List<String>
+            api.predictDocumentTopic(documentsArg) { 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)
+        }
+      }
+      run {
+        val channel = BasicMessageChannel<Any?>(binaryMessenger, "dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.generateDocumentEmbeddings$separatedMessageChannelSuffix", codec)
+        if (api != null) {
+          channel.setMessageHandler { message, reply ->
+            val args = message as List<Any?>
+            val documentsArg = args[0] as List<String>
+            api.generateDocumentEmbeddings(documentsArg) { result: Result<List<Any?>> ->
              val error = result.exceptionOrNull()
              if (error != null) {
                reply.reply(GeckoPigeonUtils.wrapError(error))
@@ -9,13 +9,23 @@ import 'package:flutter_mozilla_components/src/pigeons/gecko.g.dart';
 final _apiInstance = GeckoMlApi();

 class GeckoMlService {
-  Future<String> getContainerTopic(Set<String> titles, {int maxCount = 8}) {
+  Future<String> predictDocumentTopic(Set<String> titles, {int maxCount = 10}) {
    var selectedTitles = titles.toList();
    if (selectedTitles.length > maxCount) {
      //TODO: Randomize for now, maybe use clusters later
      selectedTitles = (selectedTitles..shuffle()).take(maxCount).toList();
    }

-    return _apiInstance.getContainerTopic(selectedTitles);
+    return _apiInstance.predictDocumentTopic(selectedTitles);
+  }
+
+  Future<List<List<double>>> generateDocumentEmbeddings(
+    List<String> documents,
+  ) async {
+    final embeddings = await _apiInstance.generateDocumentEmbeddings(documents);
+
+    return embeddings
+        .map((values) => (values! as List).cast<double>())
+        .toList();
  }
 }
@@ -3968,14 +3968,14 @@ class GeckoMlApi {

  final String pigeonVar_messageChannelSuffix;

-  Future<String> getContainerTopic(List<String> titles) async {
-    final String pigeonVar_channelName = 'dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.getContainerTopic$pigeonVar_messageChannelSuffix';
+  Future<String> predictDocumentTopic(List<String> documents) async {
+    final String pigeonVar_channelName = 'dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.predictDocumentTopic$pigeonVar_messageChannelSuffix';
    final BasicMessageChannel<Object?> pigeonVar_channel = BasicMessageChannel<Object?>(
      pigeonVar_channelName,
      pigeonChannelCodec,
      binaryMessenger: pigeonVar_binaryMessenger,
    );
-    final Future<Object?> pigeonVar_sendFuture = pigeonVar_channel.send(<Object?>[titles]);
+    final Future<Object?> pigeonVar_sendFuture = pigeonVar_channel.send(<Object?>[documents]);
    final List<Object?>? pigeonVar_replyList =
        await pigeonVar_sendFuture as List<Object?>?;
    if (pigeonVar_replyList == null) {
@@ -3995,6 +3995,34 @@ class GeckoMlApi {
      return (pigeonVar_replyList[0] as String?)!;
    }
  }
+
+  Future<List<Object?>> generateDocumentEmbeddings(List<String> documents) async {
+    final String pigeonVar_channelName = 'dev.flutter.pigeon.flutter_mozilla_components.GeckoMlApi.generateDocumentEmbeddings$pigeonVar_messageChannelSuffix';
+    final BasicMessageChannel<Object?> pigeonVar_channel = BasicMessageChannel<Object?>(
+      pigeonVar_channelName,
+      pigeonChannelCodec,
+      binaryMessenger: pigeonVar_binaryMessenger,
+    );
+    final Future<Object?> pigeonVar_sendFuture = pigeonVar_channel.send(<Object?>[documents]);
+    final List<Object?>? pigeonVar_replyList =
+        await pigeonVar_sendFuture as List<Object?>?;
+    if (pigeonVar_replyList == null) {
+      throw _createConnectionError(pigeonVar_channelName);
+    } else if (pigeonVar_replyList.length > 1) {
+      throw PlatformException(
+        code: pigeonVar_replyList[0]! as String,
+        message: pigeonVar_replyList[1] as String?,
+        details: pigeonVar_replyList[2],
+      );
+    } else if (pigeonVar_replyList[0] == null) {
+      throw PlatformException(
+        code: 'null-error',
+        message: 'Host platform returned null value for non-null return value.',
+      );
+    } else {
+      return (pigeonVar_replyList[0] as List<Object?>?)!;
+    }
+  }
 }

 class GeckoBrowserExtensionApi {
@@ -956,7 +956,9 @@ abstract class GeckoPrefApi {
@HostApi()
 abstract class GeckoMlApi {
  @async
-  String getContainerTopic(List<String> titles);
+  String predictDocumentTopic(List<String> documents);
+  @async
+  List generateDocumentEmbeddings(List<String> documents);
 }

@HostApi()