Repository: run-llama/llama_index
Pull request: Fix: ChromaVectorStore can attempt to add in excess of chromadb batch…
… size

# Description

Changed `ChromaVectorStore.add` to batch inputs at one less than chromadb's max to avoid the oversize error. Adds a class method to do chunking.

Fixes #7648 

## Type of Change

Please delete options that are not relevant.

- [X] Bug fix (non-breaking change which fixes an issue)

# How Has This Been Tested?

I have passing unit tests on my branch but have not included them. The unit test to add a large number of nodes takes ~5 secs to run. If that's acceptable, I can submit the unit tests for `add`

I noticed tests were missing for all of the llama_index.vector_stores.chroma.py. So I also implemented tests for several other methods.

Let me know if you want all the unit tests, all the tests excluding `add`, or none of them.

- [see above] Added new unit/integration tests
- [X] I stared at the code and made sure it makes sense

# Suggested Checklist:

- [X] I have performed a self-review of my own code
- [X] I have commented my code, particularly in hard-to-understand areas
- [N/A ] I have made corresponding changes to the documentation
- [X] My changes generate no new warnings
- [see above] I have added tests that prove my fix is effective or that my feature works
- [X] New and existing unit tests pass locally with my changes
- [X] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Bug]: Exceeding submit limit in ChromaVectorStore
### Bug Description

There are too many documents. When converting them into embeddings and loading them into chroma at once, chroma shows that they exceed the submit limit.

### Version

v0.8.24.post1 Latest

### Steps to Reproduce

```python
self.index = VectorStoreIndex.from_documents(
                                    documents,
                                    storage_context=storage_context,
                                    service_context=service_context,
                                    show_progress=True,
                                )
```

### Relevant Logs/Tracbacks

```shell
Batches: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.57s/it]
Generating embeddings: 100%|███████████████████████████████████████████████████████████████████████████████████████| 293144/293144 [7:01:19<00:00, 11.60it/s]
Traceback (most recent call last):
  File ".../index/index.py", line 147, in <module>
    index = MobilityBuddyIndex()
            ^^^^^^^^^^^^^^^^^^^^
  File ".../index/index.py", line 36, in __init__
    self._setup_index()
  File ".../index/index.py", line 87, in _setup_index
    self.index = VectorStoreIndex.from_documents(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/llama_index/indices/base.py", line 102, in from_documents
    return cls(
           ^^^^
  File ".../.venv/lib/python3.11/site-packages/llama_index/indices/vector_store/base.py", line 46, in __init__
    super().__init__(
  File ".../.venv/lib/python3.11/site-packages/llama_index/indices/base.py", line 71, in __init__
    index_struct = self.build_index_from_nodes(nodes)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/llama_index/indices/vector_store/base.py", line 265, in build_index_from_nodes
    return self._build_index_from_nodes(nodes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/llama_index/indices/vector_store/base.py", line 253, in _build_index_from_nodes
    self._add_nodes_to_index(
  File ".../.venv/lib/python3.11/site-packages/llama_index/indices/vector_store/base.py", line 214, in _add_nodes_to_index
    new_ids = self._vector_store.add(nodes)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/llama_index/vector_stores/chroma.py", line 146, in add
    self._collection.add(
  File ".../.venv/lib/python3.11/site-packages/chromadb/api/models/Collection.py", line 99, in add
    self._client._add(ids, self.id, embeddings, metadatas, documents)
  File ".../.venv/lib/python3.11/site-packages/chromadb/api/segment.py", line 249, in _add
    self._producer.submit_embeddings(coll["topic"], records_to_submit)
  File ".../.venv/lib/python3.11/site-packages/chromadb/db/mixins/embeddings_queue.py", line 127, in submit_embeddings
    raise ValueError(
ValueError: 
                Cannot submit more than 41,666 embeddings at once.
                Please submit your embeddings in batches of size
                41,666 or less.
```

Hints from issue discussion:
here: .../llama_index/vector_stores/chroma.py:136
```python
        for node in nodes:
            embeddings.append(node.get_embedding())
            metadatas.append(
                node_to_metadata_dict(
                    node, remove_text=True, flat_metadata=self.flat_metadata
                )
            )
            ids.append(node.node_id)
            documents.append(node.get_content(metadata_mode=MetadataMode.NONE))

        self._collection.add(
            embeddings=embeddings,
            ids=ids,
            metadatas=metadatas,
            documents=documents,
        )
```
should be use smaller chunks to add node into self._collection.
The possible solution is to use `index.insert()`:
```python

index = VectorStoreIndex.from_documents(
    [],
    storage_context=storage_context,
    service_context=service_context,
)

for doc in docs:
    index.insert(doc)
```
Thanks for raising this issue. Looks like we are missing some batching logic here. Will try to get a fix in soon!

