From a15e44a5ffd48e08a6d4bf2bb99638f2e2e14c38 Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Mon, 27 Jul 2026 00:12:08 +0200 Subject: [PATCH] refac: use MilvusClient instead of deprecated ORM-style PyMilvus APIs (#27521) * refac: use MilvusClient instead of deprecated ORM-style PyMilvus APIs PyMilvus 2.6 emits a PyMilvusDeprecationWarning for every ORM-style call (`connections.connect`, `utility.*`, `Collection` and its methods) and will remove those APIs in PyMilvus 3.1. Both Milvus backends still used them, so a running instance floods its logs with deprecation warnings during indexing and retrieval, and would break outright once PyMilvus 3.1 lands. Both vector clients now go through `MilvusClient`: - `milvus_multitenancy.py`: collection creation, index creation, has_collection, insert, search, query iteration, delete and reset. - `milvus.py`: the remaining ORM calls in `query()` (`connections.connect`, `Collection(...).load()`, `Collection.query_iterator`), plus the now-unused `FieldSchema` import. Behaviour is unchanged: same schema, same index parameters and the same two-step scalar-index fallback, same filter expressions, same result shapes. Verified against embedded Milvus (milvus-lite, pymilvus 2.6.14) with a functional harness over both clients: insert, get, query by string/int/bool metadata filters, vector search, tenant isolation, oversized-text truncation, delete by id and by filter, delete_collection and reset all return identical results before and after, while the deprecation warnings drop from 57 to 0 for the multi-tenancy client and from 16 to 0 for the standard one. One Milvus Lite nuance worth recording: `MilvusClient` sends index build parameters (`M`, `efConstruction`, `nlist`) as flat keys rather than as a nested `params` blob. A Milvus server accepts both forms, Milvus Lite only reads the nested one, so those tuning values are ignored on Lite. `MilvusClient` offers no way to send the nested form, and `milvus.py` already built its index parameters this way, so both backends are now consistent. Fixes #26978 * refac: correct the Milvus scalar-index comment The comment claimed that embedded Milvus Lite requires an explicit scalar index type. It does not: Milvus Lite rejects `create_index` on a VARCHAR field outright ("create_index only supports vector fields"), for every index type and with or without a metric type, so neither the parameterless call nor the explicit INVERTED fallback can succeed there. Filtered queries on `resource_id` still work on Lite, just unindexed. Only the accurate half is kept, which is the reason the parameterless call is deliberate rather than an omission. --- .../open_webui/retrieval/vector/dbs/milvus.py | 12 +- .../vector/dbs/milvus_multitenancy.py | 136 ++++++++---------- 2 files changed, 67 insertions(+), 81 deletions(-) diff --git a/backend/open_webui/retrieval/vector/dbs/milvus.py b/backend/open_webui/retrieval/vector/dbs/milvus.py index b0331e3eea..fa2abe85d7 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus.py @@ -25,7 +25,7 @@ from open_webui.retrieval.vector.main import ( VectorItem, ) from open_webui.retrieval.vector.utils import process_metadata -from pymilvus import Collection, DataType, FieldSchema, connections +from pymilvus import DataType from pymilvus import MilvusClient as Client from pymilvus.exceptions import MilvusException @@ -202,8 +202,6 @@ class MilvusClient(VectorDBBase): return self._result_to_search_result(result) def query(self, collection_name: str, filter: dict, limit: int = -1): - connections.connect(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB) - collection_name = collection_name.replace('-', '_') if not self.has_collection(collection_name): log.warning(f'Query attempted on non-existent collection: {self.collection_prefix}_{collection_name}') @@ -218,16 +216,16 @@ class MilvusClient(VectorDBBase): filter_string = ' && '.join(filter_expressions) - collection = Collection(f'{self.collection_prefix}_{collection_name}') - collection.load() + self.client.load_collection(collection_name=f'{self.collection_prefix}_{collection_name}') try: log.info( f"Querying collection {self.collection_prefix}_{collection_name} with filter: '{filter_string}', limit: {limit}" ) - iterator = collection.query_iterator( - expr=filter_string, + iterator = self.client.query_iterator( + collection_name=f'{self.collection_prefix}_{collection_name}', + filter=filter_string, output_fields=[ 'id', 'data', diff --git a/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py b/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py index f2f0b92817..599cb8712c 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py @@ -23,14 +23,8 @@ from open_webui.retrieval.vector.main import ( VectorDBBase, VectorItem, ) -from pymilvus import ( - Collection, - CollectionSchema, - DataType, - FieldSchema, - connections, - utility, -) +from pymilvus import DataType +from pymilvus import MilvusClient as Client from pymilvus.exceptions import MilvusException log = logging.getLogger(__name__) @@ -70,12 +64,7 @@ class MilvusClient(VectorDBBase): def __init__(self): # Milvus collection names can only contain numbers, letters, and underscores. self.collection_prefix = MILVUS_COLLECTION_PREFIX.replace('-', '_') - connections.connect( - alias='default', - uri=MILVUS_URI, - token=MILVUS_TOKEN, - db_name=MILVUS_DB, - ) + self.client = Client(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB) # Main collection types for multi-tenancy self.MEMORY_COLLECTION = f'{self.collection_prefix}_memories' @@ -116,63 +105,66 @@ class MilvusClient(VectorDBBase): return self.KNOWLEDGE_COLLECTION, resource_id def _create_shared_collection(self, mt_collection_name: str, dimension: int): - fields = [ - FieldSchema( - name='id', - dtype=DataType.VARCHAR, - is_primary=True, - auto_id=False, - max_length=36, - ), - FieldSchema(name='vector', dtype=DataType.FLOAT_VECTOR, dim=dimension), - FieldSchema(name='text', dtype=DataType.VARCHAR, max_length=65535), - FieldSchema(name='metadata', dtype=DataType.JSON), - FieldSchema(name=RESOURCE_ID_FIELD, dtype=DataType.VARCHAR, max_length=255), - ] - schema = CollectionSchema(fields, 'Shared collection for multi-tenancy') - collection = Collection(mt_collection_name, schema) + schema = self.client.create_schema(auto_id=False, description='Shared collection for multi-tenancy') + schema.add_field(field_name='id', datatype=DataType.VARCHAR, is_primary=True, max_length=36) + schema.add_field(field_name='vector', datatype=DataType.FLOAT_VECTOR, dim=dimension) + schema.add_field(field_name='text', datatype=DataType.VARCHAR, max_length=MILVUS_TEXT_MAX_LENGTH) + schema.add_field(field_name='metadata', datatype=DataType.JSON) + schema.add_field(field_name=RESOURCE_ID_FIELD, datatype=DataType.VARCHAR, max_length=255) - index_params = { - 'metric_type': MILVUS_METRIC_TYPE, - 'index_type': MILVUS_INDEX_TYPE, - 'params': {}, - } + index_build_params = {} if MILVUS_INDEX_TYPE == 'HNSW': - index_params['params'] = { + index_build_params = { 'M': MILVUS_HNSW_M, 'efConstruction': MILVUS_HNSW_EFCONSTRUCTION, } elif MILVUS_INDEX_TYPE == 'IVF_FLAT': - index_params['params'] = {'nlist': MILVUS_IVF_FLAT_NLIST} + index_build_params = {'nlist': MILVUS_IVF_FLAT_NLIST} - collection.create_index('vector', index_params) + vector_index = self.client.prepare_index_params( + field_name='vector', + index_type=MILVUS_INDEX_TYPE, + metric_type=MILVUS_METRIC_TYPE, + params=index_build_params, + ) + + self.client.create_collection(collection_name=mt_collection_name, schema=schema) + self.client.create_index(collection_name=mt_collection_name, index_params=vector_index) try: - # A Milvus server auto-selects the scalar index type; embedded - # Milvus Lite requires an explicit one. - collection.create_index(RESOURCE_ID_FIELD) + # A Milvus server auto-selects the scalar index type from a parameterless call. + self.client.create_index( + collection_name=mt_collection_name, + index_params=self.client.prepare_index_params(field_name=RESOURCE_ID_FIELD), + ) except MilvusException: try: - collection.create_index(RESOURCE_ID_FIELD, {'index_type': 'INVERTED'}) + self.client.create_index( + collection_name=mt_collection_name, + index_params=self.client.prepare_index_params(field_name=RESOURCE_ID_FIELD, index_type='INVERTED'), + ) except MilvusException as e: # The index only accelerates resource_id filters; never fail # collection creation over it. log.warning(f'Could not create {RESOURCE_ID_FIELD} index on {mt_collection_name}: {e}') log.info(f'Created shared collection: {mt_collection_name}') - return collection def _ensure_collection(self, mt_collection_name: str, dimension: int): - if not utility.has_collection(mt_collection_name): + if not self.client.has_collection(mt_collection_name): self._create_shared_collection(mt_collection_name, dimension) def has_collection(self, collection_name: str) -> bool: mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) _validate_resource_id(resource_id) - if not utility.has_collection(mt_collection): + if not self.client.has_collection(mt_collection): return False - collection = Collection(mt_collection) - collection.load() - res = collection.query(expr=f"{RESOURCE_ID_FIELD} == '{resource_id}'", limit=1) + self.client.load_collection(mt_collection) + res = self.client.query( + collection_name=mt_collection, + filter=f"{RESOURCE_ID_FIELD} == '{resource_id}'", + output_fields=['id'], + limit=1, + ) return len(res) > 0 def upsert(self, collection_name: str, items: List[VectorItem]): @@ -182,7 +174,6 @@ class MilvusClient(VectorDBBase): _validate_resource_id(resource_id) dimension = len(items[0]['vector']) self._ensure_collection(mt_collection, dimension) - collection = Collection(mt_collection) entities = [] for item in items: @@ -205,7 +196,7 @@ class MilvusClient(VectorDBBase): ) try: - collection.insert(entities) + self.client.insert(collection_name=mt_collection, data=entities) except MilvusException as e: log.error( f'Milvus insert failed (collection={mt_collection}, ' @@ -225,19 +216,18 @@ class MilvusClient(VectorDBBase): mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) _validate_resource_id(resource_id) - if not utility.has_collection(mt_collection): + if not self.client.has_collection(mt_collection): return None - collection = Collection(mt_collection) - collection.load() + self.client.load_collection(mt_collection) - search_params = {'metric_type': MILVUS_METRIC_TYPE, 'params': {}} - results = collection.search( + results = self.client.search( + collection_name=mt_collection, data=vectors, anns_field='vector', - param=search_params, + search_params={'metric_type': MILVUS_METRIC_TYPE, 'params': {}}, limit=limit, - expr=f"{RESOURCE_ID_FIELD} == '{resource_id}'", + filter=f"{RESOURCE_ID_FIELD} == '{resource_id}'", output_fields=['id', 'text', 'metadata'], ) @@ -245,10 +235,11 @@ class MilvusClient(VectorDBBase): for hits in results: batch_ids, batch_docs, batch_metadatas, batch_dists = [], [], [], [] for hit in hits: - batch_ids.append(hit.entity.get('id')) - batch_docs.append(hit.entity.get('text')) - batch_metadatas.append(hit.entity.get('metadata')) - batch_dists.append(hit.distance) + entity = hit.get('entity', {}) + batch_ids.append(entity.get('id')) + batch_docs.append(entity.get('text')) + batch_metadatas.append(entity.get('metadata')) + batch_dists.append(hit.get('distance')) ids.append(batch_ids) documents.append(batch_docs) metadatas.append(batch_metadatas) @@ -264,11 +255,9 @@ class MilvusClient(VectorDBBase): ): mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) _validate_resource_id(resource_id) - if not utility.has_collection(mt_collection): + if not self.client.has_collection(mt_collection): return - collection = Collection(mt_collection) - expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"] if ids: # Milvus expects a string list for 'in' operator @@ -280,30 +269,28 @@ class MilvusClient(VectorDBBase): _validate_metadata_key(key) expr.append(f"metadata['{key}'] == '{_escape_milvus_string(str(value))}'") - collection.delete(' and '.join(expr)) + self.client.delete(collection_name=mt_collection, filter=' and '.join(expr)) def reset(self): for collection_name in self.shared_collections: - if utility.has_collection(collection_name): - utility.drop_collection(collection_name) + if self.client.has_collection(collection_name): + self.client.drop_collection(collection_name) def delete_collection(self, collection_name: str): mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) _validate_resource_id(resource_id) - if not utility.has_collection(mt_collection): + if not self.client.has_collection(mt_collection): return - collection = Collection(mt_collection) - collection.delete(f"{RESOURCE_ID_FIELD} == '{resource_id}'") + self.client.delete(collection_name=mt_collection, filter=f"{RESOURCE_ID_FIELD} == '{resource_id}'") def query(self, collection_name: str, filter: Dict[str, Any], limit: Optional[int] = None) -> Optional[GetResult]: mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) _validate_resource_id(resource_id) - if not utility.has_collection(mt_collection): + if not self.client.has_collection(mt_collection): return None - collection = Collection(mt_collection) - collection.load() + self.client.load_collection(mt_collection) expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"] if filter: @@ -318,8 +305,9 @@ class MilvusClient(VectorDBBase): else: raise TypeError(f'Unsupported Milvus filter value type for key {key!r}: {type(value).__name__}') - iterator = collection.query_iterator( - expr=' and '.join(expr), + iterator = self.client.query_iterator( + collection_name=mt_collection, + filter=' and '.join(expr), output_fields=['id', 'text', 'metadata'], limit=limit if limit else -1, )