ai_manager_keyed_service.cc | Explore in Territory

// 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.

#include "chrome/browser/ai/ai_manager_keyed_service.h"

#include <memory>

#include "base/containers/fixed_flat_set.h"
#include "base/containers/flat_set.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/functional/bind.h"
#include "base/functional/callback_forward.h"
#include "base/memory/weak_ptr.h"
#include "base/notreached.h"
#include "base/strings/stringprintf.h"
#include "base/task/task_traits.h"
#include "base/task/thread_pool.h"
#include "base/types/pass_key.h"
#include "chrome/browser/ai/ai_context_bound_object.h"
#include "chrome/browser/ai/ai_context_bound_object_set.h"
#include "chrome/browser/ai/ai_rewriter.h"
#include "chrome/browser/ai/ai_summarizer.h"
#include "chrome/browser/ai/ai_text_session.h"
#include "chrome/browser/ai/ai_writer.h"
#include "chrome/browser/optimization_guide/optimization_guide_keyed_service.h"
#include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
#include "chrome/browser/profiles/profile.h"
#include "components/optimization_guide/core/model_execution/feature_keys.h"
#include "components/optimization_guide/core/model_util.h"
#include "components/optimization_guide/core/optimization_guide_enums.h"
#include "components/optimization_guide/core/optimization_guide_features.h"
#include "components/optimization_guide/core/optimization_guide_model_executor.h"
#include "components/optimization_guide/core/optimization_guide_switches.h"
#include "mojo/public/cpp/bindings/self_owned_receiver.h"
#include "third_party/blink/public/mojom/ai/ai_manager.mojom.h"
#include "third_party/blink/public/mojom/ai/ai_text_session_info.mojom.h"
#include "third_party/blink/public/mojom/ai/model_streaming_responder.mojom.h"
#include "third_party/blink/public/mojom/devtools/console_message.mojom-shared.h"

namespace {

// Checks if the model path configured via command line is valid.
bool IsModelPathValid(const std::string& model_path_str) { … }

blink::mojom::ModelAvailabilityCheckResult
ConvertOnDeviceModelEligibilityReasonToModelAvailabilityCheckResult(
    optimization_guide::OnDeviceModelEligibilityReason
        on_device_model_eligibility_reason) { … }

// Currently, the following errors, which are used when a model may have been
// installed but not yet loaded, are treated as waitable.
// TODO(crbug.com/361537114): Consider making the kModelToBeInstalled error
// waitable as well.
static constexpr auto kWaitableReasons = …;

// A base class for tasks which create an on-device session. See the method
// comment of `Run()` for the details.
class CreateOnDeviceSessionTask
    : public AIContextBoundObject,
      public optimization_guide::OnDeviceModelAvailabilityObserver { … };

template <typename ContextBoundObjectType,
          typename ContextBoundObjectReceiverInterface,
          typename ClientRemoteInterface>
class CreateContextBoundObjectTask : public CreateOnDeviceSessionTask { … };

}  // namespace

AIManagerKeyedService::AIManagerKeyedService(
    content::BrowserContext* browser_context)
    : … { … }

AIManagerKeyedService::~AIManagerKeyedService() { … }

void AIManagerKeyedService::AddReceiver(
    mojo::PendingReceiver<blink::mojom::AIManager> receiver,
    AIContextBoundObjectSet::ReceiverContext context) { … }

void AIManagerKeyedService::CanCreateTextSession(
    CanCreateTextSessionCallback callback) { … }

std::unique_ptr<AITextSession> AIManagerKeyedService::CreateTextSessionInternal(
    mojo::PendingReceiver<blink::mojom::AITextSession> receiver,
    const blink::mojom::AITextSessionSamplingParamsPtr& sampling_params,
    AIContextBoundObjectSet* context_bound_object_set,
    const std::optional<const AITextSession::Context>& context) { … }

void AIManagerKeyedService::CreateTextSession(
    mojo::PendingReceiver<blink::mojom::AITextSession> receiver,
    blink::mojom::AITextSessionSamplingParamsPtr sampling_params,
    const std::optional<std::string>& system_prompt,
    CreateTextSessionCallback callback) { … }

void AIManagerKeyedService::CanCreateSummarizer(
    CanCreateSummarizerCallback callback) { … }

void AIManagerKeyedService::CreateSummarizer(
    mojo::PendingRemote<blink::mojom::AIManagerCreateSummarizerClient> client,
    blink::mojom::AISummarizerOptionsPtr options,
    const std::optional<std::string>& shared_context) { … }

void AIManagerKeyedService::GetTextModelInfo(
    GetTextModelInfoCallback callback) { … }

void AIManagerKeyedService::CreateWriter(
    const std::optional<std::string>& shared_context,
    mojo::PendingRemote<blink::mojom::AIManagerCreateWriterClient> client) { … }

void AIManagerKeyedService::CreateRewriter(
    const std::optional<std::string>& shared_context,
    blink::mojom::AIRewriterTone tone,
    blink::mojom::AIRewriterLength length,
    mojo::PendingRemote<blink::mojom::AIManagerCreateRewriterClient> client) { … }

void AIManagerKeyedService::CheckModelPathOverrideCanCreateSession(
    const std::string& model_path,
    optimization_guide::ModelBasedCapabilityKey capability_) { … }

void AIManagerKeyedService::
    CanOptimizationGuideKeyedServiceCreateGenericSession(
        optimization_guide::ModelBasedCapabilityKey capability,
        CanCreateTextSessionCallback callback) { … }

void AIManagerKeyedService::CreateTextSessionForCloning(
    base::PassKey<AITextSession> pass_key,
    mojo::PendingReceiver<blink::mojom::AITextSession> receiver,
    blink::mojom::AITextSessionSamplingParamsPtr sampling_params,
    AIContextBoundObjectSet* context_bound_object_set,
    const AITextSession::Context& context,
    CreateTextSessionCallback callback) { … }

void AIManagerKeyedService::OnModelPathValidationComplete(
    const std::string& model_path,
    bool is_valid_path) { … }
chromium/chrome/browser/ai/ai_manager_keyed_service.cc
