From 5b6b510077feaaf2f31b19c91f53d9c0b2021ab9 Mon Sep 17 00:00:00 2001 From: Greg Lum Date: Thu, 16 Jul 2026 09:21:16 -1000 Subject: [PATCH] feat: add opt-in batched file uploads --- python/ai-server/README.md | 8 +- .../src/ai_server/py_client/gaas/vector.py | 4 + .../server_resources/server_client.py | 50 +++++- .../ai_server/tests/test_upload_batching.py | 166 ++++++++++++++++++ 4 files changed, 219 insertions(+), 9 deletions(-) create mode 100644 python/ai-server/src/ai_server/tests/test_upload_batching.py diff --git a/python/ai-server/README.md b/python/ai-server/README.md index 9b00cbb..e9c9e66 100644 --- a/python/ai-server/README.md +++ b/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") ``` diff --git a/python/ai-server/src/ai_server/py_client/gaas/vector.py b/python/ai-server/src/ai_server/py_client/gaas/vector.py index 5046433..19f7303 100644 --- a/python/ai-server/src/ai_server/py_client/gaas/vector.py +++ b/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 = ( diff --git a/python/ai-server/src/ai_server/server_resources/server_client.py b/python/ai-server/src/ai_server/server_resources/server_client.py index da6def7..ea9ba7a 100644 --- a/python/ai-server/src/ai_server/server_resources/server_client.py +++ b/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 diff --git a/python/ai-server/src/ai_server/tests/test_upload_batching.py b/python/ai-server/src/ai_server/tests/test_upload_batching.py new file mode 100644 index 0000000..fda77ee --- /dev/null +++ b/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()