// Copyright 2025 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/encryption_spec.proto";
import "google/cloud/aiplatform/v1beta1/env_var.proto";
import "google/cloud/aiplatform/v1beta1/service_networking.proto";
import "google/protobuf/duration.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 = "ReasoningEngineProto";
option java_package = "com.google.cloud.aiplatform.v1beta1";
option php_namespace = "Google\\Cloud\\AIPlatform\\V1beta1";
option ruby_package = "Google::Cloud::AIPlatform::V1beta1";

// ReasoningEngine configurations
message ReasoningEngineSpec {
  // User provided package spec like pickled object and package requirements.
  message PackageSpec {
    // Optional. The Cloud Storage URI of the pickled python object.
    string pickle_object_gcs_uri = 1 [(google.api.field_behavior) = OPTIONAL];

    // Optional. The Cloud Storage URI of the dependency files in tar.gz format.
    string dependency_files_gcs_uri = 2
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. The Cloud Storage URI of the `requirements.txt` file
    string requirements_gcs_uri = 3 [(google.api.field_behavior) = OPTIONAL];

    // Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11.
    // If not specified, default value is 3.10.
    string python_version = 4 [(google.api.field_behavior) = OPTIONAL];
  }

  // The specification of a Reasoning Engine deployment.
  message DeploymentSpec {
    // Optional. Environment variables to be set with the Reasoning Engine
    // deployment. The environment variables can be updated through the
    // UpdateReasoningEngine API.
    repeated EnvVar env = 1 [(google.api.field_behavior) = OPTIONAL];

    // Optional. Environment variables where the value is a secret in Cloud
    // Secret Manager.
    // To use this feature, add 'Secret Manager Secret Accessor' role
    // (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine
    // Service Agent.
    repeated SecretEnvVar secret_env = 2
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. Configuration for PSC-I.
    PscInterfaceConfig psc_interface_config = 4
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. The minimum number of application instances that will be kept
    // running at all times. Defaults to 1.
    optional int32 min_instances = 5 [(google.api.field_behavior) = OPTIONAL];

    // Optional. The maximum number of application instances that can be
    // launched to handle increased traffic. Defaults to 100.
    optional int32 max_instances = 6 [(google.api.field_behavior) = OPTIONAL];

    // Optional. Resource limits for each container. Only 'cpu' and 'memory'
    // keys are supported. Defaults to {"cpu": "4", "memory": "4Gi"}.
    //
    //   * The only supported values for CPU are '1', '2', '4', and '8'. For
    //   more information, go to
    //   https://cloud.google.com/run/docs/configuring/cpu.
    //   * For supported 'memory' values and syntax, go to
    //   https://cloud.google.com/run/docs/configuring/memory-limits
    map<string, string> resource_limits = 7
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. Concurrency for each container and agent server. Recommended
    // value: 2 * cpu + 1. Defaults to 9.
    optional int32 container_concurrency = 8
        [(google.api.field_behavior) = OPTIONAL];
  }

  // Optional. The service account that the Reasoning Engine artifact runs as.
  // It should have "roles/storage.objectViewer" for reading the user project's
  // Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If
  // not specified, the Vertex AI Reasoning Engine Service Agent in the project
  // will be used.
  optional string service_account = 1 [(google.api.field_behavior) = OPTIONAL];

  // Optional. User provided package spec of the ReasoningEngine.
  // Ignored when users directly specify a deployment image through
  // `deployment_spec.first_party_image_override`, but keeping the
  // field_behavior to avoid introducing breaking changes.
  PackageSpec package_spec = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. The specification of a Reasoning Engine deployment.
  DeploymentSpec deployment_spec = 4 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Declarations for object class methods in OpenAPI specification
  // format.
  repeated google.protobuf.Struct class_methods = 3
      [(google.api.field_behavior) = OPTIONAL];

  // Optional. The OSS agent framework used to develop the agent.
  // Currently supported values: "google-adk", "langchain", "langgraph", "ag2",
  // "llama-index", "custom".
  string agent_framework = 5 [(google.api.field_behavior) = OPTIONAL];
}

