// 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/openapi.proto";
import "google/protobuf/struct.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 = "ToolProto";
option java_package = "com.google.cloud.aiplatform.v1beta1";
option php_namespace = "Google\\Cloud\\AIPlatform\\V1beta1";
option ruby_package = "Google::Cloud::AIPlatform::V1beta1";

// Tool details that the model may use to generate response.
//
// A `Tool` is a piece of code that enables the system to interact with
// external systems to perform an action, or set of actions, outside of
// knowledge and scope of the model. A Tool object should contain exactly
// one type of Tool (e.g FunctionDeclaration, Retrieval or
// GoogleSearchRetrieval).
message Tool {
  // Optional. Function tool type.
  // One or more function declarations to be passed to the model along with the
  // current user query. Model may decide to call a subset of these functions
  // by populating [FunctionCall][content.part.function_call] in the response.
  // User should provide a [FunctionResponse][content.part.function_response]
  // for each function call in the next turn. Based on the function responses,
  // Model will generate the final response back to the user.
  // Maximum 64 function declarations can be provided.
  repeated FunctionDeclaration function_declarations = 1
      [(google.api.field_behavior) = OPTIONAL];

  // Optional. Retrieval tool type.
  // System will always execute the provided retrieval tool(s) to get external
  // knowledge to answer the prompt. Retrieval results are presented to the
  // model for generation.
  Retrieval retrieval = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. GoogleSearchRetrieval tool type.
  // Specialized retrieval tool that is powered by Google search.
  GoogleSearchRetrieval google_search_retrieval = 3
      [(google.api.field_behavior) = OPTIONAL];
}

// A single example of the tool usage.
message ToolUseExample {
  // Identifies one operation of the extension.
  message ExtensionOperation {
    // Resource name of the extension.
    string extension = 1 [(google.api.resource_reference) = {
      type: "aiplatform.googleapis.com/Extension"
    }];

    // Required. Operation ID of the extension.
    string operation_id = 2 [(google.api.field_behavior) = REQUIRED];
  }

  // Target tool to use.
  oneof Target {
    // Extension operation to call.
    ExtensionOperation extension_operation = 10;

    // Function name to call.
    string function_name = 11;
  }

  // Required. The display name for example.
  string display_name = 1 [(google.api.field_behavior) = REQUIRED];

  // Required. Query that should be routed to this tool.
  string query = 2 [(google.api.field_behavior) = REQUIRED];

  // Request parameters used for executing this tool.
  google.protobuf.Struct request_params = 3;

  // Response parameters generated by this tool.
  google.protobuf.Struct response_params = 4;

  // Summary of the tool response to the user query.
  string response_summary = 5;
}

// Structured representation of a function declaration as defined by the
// [OpenAPI 3.0 specification](https://spec.openapis.org/oas/v3.0.3). Included
// in this declaration are the function name and parameters. This
// FunctionDeclaration is a representation of a block of code that can be used
// as a `Tool` by the model and executed by the client.
message FunctionDeclaration {
  // Required. The name of the function to call.
  // Must start with a letter or an underscore.
  // Must be a-z, A-Z, 0-9, or contain underscores, dots and dashes, with a
  // maximum length of 64.
  string name = 1 [(google.api.field_behavior) = REQUIRED];

  // Optional. Description and purpose of the function.
  // Model uses it to decide how and whether to call the function.
  string description = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Describes the parameters to this function in JSON Schema Object
  // format. Reflects the Open API 3.03 Parameter Object. string Key: the name
  // of the parameter. Parameter names are case sensitive. Schema Value: the
  // Schema defining the type used for the parameter. For function with no
  // parameters, this can be left unset. Parameter names must start with a
  // letter or an underscore and must only contain chars a-z, A-Z, 0-9, or
  // underscores with a maximum length of 64. Example with 1 required and 1
  // optional parameter: type: OBJECT properties:
  //  param1:
  //    type: STRING
  //  param2:
  //    type: INTEGER
  // required:
  //  - param1
  Schema parameters = 3 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Describes the output from this function in JSON Schema format.
  // Reflects the Open API 3.03 Response Object. The Schema defines the type
  // used for the response value of the function.
  Schema response = 4 [(google.api.field_behavior) = OPTIONAL];
}

// A predicted [FunctionCall] returned from the model that contains a string
// representing the [FunctionDeclaration.name] and a structured JSON object
// containing the parameters and their values.
message FunctionCall {
  // Required. The name of the function to call.
  // Matches [FunctionDeclaration.name].
  string name = 1 [(google.api.field_behavior) = REQUIRED];

  // Optional. Required. The function parameters and values in JSON object
  // format. See [FunctionDeclaration.parameters] for parameter details.
  google.protobuf.Struct args = 2 [(google.api.field_behavior) = OPTIONAL];
}

