fix: chromadb offload async batch upsert/insert to worker thread (#7711)

Author: basnijholt

libs/agno/agno/vectordb/chroma/chromadb.py

@@ -465,7 +465,11 @@ async def async_insert(
         if self._collection is None:
             logger.warning("Collection does not exist")
         else:
-            self._batch_operation(
+            # Run the synchronous ChromaDB batch on a worker thread so the
+            # asyncio event loop stays responsive while the (potentially large)
+            # write completes.
+            await asyncio.to_thread(
+                self._batch_operation,
                 ids=ids,
                 embeddings=docs_embeddings,
                 documents=docs,
@@ -656,7 +660,11 @@ async def _async_upsert(
         if self._collection is None:
             logger.warning("Collection does not exist")
         else:
-            self._batch_operation(
+            # Run the synchronous ChromaDB batch on a worker thread so the
+            # asyncio event loop stays responsive while the (potentially large)
+            # write completes.
+            await asyncio.to_thread(
+                self._batch_operation,
                 ids=ids,
                 embeddings=docs_embeddings,
                 documents=docs,

libs/agno/tests/unit/vectordb/test_chromadb.py

@@ -1481,3 +1481,57 @@ def test_no_batch_size_no_batching_logs(chroma_db, caplog):
     assert not batch_log_found, "Expected no batch log messages, but found batch related logs"
 
     assert chroma_db.get_count() == num_documents
+
+
+@pytest.mark.asyncio
+async def test_async_insert_runs_batch_on_worker_thread(chroma_db, sample_documents):
+    """Ensure async_insert offloads the synchronous ChromaDB batch to a worker thread."""
+    from threading import get_ident
+
+    for doc in sample_documents:
+        doc.embedding = chroma_db.embedder.get_embedding(doc.content)
+
+    main_thread_id = get_ident()
+    batch_thread_ids: List[int] = []
+    original_batch = ChromaDb._batch_operation
+
+    def _record_batch(self, *args, **kwargs):
+        batch_thread_ids.append(get_ident())
+        return original_batch(self, *args, **kwargs)
+
+    with patch.object(ChromaDb, "_batch_operation", _record_batch):
+        await chroma_db.async_insert(content_hash="test_hash", documents=sample_documents)
+
+    assert batch_thread_ids, "expected _batch_operation to be invoked"
+    for thread_id in batch_thread_ids:
+        assert thread_id != main_thread_id, (
+            "_batch_operation ran on the asyncio main thread; the synchronous ChromaDB "
+            "batch must be offloaded so the event loop stays responsive during async_insert"
+        )
+
+
+@pytest.mark.asyncio
+async def test_async_upsert_runs_batch_on_worker_thread(chroma_db, sample_documents):
+    """Ensure async_upsert offloads the synchronous ChromaDB batch to a worker thread."""
+    from threading import get_ident
+
+    for doc in sample_documents:
+        doc.embedding = chroma_db.embedder.get_embedding(doc.content)
+
+    main_thread_id = get_ident()
+    batch_thread_ids: List[int] = []
+    original_batch = ChromaDb._batch_operation
+
+    def _record_batch(self, *args, **kwargs):
+        batch_thread_ids.append(get_ident())
+        return original_batch(self, *args, **kwargs)
+
+    with patch.object(ChromaDb, "_batch_operation", _record_batch):
+        await chroma_db.async_upsert(content_hash="test_hash", documents=sample_documents)
+
+    assert batch_thread_ids, "expected _batch_operation to be invoked"
+    for thread_id in batch_thread_ids:
+        assert thread_id != main_thread_id, (
+            "_batch_operation ran on the asyncio main thread; the synchronous ChromaDB "
+            "batch must be offloaded so the event loop stays responsive during async_upsert"
+        )