// ReasoningEngine provides a customizable runtime for models to determine
// which actions to take and in which order.
message ReasoningEngine {
  option (google.api.resource) = {
    type: "aiplatform.googleapis.com/ReasoningEngine"
    pattern: "projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}"
    plural: "reasoningEngines"
    singular: "reasoningEngine"
  };

  // Identifier. The resource name of the ReasoningEngine.
  // Format:
  // `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`
  string name = 1 [(google.api.field_behavior) = IDENTIFIER];

  // Required. The display name of the ReasoningEngine.
  string display_name = 2 [(google.api.field_behavior) = REQUIRED];

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

  // Optional. Configurations of the ReasoningEngine
  ReasoningEngineSpec spec = 3 [(google.api.field_behavior) = OPTIONAL];

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

  // Output only. Timestamp when this ReasoningEngine was most recently updated.
  google.protobuf.Timestamp update_time = 5
      [(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 = 6 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Configuration for how Agent Engine sub-resources should manage
  // context.
  ReasoningEngineContextSpec context_spec = 9
      [(google.api.field_behavior) = OPTIONAL];

  // Customer-managed encryption key spec for a ReasoningEngine. If set, this
  // ReasoningEngine and all sub-resources of this ReasoningEngine will be
  // secured by this key.
  EncryptionSpec encryption_spec = 11;
}

// Configuration for how Agent Engine sub-resources should manage context.
message ReasoningEngineContextSpec {
  // Specification for a Memory Bank.
  message MemoryBankConfig {
    // Configuration for automatically setting the TTL ("time-to-live") of the
    // memories in the Memory Bank.
    message TtlConfig {
      // Configuration for TTL of the memories in the Memory Bank based on the
      // action that created or updated the memory.
      message GranularTtlConfig {
        // Optional. The TTL duration for memories uploaded via CreateMemory.
        google.protobuf.Duration create_ttl = 1
            [(google.api.field_behavior) = OPTIONAL];

        // Optional. The TTL duration for memories newly generated via
        // GenerateMemories
        // ([GenerateMemoriesResponse.GeneratedMemory.Action.CREATED][google.cloud.aiplatform.v1beta1.GenerateMemoriesResponse.GeneratedMemory.Action.CREATED]).
        google.protobuf.Duration generate_created_ttl = 2
            [(google.api.field_behavior) = OPTIONAL];

        // Optional. The TTL duration for memories updated via GenerateMemories
        // ([GenerateMemoriesResponse.GeneratedMemory.Action.CREATED][google.cloud.aiplatform.v1beta1.GenerateMemoriesResponse.GeneratedMemory.Action.CREATED]).
        // In the case of an UPDATE action, the `expire_time` of the existing
        // memory will be updated to the new value (now + TTL).
        google.protobuf.Duration generate_updated_ttl = 3
            [(google.api.field_behavior) = OPTIONAL];
      }

      // Configuration for automatically setting the TTL of the memories in the
      // Memory Bank.
      oneof ttl {
        // Optional. The default TTL duration of the memories in the Memory
        // Bank. This applies to all operations that create or update a memory.
        google.protobuf.Duration default_ttl = 1
            [(google.api.field_behavior) = OPTIONAL];

        // Optional. The granular TTL configuration of the memories in the
        // Memory Bank.
        GranularTtlConfig granular_ttl_config = 2
            [(google.api.field_behavior) = OPTIONAL];
      }
    }

    // Configuration for how to generate memories.
    message GenerationConfig {
      // Required. The model used to generate memories.
      // Format:
      // `projects/{project}/locations/{location}/publishers/google/models/{model}`.
      string model = 1 [
        (google.api.field_behavior) = REQUIRED,
        (google.api.resource_reference) = {
          type: "aiplatform.googleapis.com/Endpoint"
        }
      ];
    }

    // Configuration for how to perform similarity search on memories.
    message SimilaritySearchConfig {
      // Required. The model used to generate embeddings to lookup similar
      // memories. Format:
      // `projects/{project}/locations/{location}/publishers/google/models/{model}`.
      string embedding_model = 1 [
        (google.api.field_behavior) = REQUIRED,
        (google.api.resource_reference) = {
          type: "aiplatform.googleapis.com/Endpoint"
        }
      ];
    }

    // Optional. Configuration for how to generate memories for the Memory Bank.
    GenerationConfig generation_config = 1
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. Configuration for how to perform similarity search on memories.
    // If not set, the Memory Bank will use the default embedding model
    // `text-embedding-005`.
    SimilaritySearchConfig similarity_search_config = 2
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. Configuration for automatic TTL ("time-to-live") of the
    // memories in the Memory Bank. If not set, TTL will not be applied
    // automatically. The TTL can be explicitly set by modifying the
    // `expire_time` of each Memory resource.
    TtlConfig ttl_config = 5 [(google.api.field_behavior) = OPTIONAL];
  }

  // Optional. Specification for a Memory Bank, which manages memories for the
  // Agent Engine.
  MemoryBankConfig memory_bank_config = 1
      [(google.api.field_behavior) = OPTIONAL];
}