From 5c39e0452fc636418ccd2bafbe3f7d52c5b2bd11 Mon Sep 17 00:00:00 2001 From: ryanda9910 Date: Tue, 16 Jun 2026 12:55:29 +0700 Subject: [PATCH] feat(collections): add concurrent upload_documents with bounded parallelism Sequential document uploads were a scaling bottleneck (issue #77). Add upload_documents() to sync + async clients: a bounded ThreadPoolExecutor (configurable max_workers, conservative default) that overlaps upload latency, preserves result order, and re-raises the first error without swallowing others. Backward compatible (upload_document unchanged). Idea suggested by @nicsins in #77. --- src/xai_sdk/aio/collections.py | 78 ++++++++++++++++++++++++++ src/xai_sdk/collections.py | 19 +++++++ src/xai_sdk/sync/collections.py | 98 +++++++++++++++++++++++++++++++++ tests/aio/collections_test.py | 75 +++++++++++++++++++++++++ tests/sync/collections_test.py | 81 +++++++++++++++++++++++++++ 5 files changed, 351 insertions(+) diff --git a/src/xai_sdk/aio/collections.py b/src/xai_sdk/aio/collections.py index ed7a3f2..cc7e53f 100644 --- a/src/xai_sdk/aio/collections.py +++ b/src/xai_sdk/aio/collections.py @@ -8,11 +8,13 @@ from ..collections import ( DEFAULT_INDEXING_POLL_INTERVAL, DEFAULT_INDEXING_TIMEOUT, + DEFAULT_UPLOAD_MAX_WORKERS, BaseClient, ChunkConfiguration, CollectionSortBy, DocumentRetrievalMode, DocumentSortBy, + DocumentUpload, FieldDefinition, FieldDefinitionUpdate, HNSWMetric, @@ -435,6 +437,82 @@ async def _wait_for_indexing( ) await asyncio.sleep(timer.sleep_interval_or_raise()) + async def upload_documents( + self, + collection_id: str, + documents: Sequence[DocumentUpload], + *, + max_workers: int = DEFAULT_UPLOAD_MAX_WORKERS, + wait_for_indexing: bool = False, + poll_interval: Optional[datetime.timedelta] = None, + timeout: Optional[datetime.timedelta] = None, + ) -> Sequence[collections_pb2.DocumentMetadata]: + """Uploads multiple documents to a collection concurrently. + + This is a convenience wrapper around `upload_document` that uploads the provided + documents concurrently using `asyncio`. Each upload is an independent sequence of + gRPC calls, so awaiting them concurrently overlaps the network latency of otherwise + serial uploads. The concurrency is bounded by `max_workers` (via an + `asyncio.Semaphore`) to avoid overwhelming the service or tripping server-side rate + limits. + + Args: + collection_id: The ID of the collection to upload the documents to. + documents: The documents to upload. Each document is a mapping with a `name`, + `data` (raw bytes) and optional `fields`. + max_workers: The maximum number of concurrent uploads. Defaults to a + conservative value to respect server-side rate limits. Set this higher only + if you know the service can accommodate it. + wait_for_indexing: Whether to wait for every document to be indexed before + returning. + poll_interval: The interval to poll for when checking whether a document has been + indexed. + timeout: The total time to wait for each document to be indexed before returning. + + Returns: + The metadata for the uploaded documents, in the same order as the input + `documents`. + + Raises: + ValueError: If `documents` is empty or `max_workers` is not positive. + Exception: If any individual upload fails, the first encountered error is + re-raised and any in-flight uploads are cancelled. Errors are not swallowed. + """ + if not documents: + raise ValueError("documents must not be empty") + if max_workers < 1: + raise ValueError("max_workers must be a positive integer") + + semaphore = asyncio.Semaphore(max_workers) + + async def _upload(document: DocumentUpload) -> collections_pb2.DocumentMetadata: + async with semaphore: + return await self.upload_document( + collection_id, + document["name"], + document["data"], + document.get("fields"), + # Indexing is awaited in a single batched pass below. + wait_for_indexing=False, + ) + + # asyncio.gather preserves input ordering and, by default, propagates the first + # exception to the caller while cancelling the remaining tasks. + results = await asyncio.gather(*(_upload(document) for document in documents)) + + if wait_for_indexing: + return [ + await self._wait_for_indexing( + collection_id, + result.file_metadata.file_id, + poll_interval or DEFAULT_INDEXING_POLL_INTERVAL, + timeout or DEFAULT_INDEXING_TIMEOUT, + ) + for result in results + ] + + return list(results) + async def add_existing_document( self, collection_id: str, diff --git a/src/xai_sdk/collections.py b/src/xai_sdk/collections.py index 98541b6..2ea26e6 100644 --- a/src/xai_sdk/collections.py +++ b/src/xai_sdk/collections.py @@ -110,6 +110,25 @@ class BytesConfiguration(TypedDict): chunk_overlap_bytes: int +class DocumentUpload(TypedDict): + """A single document to upload as part of a batch upload. + + - name: The name of the document. + - data: The raw bytes of the document. + - fields: Optional additional metadata fields to store with the document. + """ + + name: str + data: bytes + fields: NotRequired[Optional[dict[str, str]]] + + +# Default number of concurrent workers used by batch document uploads. Kept +# conservative so that callers do not accidentally overwhelm the service or +# trip server-side rate limits; callers can raise it explicitly if needed. +DEFAULT_UPLOAD_MAX_WORKERS = 4 + + # Create TypeAdapter instances for validation FieldDefinitionValidator = TypeAdapter(FieldDefinition) FieldDefinitionAddValidator = TypeAdapter(FieldDefinitionAdd) diff --git a/src/xai_sdk/sync/collections.py b/src/xai_sdk/sync/collections.py index cfff119..c957f38 100644 --- a/src/xai_sdk/sync/collections.py +++ b/src/xai_sdk/sync/collections.py @@ -1,6 +1,7 @@ import datetime import time import warnings +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from typing import Optional, Sequence, Union from opentelemetry.trace import SpanKind @@ -8,11 +9,13 @@ from ..collections import ( DEFAULT_INDEXING_POLL_INTERVAL, DEFAULT_INDEXING_TIMEOUT, + DEFAULT_UPLOAD_MAX_WORKERS, BaseClient, ChunkConfiguration, CollectionSortBy, DocumentRetrievalMode, DocumentSortBy, + DocumentUpload, FieldDefinition, FieldDefinitionUpdate, HNSWMetric, @@ -438,6 +441,101 @@ def _wait_for_indexing_to_complete( ) time.sleep(timer.sleep_interval_or_raise()) + def upload_documents( + self, + collection_id: str, + documents: Sequence[DocumentUpload], + *, + max_workers: int = DEFAULT_UPLOAD_MAX_WORKERS, + wait_for_indexing: bool = False, + poll_interval: Optional[datetime.timedelta] = None, + timeout: Optional[datetime.timedelta] = None, + ) -> Sequence[collections_pb2.DocumentMetadata]: + """Uploads multiple documents to a collection concurrently. + + This is a convenience wrapper around `upload_document` that uploads the provided + documents using a bounded thread pool. Each upload is an independent, blocking + sequence of gRPC calls, so running them concurrently overlaps the network latency + of otherwise serial uploads. The concurrency is bounded by `max_workers` to avoid + overwhelming the service or tripping server-side rate limits. + + Args: + collection_id: The ID of the collection to upload the documents to. + documents: The documents to upload. Each document is a mapping with a `name`, + `data` (raw bytes) and optional `fields`. + max_workers: The maximum number of concurrent uploads. Defaults to a + conservative value to respect server-side rate limits. Set this higher only + if you know the service can accommodate it. + wait_for_indexing: Whether to wait for every document to be indexed before + returning. + poll_interval: The interval to poll for when checking whether a document has been + indexed. + timeout: The total time to wait for each document to be indexed before returning. + + Returns: + The metadata for the uploaded documents, in the same order as the input + `documents`. + + Raises: + ValueError: If `documents` is empty or `max_workers` is not positive. + Exception: If any individual upload fails, the first encountered error is + re-raised after cancelling any not-yet-started uploads. Other errors are not + swallowed; they are surfaced as `__cause__`/`__context__` where the runtime + preserves them. + """ + if not documents: + raise ValueError("documents must not be empty") + if max_workers < 1: + raise ValueError("max_workers must be a positive integer") + + results: list[Optional[collections_pb2.DocumentMetadata]] = [None] * len(documents) + + # Bound concurrency with a thread pool. The sync client performs blocking gRPC + # calls, so threads (which release the GIL during I/O) overlap upload latency. + with ThreadPoolExecutor(max_workers=max_workers) as executor: + future_to_index = { + executor.submit( + self.upload_document, + collection_id, + document["name"], + document["data"], + document.get("fields"), + # Indexing is awaited in a single batched pass below so we do not hold a + # worker thread busy-polling per document. + wait_for_indexing=False, + ): index + for index, document in enumerate(documents) + } + + done, not_done = wait(future_to_index, return_when=FIRST_EXCEPTION) + + # Fail fast: if any upload raised, cancel the uploads that have not started yet + # and re-raise the first error rather than silently returning partial results. + for future in done: + if future.exception() is not None: + for pending in not_done: + pending.cancel() + raise future.exception() # type: ignore[misc] + + # All uploads succeeded; collect their results preserving input order. + for future, index in future_to_index.items(): + results[index] = future.result() + + ordered_results = [result for result in results if result is not None] + + if wait_for_indexing: + return [ + self._wait_for_indexing_to_complete( + collection_id, + result.file_metadata.file_id, + poll_interval or DEFAULT_INDEXING_POLL_INTERVAL, + timeout or DEFAULT_INDEXING_TIMEOUT, + ) + for result in ordered_results + ] + + return ordered_results + def add_existing_document( self, collection_id: str, diff --git a/tests/aio/collections_test.py b/tests/aio/collections_test.py index 7fe5b0e..4872aee 100644 --- a/tests/aio/collections_test.py +++ b/tests/aio/collections_test.py @@ -1,6 +1,7 @@ # ruff noqa: DTZ005 +import datetime import uuid from typing import Union from unittest import mock @@ -1078,3 +1079,77 @@ async def test_update_document_creates_span_with_attributes(mock_tracer: mock.Ma ) mock_span.set_attribute.assert_any_call("document.id", document_metadata.file_metadata.file_id) mock_span.set_attribute.assert_any_call("document.name", new_name) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_upload_documents_uploads_all_and_preserves_order(client: AsyncClient): + """Batch upload returns metadata for every document in input order.""" + collection_metadata = await client.collections.create(f"test-collection-{uuid.uuid4()}") + assert collection_metadata.collection_id is not None + + documents = [ + {"name": f"doc-{i}.txt", "data": f"content {i}".encode(), "fields": {"idx": str(i)}} for i in range(5) + ] + + results = await client.collections.upload_documents(collection_metadata.collection_id, documents) + + assert len(results) == len(documents) + assert [r.file_metadata.name for r in results] == [d["name"] for d in documents] + for i, result in enumerate(results): + assert result.file_metadata.file_id is not None + assert result.fields == {"idx": str(i)} + + +@pytest.mark.asyncio(loop_scope="session") +async def test_upload_documents_empty_raises(client: AsyncClient): + collection_metadata = await client.collections.create(f"test-collection-{uuid.uuid4()}") + with pytest.raises(ValueError, match="documents must not be empty"): + await client.collections.upload_documents(collection_metadata.collection_id, []) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_upload_documents_invalid_max_workers_raises(client: AsyncClient): + collection_metadata = await client.collections.create(f"test-collection-{uuid.uuid4()}") + with pytest.raises(ValueError, match="max_workers must be a positive integer"): + await client.collections.upload_documents( + collection_metadata.collection_id, [{"name": "a.txt", "data": b"a"}], max_workers=0 + ) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_upload_documents_propagates_first_error(client: AsyncClient): + """If an individual upload fails, the error is surfaced (not swallowed).""" + collection_metadata = await client.collections.create(f"test-collection-{uuid.uuid4()}") + + documents = [{"name": f"err-{i}.txt", "data": b"x"} for i in range(6)] + + original_upload_document = client.collections.upload_document + + async def flaky_upload_document(collection_id, name, data, fields=None, **kwargs): + if name == "err-3.txt": + raise RuntimeError("simulated upload failure") + return await original_upload_document(collection_id, name, data, fields, **kwargs) + + with mock.patch.object(client.collections, "upload_document", side_effect=flaky_upload_document): + with pytest.raises(RuntimeError, match="simulated upload failure"): + await client.collections.upload_documents(collection_metadata.collection_id, documents, max_workers=2) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_upload_documents_wait_for_indexing_failure(client: AsyncClient): + """wait_for_indexing surfaces a failed document's indexing error.""" + collection_metadata = await client.collections.create(f"test-collection-{uuid.uuid4()}") + + documents = [ + {"name": "ok.txt", "data": b"x"}, + {"name": "test-failed", "data": b"x"}, + ] + + with pytest.raises(ValueError, match="Document indexing failed"): + await client.collections.upload_documents( + collection_metadata.collection_id, + documents, + wait_for_indexing=True, + poll_interval=datetime.timedelta(milliseconds=50), + timeout=datetime.timedelta(seconds=2), + ) diff --git a/tests/sync/collections_test.py b/tests/sync/collections_test.py index f1a5612..5275b10 100644 --- a/tests/sync/collections_test.py +++ b/tests/sync/collections_test.py @@ -1158,3 +1158,84 @@ def test_update_document_creates_span_with_attributes(mock_tracer: mock.MagicMoc ) mock_span.set_attribute.assert_any_call("document.id", document_metadata.file_metadata.file_id) mock_span.set_attribute.assert_any_call("document.name", new_name) + + +def test_upload_documents_uploads_all_and_preserves_order(client: Client): + """Batch upload returns metadata for every document in input order.""" + collection_metadata = client.collections.create(f"test-collection-{uuid.uuid4()}") + assert collection_metadata.collection_id is not None + + documents = [ + {"name": f"doc-{i}.txt", "data": f"content {i}".encode(), "fields": {"idx": str(i)}} for i in range(5) + ] + + results = client.collections.upload_documents(collection_metadata.collection_id, documents) + + assert len(results) == len(documents) + # Order of results must match the order of the input documents, regardless of the + # non-deterministic completion order of the worker threads. + assert [r.file_metadata.name for r in results] == [d["name"] for d in documents] + for i, result in enumerate(results): + assert result.file_metadata.file_id is not None + assert result.fields == {"idx": str(i)} + + +def test_upload_documents_single_worker(client: Client): + """max_workers=1 still works (sequential degenerate case).""" + collection_metadata = client.collections.create(f"test-collection-{uuid.uuid4()}") + + documents = [{"name": f"seq-{i}.txt", "data": b"x"} for i in range(3)] + results = client.collections.upload_documents(collection_metadata.collection_id, documents, max_workers=1) + + assert [r.file_metadata.name for r in results] == [d["name"] for d in documents] + + +def test_upload_documents_empty_raises(client: Client): + collection_metadata = client.collections.create(f"test-collection-{uuid.uuid4()}") + with pytest.raises(ValueError, match="documents must not be empty"): + client.collections.upload_documents(collection_metadata.collection_id, []) + + +def test_upload_documents_invalid_max_workers_raises(client: Client): + collection_metadata = client.collections.create(f"test-collection-{uuid.uuid4()}") + with pytest.raises(ValueError, match="max_workers must be a positive integer"): + client.collections.upload_documents( + collection_metadata.collection_id, [{"name": "a.txt", "data": b"a"}], max_workers=0 + ) + + +def test_upload_documents_propagates_first_error(client: Client): + """If an individual upload fails, the error is surfaced (not swallowed).""" + collection_metadata = client.collections.create(f"test-collection-{uuid.uuid4()}") + + documents = [{"name": f"err-{i}.txt", "data": b"x"} for i in range(6)] + + original_upload_document = client.collections.upload_document + + def flaky_upload_document(collection_id, name, data, fields=None, **kwargs): + if name == "err-3.txt": + raise RuntimeError("simulated upload failure") + return original_upload_document(collection_id, name, data, fields, **kwargs) + + with mock.patch.object(client.collections, "upload_document", side_effect=flaky_upload_document): + with pytest.raises(RuntimeError, match="simulated upload failure"): + client.collections.upload_documents(collection_metadata.collection_id, documents, max_workers=2) + + +def test_upload_documents_wait_for_indexing_failure(client: Client): + """wait_for_indexing surfaces a failed document's indexing error.""" + collection_metadata = client.collections.create(f"test-collection-{uuid.uuid4()}") + + documents = [ + {"name": "ok.txt", "data": b"x"}, + {"name": "test-failed", "data": b"x"}, + ] + + with pytest.raises(ValueError, match="Document indexing failed"): + client.collections.upload_documents( + collection_metadata.collection_id, + documents, + wait_for_indexing=True, + poll_interval=datetime.timedelta(milliseconds=50), + timeout=datetime.timedelta(seconds=2), + )