diff --git a/backend/open_webui/retrieval/vector/dbs/milvus.py b/backend/open_webui/retrieval/vector/dbs/milvus.py index 9b356b1c69..1e8439cdaa 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus.py @@ -27,9 +27,15 @@ from open_webui.retrieval.vector.main import ( from open_webui.retrieval.vector.utils import process_metadata from pymilvus import Collection, DataType, FieldSchema, connections from pymilvus import MilvusClient as Client +from pymilvus.exceptions import MilvusException log = logging.getLogger(__name__) +# Milvus caps stored text length (here the chunk lives under the JSON `data` +# field). Clamp long chunks before insert so one oversized chunk can't fail the +# whole batch and leave the file with zero embeddings. +MILVUS_TEXT_MAX_LENGTH = 65535 + class MilvusClient(VectorDBBase): def __init__(self): @@ -270,18 +276,34 @@ class MilvusClient(VectorDBBase): self._create_collection(collection_name=collection_name, dimension=len(items[0]['vector'])) log.info(f'Inserting {len(items)} items into collection {self.collection_prefix}_{collection_name}.') - return self.client.insert( - collection_name=f'{self.collection_prefix}_{collection_name}', - data=[ + data = [] + for item in items: + text = item['text'] or '' + if len(text) > MILVUS_TEXT_MAX_LENGTH: + log.warning( + f'Milvus: truncating text id={item["id"]} ' + f'{len(text)}->{MILVUS_TEXT_MAX_LENGTH} chars' + ) + text = text[:MILVUS_TEXT_MAX_LENGTH] + data.append( { 'id': item['id'], 'vector': item['vector'], - 'data': {'text': item['text']}, + 'data': {'text': text}, 'metadata': process_metadata(item['metadata']), } - for item in items - ], - ) + ) + try: + return self.client.insert( + collection_name=f'{self.collection_prefix}_{collection_name}', + data=data, + ) + except MilvusException as e: + log.error( + f'Milvus insert failed for {self.collection_prefix}_{collection_name} ' + f'({len(items)} items): {e}' + ) + raise def upsert(self, collection_name: str, items: list[VectorItem]): # Update the items in the collection, if the items are not present, insert them. If the collection does not exist, it will be created. @@ -298,18 +320,34 @@ class MilvusClient(VectorDBBase): self._create_collection(collection_name=collection_name, dimension=len(items[0]['vector'])) log.info(f'Upserting {len(items)} items into collection {self.collection_prefix}_{collection_name}.') - return self.client.upsert( - collection_name=f'{self.collection_prefix}_{collection_name}', - data=[ + data = [] + for item in items: + text = item['text'] or '' + if len(text) > MILVUS_TEXT_MAX_LENGTH: + log.warning( + f'Milvus: truncating text id={item["id"]} ' + f'{len(text)}->{MILVUS_TEXT_MAX_LENGTH} chars' + ) + text = text[:MILVUS_TEXT_MAX_LENGTH] + data.append( { 'id': item['id'], 'vector': item['vector'], - 'data': {'text': item['text']}, + 'data': {'text': text}, 'metadata': process_metadata(item['metadata']), } - for item in items - ], - ) + ) + try: + return self.client.upsert( + collection_name=f'{self.collection_prefix}_{collection_name}', + data=data, + ) + except MilvusException as e: + log.error( + f'Milvus upsert failed for {self.collection_prefix}_{collection_name} ' + f'({len(items)} items): {e}' + ) + raise def delete( self,