tor-browser

commit de5612d526fd3cdb4744a429debe230fdd2648f4
parent 30b907bc2066ae5d0a2e9f618dbddbed525b5582
Author: Tzu-An Liu <tliu@mozilla.com>
Date:   Thu, 30 Oct 2025 22:29:23 +0000

Bug 1996648 - Update query intent detection model to mobilebert v0.2 r=Mardak,firefox-ai-ml-reviewers

Differential Revision: https://phabricator.services.mozilla.com/D270247

Diffstat:
M
browser/components/genai/SmartAssistEngine.sys.mjs
 | 
25
++++++++++++++++++++++---
M
browser/components/genai/tests/xpcshell/test_smart_assist_engine.js
 | 
50
++++++++++++++++++++++++++++++++++++++++++++------

2 files changed, 66 insertions(+), 9 deletions(-)
diff --git a/browser/components/genai/SmartAssistEngine.sys.mjs b/browser/components/genai/SmartAssistEngine.sys.mjs
@@ -223,14 +223,33 @@ export const SmartAssistEngine = {
       const engine = await this._createEngine({
         featureId: "smart-intent",
         modelId: "mozilla/mobilebert-query-intent-detection",
-        modelRevision: "v0.1.0",
+        modelRevision: "v0.2.0",
         taskName: "text-classification",
       });
-      const resp = await engine.run({ args: [[query]] });
-      return resp[0].label.toLowerCase();
+      const threshold = 0.6;
+      const cleanedQuery = this._preprocessQuery(query);
+      const resp = await engine.run({ args: [[cleanedQuery]] });
+      // resp example: [{ label: "chat", score: 0.95 }, { label: "search", score: 0.04 }]
+      if (
+        resp[0].label.toLowerCase() === "chat" &&
+        resp[0].score >= threshold
+      ) {
+        return "chat";
+      }
+      return "search";
     } catch (error) {
       console.error("Error using intent detection model:", error);
       throw error;
     }
   },
+
+  // Helper function for preprocessing text input
+  _preprocessQuery(query) {
+    if (typeof query !== "string") {
+      throw new TypeError(
+        `Expected a string for query preprocessing, but received ${typeof query}`
+      );
+    }
+    return query.replace(/\?/g, "").trim();
+  },
 };
diff --git a/browser/components/genai/tests/xpcshell/test_smart_assist_engine.js b/browser/components/genai/tests/xpcshell/test_smart_assist_engine.js
@@ -199,14 +199,21 @@ add_task(async function test_getPromptIntent_basic() {
         "explore",
       ];
       const formattedPrompt = query.toLowerCase();
-
-      const intent = searchKeywords.some(keyword =>
+      const isSearch = searchKeywords.some(keyword =>
         formattedPrompt.includes(keyword)
-      )
-        ? "search"
-        : "chat";
+      );
 
-      return [{ label: intent }];
+      // Simulate model confidence scores
+      if (isSearch) {
+        return [
+          { label: "search", score: 0.95 },
+          { label: "chat", score: 0.05 },
+        ];
+      }
+      return [
+        { label: "chat", score: 0.95 },
+        { label: "search", score: 0.05 },
+      ];
     },
   };
 
@@ -220,4 +227,35 @@ add_task(async function test_getPromptIntent_basic() {
       `getPromptIntent("${prompt}") should return "${expected}"`
     );
   }
+
+  sb.restore();
+});
+
+add_task(async function test_preprocessQuery_removes_question_marks() {
+  const sb = sinon.createSandbox();
+
+  // Use a minimal fake SmartAssistEngine to test only preprocessing
+  const engine = Object.create(SmartAssistEngine);
+
+  const cases = [
+    { input: "hello?", expected: "hello" },
+    { input: "?prompt", expected: "prompt" },
+    { input: "multiple???", expected: "multiple" },
+    { input: "mid?dle", expected: "middle" },
+    { input: "question? ", expected: "question" },
+    { input: " no?  spaces? ", expected: "no  spaces" },
+    { input: "???", expected: "" },
+    { input: "clean input", expected: "clean input" },
+  ];
+
+  for (const { input, expected } of cases) {
+    const result = engine._preprocessQuery(input);
+    Assert.equal(
+      result,
+      expected,
+      `Expected "${input}" to preprocess to "${expected}", got "${result}"`
+    );
+  }
+
+  sb.restore();
 });