fix: Generalize RAG files import from Google Drive

PiperOrigin-RevId: 639184120
googleapis · May 31, 2024 · 88c6a6a · 88c6a6a
1 parent ba65828
commit 88c6a6a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/tests/unit/vertex_rag/test_rag_constants.py b/tests/unit/vertex_rag/test_rag_constants.py
@@ -91,6 +91,9 @@
 TEST_DRIVE_FOLDER = (
     f"https://drive.google.com/corp/drive/folders/{TEST_DRIVE_FOLDER_ID}"
 )
+TEST_DRIVE_FOLDER_2 = (
+    f"https://drive.google.com/drive/folders/{TEST_DRIVE_FOLDER_ID}?resourcekey=0-eiOT3"
+)
 TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER = ImportRagFilesConfig()
 TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER.google_drive_source.resource_ids = [
     GoogleDriveSource.ResourceId(

diff --git a/tests/unit/vertex_rag/test_rag_data.py b/tests/unit/vertex_rag/test_rag_data.py
@@ -374,11 +374,11 @@ def test_prepare_import_files_request_list_gcs_uris(self):
         )
         import_files_request_eq(request, tc.TEST_IMPORT_REQUEST_GCS)
 
-    def test_prepare_import_files_request_drive_folders(self):
-        paths = [tc.TEST_DRIVE_FOLDER]
+    @pytest.mark.parametrize("path", [tc.TEST_DRIVE_FOLDER, tc.TEST_DRIVE_FOLDER_2])
+    def test_prepare_import_files_request_drive_folders(self, path):
         request = prepare_import_files_request(
             corpus_name=tc.TEST_RAG_CORPUS_RESOURCE_NAME,
-            paths=paths,
+            paths=[path],
             chunk_size=tc.TEST_CHUNK_SIZE,
             chunk_overlap=tc.TEST_CHUNK_OVERLAP,
         )

diff --git a/vertexai/preview/rag/utils/_gapic_utils.py b/vertexai/preview/rag/utils/_gapic_utils.py
@@ -97,10 +97,12 @@ def convert_path_to_resource_id(
         # Google Drive source
         path_list = path.split("/")
         if "file" in path_list:
-            resource_id = path_list[5]
+            index = path_list.index("file") + 2
+            resource_id = path_list[index].split("?")[0]
             resource_type = GoogleDriveSource.ResourceId.ResourceType.RESOURCE_TYPE_FILE
         elif "folders" in path_list:
-            resource_id = path_list[6]
+            index = path_list.index("folders") + 1
+            resource_id = path_list[index].split("?")[0]
             resource_type = (
                 GoogleDriveSource.ResourceId.ResourceType.RESOURCE_TYPE_FOLDER
             )