feat: add page spans in retrieved contexts from Vertex RAG Engine in aiplatform v1beta1

Google APIs · copybara-github · commit d427057ed9dd · 2025-03-27T13:47:22.000-07:00
docs: A comment for field `rag_files_count` in message `.google.cloud.aiplatform.v1beta1.RagCorpus` is changed
docs: A comment for field `model_name` in message `.google.cloud.aiplatform.v1beta1.RagFileParsingConfig` is changed
PiperOrigin-RevId: 741268484
diff --git a/google/cloud/aiplatform/v1beta1/content.proto b/google/cloud/aiplatform/v1beta1/content.proto
@@ -20,6 +20,7 @@ import "google/api/field_behavior.proto";
 import "google/api/resource.proto";
 import "google/cloud/aiplatform/v1beta1/openapi.proto";
 import "google/cloud/aiplatform/v1beta1/tool.proto";
+import "google/cloud/aiplatform/v1beta1/vertex_rag_data.proto";
 import "google/protobuf/duration.proto";
 import "google/type/date.proto";
 
@@ -631,6 +632,13 @@ message GroundingChunk {
 
   // Chunk from context retrieved by the retrieval tools.
   message RetrievedContext {
+    // Tool-specific details about the retrieved context.
+    oneof context_details {
+      // Additional context for the RAG retrieval result. This is only populated
+      // when using the RAG retrieval tool.
+      RagChunk rag_chunk = 4;
+    }
+
     // URI reference of the attribution.
     optional string uri = 1;
 
diff --git a/google/cloud/aiplatform/v1beta1/vertex_rag_data.proto b/google/cloud/aiplatform/v1beta1/vertex_rag_data.proto
@@ -254,6 +254,22 @@ message RagCorpus {
     singular: "ragCorpus"
   };
 
+  // The backend config of the RagCorpus.
+  // It can be data store and/or retrieval engine.
+  oneof backend_config {
+    // Optional. Immutable. The config for the Vector DBs.
+    RagVectorDbConfig vector_db_config = 9 [
+      (google.api.field_behavior) = OPTIONAL,
+      (google.api.field_behavior) = IMMUTABLE
+    ];
+
+    // Optional. Immutable. The config for the Vertex AI Search.
+    VertexAiSearchConfig vertex_ai_search_config = 10 [
+      (google.api.field_behavior) = OPTIONAL,
+      (google.api.field_behavior) = IMMUTABLE
+    ];
+  }
+
   // Output only. The resource name of the RagCorpus.
   string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
 
@@ -290,23 +306,8 @@ message RagCorpus {
   // Output only. RagCorpus state.
   CorpusStatus corpus_status = 8 [(google.api.field_behavior) = OUTPUT_ONLY];
 
-  // The backend config of the RagCorpus.
-  // It can be data store and/or retrieval engine.
-  oneof backend_config {
-    // Optional. Immutable. The config for the Vector DBs.
-    RagVectorDbConfig vector_db_config = 9 [
-      (google.api.field_behavior) = OPTIONAL,
-      (google.api.field_behavior) = IMMUTABLE
-    ];
-
-    // Optional. Immutable. The config for the Vertex AI Search.
-    VertexAiSearchConfig vertex_ai_search_config = 10 [
-      (google.api.field_behavior) = OPTIONAL,
-      (google.api.field_behavior) = IMMUTABLE
-    ];
-  }
-  // Output only. The number of RagFiles in the RagCorpus.
-  int32 rag_files_count = 11;
+  // Output only. Number of RagFiles in the RagCorpus.
+  int32 rag_files_count = 11 [(google.api.field_behavior) = OUTPUT_ONLY];
 }
 
 // A RagFile contains user data for chunking, embedding and indexing.
@@ -386,6 +387,25 @@ message RagFile {
   FileStatus file_status = 13 [(google.api.field_behavior) = OUTPUT_ONLY];
 }
 
+// A RagChunk includes the content of a chunk of a RagFile, and associated
+// metadata.
+message RagChunk {
+  // Represents where the chunk starts and ends in the document.
+  message PageSpan {
+    // Page where chunk starts in the document. Inclusive. 1-indexed.
+    int32 first_page = 1;
+
+    // Page where chunk ends in the document. Inclusive. 1-indexed.
+    int32 last_page = 2;
+  }
+
+  // The content of the chunk.
+  string text = 1;
+
+  // If populated, represents where the chunk starts and ends in the document.
+  optional PageSpan page_span = 2;
+}
+
 // Specifies the size and overlap of chunks for RagFiles.
 message RagFileChunkingConfig {
   // Specifies the fixed length chunking config.
@@ -445,7 +465,8 @@ message RagFileParsingConfig {
   // Specifies the advanced parsing for RagFiles.
   message LlmParser {
     // The name of a LLM model used for parsing.
-    // Format: `gemini-1.5-pro-002`
+    // Format:
+    // * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`
     string model_name = 1;
 
     // The maximum number of requests the job is allowed to make to the
diff --git a/google/cloud/aiplatform/v1beta1/vertex_rag_service.proto b/google/cloud/aiplatform/v1beta1/vertex_rag_service.proto
@@ -22,6 +22,7 @@ import "google/api/field_behavior.proto";
 import "google/api/resource.proto";
 import "google/cloud/aiplatform/v1beta1/content.proto";
 import "google/cloud/aiplatform/v1beta1/tool.proto";
+import "google/cloud/aiplatform/v1beta1/vertex_rag_data.proto";
 
 option csharp_namespace = "Google.Cloud.AIPlatform.V1Beta1";
 option go_package = "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb;aiplatformpb";
@@ -194,6 +195,9 @@ message RagContexts {
     // less relevant the context is to the query. The range is [0, 2], while 0
     // means the most relevant and 2 means the least relevant.
     optional double score = 6;
+
+    // Context of the retrieved chunk.
+    RagChunk chunk = 7;
   }
 
   // All its contexts.
@@ -323,6 +327,9 @@ message Fact {
   // relevant the fact is to the query. The range is [0, 2], while 0 means the
   // most relevant and 2 means the least relevant.
   optional double score = 6;
+
+  // If present, chunk properties.
+  optional RagChunk chunk = 7;
 }
 
 // Claim that is extracted from the input text and facts that support it.
