// Copyright 2026 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/annotations.proto";
import "google/api/client.proto";
import "google/api/field_behavior.proto";
import "google/api/resource.proto";
import "google/cloud/aiplatform/v1beta1/content.proto";
import "google/cloud/aiplatform/v1beta1/operation.proto";
import "google/cloud/aiplatform/v1beta1/tool.proto";
import "google/cloud/aiplatform/v1beta1/vertex_rag_data.proto";
import "google/longrunning/operations.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 = "VertexRagServiceProto";
option java_package = "com.google.cloud.aiplatform.v1beta1";
option php_namespace = "Google\\Cloud\\AIPlatform\\V1beta1";
option ruby_package = "Google::Cloud::AIPlatform::V1beta1";

// A service for retrieving relevant contexts.
service VertexRagService {
  option (google.api.default_host) = "aiplatform.googleapis.com";
  option (google.api.oauth_scopes) =
      "https://www.googleapis.com/auth/cloud-platform";

  // Retrieves relevant contexts for a query.
  rpc RetrieveContexts(RetrieveContextsRequest)
      returns (RetrieveContextsResponse) {
    option (google.api.http) = {
      post: "/v1beta1/{parent=projects/*/locations/*}:retrieveContexts"
      body: "*"
    };
    option (google.api.method_signature) = "parent,query";
  }

  // Given an input prompt, it returns augmented prompt from vertex rag store
  //  to guide LLM towards generating grounded responses.
  rpc AugmentPrompt(AugmentPromptRequest) returns (AugmentPromptResponse) {
    option (google.api.http) = {
      post: "/v1beta1/{parent=projects/*/locations/*}:augmentPrompt"
      body: "*"
    };
    option (google.api.method_signature) = "parent,model,vertex_rag_store";
  }

  // Given an input text, it returns a score that evaluates the factuality of
  // the text. It also extracts and returns claims from the text and provides
  // supporting facts.
  rpc CorroborateContent(CorroborateContentRequest)
      returns (CorroborateContentResponse) {
    option (google.api.http) = {
      post: "/v1beta1/{parent=projects/*/locations/*}:corroborateContent"
      body: "*"
    };
    option (google.api.method_signature) = "parent,content,facts";
  }

  // Agentic Retrieval Ask API for RAG.
  rpc AskContexts(AskContextsRequest) returns (AskContextsResponse) {
    option (google.api.http) = {
      post: "/v1beta1/{parent=projects/*/locations/*}:askContexts"
      body: "*"
    };
    option (google.api.method_signature) = "parent,query";
  }

  // Asynchronous API to retrieves relevant contexts for a query.
  rpc AsyncRetrieveContexts(AsyncRetrieveContextsRequest)
      returns (google.longrunning.Operation) {
    option (google.api.http) = {
      post: "/v1beta1/{parent=projects/*/locations/*}:asyncRetrieveContexts"
      body: "*"
    };
    option (google.api.method_signature) = "parent,query";
    option (google.longrunning.operation_info) = {
      response_type: "AsyncRetrieveContextsResponse"
      metadata_type: "AsyncRetrieveContextsOperationMetadata"
    };
  }
}

// A query to retrieve relevant contexts.
message RagQuery {
  // Configurations for hybrid search results ranking.
  message Ranking {
    // Optional. Alpha value controls the weight between dense and sparse vector
    // search results. The range is [0, 1], while 0 means sparse vector search
    // only and 1 means dense vector search only. The default value is 0.5 which
    // balances sparse and dense vector search equally.
    optional float alpha = 1 [(google.api.field_behavior) = OPTIONAL];
  }

  // The query to retrieve contexts.
  // Currently only text query is supported.
  oneof query {
    // Optional. The query in text format to get relevant contexts.
    string text = 1 [(google.api.field_behavior) = OPTIONAL];
  }

  // Optional. The number of contexts to retrieve.
  int32 similarity_top_k = 2
      [deprecated = true, (google.api.field_behavior) = OPTIONAL];

  // Optional. Configurations for hybrid search results ranking.
  Ranking ranking = 4
      [deprecated = true, (google.api.field_behavior) = OPTIONAL];

  // Optional. The retrieval config for the query.
  RagRetrievalConfig rag_retrieval_config = 6
      [(google.api.field_behavior) = OPTIONAL];
}

