// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

syntax = "proto3";

package google.cloud.aiplatform.v1beta1;

import "google/api/field_behavior.proto";
import "google/api/resource.proto";
import "google/cloud/aiplatform/v1beta1/tool.proto";
import "google/protobuf/struct.proto";
import "google/protobuf/timestamp.proto";

option csharp_namespace = "Google.Cloud.AIPlatform.V1Beta1";
option go_package = "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb;aiplatformpb";
option java_multiple_files = true;
option java_outer_classname = "ExtensionProto";
option java_package = "com.google.cloud.aiplatform.v1beta1";
option php_namespace = "Google\\Cloud\\AIPlatform\\V1beta1";
option ruby_package = "Google::Cloud::AIPlatform::V1beta1";
option (google.api.resource_definition) = {
  type: "secretmanager.googleapis.com/SecretVersion"
  pattern: "projects/{project}/secrets/{secret}/versions/{secret_version}"
  pattern: "projects/{project}/locations/{location}/secrets/{secret}/versions/{secret_version}"
};
option (google.api.resource_definition) = {
  type: "servicedirectory.googleapis.com/Service"
  pattern: "projects/{project}/locations/{location}/namespaces/{namespace}/services/{service}"
};

// Enum of location an HTTP element can be.
enum HttpElementLocation {
  HTTP_IN_UNSPECIFIED = 0;

  // Element is in the HTTP request query.
  HTTP_IN_QUERY = 1;

  // Element is in the HTTP request header.
  HTTP_IN_HEADER = 2;

  // Element is in the HTTP request path.
  HTTP_IN_PATH = 3;

  // Element is in the HTTP request body.
  HTTP_IN_BODY = 4;

  // Element is in the HTTP request cookie.
  HTTP_IN_COOKIE = 5;
}

// Type of Auth.
enum AuthType {
  AUTH_TYPE_UNSPECIFIED = 0;

  // No Auth.
  NO_AUTH = 1;

  // API Key Auth.
  API_KEY_AUTH = 2;

  // HTTP Basic Auth.
  HTTP_BASIC_AUTH = 3;

  // Google Service Account Auth.
  GOOGLE_SERVICE_ACCOUNT_AUTH = 4;

  // OAuth auth.
  OAUTH = 6;

  // OpenID Connect (OIDC) Auth.
  OIDC_AUTH = 8;
}

// Extensions are tools for large language models to access external data, run
// computations, etc.
message Extension {
  option (google.api.resource) = {
    type: "aiplatform.googleapis.com/Extension"
    pattern: "projects/{project}/locations/{location}/extensions/{extension}"
    plural: "extensions"
    singular: "extension"
  };

  // Identifier. The resource name of the Extension.
  string name = 1 [(google.api.field_behavior) = IDENTIFIER];

  // Required. The display name of the Extension.
  // The name can be up to 128 characters long and can consist of any UTF-8
  // characters.
  string display_name = 3 [(google.api.field_behavior) = REQUIRED];

  // Optional. The description of the Extension.
  string description = 4 [(google.api.field_behavior) = OPTIONAL];

  // Output only. Timestamp when this Extension was created.
  google.protobuf.Timestamp create_time = 5
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Timestamp when this Extension was most recently updated.
  google.protobuf.Timestamp update_time = 6
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Optional. Used to perform consistent read-modify-write updates. If not set,
  // a blind "overwrite" update happens.
  string etag = 7 [(google.api.field_behavior) = OPTIONAL];

  // Required. Manifest of the Extension.
  ExtensionManifest manifest = 9 [(google.api.field_behavior) = REQUIRED];

  // Output only. Supported operations.
  repeated ExtensionOperation extension_operations = 11
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Optional. Runtime config controlling the runtime behavior of this
  // Extension.
  RuntimeConfig runtime_config = 13 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Examples to illustrate the usage of the extension as a tool.
  repeated ToolUseExample tool_use_examples = 15
      [(google.api.field_behavior) = OPTIONAL];

  // Optional. The PrivateServiceConnect config for the extension.
  // If specified, the service endpoints associated with the
  // Extension should be registered with private network access in the provided
  // Service Directory
  // (https://cloud.google.com/service-directory/docs/configuring-private-network-access).
  //
  // If the service contains more than one endpoint with a network, the service
  // will arbitrarilty choose one of the endpoints to use for extension
  // execution.
  ExtensionPrivateServiceConnectConfig private_service_connect_config = 16
      [(google.api.field_behavior) = OPTIONAL];
}

