feat: add Vertex RAG service proto to v1

PiperOrigin-RevId: 698990917
googleapis · Nov 22, 2024 · 6f7652b · 6f7652b
1 parent 55c0c09
commit 6f7652b
Show file tree

Hide file tree

Showing 8 changed files with 1,375 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/v1/BUILD.bazel b/google/cloud/aiplatform/v1/BUILD.bazel
@@ -31,6 +31,7 @@ proto_library(
         "accelerator_type.proto",
         "annotation.proto",
         "annotation_spec.proto",
+        "api_auth.proto",
         "artifact.proto",
         "batch_prediction_job.proto",
         "completion_stats.proto",
@@ -136,6 +137,9 @@ proto_library(
         "unmanaged_container_model.proto",
         "user_action_reference.proto",
         "value.proto",
+        "vertex_rag_data.proto",
+        "vertex_rag_data_service.proto",
+        "vertex_rag_service.proto",
         "vizier_service.proto",
     ],
     deps = [

diff --git a/google/cloud/aiplatform/v1/aiplatform_v1.yaml b/google/cloud/aiplatform/v1/aiplatform_v1.yaml
@@ -312,6 +312,8 @@ http:
     - post: '/v1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:cancel'
     - post: '/v1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:cancel'
     - post: '/v1/{name=projects/*/locations/*/persistentResources/*/operations/*}:cancel'
+    - post: '/v1/{name=projects/*/locations/*/ragCorpora/*/operations/*}:cancel'
+    - post: '/v1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}:cancel'
     - post: '/v1/{name=projects/*/locations/*/studies/*/operations/*}:cancel'
     - post: '/v1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:cancel'
     - post: '/v1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:cancel'
@@ -399,6 +401,8 @@ http:
     - delete: '/v1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}'
     - delete: '/v1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}'
     - delete: '/v1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}'
+    - delete: '/v1/{name=projects/*/locations/*/ragCorpora/*/operations/*}'
+    - delete: '/v1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}'
     - delete: '/v1/{name=projects/*/locations/*/studies/*/operations/*}'
     - delete: '/v1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}'
     - delete: '/v1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}'
@@ -494,6 +498,8 @@ http:
     - get: '/v1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}'
     - get: '/v1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}'
     - get: '/v1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}'
+    - get: '/v1/{name=projects/*/locations/*/ragCorpora/*/operations/*}'
+    - get: '/v1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}'
     - get: '/v1/{name=projects/*/locations/*/studies/*/operations/*}'
     - get: '/v1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}'
     - get: '/v1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}'
@@ -593,6 +599,8 @@ http:
     - get: '/v1/{name=projects/*/locations/*/trainingPipelines/*}/operations'
     - get: '/v1/{name=projects/*/locations/*/persistentResources/*}/operations'
     - get: '/v1/{name=projects/*/locations/*/pipelineJobs/*}/operations'
+    - get: '/v1/{name=projects/*/locations/*/ragCorpora/*}/operations'
+    - get: '/v1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*}/operations'
     - get: '/v1/{name=projects/*/locations/*/schedules/*}/operations'
     - get: '/v1/{name=projects/*/locations/*/specialistPools/*}/operations'
     - get: '/v1/{name=projects/*/locations/*/tensorboards/*}/operations'
@@ -681,6 +689,8 @@ http:
     - post: '/v1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:wait'
     - post: '/v1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:wait'
     - post: '/v1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:wait'
+    - post: '/v1/{name=projects/*/locations/*/ragCorpora/*/operations/*}:wait'
+    - post: '/v1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}:wait'
     - post: '/v1/{name=projects/*/locations/*/studies/*/operations/*}:wait'
     - post: '/v1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:wait'
     - post: '/v1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:wait'
@@ -875,6 +885,14 @@ authentication:
       canonical_scopes: |-
         https://www.googleapis.com/auth/cloud-platform,
         https://www.googleapis.com/auth/cloud-platform.read-only
+  - selector: 'google.cloud.aiplatform.v1.VertexRagDataService.*'
+    oauth:
+      canonical_scopes: |-
+        https://www.googleapis.com/auth/cloud-platform
+  - selector: 'google.cloud.aiplatform.v1.VertexRagService.*'
+    oauth:
+      canonical_scopes: |-
+        https://www.googleapis.com/auth/cloud-platform
   - selector: 'google.cloud.aiplatform.v1.VizierService.*'
     oauth:
       canonical_scopes: |-

diff --git a/google/cloud/aiplatform/v1/api_auth.proto b/google/cloud/aiplatform/v1/api_auth.proto
@@ -0,0 +1,53 @@
+// 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.v1;
+
+import "google/api/field_behavior.proto";
+import "google/api/resource.proto";
+
+option csharp_namespace = "Google.Cloud.AIPlatform.V1";
+option go_package = "cloud.google.com/go/aiplatform/apiv1/aiplatformpb;aiplatformpb";
+option java_multiple_files = true;
+option java_outer_classname = "ApiAuthProto";
+option java_package = "com.google.cloud.aiplatform.v1";
+option php_namespace = "Google\\Cloud\\AIPlatform\\V1";
+option ruby_package = "Google::Cloud::AIPlatform::V1";
+option (google.api.resource_definition) = {
+  type: "secretmanager.googleapis.com/SecretVersion"
+  pattern: "projects/{project}/secrets/{secret}/versions/{secret_version}"
+};
+
+// The generic reusable api auth config.
+message ApiAuth {
+  // The API secret.
+  message ApiKeyConfig {
+    // Required. The SecretManager secret version resource name storing API key.
+    // e.g. projects/{project}/secrets/{secret}/versions/{version}
+    string api_key_secret_version = 1 [
+      (google.api.field_behavior) = REQUIRED,
+      (google.api.resource_reference) = {
+        type: "secretmanager.googleapis.com/SecretVersion"
+      }
+    ];
+  }
+
+  // The auth config.
+  oneof auth_config {
+    // The API secret.
+    ApiKeyConfig api_key_config = 1;
+  }
+}
diff --git a/google/cloud/aiplatform/v1/io.proto b/google/cloud/aiplatform/v1/io.proto
@@ -17,6 +17,8 @@ syntax = "proto3";
 package google.cloud.aiplatform.v1;
 
 import "google/api/field_behavior.proto";
+import "google/cloud/aiplatform/v1/api_auth.proto";
+import "google/protobuf/timestamp.proto";
 
 option csharp_namespace = "Google.Cloud.AIPlatform.V1";
 option go_package = "cloud.google.com/go/aiplatform/apiv1/aiplatformpb;aiplatformpb";
@@ -107,3 +109,142 @@ message ContainerRegistryDestination {
   // If a tag is not specified, "latest" will be used as the default tag.
   string output_uri = 1 [(google.api.field_behavior) = REQUIRED];
 }
+
+// The Google Drive location for the input content.
+message GoogleDriveSource {
+  // The type and ID of the Google Drive resource.
+  message ResourceId {
+    // The type of the Google Drive resource.
+    enum ResourceType {
+      // Unspecified resource type.
+      RESOURCE_TYPE_UNSPECIFIED = 0;
+
+      // File resource type.
+      RESOURCE_TYPE_FILE = 1;
+
+      // Folder resource type.
+      RESOURCE_TYPE_FOLDER = 2;
+    }
+
+    // Required. The type of the Google Drive resource.
+    ResourceType resource_type = 1 [(google.api.field_behavior) = REQUIRED];
+
+    // Required. The ID of the Google Drive resource.
+    string resource_id = 2 [(google.api.field_behavior) = REQUIRED];
+  }
+
+  // Required. Google Drive resource IDs.
+  repeated ResourceId resource_ids = 1 [(google.api.field_behavior) = REQUIRED];
+}
+
+// The input content is encapsulated and uploaded in the request.
+message DirectUploadSource {}
+
+// The Slack source for the ImportRagFilesRequest.
+message SlackSource {
+  // SlackChannels contains the Slack channels and corresponding access token.
+  message SlackChannels {
+    // SlackChannel contains the Slack channel ID and the time range to import.
+    message SlackChannel {
+      // Required. The Slack channel ID.
+      string channel_id = 1 [(google.api.field_behavior) = REQUIRED];
+
+      // Optional. The starting timestamp for messages to import.
+      google.protobuf.Timestamp start_time = 2
+          [(google.api.field_behavior) = OPTIONAL];
+
+      // Optional. The ending timestamp for messages to import.
+      google.protobuf.Timestamp end_time = 3
+          [(google.api.field_behavior) = OPTIONAL];
+    }
+
+    // Required. The Slack channel IDs.
+    repeated SlackChannel channels = 1 [(google.api.field_behavior) = REQUIRED];
+
+    // Required. The SecretManager secret version resource name (e.g.
+    // projects/{project}/secrets/{secret}/versions/{version}) storing the
+    // Slack channel access token that has access to the slack channel IDs.
+    // See: https://api.slack.com/tutorials/tracks/getting-a-token.
+    ApiAuth.ApiKeyConfig api_key_config = 3
+        [(google.api.field_behavior) = REQUIRED];
+  }
+
+  // Required. The Slack channels.
+  repeated SlackChannels channels = 1 [(google.api.field_behavior) = REQUIRED];
+}
+
+// The Jira source for the ImportRagFilesRequest.
+message JiraSource {
+  // JiraQueries contains the Jira queries and corresponding authentication.
+  message JiraQueries {
+    // A list of Jira projects to import in their entirety.
+    repeated string projects = 3;
+
+    // A list of custom Jira queries to import. For information about JQL (Jira
+    // Query Language), see
+    // https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/
+    repeated string custom_queries = 4;
+
+    // Required. The Jira email address.
+    string email = 5 [(google.api.field_behavior) = REQUIRED];
+
+    // Required. The Jira server URI.
+    string server_uri = 6 [(google.api.field_behavior) = REQUIRED];
+
+    // Required. The SecretManager secret version resource name (e.g.
+    // projects/{project}/secrets/{secret}/versions/{version}) storing the
+    // Jira API key. See [Manage API tokens for your Atlassian
+    // account](https://support.atlassian.com/atlassian-account/docs/manage-api-tokens-for-your-atlassian-account/).
+    ApiAuth.ApiKeyConfig api_key_config = 7
+        [(google.api.field_behavior) = REQUIRED];
+  }
+
+  // Required. The Jira queries.
+  repeated JiraQueries jira_queries = 1
+      [(google.api.field_behavior) = REQUIRED];
+}
+
+// The SharePointSources to pass to ImportRagFiles.
+message SharePointSources {
+  // An individual SharePointSource.
+  message SharePointSource {
+    // The SharePoint folder source. If not provided, uses "root".
+    oneof folder_source {
+      // The path of the SharePoint folder to download from.
+      string sharepoint_folder_path = 5;
+
+      // The ID of the SharePoint folder to download from.
+      string sharepoint_folder_id = 6;
+    }
+
+    // The SharePoint drive source.
+    oneof drive_source {
+      // The name of the drive to download from.
+      string drive_name = 7;
+
+      // The ID of the drive to download from.
+      string drive_id = 8;
+    }
+
+    // The Application ID for the app registered in Microsoft Azure Portal.
+    // The application must also be configured with MS Graph permissions
+    // "Files.ReadAll", "Sites.ReadAll" and BrowserSiteLists.Read.All.
+    string client_id = 1;
+
+    // The application secret for the app registered in Azure.
+    ApiAuth.ApiKeyConfig client_secret = 2;
+
+    // Unique identifier of the Azure Active Directory Instance.
+    string tenant_id = 3;
+
+    // The name of the SharePoint site to download from. This can be the site
+    // name or the site id.
+    string sharepoint_site_name = 4;
+
+    // Output only. The SharePoint file id. Output only.
+    string file_id = 9 [(google.api.field_behavior) = OUTPUT_ONLY];
+  }
+
+  // The SharePoint sources.
+  repeated SharePointSource share_point_sources = 1;
+}
diff --git a/google/cloud/aiplatform/v1/tool.proto b/google/cloud/aiplatform/v1/tool.proto
@@ -133,13 +133,57 @@ message Retrieval {
   oneof source {
     // Set to use data source powered by Vertex AI Search.
     VertexAISearch vertex_ai_search = 2;
+
+    // Set to use data source powered by Vertex RAG store.
+    // User data is uploaded via the VertexRagDataService.
+    VertexRagStore vertex_rag_store = 4;
   }
 
   // Optional. Deprecated. This option is no longer supported.
   bool disable_attribution = 3
       [deprecated = true, (google.api.field_behavior) = OPTIONAL];
 }
 
+// Retrieve from Vertex RAG Store for grounding.
+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. 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 = 4
+      [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. Number of top k results to return from the selected corpora.
+  optional int32 similarity_top_k = 2
+      [deprecated = true, (google.api.field_behavior) = OPTIONAL];
+
+  // Optional. Only return results with vector distance smaller than the
+  // threshold.
+  optional double vector_distance_threshold = 3
+      [deprecated = true, (google.api.field_behavior) = OPTIONAL];
+
+  // Optional. The retrieval config for the Rag query.
+  RagRetrievalConfig rag_retrieval_config = 6
+      [(google.api.field_behavior) = OPTIONAL];
+}
+
 // Retrieve from Vertex AI Search datastore for grounding.
 // See https://cloud.google.com/products/agent-builder
 message VertexAISearch {
@@ -212,3 +256,32 @@ message FunctionCallingConfig {
   repeated string allowed_function_names = 2
       [(google.api.field_behavior) = OPTIONAL];
 }
+
+// Specifies the context retrieval config.
+message RagRetrievalConfig {
+  // Config for filters.
+  message Filter {
+    // Filter contexts retrieved from the vector DB based on either vector
+    // distance or vector similarity.
+    oneof vector_db_threshold {
+      // Optional. Only returns contexts with vector distance smaller than the
+      // threshold.
+      double vector_distance_threshold = 3
+          [(google.api.field_behavior) = OPTIONAL];
+
+      // Optional. Only returns contexts with vector similarity larger than the
+      // threshold.
+      double vector_similarity_threshold = 4
+          [(google.api.field_behavior) = OPTIONAL];
+    }
+
+    // Optional. String for metadata filtering.
+    string metadata_filter = 2 [(google.api.field_behavior) = OPTIONAL];
+  }
+
+  // Optional. The number of contexts to retrieve.
+  int32 top_k = 1 [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. Config for filters.
+  Filter filter = 3 [(google.api.field_behavior) = OPTIONAL];
+}