feat: Export global quota configs in preview sdk

PiperOrigin-RevId: 770276336

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants_preview.py

@@ -485,10 +485,12 @@
 TEST_IMPORT_FILES_CONFIG_GCS = ImportRagFilesConfig(
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
     rebuild_ann_index=False,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_GCS_REBUILD_ANN_INDEX = ImportRagFilesConfig(
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
     rebuild_ann_index=True,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_GCS_REBUILD_ANN_INDEX.gcs_source.uris = [TEST_GCS_PATH]
 TEST_IMPORT_FILES_CONFIG_GCS_REBUILD_ANN_INDEX.rag_file_parsing_config.advanced_parser.use_advanced_pdf_parsing = (
@@ -517,6 +519,7 @@
 TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER = ImportRagFilesConfig(
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
     rebuild_ann_index=False,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER.google_drive_source.resource_ids = [
     GoogleDriveSource.ResourceId(
@@ -530,6 +533,7 @@
 TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING = ImportRagFilesConfig(
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
     rebuild_ann_index=False,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING.google_drive_source.resource_ids = [
     GoogleDriveSource.ResourceId(
@@ -589,6 +593,7 @@
         )
     ),
     rebuild_ann_index=False,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_DRIVE_FILE.max_embedding_requests_per_min = 800
 
@@ -603,6 +608,14 @@
     import_rag_files_config=TEST_IMPORT_FILES_CONFIG_DRIVE_FILE,
 )
 
+TEST_IMPORT_REQUEST_DRIVE_FILE_GLOBAL_QUOTA_CONTROL = ImportRagFilesRequest(
+    parent=TEST_RAG_CORPUS_RESOURCE_NAME,
+    import_rag_files_config=TEST_IMPORT_FILES_CONFIG_DRIVE_FILE,
+)
+TEST_IMPORT_REQUEST_DRIVE_FILE_GLOBAL_QUOTA_CONTROL.import_rag_files_config.global_max_embedding_requests_per_min = (
+    8000
+)
+
 TEST_IMPORT_RESPONSE = ImportRagFilesResponse(imported_rag_files_count=2)
 
 TEST_GAPIC_RAG_FILE = GapicRagFile(
@@ -649,6 +662,7 @@
     rag_file_parsing_config=TEST_RAG_FILE_PARSING_CONFIG,
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
     rebuild_ann_index=False,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_SLACK_SOURCE.slack_source.channels = [
     GapicSlackSource.SlackChannels(
@@ -703,6 +717,7 @@
     rag_file_parsing_config=TEST_RAG_FILE_PARSING_CONFIG,
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
     rebuild_ann_index=False,
+    max_embedding_requests_per_min=1000,
 )
 TEST_IMPORT_FILES_CONFIG_JIRA_SOURCE.jira_source.jira_queries = [
     GapicJiraSource.JiraQueries(
@@ -736,6 +751,7 @@
 TEST_IMPORT_FILES_CONFIG_SHARE_POINT_SOURCE = ImportRagFilesConfig(
     rag_file_parsing_config=TEST_RAG_FILE_PARSING_CONFIG,
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
+    max_embedding_requests_per_min=1000,
     share_point_sources=GapicSharePointSources(
         share_point_sources=[
             GapicSharePointSources.SharePointSource(
@@ -813,6 +829,7 @@
 TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG = LayoutParserConfig(
     processor_name="projects/test-project/locations/us/processors/abc123",
     max_parsing_requests_per_min=100,
+    global_max_parsing_requests_per_min=1000,
 )
 
 TEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH_CONFIG = LayoutParserConfig(
@@ -885,6 +902,7 @@
 
 TEST_IMPORT_FILES_CONFIG_SHARE_POINT_SOURCE_NO_FOLDERS = ImportRagFilesConfig(
     rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
+    max_embedding_requests_per_min=1000,
     share_point_sources=GapicSharePointSources(
         share_point_sources=[
             GapicSharePointSources.SharePointSource(
@@ -914,6 +932,7 @@
         layout_parser=RagFileParsingConfig.LayoutParser(
             processor_name="projects/test-project/locations/us/processors/abc123",
             max_parsing_requests_per_min=100,
+            global_max_parsing_requests_per_min=1000,
         )
     )
 )

tests/unit/vertex_rag/test_rag_data_preview.py

@@ -659,6 +659,14 @@ def import_files_request_eq(returned_request, expected_request):
         returned_request.import_rag_files_config.rebuild_ann_index
         == expected_request.import_rag_files_config.rebuild_ann_index
     )
+    assert (
+        returned_request.import_rag_files_config.max_embedding_requests_per_min
+        == expected_request.import_rag_files_config.max_embedding_requests_per_min
+    )
+    assert (
+        returned_request.import_rag_files_config.global_max_embedding_requests_per_min
+        == expected_request.import_rag_files_config.global_max_embedding_requests_per_min
+    )
 
 
 def rag_engine_config_eq(returned_config, expected_config):
@@ -1349,6 +1357,20 @@ def test_prepare_import_files_request_drive_files(self):
             request, test_rag_constants_preview.TEST_IMPORT_REQUEST_DRIVE_FILE
         )
 
+    def test_prepare_import_files_request_drive_files_with_global_quota_control(self):
+        paths = [test_rag_constants_preview.TEST_DRIVE_FILE]
+        request = prepare_import_files_request(
+            corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
+            paths=paths,
+            transformation_config=create_transformation_config(),
+            max_embedding_requests_per_min=800,
+            global_max_embedding_requests_per_min=8000,
+        )
+        import_files_request_eq(
+            request,
+            test_rag_constants_preview.TEST_IMPORT_REQUEST_DRIVE_FILE_GLOBAL_QUOTA_CONTROL,
+        )
+
     def test_prepare_import_files_request_invalid_drive_path(self):
         with pytest.raises(ValueError) as e:
             paths = ["https://drive.google.com/bslalsdfk/whichever_file/456"]

vertexai/preview/rag/rag_data.py

@@ -496,6 +496,7 @@ def import_files(
     transformation_config: Optional[TransformationConfig] = None,
     timeout: int = 600,
     max_embedding_requests_per_min: int = 1000,
+    global_max_embedding_requests_per_min: Optional[int] = None,
     use_advanced_pdf_parsing: Optional[bool] = False,
     partial_failures_sink: Optional[str] = None,
     layout_parser: Optional[LayoutParserConfig] = None,
@@ -605,6 +606,13 @@ def import_files(
             page on the project to set an appropriate value
             here. If unspecified, a default value of 1,000
             QPM would be used.
+        global_max_embedding_requests_per_min:
+            Optional. The max number of queries per minute that the indexing
+            pipeline job is allowed to make to the embedding model specified in
+            the project. Please follow the quota usage guideline of the embedding
+            model you use to set the value properly. If this value is not specified,
+            max_embedding_requests_per_min will be used by indexing pipeline job
+            as the global limit and this means parallel import jobs are not allowed.
         timeout: Default is 600 seconds.
         use_advanced_pdf_parsing: Whether to use advanced PDF
             parsing on uploaded files. This field is deprecated.
@@ -663,6 +671,7 @@ def import_files(
         chunk_overlap=chunk_overlap,
         transformation_config=transformation_config,
         max_embedding_requests_per_min=max_embedding_requests_per_min,
+        global_max_embedding_requests_per_min=global_max_embedding_requests_per_min,
         use_advanced_pdf_parsing=use_advanced_pdf_parsing,
         partial_failures_sink=partial_failures_sink,
         layout_parser=layout_parser,
@@ -686,6 +695,7 @@ async def import_files_async(
     chunk_overlap: int = 200,
     transformation_config: Optional[TransformationConfig] = None,
     max_embedding_requests_per_min: int = 1000,
+    global_max_embedding_requests_per_min: Optional[int] = None,
     use_advanced_pdf_parsing: Optional[bool] = False,
     partial_failures_sink: Optional[str] = None,
     layout_parser: Optional[LayoutParserConfig] = None,
@@ -796,6 +806,13 @@ 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.
+        global_max_embedding_requests_per_min:
+            Optional. The max number of queries per minute that the indexing
+            pipeline job is allowed to make to the embedding model specified in
+            the project. Please follow the quota usage guideline of the embedding
+            model you use to set the value properly. If this value is not specified,
+            max_embedding_requests_per_min will be used by indexing pipeline job
+            as the global limit and this means parallel import jobs are not allowed.
         use_advanced_pdf_parsing: Whether to use advanced PDF
             parsing on uploaded files.
         partial_failures_sink: Either a GCS path to store partial failures or a
@@ -852,6 +869,7 @@ async def import_files_async(
         chunk_overlap=chunk_overlap,
         transformation_config=transformation_config,
         max_embedding_requests_per_min=max_embedding_requests_per_min,
+        global_max_embedding_requests_per_min=global_max_embedding_requests_per_min,
         use_advanced_pdf_parsing=use_advanced_pdf_parsing,
         partial_failures_sink=partial_failures_sink,
         layout_parser=layout_parser,

vertexai/preview/rag/utils/_gapic_utils.py

@@ -537,6 +537,7 @@ def prepare_import_files_request(
     chunk_overlap: int = 200,
     transformation_config: Optional[TransformationConfig] = None,
     max_embedding_requests_per_min: int = 1000,
+    global_max_embedding_requests_per_min: Optional[int] = None,
     use_advanced_pdf_parsing: bool = False,
     partial_failures_sink: Optional[str] = None,
     layout_parser: Optional[LayoutParserConfig] = None,
@@ -569,8 +570,15 @@ def prepare_import_files_request(
             )
         rag_file_parsing_config.layout_parser = RagFileParsingConfig.LayoutParser(
             processor_name=layout_parser.processor_name,
-            max_parsing_requests_per_min=layout_parser.max_parsing_requests_per_min,
         )
+        if layout_parser.max_parsing_requests_per_min is not None:
+            rag_file_parsing_config.layout_parser.max_parsing_requests_per_min = (
+                layout_parser.max_parsing_requests_per_min
+            )
+        if layout_parser.global_max_parsing_requests_per_min is not None:
+            rag_file_parsing_config.layout_parser.global_max_parsing_requests_per_min = (
+                layout_parser.global_max_parsing_requests_per_min
+            )
     if llm_parser is not None:
         rag_file_parsing_config.llm_parser = RagFileParsingConfig.LlmParser(
             model_name=llm_parser.model_name
@@ -609,6 +617,10 @@ def prepare_import_files_request(
         rebuild_ann_index=rebuild_ann_index,
     )
 
+    if global_max_embedding_requests_per_min is not None:
+        import_rag_files_config.global_max_embedding_requests_per_min = (
+            global_max_embedding_requests_per_min
+        )
     if source is not None:
         gapic_source = convert_source_for_rag_import(source)
         if isinstance(gapic_source, GapicSlackSource):

vertexai/preview/rag/utils/resources.py

@@ -515,10 +515,19 @@ class LayoutParserConfig:
             https://cloud.google.com/document-ai/quotas and the Quota page for
             your project to set an appropriate value here. If unspecified, a
             default value of 120 QPM will be used.
+        global_max_parsing_requests_per_min (int):
+            The maximum number of requests the job is allowed to make to
+            the Document AI processor per minute in this project.
+            Consult https://cloud.google.com/document-ai/quotas and the
+            Quota page for your project to set an appropriate value
+            here. If this value is not specified,
+            max_parsing_requests_per_min will be used by indexing
+            pipeline as the global limit.
     """
 
     processor_name: str
     max_parsing_requests_per_min: Optional[int] = None
+    global_max_parsing_requests_per_min: Optional[int] = None
 
 
 @dataclasses.dataclass