Add opt-in multipart file upload batching

python/ai-server/README.md

@@ -91,9 +91,12 @@ from ai_server import VectorEngine
 # initialize the connection to the vector database
 vectorEngine = VectorEngine(engine_id="221a50a4-060c-4aa8-8b7c-e2bc97ee3396", insight_id=server_connection.cur_insight)
 
-# Add document(s) that have been uploaded to the insight
+# Add document(s). Upload transport remains one file per request by default.
 vectorEngine.addDocument(file_paths = ['fileName1.pdf', 'fileName2.pdf', ..., 'fileNameX.pdf'])
 
+# Opt in to repeated multipart file parts when the server supports batched uploads.
+vectorEngine.addDocument(file_paths = ['fileName1.pdf', 'fileName2.pdf', ..., 'fileNameX.pdf'], upload_batch_size = 4)
+
 # Add Vector CSV File document(s) that have been uploaded to the insight
 vectorEngine.addVectorCSVFile(file_paths = ['fileName1.csv', 'fileName2csv', ..., 'fileNameX.csv'])
 
@@ -223,6 +226,9 @@ server_connection = ServerClient(access_key=loginKeys['accessKey'], secret_key=l
 
 server_connection.upload_files(files=["path_to_local_file1", "path_to_local_file2"], project_id="your_project_id", insight_id="your_insight_id", path="path_to_upload_files_in_insight")
 
+# Opt in to four files per multipart request. The default batch_size is one.
+server_connection.upload_files(files=["path_to_local_file1", "path_to_local_file2"], insight_id="your_insight_id", batch_size=4)
+
 server_connection.download_file(file=["path_to_insight_file"], project_id="your_project_id", insight_id="your_insight_id",custom_filename="filename_for_download")
 ```
 

python/ai-server/src/ai_server/py_client/gaas/vector.py

@@ -27,6 +27,7 @@ def addDocument(
         file_paths: List[str],
         param_dict: Optional[Dict] = {},
         insight_id: Optional[str] = None,
+        upload_batch_size: int = 1,
     ) -> Union[bool, List[Dict]]:
         """Adds documents to the vector database.
 
@@ -35,6 +36,8 @@ def addDocument(
             param_dict: Optional; A dictionary of additional parameters for processing the documents.
             insight_id: Optional; The unique identifier for the temporal workspace.
                         If None, the session's default insight_id is used.
+            upload_batch_size: Optional; Number of files sent in each multipart upload request.
+                               Defaults to one for backward compatibility.
 
         Returns:
             Union[bool, List[Dict]]:  List of dicts with metadata around the state of uploading each document provided.
@@ -53,6 +56,7 @@ def addDocument(
         insight_files = self.server.upload_files(
             files=file_paths,
             insight_id=insight_id,
+            batch_size=upload_batch_size,
         )
 
         optionalParams = (

python/ai-server/src/ai_server/server_resources/server_client.py

@@ -6,6 +6,7 @@
 import logging
 from urllib.parse import urlparse, unquote
 from pathlib import Path
+from contextlib import ExitStack
 import os
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -550,6 +551,7 @@ def upload_files(
         project_id: Optional[str] = None,
         insight_id: Optional[str] = None,
         path: Optional[str] = None,
+        batch_size: int = 1,
     ) -> List[str]:
         """
         Uploads files from the local device to the server.
@@ -563,22 +565,37 @@ def upload_files(
                 Given project/app unique identifier
             path (Optional[`str`]):
                 Specific upload path
+            batch_size (`int`):
+                Number of files to include in each multipart request. Defaults to one for
+                backward compatibility.
 
         Returns (`List[str]`):
             List of file names that have been successfully uploaded
         """
+        if isinstance(files, str):
+            files = [files]
         if not files:
-            raise Exception("Must provide atleast one file to upload")
+            raise ValueError("Must provide at least one file to upload")
+        if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size < 1:
+            raise ValueError("batch_size must be a positive integer")
+
+        validated_files = []
+        for filepath in files:
+            try:
+                normalized = os.fspath(filepath)
+            except TypeError as error:
+                raise TypeError("Every upload file must be a path-like value") from error
+            if not Path(normalized).is_file():
+                raise FileNotFoundError(f"Upload file does not exist or is not a file: {normalized}")
+            validated_files.append(normalized)
 
         # .../Monolith/api/uploadFile/baseUpload?insightId=de43ce0d-db2e-4ab9-a807-336bb86c4ea0&projectId=4c14bc58-973f-4293-87ed-a5d32c24f418&path=version/assets/
-        if isinstance(files, str):
-            files = [files]
 
         param = ""
         path = path or ""
 
         if insight_id or project_id or path:
-            if insight_id == None:
+            if insight_id is None:
                 insight_id = self.cur_insight
 
             param += f"insightId={insight_id}"
@@ -602,15 +619,32 @@ def upload_files(
         logger.info("The upload url is " + upload_post_request)
 
         insight_file_paths = []
-        for filepath in files:
-            with open(filepath, "rb") as fobj:
+        for start in range(0, len(validated_files), batch_size):
+            group = validated_files[start : start + batch_size]
+            with ExitStack() as stack:
+                multipart_files = [
+                    ("file", stack.enter_context(open(filepath, "rb"))) for filepath in group
+                ]
                 response = requests.post(
                     upload_post_request,
                     cookies=self.cookies,
-                    files={"file": fobj},
+                    files=multipart_files,
                     headers=self.required_headers.copy(),
                 )
-                insight_file_paths.append(response.json()[0]["fileName"])
+                response.raise_for_status()
+                payload = response.json()
+                if not isinstance(payload, list) or len(payload) != len(group):
+                    raise ValueError(
+                        "Upload response must be an ordered list matching the submitted file count"
+                    )
+                for index, item in enumerate(payload):
+                    if (
+                        not isinstance(item, dict)
+                        or not isinstance(item.get("fileName"), str)
+                        or not item["fileName"]
+                    ):
+                        raise ValueError(f"Upload response item {index} has no valid fileName")
+                    insight_file_paths.append(item["fileName"])
 
         return insight_file_paths
 

python/ai-server/src/ai_server/tests/test_upload_batching.py

@@ -0,0 +1,166 @@
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+import requests
+
+from ai_server.py_client.gaas.vector import VectorEngine
+from ai_server.server_resources.server_client import ServerClient
+
+
+class FakeResponse:
+    def __init__(self, payload, error=None):
+        self.payload = payload
+        self.error = error
+
+    def raise_for_status(self):
+        if self.error is not None:
+            raise self.error
+
+    def json(self):
+        return self.payload
+
+
+class UploadBatchingTests(unittest.TestCase):
+    def setUp(self):
+        self.client = object.__new__(ServerClient)
+        self.client.main_url = "https://example.test/Monolith/api"
+        self.client.cur_insight = "insight-current"
+        self.client.cookies = {}
+        self.client.required_headers = {"X-Test": "true"}
+        self.temporary_directory = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temporary_directory.cleanup)
+
+    def files(self, count):
+        paths = []
+        for index in range(count):
+            path = Path(self.temporary_directory.name) / f"file-{index}.txt"
+            path.write_text(f"file {index}", encoding="utf-8")
+            paths.append(str(path))
+        return paths
+
+    def test_upload_files_batches_four_and_preserves_response_order(self):
+        paths = self.files(9)
+        request_groups = []
+        opened_handles = []
+
+        def post(url, *, cookies, files, headers):
+            self.assertIn("insightId=insight-current", url)
+            self.assertEqual(cookies, {})
+            self.assertEqual(headers, {"X-Test": "true"})
+            self.assertTrue(all(field == "file" for field, _handle in files))
+            handles = [handle for _field, handle in files]
+            self.assertTrue(all(not handle.closed for handle in handles))
+            request_groups.append([Path(handle.name).name for handle in handles])
+            opened_handles.extend(handles)
+            return FakeResponse([{"fileName": f"remote/{Path(handle.name).name}"} for handle in handles])
+
+        with patch("ai_server.server_resources.server_client.requests.post", side_effect=post):
+            uploaded = self.client.upload_files(paths, batch_size=4)
+
+        self.assertEqual([len(group) for group in request_groups], [4, 4, 1])
+        self.assertEqual(uploaded, [f"remote/file-{index}.txt" for index in range(9)])
+        self.assertTrue(all(handle.closed for handle in opened_handles))
+
+    def test_upload_files_keeps_one_file_requests_by_default(self):
+        paths = self.files(3)
+        group_sizes = []
+
+        def post(url, *, cookies, files, headers):
+            del url, cookies, headers
+            group_sizes.append(len(files))
+            return FakeResponse([{"fileName": Path(files[0][1].name).name}])
+
+        with patch("ai_server.server_resources.server_client.requests.post", side_effect=post):
+            uploaded = self.client.upload_files(paths)
+
+        self.assertEqual(group_sizes, [1, 1, 1])
+        self.assertEqual(uploaded, ["file-0.txt", "file-1.txt", "file-2.txt"])
+
+    def test_upload_files_validates_all_inputs_before_first_request(self):
+        existing = self.files(1)[0]
+        missing = str(Path(self.temporary_directory.name) / "missing.txt")
+        post = Mock()
+
+        with patch("ai_server.server_resources.server_client.requests.post", post):
+            with self.assertRaises(FileNotFoundError):
+                self.client.upload_files([existing, missing], batch_size=4)
+            with self.assertRaises(ValueError):
+                self.client.upload_files([existing], batch_size=0)
+            with self.assertRaises(ValueError):
+                self.client.upload_files([existing], batch_size=True)
+
+        post.assert_not_called()
+
+    def test_upload_files_raises_for_http_failure_and_closes_handles(self):
+        paths = self.files(4)
+        opened_handles = []
+
+        def post(url, *, cookies, files, headers):
+            del url, cookies, headers
+            opened_handles.extend(handle for _field, handle in files)
+            return FakeResponse([], requests.HTTPError("upload failed"))
+
+        with patch("ai_server.server_resources.server_client.requests.post", side_effect=post):
+            with self.assertRaises(requests.HTTPError):
+                self.client.upload_files(paths, batch_size=4)
+
+        self.assertTrue(all(handle.closed for handle in opened_handles))
+
+    def test_upload_files_rejects_malformed_responses_and_closes_handles(self):
+        paths = self.files(2)
+        malformed_payloads = ([{"fileName": "only-one"}], [{"fileName": "one"}, {}])
+        for payload in malformed_payloads:
+            with self.subTest(payload=payload):
+                opened_handles = []
+
+                def post(url, *, cookies, files, headers):
+                    del url, cookies, headers
+                    opened_handles.extend(handle for _field, handle in files)
+                    return FakeResponse(payload)
+
+                with patch("ai_server.server_resources.server_client.requests.post", side_effect=post):
+                    with self.assertRaises(ValueError):
+                        self.client.upload_files(paths, batch_size=2)
+                self.assertTrue(all(handle.closed for handle in opened_handles))
+
+    def test_vector_add_document_forwards_upload_batch_size(self):
+        class FakeServer:
+            cur_insight = "insight-current"
+
+            def __init__(self):
+                self.upload_arguments = None
+
+            def upload_files(self, **kwargs):
+                self.upload_arguments = kwargs
+                return ["remote/one.txt", "remote/two.txt"]
+
+            def run_pixel(self, **kwargs):
+                return {"pixelReturn": [{"operationType": ["OPERATION"], "output": True}]}
+
+        fake_server = FakeServer()
+        original = ServerClient.da_server
+        ServerClient.da_server = fake_server
+        self.addCleanup(setattr, ServerClient, "da_server", original)
+        engine = VectorEngine("vector-1", insight_id="insight-current")
+
+        result = engine.addDocument(
+            ["one.txt", "two.txt"],
+            insight_id="insight-current",
+            upload_batch_size=4,
+        )
+
+        self.assertTrue(result)
+        self.assertEqual(
+            fake_server.upload_arguments,
+            {
+                "files": ["one.txt", "two.txt"],
+                "insight_id": "insight-current",
+                "batch_size": 4,
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()