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

Google APIs · copybara-github · commit 5b04a07f729c · 2025-03-27T12:00:29.000-07:00
PiperOrigin-RevId: 741233157
diff --git a/google/cloud/aiplatform/v1/content.proto b/google/cloud/aiplatform/v1/content.proto
@@ -19,6 +19,7 @@ package google.cloud.aiplatform.v1;
 import "google/api/field_behavior.proto";
 import "google/cloud/aiplatform/v1/openapi.proto";
 import "google/cloud/aiplatform/v1/tool.proto";
+import "google/cloud/aiplatform/v1/vertex_rag_data.proto";
 import "google/protobuf/duration.proto";
 import "google/type/date.proto";
 
@@ -538,6 +539,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/v1/vertex_rag_data.proto b/google/cloud/aiplatform/v1/vertex_rag_data.proto
@@ -186,6 +186,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];
 
@@ -207,22 +223,6 @@ 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
-    ];
-  }
 }
 
 // A RagFile contains user data for chunking, embedding and indexing.
@@ -284,6 +284,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.
diff --git a/google/cloud/aiplatform/v1/vertex_rag_service.proto b/google/cloud/aiplatform/v1/vertex_rag_service.proto
@@ -22,6 +22,7 @@ import "google/api/field_behavior.proto";
 import "google/api/resource.proto";
 import "google/cloud/aiplatform/v1/content.proto";
 import "google/cloud/aiplatform/v1/tool.proto";
+import "google/cloud/aiplatform/v1/vertex_rag_data.proto";
 
 option csharp_namespace = "Google.Cloud.AIPlatform.V1";
 option go_package = "cloud.google.com/go/aiplatform/apiv1/aiplatformpb;aiplatformpb";
@@ -164,6 +165,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.
@@ -293,6 +297,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.