// Manifest spec of an Extension needed for runtime execution.
message ExtensionManifest {
  // The API specification shown to the LLM.
  message ApiSpec {
    oneof api_spec {
      // The API spec in Open API standard and YAML format.
      string open_api_yaml = 1;

      // Cloud Storage URI pointing to the OpenAPI spec.
      string open_api_gcs_uri = 2;
    }
  }

  // Required. Extension name shown to the LLM.
  // The name can be up to 128 characters long.
  string name = 1 [(google.api.field_behavior) = REQUIRED];

  // Required. The natural language description shown to the LLM.
  // It should describe the usage of the extension, and is essential for the LLM
  // to perform reasoning. e.g., if the extension is a data store, you can let
  // the LLM know what data it contains.
  string description = 2 [(google.api.field_behavior) = REQUIRED];

  // Required. Immutable. The API specification shown to the LLM.
  ApiSpec api_spec = 3 [
    (google.api.field_behavior) = IMMUTABLE,
    (google.api.field_behavior) = REQUIRED
  ];

  // Required. Immutable. Type of auth supported by this extension.
  AuthConfig auth_config = 5 [
    (google.api.field_behavior) = IMMUTABLE,
    (google.api.field_behavior) = REQUIRED
  ];
}

// Operation of an extension.
message ExtensionOperation {
  // Operation ID that uniquely identifies the operations among the extension.
  // See: "Operation Object" in https://swagger.io/specification/.
  //
  // This field is parsed from the OpenAPI spec. For HTTP extensions, if it does
  // not exist in the spec, we will generate one from the HTTP method and path.
  string operation_id = 1;

  // Output only. Structured representation of a function declaration as defined
  // by the OpenAPI Spec.
  FunctionDeclaration function_declaration = 3
      [(google.api.field_behavior) = OUTPUT_ONLY];
}

// Auth configuration to run the extension.
message AuthConfig {
  // Config for authentication with API key.
  message ApiKeyConfig {
    // Required. The parameter name of the API key.
    // E.g. If the API request is "https://example.com/act?api_key=<API KEY>",
    // "api_key" would be the parameter name.
    string name = 1 [(google.api.field_behavior) = REQUIRED];

    // Required. The name of the SecretManager secret version resource storing
    // the API key. Format:
    // `projects/{project}/secrets/{secrete}/versions/{version}`
    //
    // - If specified, the `secretmanager.versions.access` permission should be
    // granted to Vertex AI Extension Service Agent
    // (https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents)
    // on the specified resource.
    string api_key_secret = 2 [
      (google.api.field_behavior) = REQUIRED,
      (google.api.resource_reference) = {
        type: "secretmanager.googleapis.com/SecretVersion"
      }
    ];

    // Required. The location of the API key.
    HttpElementLocation http_element_location = 3
        [(google.api.field_behavior) = REQUIRED];
  }

  // Config for HTTP Basic Authentication.
  message HttpBasicAuthConfig {
    // Required. The name of the SecretManager secret version resource storing
    // the base64 encoded credentials. Format:
    // `projects/{project}/secrets/{secrete}/versions/{version}`
    //
    // - If specified, the `secretmanager.versions.access` permission should be
    // granted to Vertex AI Extension Service Agent
    // (https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents)
    // on the specified resource.
    string credential_secret = 2 [
      (google.api.field_behavior) = REQUIRED,
      (google.api.resource_reference) = {
        type: "secretmanager.googleapis.com/SecretVersion"
      }
    ];
  }

  // Config for Google Service Account Authentication.
  message GoogleServiceAccountConfig {
    // Optional. The service account that the extension execution service runs
    // as.
    //
    // - If the service account is specified,
    // the `iam.serviceAccounts.getAccessToken` permission should be granted to
    // Vertex AI Extension Service Agent
    // (https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents)
    // on the specified service account.
    //
    // - If not specified, the Vertex AI Extension Service Agent
    // will be used to execute the Extension.
    string service_account = 1 [(google.api.field_behavior) = OPTIONAL];
  }

  // Config for user oauth.
  message OauthConfig {
    oneof oauth_config {
      // Access token for extension endpoint.
      // Only used to propagate token from
      // [[ExecuteExtensionRequest.runtime_auth_config]] at request time.
      string access_token = 1;

      // The service account used to generate access tokens for executing the
      // Extension.
      //
      // - If the service account is specified,
      // the `iam.serviceAccounts.getAccessToken` permission should be granted
      // to Vertex AI Extension Service Agent
      // (https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents)
      // on the provided service account.
      string service_account = 2;
    }
  }

  // Config for user OIDC auth.
  message OidcConfig {
    oneof oidc_config {
      // OpenID Connect formatted ID token for extension endpoint.
      // Only used to propagate token from
      // [[ExecuteExtensionRequest.runtime_auth_config]] at request time.
      string id_token = 1;

      // The service account used to generate an OpenID Connect
      // (OIDC)-compatible JWT token signed by the Google OIDC Provider
      // (accounts.google.com) for extension endpoint
      // (https://cloud.google.com/iam/docs/create-short-lived-credentials-direct#sa-credentials-oidc).
      //
      // - The audience for the token will be set to the URL in the server url
      // defined in the OpenApi spec.
      //
      // - If the service account is provided, the service account should grant
      // `iam.serviceAccounts.getOpenIdToken` permission to Vertex AI Extension
      // Service Agent
      // (https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents).
      string service_account = 2;
    }
  }

  oneof auth_config {
    // Config for API key auth.
    ApiKeyConfig api_key_config = 2;

    // Config for HTTP Basic auth.
    HttpBasicAuthConfig http_basic_auth_config = 3;

    // Config for Google Service Account auth.
    GoogleServiceAccountConfig google_service_account_config = 4;

    // Config for user oauth.
    OauthConfig oauth_config = 5;

    // Config for user OIDC auth.
    OidcConfig oidc_config = 7;
  }

  // Type of auth scheme.
  AuthType auth_type = 101;
}