Commit message: Fix: ChromaVectorStore can attempt to add in excess of chromadb batch size

Patch to review:
diff --git a/llama_index/vector_stores/chroma.py b/llama_index/vector_stores/chroma.py
index d2373fe206..3cf9f2b35d 100644
--- a/llama_index/vector_stores/chroma.py
+++ b/llama_index/vector_stores/chroma.py
@@ -1,7 +1,7 @@
 """Chroma vector store."""
 import logging
 import math
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, Generator, List, Optional, cast
 
 from llama_index.bridge.pydantic import Field, PrivateAttr
 from llama_index.schema import BaseNode, MetadataMode, TextNode
@@ -32,6 +32,22 @@ def _to_chroma_filter(standard_filters: MetadataFilters) -> dict:
 import_err_msg = "`chromadb` package not found, please run `pip install chromadb`"
 
 
+def chunk_list(
+    lst: List[BaseNode], max_chunk_size: int
+) -> Generator[List[BaseNode], None, None]:
+    """Yield successive max_chunk_size-sized chunks from lst.
+
+    Args:
+        lst (List[BaseNode]): list of nodes with embeddings
+        max_chunk_size (int): max chunk size
+
+    Yields:
+        Generator[List[BaseNode], None, None]: list of nodes with embeddings
+    """
+    for i in range(0, len(lst), max_chunk_size):
+        yield lst[i : i + max_chunk_size]
+
+
 class ChromaVectorStore(BasePydanticVectorStore):
     """Chroma vector store.
 
@@ -129,27 +145,34 @@ def add(self, nodes: List[BaseNode]) -> List[str]:
         if not self._collection:
             raise ValueError("Collection not initialized")
 
-        embeddings = []
-        metadatas = []
-        ids = []
-        documents = []
-        for node in nodes:
-            embeddings.append(node.get_embedding())
-            metadatas.append(
-                node_to_metadata_dict(
-                    node, remove_text=True, flat_metadata=self.flat_metadata
+        max_chunk_size = 41665  # One less than the max chunk size for ChromaDB
+        node_chunks = chunk_list(nodes, max_chunk_size)
+
+        all_ids = []
+        for node_chunk in node_chunks:
+            embeddings = []
+            metadatas = []
+            ids = []
+            documents = []
+            for node in node_chunk:
+                embeddings.append(node.get_embedding())
+                metadatas.append(
+                    node_to_metadata_dict(
+                        node, remove_text=True, flat_metadata=self.flat_metadata
+                    )
                 )
+                ids.append(node.node_id)
+                documents.append(node.get_content(metadata_mode=MetadataMode.NONE))
+
+            self._collection.add(
+                embeddings=embeddings,
+                ids=ids,
+                metadatas=metadatas,
+                documents=documents,
             )
-            ids.append(node.node_id)
-            documents.append(node.get_content(metadata_mode=MetadataMode.NONE))
-
-        self._collection.add(
-            embeddings=embeddings,
-            ids=ids,
-            metadatas=metadatas,
-            documents=documents,
-        )
-        return ids
+            all_ids.extend(ids)
+
+        return all_ids
 
     def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
         """


Code review comments:
diff_hunk: @@ -129,27 +145,34 @@ def add(self, nodes: List[BaseNode]) -> List[str]:
         if not self._collection:
             raise ValueError("Collection not initialized")
 
-        embeddings = []
-        metadatas = []
-        ids = []
-        documents = []
-        for node in nodes:
-            embeddings.append(node.get_embedding())
-            metadatas.append(
-                node_to_metadata_dict(
-                    node, remove_text=True, flat_metadata=self.flat_metadata
+        max_chunk_size = 41665  # One less than the max chunk size for ChromaDB
path: llama_index/vector_stores/chroma.py
text: @user1:
very nit: could we define this as `MAX_CHUNK_SIZE` constant at the top of the file?

@author:
100% legit, I should have done that at the start. Apologies. Change pushed.