// Copyright 2023 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 {
  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];
}

// Request message for
// [VertexRagService.RetrieveContexts][google.cloud.aiplatform.v1beta1.VertexRagService.RetrieveContexts].
message RetrieveContextsRequest {
  // The data source for Vertex RagStore.
  message VertexRagStore {
    // Required. RagCorpora resource name.
    // Format:
    // `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`
    // Currently only one corpus is allowed.
    // In the future we may open up multiple corpora support. However, they
    // should be from the same project and location.
    repeated string rag_corpora = 1 [(google.api.field_behavior) = REQUIRED];
  }

  // 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 vector and the context text vector.
    double distance = 3;
  }

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