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

// MatchService is a Google managed service for efficient vector similarity
// search at scale.
service MatchService {
  option (google.api.default_host) = "aiplatform.googleapis.com";
  option (google.api.oauth_scopes) =
      "https://www.googleapis.com/auth/cloud-platform";

  // Finds the nearest neighbors of each vector within the request.
  rpc FindNeighbors(FindNeighborsRequest) returns (FindNeighborsResponse) {
    option (google.api.http) = {
      post: "/v1beta1/{index_endpoint=projects/*/locations/*/indexEndpoints/*}:findNeighbors"
      body: "*"
    };
  }

  // Reads the datapoints/vectors of the given IDs.
  // A maximum of 1000 datapoints can be retrieved in a batch.
  rpc ReadIndexDatapoints(ReadIndexDatapointsRequest)
      returns (ReadIndexDatapointsResponse) {
    option (google.api.http) = {
      post: "/v1beta1/{index_endpoint=projects/*/locations/*/indexEndpoints/*}:readIndexDatapoints"
      body: "*"
    };
  }
}

// The request message for
// [MatchService.FindNeighbors][google.cloud.aiplatform.v1beta1.MatchService.FindNeighbors].
message FindNeighborsRequest {
  // A query to find a number of the nearest neighbors (most similar vectors)
  // of a vector.
  message Query {
    // Parameters for RRF algorithm that combines search results.
    message RRF {
      // Required. Users can provide an alpha value to give more weight to dense
      // vs sparse results. For example, if the alpha is 0, we only return
      // sparse and if the alpha is 1, we only return dense.
      float alpha = 1 [(google.api.field_behavior) = REQUIRED];
    }

    oneof ranking {
      // Optional. Represents RRF algorithm that combines search results.
      RRF rrf = 6 [(google.api.field_behavior) = OPTIONAL];
    }

    // Required. The datapoint/vector whose nearest neighbors should be searched
    // for.
    IndexDatapoint datapoint = 1 [(google.api.field_behavior) = REQUIRED];

    // The number of nearest neighbors to be retrieved from database for each
    // query. If not set, will use the default from the service configuration
    // (https://cloud.google.com/vertex-ai/docs/matching-engine/configuring-indexes#nearest-neighbor-search-config).
    int32 neighbor_count = 2;

    // Crowding is a constraint on a neighbor list produced by nearest neighbor
    // search requiring that no more than some value k' of the k neighbors
    // returned have the same value of crowding_attribute.
    // It's used for improving result diversity.
    // This field is the maximum number of matches with the same crowding tag.
    int32 per_crowding_attribute_neighbor_count = 3;

    // The number of neighbors to find via approximate search before
    // exact reordering is performed. If not set, the default value from scam
    // config is used; if set, this value must be > 0.
    int32 approximate_neighbor_count = 4;

    // The fraction of the number of leaves to search, set at query time allows
    // user to tune search performance. This value increase result in both
    // search accuracy and latency increase. The value should be between 0.0
    // and 1.0. If not set or set to 0.0, query uses the default value specified
    // in
    // NearestNeighborSearchConfig.TreeAHConfig.fraction_leaf_nodes_to_search.
    double fraction_leaf_nodes_to_search_override = 5;
  }

  // Required. The name of the index endpoint.
  // Format:
  // `projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}`
  string index_endpoint = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "aiplatform.googleapis.com/IndexEndpoint"
    }
  ];

  // The ID of the DeployedIndex that will serve the request. This request is
  // sent to a specific IndexEndpoint, as per the IndexEndpoint.network. That
  // IndexEndpoint also has IndexEndpoint.deployed_indexes, and each such index
  // has a DeployedIndex.id field.
  // The value of the field below must equal one of the DeployedIndex.id
  // fields of the IndexEndpoint that is being called for this request.
  string deployed_index_id = 2;

  // The list of queries.
  repeated Query queries = 3;

  // If set to true, the full datapoints (including all vector values and
  // restricts) of the nearest neighbors are returned.
  // Note that returning full datapoint will significantly increase the
  // latency and cost of the query.
  bool return_full_datapoint = 4;
}

// The response message for
// [MatchService.FindNeighbors][google.cloud.aiplatform.v1beta1.MatchService.FindNeighbors].
message FindNeighborsResponse {
  // A neighbor of the query vector.
  message Neighbor {
    // The datapoint of the neighbor.
    // Note that full datapoints are returned only when "return_full_datapoint"
    // is set to true. Otherwise, only the "datapoint_id" and "crowding_tag"
    // fields are populated.
    IndexDatapoint datapoint = 1;

    // The distance between the neighbor and the dense embedding query.
    double distance = 2;

    // The distance between the neighbor and the query sparse_embedding.
    double sparse_distance = 3;
  }

  // Nearest neighbors for one query.
  message NearestNeighbors {
    // The ID of the query datapoint.
    string id = 1;

    // All its neighbors.
    repeated Neighbor neighbors = 2;
  }

  // The nearest neighbors of the query datapoints.
  repeated NearestNeighbors nearest_neighbors = 1;
}

// The request message for
// [MatchService.ReadIndexDatapoints][google.cloud.aiplatform.v1beta1.MatchService.ReadIndexDatapoints].
message ReadIndexDatapointsRequest {
  // Required. The name of the index endpoint.
  // Format:
  // `projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}`
  string index_endpoint = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "aiplatform.googleapis.com/IndexEndpoint"
    }
  ];

  // The ID of the DeployedIndex that will serve the request.
  string deployed_index_id = 2;

  // IDs of the datapoints to be searched for.
  repeated string ids = 3;
}

// The response message for
// [MatchService.ReadIndexDatapoints][google.cloud.aiplatform.v1beta1.MatchService.ReadIndexDatapoints].
message ReadIndexDatapointsResponse {
  // The result list of datapoints.
  repeated IndexDatapoint datapoints = 1;
}