components/history_embeddings/ml_intent_classifier.h

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef COMPONENTS_HISTORY_EMBEDDINGS_ML_INTENT_CLASSIFIER_H_
#define COMPONENTS_HISTORY_EMBEDDINGS_ML_INTENT_CLASSIFIER_H_

#include <memory>

#include "base/memory/raw_ptr.h"
#include "base/task/sequenced_task_runner.h"
#include "components/history_embeddings/intent_classifier.h"
#include "components/history_embeddings/mock_intent_classifier.h"
#include "components/optimization_guide/core/model_execution/on_device_capability.h"

namespace history_embeddings {

using optimization_guide::OnDeviceCapability;

class MlIntentClassifier : public IntentClassifier {
 public:
  class Execution;
  explicit MlIntentClassifier(OnDeviceCapability* model_executor);
  ~MlIntentClassifier() override;

  int64_t GetModelVersion() override;

  void ComputeQueryIntent(std::string query,
                          ComputeQueryIntentCallback callback) override;

 private:
  // Guaranteed to outlive `this`, since
  // `model_executor_` is owned by OptimizationGuideKeyedServiceFactory,
  // which HistoryEmbeddingsServiceFactory depends on.
  raw_ptr<OnDeviceCapability> model_executor_;
  // The current execution, cancelled if another is started.
  std::unique_ptr<Execution> execution_;
};

}  // namespace history_embeddings

#endif  // COMPONENTS_HISTORY_EMBEDDINGS_ML_INTENT_CLASSIFIER_H_