feat: RAG - Add support for max_embedding_requests_per_min in upload_file request

Aligning the default to import_files and the backend of 1000

PiperOrigin-RevId: 797149402

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_data.py

@@ -756,6 +756,30 @@ def test_upload_file_success_with_transformation_config(
 
         rag_file_eq(rag_file, test_rag_constants.TEST_RAG_FILE)
 
+    @pytest.mark.usefixtures("open_file_mock")
+    def test_upload_file_success_with_max_embedding_requests_per_min(
+        self,
+        upload_file_with_upload_config_mock,
+    ):
+        aiplatform.init(
+            project=test_rag_constants.TEST_PROJECT,
+            location=test_rag_constants.TEST_REGION,
+        )
+        rag_file = rag.upload_file(
+            corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+            path=test_rag_constants.TEST_PATH,
+            display_name=test_rag_constants.TEST_FILE_DISPLAY_NAME,
+            transformation_config=create_transformation_config(),
+            max_embedding_requests_per_min=666,
+        )
+
+        upload_file_with_upload_config_mock.assert_called_once()
+        _, mock_kwargs = upload_file_with_upload_config_mock.call_args
+        assert mock_kwargs["url"] == test_rag_constants.TEST_UPLOAD_REQUEST_URI
+        assert mock_kwargs["headers"] == test_rag_constants.TEST_HEADERS
+
+        rag_file_eq(rag_file, test_rag_constants.TEST_RAG_FILE)
+
     @pytest.mark.usefixtures("rag_data_client_mock_exception", "open_file_mock")
     def test_upload_file_failure(self):
         with pytest.raises(RuntimeError) as e:

vertexai/rag/rag_data.py

@@ -343,9 +343,9 @@ def upload_file(
     description: Optional[str] = None,
     transformation_config: Optional[TransformationConfig] = None,
     timeout: int = 600,
+    max_embedding_requests_per_min: int = 1000,
 ) -> RagFile:
-    """
-    Synchronous file upload to an existing RagCorpus.
+    """Synchronous file upload to an existing RagCorpus.
 
     Example usage:
 
@@ -375,12 +375,17 @@ def upload_file(
         corpus_name: The name of the RagCorpus resource into which to upload the file.
             Format: ``projects/{project}/locations/{location}/ragCorpora/{rag_corpus}``
             or ``{rag_corpus}``.
-        path: A local file path. For example,
-            "usr/home/my_file.txt".
+        path: A local file path. For example, "usr/home/my_file.txt".
         display_name: The display name of the data file.
         description: The description of the RagFile.
         transformation_config: The config for transforming the RagFile, like chunking.
         timeout: Default is 600 seconds.
+        max_embedding_requests_per_min: Optional. The max number of queries per
+            minute that this job is allowed to make to the embedding model specified
+            on the corpus. This value is specific to this job and not shared across
+            other import jobs. Consult the Quotas page on the project to set an
+            appropriate value here. If unspecified, a default value of 1,000 QPM
+            would be used.
 
     Returns:
         RagFile.
@@ -421,7 +426,8 @@ def upload_file(
                         "chunk_overlap": chunk_overlap,
                     }
                 }
-            }
+            },
+            "max_embedding_requests_per_min": max_embedding_requests_per_min,
         }
 
     files = {