Quick Answers text annotator integration.

This CL includes:
- Generate user intent based on text annotation result. Fallbacks to
generate translation intent based on language detection result if no
other intent is generated.
- Only send query to server if there is a supported intent generated.
- Rewrite query based on generated intent.
- Notify client with preprocess output. The output will be used to
conditionally trigger Quick Answers view in a follow up CL.

Bug: b/150249535
Test: unit tests
Change-Id: I702f1f65dd73bb8ff5d36fafb65e8409230d3028
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2198115Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Commit-Queue: Li Lin <llin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#769062}

chromeos/components/quick_answers/BUILD.gn

@@ -39,6 +39,8 @@ source_set("quick_answers") {
     "//base",
     "//chromeos/components/quick_answers/public/cpp:prefs",
     "//chromeos/constants:constants",
+    "//chromeos/services/machine_learning/public/cpp",
+    "//chromeos/services/machine_learning/public/mojom",
     "//components/prefs:prefs",
     "//net:net",
     "//services/data_decoder/public/cpp",
@@ -74,6 +76,8 @@ source_set("unit_tests") {
     "//base/test:test_support",
     "//chromeos/components/quick_answers/public/cpp:prefs",
     "//chromeos/constants:constants",
+    "//chromeos/services/machine_learning/public/cpp:test_support",
+    "//chromeos/services/machine_learning/public/mojom",
     "//components/prefs:prefs",
     "//components/prefs:test_support",
     "//services/data_decoder/public/cpp:test_support",

chromeos/components/quick_answers/quick_answers_client.cc

@@ -8,7 +8,6 @@
 
 #include "base/strings/stringprintf.h"
 #include "chromeos/components/quick_answers/quick_answers_model.h"
-#include "chromeos/components/quick_answers/understanding/intent_generator.h"
 #include "chromeos/components/quick_answers/utils/quick_answers_metrics.h"
 #include "chromeos/constants/chromeos_features.h"
 #include "third_party/icu/source/common/unicode/locid.h"
@@ -20,30 +19,38 @@ namespace {
 using network::mojom::URLLoaderFactory;
 
 constexpr char kDictionaryQueryRewriteTemplate[] = "Define:%s";
+constexpr char kTranslationQueryRewriteTemplate[] = "Translate:%s";
 
 QuickAnswersClient::ResultLoaderFactoryCallback*
     g_testing_result_factory_callback = nullptr;
 
-const QuickAnswersRequest PreprocessRequest(const QuickAnswersRequest& request,
-                                            const std::string& intent_text,
-                                            IntentType intent_type) {
-  QuickAnswersRequest processed_request = request;
+QuickAnswersClient::IntentGeneratorFactoryCallback*
+    g_testing_intent_generator_factory_callback = nullptr;
+
+const PreprocessedOutput PreprocessRequest(const QuickAnswersRequest& request,
+                                           const std::string& intent_text,
+                                           IntentType intent_type) {
+  PreprocessedOutput processed_output;
+  processed_output.intent_text = intent_text;
+  processed_output.query = intent_text;
+  processed_output.intent_type = intent_type;
 
   switch (intent_type) {
     case IntentType::kUnit:
-      processed_request.selected_text = intent_text;
       break;
     case IntentType::kDictionary:
-      processed_request.selected_text = base::StringPrintf(
+      processed_output.query = base::StringPrintf(
           kDictionaryQueryRewriteTemplate, intent_text.c_str());
       break;
     case IntentType::kTranslation:
+      processed_output.query = base::StringPrintf(
+          kTranslationQueryRewriteTemplate, intent_text.c_str());
       break;
     case IntentType::kUnknown:
       // TODO(llin): Update to NOTREACHED after integrating with TCLib.
       break;
   }
-  return processed_request;
+  return processed_output;
 }
 
 }  // namespace
@@ -54,6 +61,11 @@ void QuickAnswersClient::SetResultLoaderFactoryForTesting(
   g_testing_result_factory_callback = factory;
 }
 
+void QuickAnswersClient::SetIntentGeneratorFactoryForTesting(
+    IntentGeneratorFactoryCallback* factory) {
+  g_testing_intent_generator_factory_callback = factory;
+}
+
 QuickAnswersClient::QuickAnswersClient(URLLoaderFactory* url_loader_factory,
                                        ash::AssistantState* assistant_state,
                                        QuickAnswersDelegate* delegate)
@@ -115,10 +127,8 @@ void QuickAnswersClient::SendRequest(
   RecordSelectedTextLength(quick_answers_request.selected_text.length());
 
   // Generate intent from |quick_answers_request|.
-  IntentGenerator(base::BindOnce(&QuickAnswersClient::IntentGeneratorCallback,
-                                 weak_factory_.GetWeakPtr(),
-                                 quick_answers_request))
-      .GenerateIntent(quick_answers_request);
+  intent_generator_ = CreateIntentGenerator(quick_answers_request);
+  intent_generator_->GenerateIntent(quick_answers_request);
 }
 
 void QuickAnswersClient::OnQuickAnswerClick(ResultType result_type) {
@@ -154,6 +164,15 @@ std::unique_ptr<ResultLoader> QuickAnswersClient::CreateResultLoader(
   return ResultLoader::Create(intent_type, url_loader_factory_, this);
 }
 
+std::unique_ptr<IntentGenerator> QuickAnswersClient::CreateIntentGenerator(
+    const QuickAnswersRequest& request) {
+  if (g_testing_intent_generator_factory_callback)
+    return g_testing_intent_generator_factory_callback->Run();
+  return std::make_unique<IntentGenerator>(
+      base::BindOnce(&QuickAnswersClient::IntentGeneratorCallback,
+                     weak_factory_.GetWeakPtr(), request));
+}
+
 void QuickAnswersClient::OnNetworkError() {
   DCHECK(delegate_);
   delegate_->OnNetworkError();
@@ -170,17 +189,25 @@ void QuickAnswersClient::IntentGeneratorCallback(
     const QuickAnswersRequest& quick_answers_request,
     const std::string& intent_text,
     IntentType intent_type) {
+  DCHECK(delegate_);
+
   // Preprocess the request.
-  auto& processed_request =
+  QuickAnswersRequest processed_request = quick_answers_request;
+  processed_request.preprocessed_output =
       PreprocessRequest(quick_answers_request, intent_text, intent_type);
 
   delegate_->OnRequestPreprocessFinish(processed_request);
 
-  // TODO(llin): Only fetch answer if there is a intent generated after
-  // integrating with TCLib.
+  if (features::IsQuickAnswersTextAnnotatorEnabled() &&
+      processed_request.preprocessed_output.intent_type ==
+          IntentType::kUnknown) {
+    // Don't fetch answer if no intent is generated.
+    return;
+  }
+
   result_loader_ = CreateResultLoader(intent_type);
   // Load and parse search result.
-  result_loader_->Fetch(processed_request.selected_text);
+  result_loader_->Fetch(processed_request.preprocessed_output.query);
 }
 
 base::TimeDelta QuickAnswersClient::GetImpressionDuration() const {

chromeos/components/quick_answers/quick_answers_client.h

@@ -10,6 +10,7 @@
 
 #include "ash/public/cpp/assistant/assistant_state.h"
 #include "chromeos/components/quick_answers/result_loader.h"
+#include "chromeos/components/quick_answers/understanding/intent_generator.h"
 
 namespace network {
 namespace mojom {
@@ -60,6 +61,11 @@ class QuickAnswersClient : public ash::AssistantStateObserver,
   using ResultLoaderFactoryCallback =
       base::RepeatingCallback<std::unique_ptr<ResultLoader>()>;
 
+  // Method that can be used in tests to change the intent generator returned by
+  // |CreateResultLoader| in tests.
+  using IntentGeneratorFactoryCallback =
+      base::RepeatingCallback<std::unique_ptr<IntentGenerator>()>;
+
   QuickAnswersClient(network::mojom::URLLoaderFactory* url_loader_factory,
                      ash::AssistantState* assistant_state,
                      QuickAnswersDelegate* delegate);
@@ -95,10 +101,23 @@ class QuickAnswersClient : public ash::AssistantStateObserver,
   static void SetResultLoaderFactoryForTesting(
       ResultLoaderFactoryCallback* factory);
 
+  static void SetIntentGeneratorFactoryForTesting(
+      IntentGeneratorFactoryCallback* factory);
+
  private:
+  FRIEND_TEST_ALL_PREFIXES(QuickAnswersClientTest, SendRequest);
+  FRIEND_TEST_ALL_PREFIXES(QuickAnswersClientTest,
+                           NotSendRequestForUnknownIntent);
+  FRIEND_TEST_ALL_PREFIXES(QuickAnswersClientTest, PreprocessDefinitionIntent);
+  FRIEND_TEST_ALL_PREFIXES(QuickAnswersClientTest, PreprocessTranslationIntent);
+
   // Creates a |ResultLoader| instance.
   std::unique_ptr<ResultLoader> CreateResultLoader(IntentType intent_type);
 
+  // Creates an |IntentGenerator| instance.
+  std::unique_ptr<IntentGenerator> CreateIntentGenerator(
+      const QuickAnswersRequest& request);
+
   void NotifyEligibilityChanged();
   void IntentGeneratorCallback(const QuickAnswersRequest& quick_answers_request,
                                const std::string& intent_text,
@@ -109,6 +128,7 @@ class QuickAnswersClient : public ash::AssistantStateObserver,
   ash::AssistantState* assistant_state_ = nullptr;
   QuickAnswersDelegate* delegate_ = nullptr;
   std::unique_ptr<ResultLoader> result_loader_;
+  std::unique_ptr<IntentGenerator> intent_generator_;
   bool assistant_enabled_ = false;
   bool assistant_context_enabled_ = false;
   bool quick_answers_settings_enabled_ = false;

chromeos/components/quick_answers/quick_answers_client_unittest.cc

@@ -50,6 +50,28 @@ class MockResultLoader : public TestResultLoader {
   MOCK_METHOD1(Fetch, void(const std::string&));
 };
 
+MATCHER_P(QuickAnswersRequestWithOutputEqual, quick_answers_request, "") {
+  return (arg.selected_text == quick_answers_request.selected_text &&
+          arg.preprocessed_output.intent_type ==
+              quick_answers_request.preprocessed_output.intent_type &&
+          arg.preprocessed_output.intent_text ==
+              quick_answers_request.preprocessed_output.intent_text &&
+          arg.preprocessed_output.query ==
+              quick_answers_request.preprocessed_output.query);
+}
+
+class MockIntentGenerator : public IntentGenerator {
+ public:
+  MockIntentGenerator(IntentGeneratorCallback complete_callback)
+      : IntentGenerator(std::move(complete_callback)) {}
+
+  MockIntentGenerator(const MockIntentGenerator&) = delete;
+  MockIntentGenerator& operator=(const MockIntentGenerator&) = delete;
+
+  // IntentGenerator:
+  MOCK_METHOD1(GenerateIntent, void(const QuickAnswersRequest&));
+};
+
 }  // namespace
 
 class QuickAnswersClientTest : public testing::Test {
@@ -70,13 +92,23 @@ class QuickAnswersClientTest : public testing::Test {
 
     result_loader_factory_callback_ = base::BindRepeating(
         &QuickAnswersClientTest::CreateResultLoader, base::Unretained(this));
+
+    intent_generator_factory_callback_ = base::BindRepeating(
+        &QuickAnswersClientTest::CreateIntentGenerator, base::Unretained(this));
+
+    mock_intent_generator_ = std::make_unique<MockIntentGenerator>(
+        base::BindOnce(&QuickAnswersClientTest::IntentGeneratorTestCallback,
+                       base::Unretained(this)));
   }
 
   void TearDown() override {
     QuickAnswersClient::SetResultLoaderFactoryForTesting(nullptr);
+    QuickAnswersClient::SetIntentGeneratorFactoryForTesting(nullptr);
     client_.reset();
   }
 
+  void IntentGeneratorTestCallback(const std::string& text, IntentType type) {}
+
  protected:
   void NotifyAssistantStateChange(
       bool setting_enabled,
@@ -95,14 +127,22 @@ class QuickAnswersClientTest : public testing::Test {
     return std::move(mock_result_loader_);
   }
 
+  std::unique_ptr<IntentGenerator> CreateIntentGenerator() {
+    return std::move(mock_intent_generator_);
+  }
+
   std::unique_ptr<QuickAnswersClient> client_;
   std::unique_ptr<MockQuickAnswersDelegate> mock_delegate_;
   std::unique_ptr<MockResultLoader> mock_result_loader_;
   std::unique_ptr<ash::AssistantState> assistant_state_;
+  std::unique_ptr<MockIntentGenerator> mock_intent_generator_;
   base::test::SingleThreadTaskEnvironment task_environment_;
   network::TestURLLoaderFactory test_url_loader_factory_;
   QuickAnswersClient::ResultLoaderFactoryCallback
       result_loader_factory_callback_;
+  QuickAnswersClient::IntentGeneratorFactoryCallback
+      intent_generator_factory_callback_;
+  IntentGenerator::IntentGeneratorCallback intent_generator_callback_;
 };
 
 TEST_F(QuickAnswersClientTest, FeatureEligible) {
@@ -248,15 +288,21 @@ TEST_F(QuickAnswersClientTest, SendRequest) {
       std::make_unique<QuickAnswersRequest>();
   quick_answers_request->selected_text = "sel";
 
+  // Verify that |GenerateIntent| is called.
+  EXPECT_CALL(*mock_intent_generator_,
+              GenerateIntent(QuickAnswersRequestEqual(*quick_answers_request)));
+  QuickAnswersClient::SetIntentGeneratorFactoryForTesting(
+      &intent_generator_factory_callback_);
+
   mock_result_loader_ =
       std::make_unique<MockResultLoader>(&test_url_loader_factory_, nullptr);
-  EXPECT_CALL(*mock_result_loader_, Fetch(::testing::Eq("sel")));
+  EXPECT_CALL(*mock_result_loader_, Fetch(::testing::Eq("Define:sel")));
   QuickAnswersClient::SetResultLoaderFactoryForTesting(
       &result_loader_factory_callback_);
-  EXPECT_CALL(*mock_delegate_,
-              OnRequestPreprocessFinish(
-                  QuickAnswersRequestEqual(*quick_answers_request)));
+
   client_->SendRequest(*quick_answers_request);
+  client_->IntentGeneratorCallback(*quick_answers_request, "sel",
+                                   IntentType::kDictionary);
 
   std::unique_ptr<QuickAnswer> quick_answer = std::make_unique<QuickAnswer>();
   quick_answer->primary_answer = "answer";
@@ -265,5 +311,68 @@ TEST_F(QuickAnswersClientTest, SendRequest) {
   client_->OnQuickAnswerReceived(std::move(quick_answer));
 }
 
+TEST_F(QuickAnswersClientTest, NotSendRequestForUnknownIntent) {
+  base::test::ScopedFeatureList scoped_feature_list;
+  scoped_feature_list.InitAndEnableFeature(
+      chromeos::features::kQuickAnswersTextAnnotator);
+
+  std::unique_ptr<QuickAnswersRequest> quick_answers_request =
+      std::make_unique<QuickAnswersRequest>();
+  quick_answers_request->selected_text = "sel";
+
+  mock_result_loader_ =
+      std::make_unique<MockResultLoader>(&test_url_loader_factory_, nullptr);
+  EXPECT_CALL(*mock_result_loader_, Fetch(::testing::_)).Times(0);
+  QuickAnswersClient::SetResultLoaderFactoryForTesting(
+      &result_loader_factory_callback_);
+
+  client_->IntentGeneratorCallback(*quick_answers_request, "sel",
+                                   IntentType::kUnknown);
+}
+
+TEST_F(QuickAnswersClientTest, PreprocessDefinitionIntent) {
+  std::unique_ptr<QuickAnswersRequest> quick_answers_request =
+      std::make_unique<QuickAnswersRequest>();
+  quick_answers_request->selected_text = "unfathomable";
+
+  // Verify that |OnRequestPreprocessFinish| is called.
+  std::unique_ptr<QuickAnswersRequest> processed_request =
+      std::make_unique<QuickAnswersRequest>();
+  processed_request->selected_text = "unfathomable";
+  PreprocessedOutput expected_processed_output;
+  expected_processed_output.intent_text = "unfathomable";
+  expected_processed_output.query = "Define:unfathomable";
+  expected_processed_output.intent_type = IntentType::kDictionary;
+  processed_request->preprocessed_output = expected_processed_output;
+  EXPECT_CALL(*mock_delegate_,
+              OnRequestPreprocessFinish(
+                  QuickAnswersRequestWithOutputEqual(*processed_request)));
+
+  client_->IntentGeneratorCallback(*quick_answers_request, "unfathomable",
+                                   IntentType::kDictionary);
+}
+
+TEST_F(QuickAnswersClientTest, PreprocessTranslationIntent) {
+  std::unique_ptr<QuickAnswersRequest> quick_answers_request =
+      std::make_unique<QuickAnswersRequest>();
+  quick_answers_request->selected_text = "sel";
+
+  // Verify that |OnRequestPreprocessFinish| is called.
+  std::unique_ptr<QuickAnswersRequest> processed_request =
+      std::make_unique<QuickAnswersRequest>();
+  processed_request->selected_text = "sel";
+  PreprocessedOutput expected_processed_output;
+  expected_processed_output.intent_text = "intent text";
+  expected_processed_output.query = "Translate:intent text";
+  expected_processed_output.intent_type = IntentType::kTranslation;
+  processed_request->preprocessed_output = expected_processed_output;
+  EXPECT_CALL(*mock_delegate_,
+              OnRequestPreprocessFinish(
+                  QuickAnswersRequestWithOutputEqual(*processed_request)));
+
+  client_->IntentGeneratorCallback(*quick_answers_request, "intent text",
+                                   IntentType::kTranslation);
+}
+
 }  // namespace quick_answers
 }  // namespace chromeos

chromeos/components/quick_answers/quick_answers_model.cc

@@ -8,5 +8,10 @@ namespace chromeos {
 namespace quick_answers {
 QuickAnswer::QuickAnswer() = default;
 QuickAnswer::~QuickAnswer() = default;
+
+PreprocessedOutput::PreprocessedOutput() = default;
+PreprocessedOutput::PreprocessedOutput(const PreprocessedOutput& other) =
+    default;
+PreprocessedOutput::~PreprocessedOutput() = default;
 }  // namespace quick_answers
 }  // namespace chromeos

chromeos/components/quick_answers/quick_answers_model.h

@@ -122,6 +122,22 @@ struct DeviceProperties {
   std::string language;
 };
 
+// Extract information generated from |QuickAnswersRequest|.
+struct PreprocessedOutput {
+  PreprocessedOutput();
+  PreprocessedOutput(const PreprocessedOutput& other);
+  ~PreprocessedOutput();
+
+  // Predicted intent.
+  IntentType intent_type = IntentType::kUnknown;
+
+  // The text extracted from the selected_text associated with the intent.
+  std::string intent_text;
+
+  // Rewritten query based on |intent_type| and |intent_text|.
+  std::string query;
+};
+
 // Structure to describe an quick answer request including selected content and
 // context.
 struct QuickAnswersRequest {
@@ -131,6 +147,9 @@ struct QuickAnswersRequest {
   // Device specific properties.
   DeviceProperties device_properties;
 
+  // Output of processed result.
+  PreprocessedOutput preprocessed_output;
+
   // TODO(llin): Add context and other targeted objects (e.g: images, links,
   // etc).
 };

chromeos/components/quick_answers/understanding/intent_generator.cc

@@ -6,10 +6,76 @@
 
 #include <map>
 
+#include "base/no_destructor.h"
 #include "chromeos/components/quick_answers/quick_answers_model.h"
+#include "chromeos/components/quick_answers/utils/language_detector.h"
+#include "chromeos/constants/chromeos_features.h"
+#include "chromeos/services/machine_learning/public/cpp/service_connection.h"
+#include "chromeos/services/machine_learning/public/mojom/machine_learning_service.mojom.h"
 
 namespace chromeos {
 namespace quick_answers {
+namespace {
+
+using chromeos::machine_learning::mojom::LoadModelResult;
+using machine_learning::mojom::TextAnnotationPtr;
+using machine_learning::mojom::TextAnnotationRequestPtr;
+using machine_learning::mojom::TextClassifier;
+
+// TODO(llin): Finalize on the threshold based on user feedback.
+constexpr int kUnitConversionIntentAndSelectionLengthDiffThreshold = 5;
+constexpr int kTranslationTextLengthThreshold = 50;
+constexpr int kDefinitionIntentAndSelectionLengthDiffThreshold = 0;
+
+const std::map<std::string, IntentType>& GetIntentTypeMap() {
+  static base::NoDestructor<std::map<std::string, IntentType>> kIntentTypeMap(
+      {{"unit", IntentType::kUnit}, {"dictionary", IntentType::kDictionary}});
+  return *kIntentTypeMap;
+}
+
+bool ExtractEntity(const std::string& selected_text,
+                   const std::vector<TextAnnotationPtr>& annotations,
+                   std::string* entity_str,
+                   std::string* type) {
+  for (auto& annotation : annotations) {
+    *entity_str =
+        selected_text.substr(annotation->start_offset,
+                             annotation->end_offset - annotation->start_offset);
+
+    // Use the first entity type.
+    auto intent_type_map = GetIntentTypeMap();
+    for (const auto& entity : annotation->entities) {
+      if (intent_type_map.find(entity->name) != intent_type_map.end()) {
+        *type = entity->name;
+        return true;
+      }
+    }
+  }
+
+  return false;
+}
+
+IntentType RewriteIntent(const std::string& selected_text,
+                         const std::string& entity_str,
+                         const IntentType intent) {
+  int intent_and_selection_length_diff =
+      selected_text.length() - entity_str.length();
+
+  if ((intent == IntentType::kUnit &&
+       intent_and_selection_length_diff >
+           kUnitConversionIntentAndSelectionLengthDiffThreshold) ||
+      (intent == IntentType::kDictionary &&
+       intent_and_selection_length_diff >
+           kDefinitionIntentAndSelectionLengthDiffThreshold)) {
+    // Override intent type to |kUnknown| if length diff between intent
+    // text and selection text is above the threshold.
+    return IntentType::kUnknown;
+  }
+
+  return intent;
+}
+
+}  // namespace
 
 IntentGenerator::IntentGenerator(IntentGeneratorCallback complete_callback)
     : complete_callback_(std::move(complete_callback)) {
@@ -22,15 +88,75 @@ IntentGenerator::~IntentGenerator() {
 }
 
 void IntentGenerator::GenerateIntent(const QuickAnswersRequest& request) {
-  // TODO(llin): Generate intent based on text annotation result first and
-  // fallback to language detection for generating translation intent. Text
-  // annotation will be a async call.
+  if (!features::IsQuickAnswersTextAnnotatorEnabled()) {
+    std::move(complete_callback_)
+        .Run(request.selected_text, IntentType::kUnknown);
+    return;
+  }
 
-  // Check if it is |kTranslation| intent.
+  // Load text classifier.
+  chromeos::machine_learning::ServiceConnection::GetInstance()
+      ->LoadTextClassifier(text_classifier_.BindNewPipeAndPassReceiver(),
+                           base::BindOnce(&IntentGenerator::LoadModelCallback,
+                                          weak_factory_.GetWeakPtr(), request));
+}
 
-  // Don't do language detection if no device language is provided. Returns
-  // unknown intent type.
-  if (request.device_properties.language.empty()) {
+void IntentGenerator::LoadModelCallback(const QuickAnswersRequest& request,
+                                        LoadModelResult result) {
+  if (result != LoadModelResult::OK) {
+    LOG(ERROR) << "Failed to load TextClassifier.";
+    std::move(complete_callback_)
+        .Run(request.selected_text, IntentType::kUnknown);
+    return;
+  }
+
+  if (text_classifier_) {
+    TextAnnotationRequestPtr text_annotation_request =
+        machine_learning::mojom::TextAnnotationRequest::New();
+    text_annotation_request->text = request.selected_text;
+    text_annotation_request->default_locales =
+        request.device_properties.language;
+
+    text_classifier_->Annotate(
+        std::move(text_annotation_request),
+        base::BindOnce(&IntentGenerator::AnnotationCallback,
+                       weak_factory_.GetWeakPtr(), request));
+  }
+}
+
+void IntentGenerator::AnnotationCallback(
+    const QuickAnswersRequest& request,
+    std::vector<TextAnnotationPtr> annotations) {
+  std::string entity_str;
+  std::string type;
+
+  if (ExtractEntity(request.selected_text, annotations, &entity_str, &type)) {
+    auto intent_type_map = GetIntentTypeMap();
+    auto it = intent_type_map.find(type);
+    if (it != intent_type_map.end()) {
+      std::move(complete_callback_)
+          .Run(entity_str,
+               RewriteIntent(request.selected_text, entity_str, it->second));
+      return;
+    }
+  }
+  // Fallback to language detection for generating translation intent.
+  MaybeGenerateTranslationIntent(request);
+}
+
+void IntentGenerator::SetLanguageDetectorForTesting(
+    std::unique_ptr<LanguageDetector> language_detector) {
+  language_detector_ = std::move(language_detector);
+}
+
+void IntentGenerator::MaybeGenerateTranslationIntent(
+    const QuickAnswersRequest& request) {
+  DCHECK(complete_callback_);
+
+  // Don't do language detection if no device language is provided or the length
+  // of selected text is above the threshold. Returns unknown intent type.
+  if (request.device_properties.language.empty() ||
+      request.selected_text.length() > kTranslationTextLengthThreshold) {
     std::move(complete_callback_)
         .Run(request.selected_text, IntentType::kUnknown);
     return;
@@ -41,20 +167,12 @@ void IntentGenerator::GenerateIntent(const QuickAnswersRequest& request) {
   auto detected_language =
       language_detector_->DetectLanguage(request.selected_text);
 
+  auto intent_type = IntentType::kUnknown;
   if (!detected_language.empty() &&
       detected_language != request.device_properties.language) {
-    // Detected language is different from device language.
-    std::move(complete_callback_)
-        .Run(request.selected_text, IntentType::kTranslation);
-  } else {
-    std::move(complete_callback_)
-        .Run(request.selected_text, IntentType::kUnknown);
+    intent_type = IntentType::kTranslation;
   }
-}
-
-void IntentGenerator::SetLanguageDetectorForTesting(
-    std::unique_ptr<LanguageDetector> language_detector) {
-  language_detector_ = std::move(language_detector);
+  std::move(complete_callback_).Run(request.selected_text, intent_type);
 }
 
 }  // namespace quick_answers

chromeos/components/quick_answers/understanding/intent_generator.h

@@ -10,6 +10,9 @@
 
 #include "base/callback.h"
 #include "chromeos/components/quick_answers/utils/language_detector.h"
+#include "chromeos/services/machine_learning/public/mojom/machine_learning_service.mojom.h"
+#include "chromeos/services/machine_learning/public/mojom/text_classifier.mojom.h"
+#include "mojo/public/cpp/bindings/remote.h"
 
 namespace chromeos {
 namespace quick_answers {
@@ -29,17 +32,36 @@ class IntentGenerator {
   IntentGenerator(const IntentGenerator&) = delete;
   IntentGenerator& operator=(const IntentGenerator&) = delete;
 
-  ~IntentGenerator();
+  virtual ~IntentGenerator();
 
-  // Generate intent from the |request|.
-  void GenerateIntent(const QuickAnswersRequest& request);
+  // Generate intent from the |request|. Virtual for testing.
+  virtual void GenerateIntent(const QuickAnswersRequest& request);
 
   void SetLanguageDetectorForTesting(
       std::unique_ptr<LanguageDetector> language_detector);
 
  private:
+  FRIEND_TEST_ALL_PREFIXES(IntentGeneratorTest,
+                           TextAnnotationIntentNoAnnotation);
+  FRIEND_TEST_ALL_PREFIXES(IntentGeneratorTest, TextAnnotationIntentNoEntity);
+  FRIEND_TEST_ALL_PREFIXES(IntentGeneratorTest,
+                           TextAnnotationIntentUnSupportedEntity);
+
+  void LoadModelCallback(
+      const QuickAnswersRequest& request,
+      chromeos::machine_learning::mojom::LoadModelResult result);
+  void AnnotationCallback(
+      const QuickAnswersRequest& request,
+      std::vector<machine_learning::mojom::TextAnnotationPtr> annotations);
+
+  void MaybeGenerateTranslationIntent(const QuickAnswersRequest& request);
+
   IntentGeneratorCallback complete_callback_;
   std::unique_ptr<LanguageDetector> language_detector_;
+  mojo::Remote<::chromeos::machine_learning::mojom::TextClassifier>
+      text_classifier_;
+
+  base::WeakPtrFactory<IntentGenerator> weak_factory_{this};
 };
 
 }  // namespace quick_answers