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

module blink.mojom;

import "mojo/public/mojom/base/values.mojom";
import "services/on_device_model/public/mojom/on_device_model.mojom";
import "skia/public/mojom/bitmap.mojom";
import "third_party/blink/public/mojom/ai/ai_common.mojom";
import "third_party/blink/public/mojom/ai/model_streaming_responder.mojom";

// Most of the mojo definition in this file can be viewed as a representation
// of the corresponding built-in API prompt API. Please note that the prompt
// text string below in various methods are directly passed from the JavaScript
// API, so they could contain arbitrary content provided from a untrustworthy
// source.
// See the explainer (https://github.com/explainers-by-googlers/prompt-api)
// for more information.

// This struct contains the parameters that control the sampling behavior of
// the model execution behind the session.
struct AILanguageModelSamplingParams {
  uint32 top_k;
  float temperature;
};

// This struct contains the information describing the default and max
// sampling params for session creation.
struct AILanguageModelParams {
  AILanguageModelSamplingParams default_sampling_params;
  AILanguageModelSamplingParams max_sampling_params;
};

// The information of an AILanguageModel instance, this should be returned
// from any method that creates a new AILanguageModel.
struct AILanguageModelInstanceInfo {
  // The maximum number of tokens that the AILanguageModel can hold in the
  // context. When the new prompt or response is being stored in the context,
  // the AILanguageModel will ensure the context stores no more than
  // `input_quota`, by optionally evicting some oldest entries.
  uint64 input_quota;
  // The number of tokens that are already in the context.
  uint64 input_usage;

  // The sampling params provided by the caller when creating the instance.
  AILanguageModelSamplingParams sampling_params;

  // Supported prompt input types for the session.
  array<AILanguageModelPromptType>? input_types;

  // Multimodal input format parameters for preprocessing.
  int32? audio_sample_rate_hz;
  int32? audio_channel_count;

  // The sampling mode provided by the caller when creating the instance.
  AILanguageModelSamplingMode? sampling_mode;
};

// The role for a given prompt. See LanguageModelPromptRole in IDL.
enum AILanguageModelPromptRole {
  kSystem = 0,
  kUser = 1,
  kAssistant = 2,
};

// The type of data for a given prompt. See LanguageModelMessageType in IDL.
enum AILanguageModelPromptType {
  kText = 0,
  kImage = 1,
  kAudio = 2,
  kToolCall = 3,
  kToolResponse = 4,
};

// Categorical sampling modes for ongoing parameter control explorations.
// See LanguageModelSamplingMode in IDL.
enum AILanguageModelSamplingMode {
  kMostPredictable = 0,
  kPredictable = 1,
  kBalanced = 2,
  kCreative = 3,
  kMostCreative = 4,
  kSlightlyPredictable = 5,
  kSlightlyCreative = 6,
};

// The expected data type and languages specified with session creation.
// See LanguageModelExpected in IDL.
struct AILanguageModelExpected {
  AILanguageModelPromptType type;
  array<AILanguageCode>? languages;
};

// The supported prompt input content types. These each contain arbitrary data
// provided via the JavaScript API. See LanguageModelPromptContent in IDL.
union AILanguageModelPromptContent {
  // Text provided as input.
  string text;

  // Bitmap provided as input.
  skia.mojom.BitmapN32 bitmap;

  // Audio provided as input.
  on_device_model.mojom.AudioData audio;

  // Tool call provided as input (from JavaScript LanguageModelToolCall).
  // Format: {"callID": "...", "name": "...", "arguments": {...}}
  mojo_base.mojom.DictionaryValue tool_call;

  // Tool response provided as input (from JavaScript LanguageModelToolSuccess
  // or LanguageModelToolError).
  // Success format: {"callID": "...", "name": "...", "result": [...]}
  // Error format: {"callID": "...", "name": "...", "errorMessage": "..."}
  mojo_base.mojom.DictionaryValue tool_response;
};

// A prompt includes a role and content. See LanguageModelMessage in IDL.
struct AILanguageModelPrompt {
  AILanguageModelPromptRole role;
  array<AILanguageModelPromptContent> content;
  // Whether this prompt is an assistant response prefix.
  bool is_prefix = false;
};

// Declares a tool that can be used by the language model.
// See LanguageModelToolDeclaration in IDL.
struct AILanguageModelToolDeclaration {
  // Unique name identifying this tool.
  string name;
  // Human-readable description of what the tool does.
  string description;
  // JSON Schema defining the tool's input parameters as structured data.
  mojo_base.mojom.DictionaryValue input_schema;
};

// This is used when creating a new AILanguageModel.
struct AILanguageModelCreateOptions {
  AILanguageModelSamplingParams? sampling_params;
  array<AILanguageModelPrompt> initial_prompts;
  array<AILanguageModelExpected>? expected_inputs;
  array<AILanguageModelExpected>? expected_outputs;
  array<AILanguageModelToolDeclaration>? tools;
  AILanguageModelSamplingMode? sampling_mode;
};

// The client interface that receives an AILanguageModel or errors from an
// AIManager.
interface AIManagerCreateLanguageModelClient {
  // Called with a created language model's mojo interface as a result for the
  // CreateLanguageModel() method of the AIManager. When failed to create an
  // AILanguageModel, this method is called with a null pending remote.
  // It also returns the language model info if the language model is created
  // successfully.
  OnResult(pending_remote<blink.mojom.AILanguageModel> language_model_remote,
           AILanguageModelInstanceInfo info);

  // Called when the manager cannot create an AILanguageModel.
  // `quota_error_info` may provide additional info for kInitialInputTooLarge.
  OnError(AIManagerCreateClientError error, QuotaErrorInfo? quota_error_info);
};

// A session for prompting a model with inputs to obtain streaming responses.
interface AILanguageModel {
  // Prompts the model with arbitrary content provided via the JavaScript API.
  // An optional JSON schema or regex can be provided to define structured
  // output requirements for the response. Note that only one of the JSON schema
  // or regex should be provided.
  Prompt(array<AILanguageModelPrompt> prompts,
         on_device_model.mojom.ResponseConstraint? constraint,
         pending_remote<ModelStreamingResponder> pending_responder);

  // Appends the prompts to the context, but they won't be processed yet until
  // `Prompt()` is called. This method can be called multiple times before the
  // `Prompt()`.
  Append(array<AILanguageModelPrompt> prompts,
         pending_remote<ModelStreamingResponder> client);

  // Creates a new session with the same configuration and existing context
  // as the current session.
  // Note that this method is called by the `AILanguageModel.clone()` method
  // from blink, but it is named as `Fork` instead of `Clone` to avoid confusion
  // with the "clone" term in mojo context (which indicates binding another
  // client to the same instance).
  // Similar to `AIManager.CreateLanguageModel()`, this method returns the
  // information of the created session.
  Fork(pending_remote<AIManagerCreateLanguageModelClient> client);

  // Destroys the session.
  Destroy();

  // Counts the number of token in the given prompt `input` which can
  // be arbitrary strings or multimodal data from the JavaScript API. If it's
  // unable to measure, the method will return null.
  MeasureInputUsage(array<AILanguageModelPrompt> input)
      => (uint32? number_of_tokens);
};