File: ai_create_on_device_session_task.h

package info (click to toggle)
chromium 139.0.7258.127-1
  • links: PTS, VCS
  • area: main
  • in suites:
  • size: 6,122,068 kB
  • sloc: cpp: 35,100,771; ansic: 7,163,530; javascript: 4,103,002; python: 1,436,920; asm: 946,517; xml: 746,709; pascal: 187,653; perl: 88,691; sh: 88,436; objc: 79,953; sql: 51,488; cs: 44,583; fortran: 24,137; makefile: 22,147; tcl: 15,277; php: 13,980; yacc: 8,984; ruby: 7,485; awk: 3,720; lisp: 3,096; lex: 1,327; ada: 727; jsp: 228; sed: 36
file content (101 lines) | stat: -rw-r--r-- 3,829 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
// 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 CHROME_BROWSER_AI_AI_CREATE_ON_DEVICE_SESSION_TASK_H_
#define CHROME_BROWSER_AI_AI_CREATE_ON_DEVICE_SESSION_TASK_H_

#include "base/memory/raw_ptr.h"
#include "base/state_transitions.h"
#include "chrome/browser/ai/ai_context_bound_object.h"
#include "chrome/browser/optimization_guide/optimization_guide_keyed_service.h"
#include "components/optimization_guide/core/optimization_guide_model_executor.h"
#include "services/on_device_model/public/cpp/capabilities.h"

// A base class for tasks which create an on-device session.
class CreateOnDeviceSessionTask
    : public AIContextBoundObject,
      public optimization_guide::OnDeviceModelAvailabilityObserver {
 public:
  CreateOnDeviceSessionTask(
      AIContextBoundObjectSet& context_bound_object_set,
      content::BrowserContext* browser_context,
      optimization_guide::ModelBasedCapabilityKey feature);
  ~CreateOnDeviceSessionTask() override;
  CreateOnDeviceSessionTask(const CreateOnDeviceSessionTask&) = delete;
  CreateOnDeviceSessionTask& operator=(const CreateOnDeviceSessionTask&) =
      delete;

  bool IsPending() const { return state_ == State::kPending; }

  // Starts the process of creating an on-device model session.
  // It may succeed or fail immediately, or it may move into the `kPending`
  // state if it needs to wait for the on-device model availability changes.
  // See `kWaitableReasons` for more details.
  void Start();

 protected:
  // Cancels the creation task, and deletes itself.
  void Cancel();

  virtual void OnFinish(
      std::unique_ptr<
          optimization_guide::OptimizationGuideModelExecutor::Session>
          session) = 0;

  virtual void UpdateSessionConfigParams(
      optimization_guide::SessionConfigParams* config_params) {}

 private:
  // The state of `CreateOnDeviceSessionTask`.
  // The possible transitions of state are:
  // - kNotStarted -> kFinished
  // - kNotStarted -> kPending
  // - kPending -> kFinished
  // - kPending -> kCancelled
  enum class State {
    // The task is created but not started yet.
    kNotStarted,
    // The task has started, but the on-device model is not readily available
    // and is expected to be ready soon.
    // When the task is in `kPending` state, it should be kept alive by the
    // creator as it needs to keep observing the on-device model availability
    // changes.
    // See `kWaitableReasons` for more details.
    kPending,
    // The task is finished, but it's not guaranteed that the session has been
    // created successfully.
    kFinished,
    // The task is cancelled before finishing.
    kCancelled,
  };

  friend base::StateTransitions<State>;
  friend std::ostream& operator<<(std::ostream& os, State state);

  void SetState(State state);

  // optimization_guide::OnDeviceModelAvailabilityObserver
  void OnDeviceModelAvailabilityChanged(
      optimization_guide::ModelBasedCapabilityKey feature,
      optimization_guide::OnDeviceModelEligibilityReason reason) override;

  void OnGetEligibility(
      optimization_guide::OnDeviceModelEligibilityReason eligibility);

  std::unique_ptr<optimization_guide::OptimizationGuideModelExecutor::Session>
  StartSession();

  void Finish(
      std::unique_ptr<
          optimization_guide::OptimizationGuideModelExecutor::Session> session);

  OptimizationGuideKeyedService* GetOptimizationGuideService();

  const raw_ptr<content::BrowserContext> browser_context_;
  const optimization_guide::ModelBasedCapabilityKey feature_;
  State state_ = CreateOnDeviceSessionTask::State::kNotStarted;
  base::WeakPtrFactory<CreateOnDeviceSessionTask> weak_factory_{this};
};

#endif  // CHROME_BROWSER_AI_AI_CREATE_ON_DEVICE_SESSION_TASK_H_