feat(collections): concurrent upload_documents with bounded parallelism

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,

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)

src/xai_sdk/sync/collections.py

@@ -1,18 +1,21 @@
 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
 
 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,

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),
+        )