diff --git a/backend/open_webui/retrieval/vector/dbs/milvus.py b/backend/open_webui/retrieval/vector/dbs/milvus.py index 0bf4a5955d..2eefb8f67f 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus.py @@ -86,6 +86,11 @@ def _truncate_text(text: str) -> str: return text.encode()[:MILVUS_TEXT_MAX_LENGTH].decode(errors='ignore') +def _write_in_batches(write: Callable[..., Any], collection_name: str, rows: list[dict]) -> None: + for start in range(0, len(rows), BM25_BACKFILL_BATCH_SIZE): + write(collection_name=collection_name, data=rows[start : start + BM25_BACKFILL_BATCH_SIZE]) + + def _bm25_rows(rows: list[dict]) -> list[dict]: return [ { @@ -580,10 +585,7 @@ class MilvusClient(VectorDBBase): row['data'] = {'text': text} data.append(row) try: - return self.client.insert( - collection_name=f'{self.collection_prefix}_{collection_name}', - data=data, - ) + _write_in_batches(self.client.insert, f'{self.collection_prefix}_{collection_name}', data) except MilvusException as e: log.error(f'Milvus insert failed for {self.collection_prefix}_{collection_name} ({len(items)} items): {e}') raise @@ -626,10 +628,7 @@ class MilvusClient(VectorDBBase): row['data'] = {'text': text} data.append(row) try: - return self.client.upsert( - collection_name=f'{self.collection_prefix}_{collection_name}', - data=data, - ) + _write_in_batches(self.client.upsert, f'{self.collection_prefix}_{collection_name}', data) except MilvusException as e: log.error(f'Milvus upsert failed for {self.collection_prefix}_{collection_name} ({len(items)} items): {e}') raise diff --git a/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py b/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py index d0924fa9ae..b6306b88a5 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py @@ -26,6 +26,7 @@ from open_webui.retrieval.vector.dbs.milvus import ( _metadata_exprs, _supports_bm25, _truncate_text, + _write_in_batches, ) from open_webui.retrieval.vector.main import ( GetResult, @@ -237,7 +238,7 @@ class MilvusClient(VectorDBBase): ) try: - self.client.insert(collection_name=mt_collection, data=entities) + _write_in_batches(self.client.insert, mt_collection, entities) except MilvusException as e: log.error( f'Milvus insert failed (collection={mt_collection}, '