// 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/annotations.proto";
import "google/api/client.proto";
import "google/api/field_behavior.proto";
import "google/api/resource.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";
  }
}

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

  // Optional. Configurations for hybrid search results ranking.
  Ranking ranking = 4 [(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
        [(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 {
    // For vertex RagStore, 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 text chunk.
    string text = 2;

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

    // The distance between the query sparse embedding vector and the context
    // text vector.
    double sparse_distance = 4;
  }

  // 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;
}