// Request message for
// [VertexRagService.RetrieveContexts][google.cloud.aiplatform.v1beta1.VertexRagService.RetrieveContexts].
message RetrieveContextsRequest {
  // The data source for Vertex RagStore.
  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 to specify the data
    // source.
    repeated string rag_corpora = 1
        [deprecated = true, (google.api.field_behavior) = OPTIONAL];

    // 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 = 3
        [(google.api.field_behavior) = OPTIONAL];

    // Optional. Only return contexts with vector distance smaller than the
    // threshold.
    optional double vector_distance_threshold = 2
        [deprecated = true, (google.api.field_behavior) = OPTIONAL];
  }

  // Data Source to retrieve contexts.
  oneof data_source {
    // The data source for Vertex RagStore.
    VertexRagStore vertex_rag_store = 2;
  }

  // Required. The resource name of the Location from which to retrieve
  // RagContexts. The users must have permission to make a call in the project.
  // Format:
  // `projects/{project}/locations/{location}`.
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "locations.googleapis.com/Location"
    }
  ];

  // Required. Single RAG retrieve query.
  RagQuery query = 3 [(google.api.field_behavior) = REQUIRED];
}

// Relevant contexts for one query.
message RagContexts {
  // A context of the query.
  message Context {
    // If the file is imported from Cloud Storage or Google Drive, source_uri
    // will be original file URI in Cloud Storage or Google Drive; if file is
    // uploaded, source_uri will be file display name.
    string source_uri = 1;

    // The file display name.
    string source_display_name = 5;

    // The text chunk.
    string text = 2;

    // The distance between the query dense embedding vector and the context
    // text vector.
    double distance = 3 [deprecated = true];

    // The distance between the query sparse embedding vector and the context
    // text vector.
    double sparse_distance = 4 [deprecated = true];

    // According to the underlying Vector DB and the selected metric type, the
    // score can be either the distance or the similarity between the query and
    // the context and its range depends on the metric type.
    //
    // For example, if the metric type is COSINE_DISTANCE, it represents the
    // distance between the query and the context. The larger the distance, the
    // less relevant the context is to the query. The range is [0, 2], while 0
    // means the most relevant and 2 means the least relevant.
    optional double score = 6;

    // Context of the retrieved chunk.
    RagChunk chunk = 7;
  }

  // All its contexts.
  repeated Context contexts = 1;
}

// Response message for
// [VertexRagService.RetrieveContexts][google.cloud.aiplatform.v1beta1.VertexRagService.RetrieveContexts].
message RetrieveContextsResponse {
  // The contexts of the query.
  RagContexts contexts = 1;
}

// Request message for AugmentPrompt.
message AugmentPromptRequest {
  // Metadata of the backend deployed model.
  message Model {
    // Optional. The model that the user will send the augmented prompt for
    // content generation.
    string model = 1 [(google.api.field_behavior) = OPTIONAL];

    // Optional. The model version of the backend deployed model.
    string model_version = 2 [(google.api.field_behavior) = OPTIONAL];
  }

  // The data source for retrieving contexts.
  oneof data_source {
    // Optional. Retrieves contexts from the Vertex RagStore.
    VertexRagStore vertex_rag_store = 4
        [(google.api.field_behavior) = OPTIONAL];
  }

  // Required. The resource name of the Location from which to augment prompt.
  // The users must have permission to make a call in the project.
  // Format:
  // `projects/{project}/locations/{location}`.
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "locations.googleapis.com/Location"
    }
  ];

  // Optional. Input content to augment, only text format is supported for now.
  repeated Content contents = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Metadata of the backend deployed model.
  Model model = 3 [(google.api.field_behavior) = OPTIONAL];
}

// Response message for AugmentPrompt.
message AugmentPromptResponse {
  // Augmented prompt, only text format is supported for now.
  repeated Content augmented_prompt = 1;

  // Retrieved facts from RAG data sources.
  repeated Fact facts = 2;
}

