feat: Add advanced PDF parsing option for RAG file import

PiperOrigin-RevId: 663391146

Author: speedstorm1

tests/unit/vertex_rag/test_rag_constants.py

@@ -22,6 +22,7 @@
 from google.cloud.aiplatform_v1beta1 import (
     GoogleDriveSource,
     RagFileChunkingConfig,
+    RagFileParsingConfig,
     ImportRagFilesConfig,
     ImportRagFilesRequest,
     ImportRagFilesResponse,
@@ -93,6 +94,7 @@
 # GCS
 TEST_IMPORT_FILES_CONFIG_GCS = ImportRagFilesConfig()
 TEST_IMPORT_FILES_CONFIG_GCS.gcs_source.uris = [TEST_GCS_PATH]
+TEST_IMPORT_FILES_CONFIG_GCS.rag_file_parsing_config.use_advanced_pdf_parsing = False
 TEST_IMPORT_REQUEST_GCS = ImportRagFilesRequest(
     parent=TEST_RAG_CORPUS_RESOURCE_NAME,
     import_rag_files_config=TEST_IMPORT_FILES_CONFIG_GCS,
@@ -112,18 +114,36 @@
         resource_type=GoogleDriveSource.ResourceId.ResourceType.RESOURCE_TYPE_FOLDER,
     )
 ]
+TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER.rag_file_parsing_config.use_advanced_pdf_parsing = (
+    False
+)
+TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING = ImportRagFilesConfig()
+TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING.google_drive_source.resource_ids = [
+    GoogleDriveSource.ResourceId(
+        resource_id=TEST_DRIVE_FOLDER_ID,
+        resource_type=GoogleDriveSource.ResourceId.ResourceType.RESOURCE_TYPE_FOLDER,
+    )
+]
+TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING.rag_file_parsing_config.use_advanced_pdf_parsing = (
+    True
+)
 TEST_IMPORT_REQUEST_DRIVE_FOLDER = ImportRagFilesRequest(
     parent=TEST_RAG_CORPUS_RESOURCE_NAME,
     import_rag_files_config=TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER,
 )
+TEST_IMPORT_REQUEST_DRIVE_FOLDER_PARSING = ImportRagFilesRequest(
+    parent=TEST_RAG_CORPUS_RESOURCE_NAME,
+    import_rag_files_config=TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING,
+)
 # Google Drive files
 TEST_DRIVE_FILE_ID = "456"
 TEST_DRIVE_FILE = f"https://drive.google.com/file/d/{TEST_DRIVE_FILE_ID}"
 TEST_IMPORT_FILES_CONFIG_DRIVE_FILE = ImportRagFilesConfig(
     rag_file_chunking_config=RagFileChunkingConfig(
         chunk_size=TEST_CHUNK_SIZE,
         chunk_overlap=TEST_CHUNK_OVERLAP,
-    )
+    ),
+    rag_file_parsing_config=RagFileParsingConfig(use_advanced_pdf_parsing=False),
 )
 TEST_IMPORT_FILES_CONFIG_DRIVE_FILE.max_embedding_requests_per_min = 800
 

tests/unit/vertex_rag/test_rag_data.py

@@ -166,6 +166,10 @@ def import_files_request_eq(returned_request, expected_request):
         returned_request.import_rag_files_config.jira_source.jira_queries
         == expected_request.import_rag_files_config.jira_source.jira_queries
     )
+    assert (
+        returned_request.import_rag_files_config.rag_file_parsing_config
+        == expected_request.import_rag_files_config.rag_file_parsing_config
+    )
 
 
 @pytest.mark.usefixtures("google_auth_mock")
@@ -396,6 +400,17 @@ def test_prepare_import_files_request_drive_folders(self, path):
         )
         import_files_request_eq(request, tc.TEST_IMPORT_REQUEST_DRIVE_FOLDER)
 
