[ENHANCEMENT] Sub-optimal sdk implementation for uploading collections

[gauravagerwala]

Confirmation of Request Source

• I confirm that this feature request is for the xAI SDK Python library (e.g., client features, SDK methods, or documentation) and not for the xAI API itself (e.g., model capabilities or API changes).

Describe the feature you'd like

We built a design review tool for the X AI hackathon.
The tool found scaling bottlenecks in the python sdk for collections management. The tool relies on artifacts it had generated (attached [here)]([url](

workflows.json
design-workflow-3-collections-management.md

))

Scaling Bottlenecks in the Third Workflow: Collections Management

The third workflow, as documented in .exp/design-workflow-3-collections-management.md and workflows.json, focuses on managing vector collections for document storage, embedding, indexing, searching, and retrieval. It involves operations like creating collections, uploading/indexing documents, and searching via gRPC calls to management and documents services. Exploration of the codebase (src/xai_sdk/collections.py, src/xai_sdk/sync/collections.py, src/xai_sdk/files.py, src/xai_sdk/poll_timer.py, proto files, and examples/sync/collection.py) reveals several client-side scaling bottlenecks, particularly for high-volume scenarios (e.g., ingesting thousands of documents or handling large files). These limit efficiency, increase latency, memory usage, and API call overhead:

1. Lack of Batch Operations for Document Addition (Major Bottleneck)

• The underlying protobuf API (in proto/v5/collections_pb2_grpc.py and v6) supports BatchAddDocumentToCollection(BatchAddDocumentToCollectionRequest) RPC for adding multiple documents in a single call.
• However, the SDK client (sync/collections.py) only exposes single-document methods: add_existing_document (single AddDocumentToCollection RPC) and upload_document (which internally calls single add after file upload).
• Impact for Scaling: For bulk ingestion (common in RAG/knowledge base workflows), users must loop over individual calls, resulting in N gRPC RPCs for N documents. This amplifies network latency (each RPC has overhead), risks server-side rate limiting/throttling, and serializes processing. Examples show sequential small uploads, but real-world scale (e.g., 10k+ docs) would be prohibitively slow without user-implemented parallelism (threads/asyncio).
• Evidence: No batch_add_documents method in client; proto grep confirms RPC exists but unimplemented. Design doc mentions "batch operations," but client only implements batch_get_documents.

2. Memory Inefficiency in Document Upload for Large Files

• upload_document(..., data: bytes) in sync/collections.py requires loading entire document into memory to compute len(data) and slice chunks via _chunk_file_data (in files.py:93-129), which yields fixed-size chunks ( _CHUNK_SIZE, likely ~4MB) for streaming gRPC upload.
• While upload itself streams (good), initial bytes load is client-side burden.
• Impact for Scaling: Large documents (e.g., PDFs >100MB, datasets) cause high memory spikes per upload, risking OOM in batch loops or low-memory envs. No overload for path: str or file objects in upload_document; users must detour via client.files.upload(path=...) (which streams from disk via open("rb").read(_CHUNK_SIZE) in files.py:130-164), get file_id, then add_existing_document. This adds complexity and still requires single adds for batch.
• Evidence: Examples use small b"""...""" strings; files.upload(path) handles streaming well, but collections workflow doesn't integrate it seamlessly.

3. Inefficient Per-Document Polling for Async Indexing Status

• Document indexing (chunking, embedding, HNSW indexing) is server-side async after AddDocumentToCollection.
• upload_document(wait_for_indexing=True) or manual get_document loops poll status via PollTimer (poll_timer.py), checking DocumentStatus (PROCESSING → PROCESSED/FAILED) with default 10s intervals (DEFAULT_INDEXING_POLL_INTERVAL).
• No batch status method (e.g., poll multiple via batch_get_documents is possible but manual and still per-batch-loop).
• Impact for Scaling: In bulk workflows, waiting for many docs requires many polls (e.g., 10s intervals × N docs = excessive gRPC traffic, even with configurable intervals). Busy-waiting wastes resources; no pub/sub or webhook alternatives. Timeouts (default 2min) may fail under load if server backlog.
• Evidence: _wait_for_indexing_to_complete in sync/collections.py:319-356 is single-doc loop with time.sleep; design doc notes "optional polling," but no optimized batch/multi-wait.

4. Synchronous Blocking and Lack of Built-in Concurrency

• Sync client (sync/collections.py) blocks on each RPC (upload, add, poll, search), making loops inherently sequential.
• No high-level batch/parallel methods (e.g., upload_documents(paths: list[str], parallel=True) with progress or threading).
• Async client (aio/collections.py) exists for concurrency but mirrors issues (single ops, no batch).
• Impact for Scaling: High-throughput ingestion (e.g., processing directories of files) requires user-side concurrency (e.g., concurrent.futures), increasing code complexity and potential for race conditions (e.g., adding before upload complete). Telemetry/interceptors add minor per-call overhead under load.
• Evidence: Examples (examples/sync/collection.py) run ops sequentially; BaseClient reuses channels (good), but no parallelism abstraction.

5. Minor Overhead in Repeated Validations and Conversions

• Every create/update involves Pydantic TypeAdapter validations (FieldDefinitionValidator, ChunkConfigurationValidator) and dict-to-pb conversions (e.g., _field_definition_to_pb, _chunk_configuration_to_pb).
• Impact for Scaling: Negligible for few ops, but in loops for many collections/fields, adds CPU cycles. Proto versioning (v5/v6 stubs) requires correct channel selection, potential misconfig under scale.
• Evidence: collections.py converters called per-call; good type safety but runtime cost.

Recommendations for Mitigation (Not Implemented Changes)

• Expose batch_add_documents(file_ids: list[str], fields: list[dict]) wrapping proto RPC.
• Add upload_documents overloads for paths/files with optional concurrency and batch add.
• Implement batch_wait_for_indexing(file_ids: list[str]) using batch_get_documents loops or proto batch status if available.
• Integrate files streaming directly in collections for path-based uploads.
• Add optional async wrappers or concurrency helpers in docs/examples for scale.

These issues primarily affect client-side efficiency for large-scale document management, aligning with the workflow's emphasis on "document storage, embedding, indexing" at volume. Server-side (HNSW search, embedding) scales via design (e.g., approximate NN), but SDK gaps force suboptimal usage. For current small-scale (as in examples), fine; for production RAG pipelines, significant hurdles.

Additional Context

No response

[nicsins]

gemini and i thought this may help ...
Here are the additions for xai-sdk-python/src/xai_sdk/sync/collections.py:

At the top of the file, add the following import:

1 from concurrent.futures import ThreadPoolExecutor