// Runtime configuration to run the extension.
message RuntimeConfig {
  message CodeInterpreterRuntimeConfig {
    // Optional. The Cloud Storage bucket for file input of this Extension.
    // If specified, support input from the Cloud Storage bucket.
    // Vertex Extension Custom Code Service Agent should be granted
    // file reader to this bucket.
    // If not specified, the extension will only accept file contents from
    // request body and reject Cloud Storage file inputs.
    string file_input_gcs_bucket = 1 [(google.api.field_behavior) = OPTIONAL];

    // Optional. The Cloud Storage bucket for file output of this Extension.
    // If specified, write all output files to the Cloud Storage bucket.
    // Vertex Extension Custom Code Service Agent should be granted
    // file writer to this bucket.
    // If not specified, the file content will be output in response body.
    string file_output_gcs_bucket = 2 [(google.api.field_behavior) = OPTIONAL];
  }

  message VertexAISearchRuntimeConfig {
    // Optional. Vertex AI Search serving config name. Format:
    // `projects/{project}/locations/{location}/collections/{collection}/engines/{engine}/servingConfigs/{serving_config}`
    string serving_config_name = 1 [(google.api.field_behavior) = OPTIONAL];

    // Optional. Vertex AI Search engine ID. This is used to construct the
    // search request. By setting this engine_id, API will construct the serving
    // config using the default value to call search API for the user. The
    // engine_id and serving_config_name cannot both be empty at the same time.
    string engine_id = 2 [(google.api.field_behavior) = OPTIONAL];
  }

  // Runtime configurations for Google first party extensions.
  oneof GoogleFirstPartyExtensionConfig {
    // Code execution runtime configurations for code interpreter extension.
    CodeInterpreterRuntimeConfig code_interpreter_runtime_config = 2;

    // Runtime configuration for Vertex AI Search extension.
    VertexAISearchRuntimeConfig vertex_ai_search_runtime_config = 6;
  }

  // Optional. Default parameters that will be set for all the execution of this
  // extension. If specified, the parameter values can be overridden by values
  // in [[ExecuteExtensionRequest.operation_params]] at request time.
  //
  // The struct should be in a form of map with param name as the key and actual
  // param value as the value.
  // E.g. If this operation requires a param "name" to be set to "abc". you can
  // set this to something like {"name": "abc"}.
  google.protobuf.Struct default_params = 4
      [(google.api.field_behavior) = OPTIONAL];
}

// PrivateExtensionConfig configuration for the extension.
message ExtensionPrivateServiceConnectConfig {
  // Required. The Service Directory resource name in which the service
  // endpoints associated to the extension are registered. Format:
  // `projects/{project_id}/locations/{location_id}/namespaces/{namespace_id}/services/{service_id}`
  //
  // - The Vertex AI Extension Service Agent
  // (https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents)
  // should be granted `servicedirectory.viewer` and
  // `servicedirectory.pscAuthorizedService` roles on the resource.
  string service_directory = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "servicedirectory.googleapis.com/Service"
    }
  ];
}