Bug 1932671 - MLSuggest handle city/state in subjects correctly r=adw,urlbar-reviewers

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

browser/components/urlbar/private/MLSuggest.sys.mjs

@@ -310,35 +310,27 @@ class _MLSuggest {
     if (!location || (!location.city && !location.state)) {
       return query;
     }
-    // Remove the city and state from the query
-    let subjectWithoutLocation = query;
-    if (location.city) {
-      subjectWithoutLocation = subjectWithoutLocation
-        .replace(location.city, "")
-        .trim();
-    }
-    if (location.state) {
-      subjectWithoutLocation = subjectWithoutLocation
-        .replace(location.state, "")
-        .trim();
-    }
-    // Remove leftover commas, trailing whitespace, and unnecessary punctuation
-    subjectWithoutLocation = subjectWithoutLocation
-      .replaceAll(",", "")
-      .replace(/\s+/g, " ")
-      .trim();
 
-    return this.#cleanSubject(subjectWithoutLocation);
+    // Remove the city and state values from the query
+    let locValues = Object.values(location).filter(v => !!v);
+    let words = query
+      .trim()
+      .split(/\s+|,/)
+      .filter(w => !!w && !locValues.includes(w));
+
+    let subjectWords = this.#cleanSubject(words);
+    return subjectWords.join(" ");
   }
 
-  #cleanSubject(subject) {
-    let end = PREPOSITIONS.find(
-      p => subject === p || subject.endsWith(" " + p)
-    );
-    if (end) {
-      subject = subject.substring(0, subject.length - end.length).trimEnd();
+  #cleanSubject(words) {
+    // Remove trailing prepositions from the list of words
+    while (
+      words.length > 0 &&
+      PREPOSITIONS.includes(words[words.length - 1])
+    ) {
+      words.pop();
     }
-    return subject;
+    return words;
   }
 }
 

browser/components/urlbar/tests/browser/browser_MLSuggest_integration.js

@@ -136,6 +136,34 @@ let nerResultsMap = {
       word: "ca",
     },
   ],
+  "ramen ra": [
+    {
+      entity: "B-CITY",
+      score: 0.6767462491989136,
+      index: 3,
+      word: "ra",
+    },
+  ],
+  "plumbers in seattle,wa": [
+    {
+      entity: "B-CITYSTATE",
+      index: 4,
+      score: 0.99997478723526,
+      word: "seattle",
+    },
+    {
+      entity: "I-CITYSTATE",
+      index: 5,
+      score: 0.9999989867210388,
+      word: ","
+    },
+    {
+      entity: "I-CITYSTATE",
+      index: 6,
+      score: 0.9999985098838806,
+      word: "wa",
+    },
+  ],
 };
 
 add_setup(async function () {
@@ -263,6 +291,7 @@ add_task(async function test_MLSuggest() {
   );
 
   await testSuggestion("dumplings in ca", null, "ca", remoteClients);
+  await testSuggestion("plumbers in seattle,wa", "seattle", "wa", remoteClients);
 
   Assert.strictEqual(
     Services.prefs.getFloatPref("browser.urlbar.nerThreshold"),
@@ -331,7 +360,7 @@ add_task(async function test_MLSuggest_restart_after_failure() {
 /**
  * For Mocking MLEngine with low score
  */
-class MLEngineWithLowScore {
+class MLEngineWithLowYelpIntent {
   // prefix with _query avoids lint error
   async run(_query) {
     return [
@@ -348,7 +377,7 @@ add_task(async function test_MLSuggest_low_intent_threshold() {
   sinon.restore();
 
   sinon.stub(MLSuggest, "createEngine").callsFake(() => {
-    return new MLEngineWithLowScore();
+    return new MLEngineWithLowYelpIntent();
   });
   const { cleanup } = await setup();
 
@@ -364,3 +393,50 @@ add_task(async function test_MLSuggest_low_intent_threshold() {
   await cleanup();
   sinon.restore();
 });
+
+
+/**
+ * For Mocking MLEngine with positive yelp itent
+ */
+class MLEngineWithHighYelpIntent {
+  // prefix with _query avoids lint error
+  async run(_query) {
+    return [
+      {
+        label: "yelp_intent",
+        score: 0.9,
+      },
+    ];
+  }
+}
+
+add_task(async function test_MLSuggest_city_dup_in_subject() {
+  // Restore any previous stubs
+  sinon.restore();
+
+  sinon.stub(MLSuggest, "createEngine").callsFake(() => {
+    return new MLEngineWithHighYelpIntent();
+  });
+  sinon.stub(MLSuggest, "_findNER").callsFake(query => {
+    return nerResultsMap[query] || [];
+  });
+  const { cleanup } = await setup();
+
+  await MLSuggest.initialize();
+
+  let suggestion = await MLSuggest.makeSuggestions("ramen ra");
+  Assert.ok(suggestion, "Suggestion should be good");
+  const expected = {
+    intent: "yelp_intent",
+    location: { city: "ra", state: null },
+    subject: "ramen"
+  };
+  Assert.deepEqual(suggestion.intent, expected.intent);
+  Assert.deepEqual(suggestion.location, expected.location);
+  Assert.deepEqual(suggestion.subject, expected.subject);
+
+  await MLSuggest.shutdown();
+  await EngineProcess.destroyMLEngine();
+  await cleanup();
+  sinon.restore();
+});