+    @pytest.mark.parametrize("path", [tc.TEST_DRIVE_FOLDER, tc.TEST_DRIVE_FOLDER_2])
+    def test_prepare_import_files_request_drive_folders_with_pdf_parsing(self, path):
+        request = prepare_import_files_request(
+            corpus_name=tc.TEST_RAG_CORPUS_RESOURCE_NAME,
+            paths=[path],
+            chunk_size=tc.TEST_CHUNK_SIZE,
+            chunk_overlap=tc.TEST_CHUNK_OVERLAP,
+            use_advanced_pdf_parsing=True,
+        )
+        import_files_request_eq(request, tc.TEST_IMPORT_REQUEST_DRIVE_FOLDER_PARSING)
+
     def test_prepare_import_files_request_drive_files(self):
         paths = [tc.TEST_DRIVE_FILE]
         request = prepare_import_files_request(

vertexai/preview/rag/rag_data.py

@@ -281,6 +281,7 @@ def import_files(
     chunk_overlap: int = 200,
     timeout: int = 600,
     max_embedding_requests_per_min: int = 1000,
+    use_advanced_pdf_parsing: Optional[bool] = False,
 ) -> ImportRagFilesResponse:
     """
     Import files to an existing RagCorpus, wait until completion.
@@ -364,6 +365,8 @@ def import_files(
             here. If unspecified, a default value of 1,000
             QPM would be used.
         timeout: Default is 600 seconds.
+        use_advanced_pdf_parsing: Whether to use advanced PDF
+            parsing on uploaded files.
     Returns:
         ImportRagFilesResponse.
     """
@@ -379,6 +382,7 @@ def import_files(
         chunk_size=chunk_size,
         chunk_overlap=chunk_overlap,
         max_embedding_requests_per_min=max_embedding_requests_per_min,
+        use_advanced_pdf_parsing=use_advanced_pdf_parsing,
     )
     client = _gapic_utils.create_rag_data_service_client()
     try:
@@ -396,6 +400,7 @@ async def import_files_async(
     chunk_size: int = 1024,
     chunk_overlap: int = 200,
     max_embedding_requests_per_min: int = 1000,
+    use_advanced_pdf_parsing: Optional[bool] = False,
 ) -> operation_async.AsyncOperation:
     """
     Import files to an existing RagCorpus asynchronously.
@@ -479,6 +484,8 @@ async def import_files_async(
             page on the project to set an appropriate value
             here. If unspecified, a default value of 1,000
             QPM would be used.
+        use_advanced_pdf_parsing: Whether to use advanced PDF
+            parsing on uploaded files.
     Returns:
         operation_async.AsyncOperation.
     """
@@ -494,6 +501,7 @@ async def import_files_async(
         chunk_size=chunk_size,
         chunk_overlap=chunk_overlap,
         max_embedding_requests_per_min=max_embedding_requests_per_min,
+        use_advanced_pdf_parsing=use_advanced_pdf_parsing,
     )
     async_client = _gapic_utils.create_rag_data_service_async_client()
     try:

vertexai/preview/rag/utils/_gapic_utils.py

@@ -23,6 +23,7 @@
     ImportRagFilesConfig,
     ImportRagFilesRequest,
     RagFileChunkingConfig,
+    RagFileParsingConfig,
     RagCorpus as GapicRagCorpus,
     RagFile as GapicRagFile,
     SlackSource as GapicSlackSource,
@@ -217,19 +218,24 @@ def prepare_import_files_request(
     chunk_size: int = 1024,
     chunk_overlap: int = 200,
     max_embedding_requests_per_min: int = 1000,
+    use_advanced_pdf_parsing: bool = False,
 ) -> ImportRagFilesRequest:
     if len(corpus_name.split("/")) != 6:
         raise ValueError(
             "corpus_name must be of the format `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`"
         )
 
+    rag_file_parsing_config = RagFileParsingConfig(
+        use_advanced_pdf_parsing=use_advanced_pdf_parsing,
+    )
     rag_file_chunking_config = RagFileChunkingConfig(
         chunk_size=chunk_size,
         chunk_overlap=chunk_overlap,
     )
     import_rag_files_config = ImportRagFilesConfig(
         rag_file_chunking_config=rag_file_chunking_config,
         max_embedding_requests_per_min=max_embedding_requests_per_min,
+        rag_file_parsing_config=rag_file_parsing_config,
     )
 
     if source is not None: