fix: batch insertions into Milvus vector store (run-llama#12837)

llama-index-integrations/vector_stores/llama-index-vector-stores-milvus/llama_index/vector_stores/milvus/base.py

@@ -10,6 +10,7 @@
 import pymilvus  # noqa
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.schema import BaseNode, TextNode
+from llama_index.core.utils import iter_batch
 from llama_index.core.vector_stores.types import (
     BasePydanticVectorStore,
     MetadataFilters,
@@ -27,6 +28,7 @@
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_BATCH_SIZE = 100
 MILVUS_ID_FIELD = "id"
 
 
@@ -118,6 +120,7 @@ class MilvusVectorStore(BasePydanticVectorStore):
     output_fields: List[str] = Field(default_factory=list)
     index_config: Optional[dict]
     search_config: Optional[dict]
+    batch_size: int = DEFAULT_BATCH_SIZE
 
     _milvusclient: MilvusClient = PrivateAttr()
     _collection: Any = PrivateAttr()
@@ -137,6 +140,7 @@ def __init__(
         output_fields: Optional[List[str]] = None,
         index_config: Optional[dict] = None,
         search_config: Optional[dict] = None,
+        batch_size: int = DEFAULT_BATCH_SIZE,
         **kwargs: Any,
     ) -> None:
         """Init params."""
@@ -151,6 +155,7 @@ def __init__(
             output_fields=output_fields or [],
             index_config=index_config if index_config else {},
             search_config=search_config if search_config else {},
+            batch_size=batch_size,
         )
 
         # Select the similarity metric
@@ -225,7 +230,8 @@ def add(self, nodes: List[BaseNode], **add_kwargs: Any) -> List[str]:
             insert_list.append(entry)
 
         # Insert the data into milvus
-        self._collection.insert(insert_list)
+        for insert_batch in iter_batch(insert_list, self.batch_size):
+            self._collection.insert(insert_batch)
         if add_kwargs.get("force_flush", False):
             self._collection.flush()
         self._create_index_if_required()

llama-index-integrations/vector_stores/llama-index-vector-stores-milvus/pyproject.toml

@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-milvus"
 readme = "README.md"
-version = "0.1.9"
+version = "0.1.10"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"