// Request message for CorroborateContent.
message CorroborateContentRequest {
  // Parameters that can be overrided per request.
  message Parameters {
    // Optional. Only return claims with citation score larger than the
    // threshold.
    double citation_threshold = 1 [(google.api.field_behavior) = OPTIONAL];
  }

  // Required. The resource name of the Location from which to corroborate text.
  // The users must have permission to make a call in the project.
  // Format:
  // `projects/{project}/locations/{location}`.
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "locations.googleapis.com/Location"
    }
  ];

  // Optional. Input content to corroborate, only text format is supported for
  // now.
  optional Content content = 2 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Facts used to generate the text can also be used to corroborate
  // the text.
  repeated Fact facts = 3 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Parameters that can be set to override default settings per
  // request.
  Parameters parameters = 4 [(google.api.field_behavior) = OPTIONAL];
}

// Response message for CorroborateContent.
message CorroborateContentResponse {
  // Confidence score of corroborating content. Value is [0,1] with 1 is the
  // most confidence.
  optional float corroboration_score = 1;

  // Claims that are extracted from the input content and facts that support the
  // claims.
  repeated Claim claims = 2;
}

// The fact used in grounding.
message Fact {
  // Query that is used to retrieve this fact.
  optional string query = 1;

  // If present, it refers to the title of this fact.
  optional string title = 2;

  // If present, this uri links to the source of the fact.
  optional string uri = 3;

  // If present, the summary/snippet of the fact.
  optional string summary = 4;

  // If present, the distance between the query vector and this fact vector.
  optional double vector_distance = 5 [deprecated = true];

  // If present, according to the underlying Vector DB and the selected metric
  // type, the score can be either the distance or the similarity between the
  // query and the fact and its range depends on the metric type.
  //
  // For example, if the metric type is COSINE_DISTANCE, it represents the
  // distance between the query and the fact. The larger the distance, the less
  // relevant the fact is to the query. The range is [0, 2], while 0 means the
  // most relevant and 2 means the least relevant.
  optional double score = 6;

  // If present, chunk properties.
  optional RagChunk chunk = 7;
}

// Claim that is extracted from the input text and facts that support it.
message Claim {
  // Index in the input text where the claim starts (inclusive).
  optional int32 start_index = 1;

  // Index in the input text where the claim ends (exclusive).
  optional int32 end_index = 2;

  // Indexes of the facts supporting this claim.
  repeated int32 fact_indexes = 3;

  // Confidence score of this corroboration.
  optional float score = 4;
}

// Agentic Retrieval Ask API for RAG.
// Request message for
// [VertexRagService.AskContexts][google.cloud.aiplatform.v1beta1.VertexRagService.AskContexts].
message AskContextsRequest {
  // Required. The resource name of the Location from which to retrieve
  // RagContexts. The users must have permission to make a call in the project.
  // Format:
  // `projects/{project}/locations/{location}`.
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "locations.googleapis.com/Location"
    }
  ];

  // Required. Single RAG retrieve query.
  RagQuery query = 2 [(google.api.field_behavior) = REQUIRED];

  // Optional. The tools to use for AskContexts.
  repeated Tool tools = 3 [(google.api.field_behavior) = OPTIONAL];
}

// Response message for
// [VertexRagService.AskContexts][google.cloud.aiplatform.v1beta1.VertexRagService.AskContexts].
message AskContextsResponse {
  // The Retrieval Response.
  string response = 1;

  // The contexts of the query.
  RagContexts contexts = 2;
}

// Request message for
// [VertexRagService.AsyncRetrieveContexts][google.cloud.aiplatform.v1beta1.VertexRagService.AsyncRetrieveContexts].
message AsyncRetrieveContextsRequest {
  // Required. The resource name of the Location from which to retrieve
  // RagContexts. The users must have permission to make a call in the project.
  // Format:
  // `projects/{project}/locations/{location}`.
  string parent = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "locations.googleapis.com/Location"
    }
  ];

  // Required. Single RAG retrieve query.
  RagQuery query = 2 [(google.api.field_behavior) = REQUIRED];

  // Optional. The tools to use for AskContexts.
  repeated Tool tools = 3 [(google.api.field_behavior) = OPTIONAL];
}

// Response message for
// [VertexRagService.AsyncRetrieveContexts][google.cloud.aiplatform.v1beta1.VertexRagService.AsyncRetrieveContexts].
message AsyncRetrieveContextsResponse {
  // The contexts of the query.
  RagContexts contexts = 1;
}

// Metadata for AsyncRetrieveContextsOperation.
message AsyncRetrieveContextsOperationMetadata {
  // The operation generic information.
  GenericOperationMetadata generic_metadata = 1;
}