Then, add the following methods to the Client class:

 1     def upload_document_from_path(
 2         self,
 3         collection_id: str,
 4         file_path: str,
 5         fields: Optional[dict[str, str]] = None,
 6         *,
 7         wait_for_indexing: bool = False,
 8         poll_interval: Optional[datetime.timedelta] = None,
 9         timeout: Optional[datetime.timedelta] = None,
10     ) -> collections_pb2.DocumentMetadata:
11         """Uploads a document to a collection from a file path.
12         Args:
13             collection_id: The ID of the collection to upload the document to.
14             file_path: The path to the file to upload.
15             fields: Additional metadata fields to store with the document.
16             wait_for_indexing: Whether to wait for the document to be indexed.
17             poll_interval: The interval to poll for when checking whether the document has been 
   indexed.
18             timeout: The total time to wait for the document to be indexed before returning.
19         Returns:
20             The metadata for the uploaded document.
21         """
22         from ..files import _chunk_file_from_path
23         # Upload the raw bytes via the streaming Files API, then attach to the collection.
24         upload_chunks = _chunk_file_from_path(file_path=file_path)
25         uploaded_file = self._files_stub.UploadFile(upload_chunks)
26         # Attach the uploaded file to the target collection as a document.
27         self._collections_stub.AddDocumentToCollection(
28             collections_pb2.AddDocumentToCollectionRequest(
29                 collection_id=collection_id,
30                 file_id=uploaded_file.id,
31                 fields=fields,
32             )
33         )
34         # Either wait for indexing to complete or return the current metadata.
35         if wait_for_indexing:
36             return self._wait_for_indexing_to_complete(
37                 collection_id,
38                 uploaded_file.id,
39                 poll_interval or DEFAULT_INDEXING_POLL_INTERVAL,
40                 timeout or DEFAULT_INDEXING_TIMEOUT,
41             )
42         return self.get_document(
43             uploaded_file.id,
44             collection_id,
45         )
46     def upload_documents(
47         self,
48         collection_id: str,
49         file_paths: Sequence[str],
50         max_workers: int = 4,
51         *,
52         wait_for_indexing: bool = False,
53         poll_interval: Optional[datetime.timedelta] = None,
54         timeout: Optional[datetime.timedelta] = None,
55     ) -> Sequence[collections_pb2.DocumentMetadata]:
56         """Uploads multiple documents to a collection from a list of file paths.
57         Args:
58             collection_id: The ID of the collection to upload the documents to.
59             file_paths: A list of paths to the files to upload.
60             max_workers: The maximum number of concurrent uploads.
61             wait_for_indexing: Whether to wait for all documents to be indexed.
62             poll_interval: The interval to poll for when checking whether documents have been 
   indexed.
63             timeout: The total time to wait for the documents to be indexed before returning.
64         Returns:
65             A list of metadata for the uploaded documents.
66         """
67         with ThreadPoolExecutor(max_workers=max_workers) as executor:
68             futures = [
69                 executor.submit(
70                     self.upload_document_from_path,
71                     collection_id=collection_id,
72                     file_path=file_path,
73                     wait_for_indexing=False,
74                 )
75                 for file_path in file_paths
76             ]
77             results = [future.result() for future in futures]
78         if wait_for_indexing:
79             return self._wait_for_indexing_to_complete_batch(
80                 collection_id,
81                 [result.file_id for result in results],
82                 poll_interval or DEFAULT_INDEXING_POLL_INTERVAL,
83                 timeout or DEFAULT_INDEXING_TIMEOUT,
84             )
85         return results
86     def _wait_for_indexing_to_complete_batch(
87         self,
88         collection_id: str,
89         file_ids: Sequence[str],
90         poll_interval: datetime.timedelta,
91         timeout: datetime.timedelta,
92     ) -> Sequence[collections_pb2.DocumentMetadata]:
93         """Waits for a batch of documents to be indexed.
94         Args:
95             collection_id: The ID of the collection containing the documents.
96             file_ids: The IDs of the documents to wait for.
97             poll_interval: The interval to poll for when checking whether documents have been 
   indexed.
98             timeout: The total time to wait for the documents to be indexed before returning.
99         Returns:

100 A list of metadata for the documents.
101 Raises:
102 ValueError: If document indexing fails for any document.
103 ValueError: If any document status is unknown.
104 TimeoutError: If polling times out before all documents are processed.
105 """
106 timer = PollTimer(timeout, poll_interval)
107 unprocessed_file_ids = set(file_ids)
108 while unprocessed_file_ids:
109 batch_metadata = self.batch_get_documents(
110 collection_id=collection_id,
111 file_ids=list(unprocessed_file_ids),
112 )
113 for document_metadata in batch_metadata.documents:
114 match document_metadata.status:
115 case collections_pb2.DocumentStatus.DOCUMENT_STATUS_PROCESSED:
116 unprocessed_file_ids.remove(document_metadata.file_id)
117 case collections_pb2.DocumentStatus.DOCUMENT_STATUS_PROCESSING:
118 continue
119 case collections_pb2.DocumentStatus.DOCUMENT_STATUS_FAILED:
120 raise ValueError(f"Document indexing failed for file_id
{document_metadata.file_id}: {document_metadata.error_message}")
121 case unknown_status:
122 raise ValueError(f"Unknown document status for file_id
{document_metadata.file_id}: {unknown_status}")
123 if unprocessed_file_ids:
124 time.sleep(timer.sleep_interval_or_raise())
125 return self.batch_get_documents(collection_id=collection_id, file_ids=file_ids).documents

Here are the additions for xai-sdk-python/src/xai_sdk/aio/collections.py:

Add the following methods to the Client class:

 1     async def upload_document_from_path(
 2         self,
 3         collection_id: str,
 4         file_path: str,
 5         fields: Optional[dict[str, str]] = None,
 6         *,
 7         wait_for_indexing: bool = False,
 8         poll_interval: Optional[datetime.timedelta] = None,
 9         timeout: Optional[datetime.timedelta] = None,
10     ) -> collections_pb2.DocumentMetadata:
11         """Uploads a document to a collection from a file path.
12 
13         Args:
14             collection_id: The ID of the collection to upload the document to.
15             file_path: The path to the file to upload.
16             fields: Additional metadata fields to store with the document.
17             wait_for_indexing: Whether to wait for the document to be indexed.
18             poll_interval: The interval to poll for when checking whether the document has been 
   indexed.
19             timeout: The total time to wait for the document to be indexed before returning.
20 
21         Returns:
22             The metadata for the uploaded document.
23         """
24         from ..files import _async_chunk_file_from_path
25 
26         # Upload the raw bytes via the streaming Files API, then attach to the collection.
27         upload_chunks = _async_chunk_file_from_path(file_path=file_path)
28 
29         uploaded_file = await self._files_stub.UploadFile(upload_chunks)
30 
31         # Attach the uploaded file to the target collection as a document.
32         await self._collections_stub.AddDocumentToCollection(
33             collections_pb2.AddDocumentToCollectionRequest(
34                 collection_id=collection_id,
35                 file_id=uploaded_file.id,
36                 fields=fields,
37             )
38         )
39 
40         # Either wait for indexing to complete or return the current metadata.
41         if wait_for_indexing:
42             return await self._wait_for_indexing(
43                 collection_id,
44                 uploaded_file.id,
45                 poll_interval or DEFAULT_INDEXING_POLL_INTERVAL,
46                 timeout or DEFAULT_INDEXING_TIMEOUT,
47             )
48 
49         return await self.get_document(
50             uploaded_file.id,
51             collection_id,
52         )
53 
54     async def upload_documents(
55         self,
56         collection_id: str,
57         file_paths: Sequence[str],
58         *,
59         wait_for_indexing: bool = False,
60         poll_interval: Optional[datetime.timedelta] = None,
61         timeout: Optional[datetime.timedelta] = None,
62     ) -> Sequence[collections_pb2.DocumentMetadata]:
63         """Uploads multiple documents to a collection from a list of file paths.
64 
65         Args:
66             collection_id: The ID of the collection to upload the documents to.
67             file_paths: A list of paths to the files to upload.
68             wait_for_indexing: Whether to wait for all documents to be indexed.
69             poll_interval: The interval to poll for when checking whether documents have been 
   indexed.
70             timeout: The total time to wait for the documents to be indexed before returning.
71 
72         Returns:
73             A list of metadata for the uploaded documents.
74         """
75         tasks = [
76             self.upload_document_from_path(
77                 collection_id=collection_id,
78                 file_path=file_path,
79                 wait_for_indexing=False,
80             )
81             for file_path in file_paths
82         ]
83         results = await asyncio.gather(*tasks)
84 
85         if wait_for_indexing:
86             return await self._wait_for_indexing_batch(
87                 collection_id,
88                 [result.file_id for result in results],
89                 poll_interval or DEFAULT_INDEXING_POLL_INTERVAL,
90                 timeout or DEFAULT_INDEXING_TIMEOUT,
91             )
92 
93         return results
94 
95     async def _wait_for_indexing_batch(
96         self,
97         collection_id: str,
98         file_ids: Sequence[str],
99         poll_interval: datetime.timedelta,

100 timeout: datetime.timedelta,
101 ) -> Sequence[collections_pb2.DocumentMetadata]:
102 """Waits for a batch of documents to be indexed.
103
104 Args:
105 collection_id: The ID of the collection containing the documents.
106 file_ids: The IDs of the documents to wait for.
107 poll_interval: The interval to poll for when checking whether documents have been
indexed.
108 timeout: The total time to wait for the documents to be indexed before returning.
109
110 Returns:
111 A list of metadata for the documents.
112
113 Raises:
114 ValueError: If document indexing fails for any document.
115 ValueError: If any document status is unknown.
116 TimeoutError: If polling times out before all documents are processed.
117 """
118 timer = PollTimer(timeout, poll_interval)
119 unprocessed_file_ids = set(file_ids)
120
121 while unprocessed_file_ids:
122 batch_metadata = await self.batch_get_documents(
123 collection_id=collection_id,
124 file_ids=list(unprocessed_file_ids),
125 )
126 for document_metadata in batch_metadata.documents:
127 match document_metadata.status:
128 case collections_pb2.DocumentStatus.DOCUMENT_STATUS_PROCESSED:
129 unprocessed_file_ids.remove(document_metadata.file_id)
130 case collections_pb2.DocumentStatus.DOCUMENT_STATUS_PROCESSING:
131 continue
132 case collections_pb2.DocumentStatus.DOCUMENT_STATUS_FAILED:
133 raise ValueError(f"Document indexing failed for file_id
{document_metadata.file_id}: {document_metadata.error_message}")
134 case unknown_status:
135 raise ValueError(f"Unknown document status for file_id
{document_metadata.file_id}: {unknown_status}")
136
137 if unprocessed_file_ids:
138 await asyncio.sleep(timer.sleep_interval_or_raise())
139
140 return (await self.batch_get_documents(collection_id=collection_id,
file_ids=file_ids)).documents

[yuujirohanma261-lang]

#/ @

design-workflow-3-collections-management.md