// The result output from a [FunctionCall] that contains a string representing
// the [FunctionDeclaration.name] and a structured JSON object containing any
// output from the function is used as context to the model. This should contain
// the result of a [FunctionCall] made based on model prediction.
message FunctionResponse {
  // Required. The name of the function to call.
  // Matches [FunctionDeclaration.name] and [FunctionCall.name].
  string name = 1 [(google.api.field_behavior) = REQUIRED];

  // Required. The function response in JSON object format.
  google.protobuf.Struct response = 2 [(google.api.field_behavior) = REQUIRED];
}

// Defines a retrieval tool that model can call to access external knowledge.
message Retrieval {
  // The source of the retrieval.
  oneof source {
    // Set to use data source powered by Vertex AI Search.
    VertexAISearch vertex_ai_search = 2;

    // Set to use data source powered by Vertex RAG store.
    // User data is uploaded via the VertexRagDataService.
    VertexRagStore vertex_rag_store = 4;
  }

  // Optional. Deprecated. This option is no longer supported.
  bool disable_attribution = 3
      [deprecated = true, (google.api.field_behavior) = OPTIONAL];
}

// Retrieve from Vertex RAG Store for grounding.
message VertexRagStore {
  // The definition of the Rag resource.
  message RagResource {
    // Optional. RagCorpora resource name.
    // Format:
    // `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`
    string rag_corpus = 1 [
      (google.api.field_behavior) = OPTIONAL,
      (google.api.resource_reference) = {
        type: "aiplatform.googleapis.com/RagCorpus"
      }
    ];

    // Optional. rag_file_id. The files should be in the same rag_corpus set in
    // rag_corpus field.
    repeated string rag_file_ids = 2 [(google.api.field_behavior) = OPTIONAL];
  }

  // Optional. Deprecated. Please use rag_resources instead.
  repeated string rag_corpora = 1 [
    deprecated = true,
    (google.api.field_behavior) = OPTIONAL,
    (google.api.resource_reference) = {
      type: "aiplatform.googleapis.com/RagCorpus"
    }
  ];

  // Optional. The representation of the rag source. It can be used to specify
  // corpus only or ragfiles. Currently only support one corpus or multiple
  // files from one corpus. In the future we may open up multiple corpora
  // support.
  repeated RagResource rag_resources = 4
      [(google.api.field_behavior) = OPTIONAL];

  // Optional. Number of top k results to return from the selected corpora.
  optional int32 similarity_top_k = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Only return results with vector distance smaller than the
  // threshold.
  optional double vector_distance_threshold = 3
      [(google.api.field_behavior) = OPTIONAL];
}

// Retrieve from Vertex AI Search datastore for grounding.
// See https://cloud.google.com/vertex-ai-search-and-conversation
message VertexAISearch {
  // Required. Fully-qualified Vertex AI Search's datastore resource ID.
  // Format:
  // `projects/{project}/locations/{location}/collections/{collection}/dataStores/{dataStore}`
  string datastore = 1 [(google.api.field_behavior) = REQUIRED];
}

// Tool to retrieve public web data for grounding, powered by Google.
message GoogleSearchRetrieval {}

// Tool config. This config is shared for all tools provided in the request.
message ToolConfig {
  // Optional. Function calling config.
  FunctionCallingConfig function_calling_config = 1
      [(google.api.field_behavior) = OPTIONAL];
}

// Function calling config.
message FunctionCallingConfig {
  // Function calling mode.
  enum Mode {
    // Unspecified function calling mode. This value should not be used.
    MODE_UNSPECIFIED = 0;

    // Default model behavior, model decides to predict either a function call
    // or a natural language response.
    AUTO = 1;

    // Model is constrained to always predicting a function call only.
    // If "allowed_function_names" are set, the predicted function call will be
    // limited to any one of "allowed_function_names", else the predicted
    // function call will be any one of the provided "function_declarations".
    ANY = 2;

    // Model will not predict any function call. Model behavior is same as when
    // not passing any function declarations.
    NONE = 3;
  }

  // Optional. Function calling mode.
  Mode mode = 1 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Function names to call. Only set when the Mode is ANY. Function
  // names should match [FunctionDeclaration.name]. With mode set to ANY, model
  // will predict a function call from the set of function names provided.
  repeated string allowed_function_names = 2
      [(google.api.field_behavior) = OPTIONAL];
}