tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

SuggestBackendMl.sys.mjs (3641B)

      1 /* This Source Code Form is subject to the terms of the Mozilla Public
      2 * License, v. 2.0. If a copy of the MPL was not distributed with this
      3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
      4 
      5 import { SuggestBackend } from "moz-src:///browser/components/urlbar/private/SuggestFeature.sys.mjs";
      6 
      7 const lazy = {};
      8 
      9 ChromeUtils.defineESModuleGetters(lazy, {
     10  MLSuggest: "moz-src:///browser/components/urlbar/private/MLSuggest.sys.mjs",
     11  QuickSuggest: "moz-src:///browser/components/urlbar/QuickSuggest.sys.mjs",
     12  SkippableTimer: "moz-src:///browser/components/urlbar/UrlbarUtils.sys.mjs",
     13  UrlbarPrefs: "moz-src:///browser/components/urlbar/UrlbarPrefs.sys.mjs",
     14 });
     15 
     16 /**
     17 * The Suggest ML backend. Both the ML and Rust backends can be enabled at the
     18 * same time. Features can support both backends and decide which one to use per
     19 * query.
     20 */
     21 export class SuggestBackendMl extends SuggestBackend {
     22  get enablingPreferences() {
     23    return ["quickSuggestMlEnabled", "browser.ml.enable"];
     24  }
     25 
     26  enable(enabled) {
     27    if (enabled) {
     28      this.#init();
     29    } else {
     30      this.#uninit();
     31    }
     32  }
     33 
     34  /**
     35   * Queries `MLSuggest` and returns the matching suggestion.
     36   *
     37   * @param {string} searchString
     38   *   The search string.
     39   * @param {object} options
     40   *   Options object.
     41   * @param {UrlbarQueryContext} options.queryContext
     42   *   The query context.
     43   * @returns {Promise<Array>}
     44   *   An array of matching suggestions. `MLSuggest` returns at most one
     45   *   suggestion.
     46   */
     47  async query(searchString, { queryContext }) {
     48    // `MLSuggest` requires the query to be trimmed and lowercase, which
     49    // the original `searchString` isn't necessarily.
     50    searchString = queryContext.trimmedLowerCaseSearchString;
     51 
     52    this.logger.debug("Handling query", { searchString });
     53 
     54    // Don't waste time calling into `MLSuggest` if no ML intents are enabled.
     55    if (
     56      lazy.QuickSuggest.mlFeatures
     57        .values()
     58        .every(f => !f.isEnabled || !f.isMlIntentEnabled)
     59    ) {
     60      this.logger.debug("No ML intents enabled, ignoring query");
     61      return [];
     62    }
     63 
     64    let suggestion = await lazy.MLSuggest.makeSuggestions(searchString);
     65    this.logger.debug("Got suggestion", suggestion);
     66 
     67    if (suggestion?.intent) {
     68      // `MLSuggest` doesn't have a way to return only enabled intents, so it
     69      // can return disabled ones and even ones we don't recognize. Discard the
     70      // suggestion in those cases.
     71      let feature = lazy.QuickSuggest.getFeatureByMlIntent(suggestion.intent);
     72      if (!feature?.isEnabled || !feature?.isMlIntentEnabled) {
     73        this.logger.debug("No ML feature for suggestion, ignoring query");
     74        return [];
     75      }
     76      suggestion.source = "ml";
     77      suggestion.provider = suggestion.intent;
     78      return [suggestion];
     79    }
     80 
     81    return [];
     82  }
     83 
     84  #init() {
     85    if (this.#initTimer) {
     86      return;
     87    }
     88 
     89    // Like all Suggest features, when this feature is enabled it's typically
     90    // enabled at startup. Initializing `MLSuggest` loads MB's worth of data,
     91    // which may slow down the system, so do it on a timer with a configurable
     92    // timeout.
     93    this.#initTimer = new lazy.SkippableTimer({
     94      name: `${this.name} init timer`,
     95      time: 1000 * lazy.UrlbarPrefs.get("quickSuggestMlInitDelaySeconds"),
     96      logger: this.logger,
     97      callback: async () => {
     98        this.logger.info("Init delay timer fired, initializing MLSuggest");
     99        await lazy.MLSuggest.initialize();
    100      },
    101    });
    102  }
    103 
    104  async #uninit() {
    105    this.#initTimer?.cancel();
    106    this.#initTimer = null;
    107    await lazy.MLSuggest.shutdown();
    108  }
    109 
    110  #initTimer;
    111 }