// 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.discoveryengine.v1beta;

import "google/api/annotations.proto";
import "google/api/client.proto";
import "google/api/field_behavior.proto";
import "google/api/resource.proto";
import "google/cloud/discoveryengine/v1beta/import_config.proto";
import "google/longrunning/operations.proto";
import "google/protobuf/timestamp.proto";
import "google/rpc/status.proto";

option csharp_namespace = "Google.Cloud.DiscoveryEngine.V1Beta";
option go_package = "cloud.google.com/go/discoveryengine/apiv1beta/discoveryenginepb;discoveryenginepb";
option java_multiple_files = true;
option java_outer_classname = "SearchTuningServiceProto";
option java_package = "com.google.cloud.discoveryengine.v1beta";
option objc_class_prefix = "DISCOVERYENGINE";
option php_namespace = "Google\\Cloud\\DiscoveryEngine\\V1beta";
option ruby_package = "Google::Cloud::DiscoveryEngine::V1beta";

// Service for search tuning.
service SearchTuningService {
  option (google.api.default_host) = "discoveryengine.googleapis.com";
  option (google.api.oauth_scopes) =
      "https://www.googleapis.com/auth/cloud-platform";

  // Trains a custom model.
  rpc TrainCustomModel(TrainCustomModelRequest)
      returns (google.longrunning.Operation) {
    option (google.api.http) = {
      post: "/v1beta/{data_store=projects/*/locations/*/collections/*/dataStores/*}:trainCustomModel"
      body: "*"
    };
    option (google.longrunning.operation_info) = {
      response_type: "google.cloud.discoveryengine.v1beta.TrainCustomModelResponse"
      metadata_type: "google.cloud.discoveryengine.v1beta.TrainCustomModelMetadata"
    };
  }
}

// Request message for
// [SearchTuningService.TrainCustomModel][google.cloud.discoveryengine.v1beta.SearchTuningService.TrainCustomModel]
// method.
message TrainCustomModelRequest {
  // Cloud Storage training data input.
  message GcsTrainingInput {
    // The Cloud Storage corpus data which could be associated in train data.
    // The data path format is `gs://<bucket_to_data>/<jsonl_file_name>`.
    // A newline delimited jsonl/ndjson file.
    //
    // For search-tuning model, each line should have the _id, title
    // and text. Example:
    // `{"_id": "doc1", title: "relevant doc", "text": "relevant text"}`
    string corpus_data_path = 1;

    // The gcs query data which could be associated in train data.
    // The data path format is `gs://<bucket_to_data>/<jsonl_file_name>`.
    // A newline delimited jsonl/ndjson file.
    //
    // For search-tuning model, each line should have the _id
    // and text. Example: {"_id": "query1",  "text": "example query"}
    string query_data_path = 2;

    // Cloud Storage training data path whose format should be
    // `gs://<bucket_to_data>/<tsv_file_name>`. The file should be in tsv
    // format. Each line should have the doc_id and query_id and score (number).
    //
    // For search-tuning model, it should have the query-id corpus-id
    // score as tsv file header. The score should be a number in `[0, inf+)`.
    // The larger the number is, the more relevant the pair is. Example:
    //
    // * `query-id\tcorpus-id\tscore`
    // * `query1\tdoc1\t1`
    string train_data_path = 3;

    // Cloud Storage test data. Same format as train_data_path. If not provided,
    // a random 80/20 train/test split will be performed on train_data_path.
    string test_data_path = 4;
  }

  // Model training input.
  oneof training_input {
    // Cloud Storage training input.
    GcsTrainingInput gcs_training_input = 2;
  }

  // Required. The resource name of the Data Store, such as
  // `projects/*/locations/global/collections/default_collection/dataStores/default_data_store`.
  // This field is used to identify the data store where to train the models.
  string data_store = 1 [
    (google.api.field_behavior) = REQUIRED,
    (google.api.resource_reference) = {
      type: "discoveryengine.googleapis.com/DataStore"
    }
  ];

  // Model to be trained. Supported values are:
  //
  //  * **search-tuning**: Fine tuning the search system based on data provided.
  string model_type = 3;

  // The desired location of errors incurred during the data ingestion and
  // training.
  ImportErrorConfig error_config = 4;
}

// Response of the
// [TrainCustomModelRequest][google.cloud.discoveryengine.v1beta.TrainCustomModelRequest].
// This message is returned by the google.longrunning.Operations.response field.
message TrainCustomModelResponse {
  // A sample of errors encountered while processing the data.
  repeated google.rpc.Status error_samples = 1;

  // Echoes the destination for the complete errors in the request if set.
  ImportErrorConfig error_config = 2;

  // The trained model status. Possible values are:
  //
  //  * **bad-data**: The training data quality is bad.
  //  * **no-improvement**: Tuning didn't improve performance. Won't deploy.
  //  * **in-progress**: Model training job creation is in progress.
  //  * **training**: Model is actively training.
  //  * **evaluating**: The model is evaluating trained metrics.
  //  * **indexing**: The model trained metrics are indexing.
  //  * **ready**: The model is ready for serving.
  string model_status = 3;

  // The metrics of the trained model.
  map<string, double> metrics = 4;
}

// Metadata related to the progress of the TrainCustomModel operation. This is
// returned by the google.longrunning.Operation.metadata field.
message TrainCustomModelMetadata {
  // Operation create time.
  google.protobuf.Timestamp create_time = 1;

  // Operation last update time. If the operation is done, this is also the
  // finish time.
  google.protobuf.Timestamp update_time = 2;
}