feat: Support RagFileTransformationConfig in upload_file function in V1 sdk

PiperOrigin-RevId: 700756555

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants.py

@@ -150,15 +150,31 @@
     TEST_RAG_CORPUS_RESOURCE_NAME + f"/ragFiles/{TEST_RAG_FILE_ID}"
 )
 TEST_UPLOAD_RAG_FILE_RESPONSE_CONTENT = ""
+TEST_CHUNK_SIZE = 512
+TEST_CHUNK_OVERLAP = 100
 TEST_RAG_FILE_JSON = {
     "ragFile": {
         "name": TEST_RAG_FILE_RESOURCE_NAME,
         "displayName": TEST_FILE_DISPLAY_NAME,
     }
 }
+TEST_RAG_FILE_JSON_WITH_UPLOAD_CONFIG = {
+    "ragFile": {
+        "name": TEST_RAG_FILE_RESOURCE_NAME,
+        "displayName": TEST_FILE_DISPLAY_NAME,
+    },
+    "rag_file_transformation_config": {
+        "rag_file_transformation_config": {
+            "rag_file_chunking_config": {
+                "fixed_length_chunking": {
+                    "chunk_size": TEST_CHUNK_SIZE,
+                    "chunk_overlap": TEST_CHUNK_OVERLAP,
+                }
+            }
+        }
+    },
+}
 TEST_RAG_FILE_JSON_ERROR = {"error": {"code": 13}}
-TEST_CHUNK_SIZE = 512
-TEST_CHUNK_OVERLAP = 100
 TEST_RAG_FILE_TRANSFORMATION_CONFIG = RagFileTransformationConfig(
     rag_file_chunking_config=RagFileChunkingConfig(
         fixed_length_chunking=RagFileChunkingConfig.FixedLengthChunking(

tests/unit/vertex_rag/test_rag_data.py

@@ -224,6 +224,15 @@ def upload_file_mock(authorized_session_mock):
         yield mock_post
 
 
+@pytest.fixture
+def upload_file_with_upload_config_mock(authorized_session_mock):
+    with patch.object(authorized_session_mock, "post") as mock_post:
+        mock_post.return_value = MockResponse(
+            test_rag_constants.TEST_RAG_FILE_JSON_WITH_UPLOAD_CONFIG, 200
+        )
+        yield mock_post
+
+
 @pytest.fixture
 def upload_file_not_found_mock(authorized_session_mock):
     with patch.object(authorized_session_mock, "post") as mock_post:
@@ -571,6 +580,29 @@ def test_upload_file_success(
 
         rag_file_eq(rag_file, test_rag_constants.TEST_RAG_FILE)
 
+    @pytest.mark.usefixtures("open_file_mock")
+    def test_upload_file_success_with_transformation_config(
+        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(),
+        )
+
+        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

@@ -292,6 +292,7 @@ def upload_file(
     path: Union[str, Sequence[str]],
     display_name: Optional[str] = None,
     description: Optional[str] = None,
+    transformation_config: Optional[TransformationConfig] = None,
 ) -> RagFile:
     """
     Synchronous file upload to an existing RagCorpus.
@@ -304,10 +305,19 @@ def upload_file(
 
     vertexai.init(project="my-project")
 
+    // Optional.
+    transformation_config = TransformationConfig(
+        chunking_config=ChunkingConfig(
+            chunk_size=1024,
+            chunk_overlap=200,
+        ),
+    )
+
     rag_file = rag.upload_file(
         corpus_name="projects/my-project/locations/us-central1/ragCorpora/my-corpus-1",
         display_name="my_file.txt",
         path="usr/home/my_file.txt",
+        transformation_config=transformation_config,
     )
     ```
 
@@ -319,6 +329,8 @@ def upload_file(
             "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.
+
     Returns:
         RagFile.
     Raises:
@@ -337,12 +349,25 @@ def upload_file(
         aiplatform.constants.base.API_BASE_PATH,
         corpus_name,
     )
+    js_rag_file = {"rag_file": {"display_name": display_name}}
+
     if description:
-        js_rag_file = {
-            "rag_file": {"display_name": display_name, "description": description}
+        js_rag_file["rag_file"]["description"] = description
+
+    if transformation_config and transformation_config.chunking_config:
+        chunk_size = transformation_config.chunking_config.chunk_size
+        chunk_overlap = transformation_config.chunking_config.chunk_overlap
+        js_rag_file["upload_rag_file_config"] = {
+            "rag_file_transformation_config": {
+                "rag_file_chunking_config": {
+                    "fixed_length_chunking": {
+                        "chunk_size": chunk_size,
+                        "chunk_overlap": chunk_overlap,
+                    }
+                }
+            }
         }
-    else:
-        js_rag_file = {"rag_file": {"display_name": display_name}}
+
     files = {
         "metadata": (None, str(js_rag_file)),
         "file": open(path, "rb"),