bhetherman/open-webui
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac: use MilvusClient instead of deprecated ORM-style PyMilvus APIs (#27521)

Browse Source 
* 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.
This commit is contained in:
Classic298 authored and GitHub committed 2026-07-26 18:12:08 -04:00
1 parent f2ff310b2a
commit a15e44a5ff
2 files changed
+67 -81
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/open_webui/retrieval/vector/dbs/milvus.py
+5 -7
View File
@@ -25,7 +25,7 @@ from open_webui.retrieval.vector.main import (
VectorItem, VectorItem,
) )
from open_webui.retrieval.vector.utils import process_metadata 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 import MilvusClient as Client
from pymilvus.exceptions import MilvusException from pymilvus.exceptions import MilvusException
@@ -202,8 +202,6 @@ class MilvusClient(VectorDBBase):
return self._result_to_search_result(result) return self._result_to_search_result(result)
def query(self, collection_name: str, filter: dict, limit: int = -1): 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('-', '_') collection_name = collection_name.replace('-', '_')
if not self.has_collection(collection_name): if not self.has_collection(collection_name):
log.warning(f'Query attempted on non-existent collection: {self.collection_prefix}_{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) filter_string = ' && '.join(filter_expressions)
collection = Collection(f'{self.collection_prefix}_{collection_name}') self.client.load_collection(collection_name=f'{self.collection_prefix}_{collection_name}')
collection.load()
try: try:
log.info( log.info(
f"Querying collection {self.collection_prefix}_{collection_name} with filter: '{filter_string}', limit: {limit}" f"Querying collection {self.collection_prefix}_{collection_name} with filter: '{filter_string}', limit: {limit}"
) )
iterator = collection.query_iterator( iterator = self.client.query_iterator(
expr=filter_string, collection_name=f'{self.collection_prefix}_{collection_name}',
filter=filter_string,
output_fields=[ output_fields=[
'id', 'id',
'data', 'data',
backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py
+62 -74
View File
@@ -23,14 +23,8 @@ from open_webui.retrieval.vector.main import (
VectorDBBase, VectorDBBase,
VectorItem, VectorItem,
) )
from pymilvus import ( from pymilvus import DataType
Collection, from pymilvus import MilvusClient as Client
CollectionSchema,
DataType,
FieldSchema,
connections,
utility,
)
from pymilvus.exceptions import MilvusException from pymilvus.exceptions import MilvusException
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@@ -70,12 +64,7 @@ class MilvusClient(VectorDBBase):
def __init__(self): def __init__(self):
# Milvus collection names can only contain numbers, letters, and underscores. # Milvus collection names can only contain numbers, letters, and underscores.
self.collection_prefix = MILVUS_COLLECTION_PREFIX.replace('-', '_') self.collection_prefix = MILVUS_COLLECTION_PREFIX.replace('-', '_')
connections.connect( self.client = Client(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB)
alias='default',
uri=MILVUS_URI,
token=MILVUS_TOKEN,
db_name=MILVUS_DB,
)
# Main collection types for multi-tenancy # Main collection types for multi-tenancy
self.MEMORY_COLLECTION = f'{self.collection_prefix}_memories' self.MEMORY_COLLECTION = f'{self.collection_prefix}_memories'
@@ -116,63 +105,66 @@ class MilvusClient(VectorDBBase):
return self.KNOWLEDGE_COLLECTION, resource_id return self.KNOWLEDGE_COLLECTION, resource_id
def _create_shared_collection(self, mt_collection_name: str, dimension: int): def _create_shared_collection(self, mt_collection_name: str, dimension: int):
fields = [ schema = self.client.create_schema(auto_id=False, description='Shared collection for multi-tenancy')
FieldSchema( schema.add_field(field_name='id', datatype=DataType.VARCHAR, is_primary=True, max_length=36)
name='id', schema.add_field(field_name='vector', datatype=DataType.FLOAT_VECTOR, dim=dimension)
dtype=DataType.VARCHAR, schema.add_field(field_name='text', datatype=DataType.VARCHAR, max_length=MILVUS_TEXT_MAX_LENGTH)
is_primary=True, schema.add_field(field_name='metadata', datatype=DataType.JSON)
auto_id=False, schema.add_field(field_name=RESOURCE_ID_FIELD, datatype=DataType.VARCHAR, max_length=255)
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)
index_params = { index_build_params = {}
'metric_type': MILVUS_METRIC_TYPE,
'index_type': MILVUS_INDEX_TYPE,
'params': {},
}
if MILVUS_INDEX_TYPE == 'HNSW': if MILVUS_INDEX_TYPE == 'HNSW':
index_params['params'] = { index_build_params = {
'M': MILVUS_HNSW_M, 'M': MILVUS_HNSW_M,
'efConstruction': MILVUS_HNSW_EFCONSTRUCTION, 'efConstruction': MILVUS_HNSW_EFCONSTRUCTION,
} }
elif MILVUS_INDEX_TYPE == 'IVF_FLAT': 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: try:
# A Milvus server auto-selects the scalar index type; embedded # A Milvus server auto-selects the scalar index type from a parameterless call.
# Milvus Lite requires an explicit one. self.client.create_index(
collection.create_index(RESOURCE_ID_FIELD) collection_name=mt_collection_name,
index_params=self.client.prepare_index_params(field_name=RESOURCE_ID_FIELD),
)
except MilvusException: except MilvusException:
try: 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: except MilvusException as e:
# The index only accelerates resource_id filters; never fail # The index only accelerates resource_id filters; never fail
# collection creation over it. # collection creation over it.
log.warning(f'Could not create {RESOURCE_ID_FIELD} index on {mt_collection_name}: {e}') log.warning(f'Could not create {RESOURCE_ID_FIELD} index on {mt_collection_name}: {e}')
log.info(f'Created shared collection: {mt_collection_name}') log.info(f'Created shared collection: {mt_collection_name}')
return collection
def _ensure_collection(self, mt_collection_name: str, dimension: int): 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) self._create_shared_collection(mt_collection_name, dimension)
def has_collection(self, collection_name: str) -> bool: def has_collection(self, collection_name: str) -> bool:
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
_validate_resource_id(resource_id) _validate_resource_id(resource_id)
if not utility.has_collection(mt_collection): if not self.client.has_collection(mt_collection):
return False return False
collection = Collection(mt_collection) self.client.load_collection(mt_collection)
collection.load() res = self.client.query(
res = collection.query(expr=f"{RESOURCE_ID_FIELD} == '{resource_id}'", limit=1) collection_name=mt_collection,
filter=f"{RESOURCE_ID_FIELD} == '{resource_id}'",
output_fields=['id'],
limit=1,
)
return len(res) > 0 return len(res) > 0
def upsert(self, collection_name: str, items: List[VectorItem]): def upsert(self, collection_name: str, items: List[VectorItem]):
@@ -182,7 +174,6 @@ class MilvusClient(VectorDBBase):
_validate_resource_id(resource_id) _validate_resource_id(resource_id)
dimension = len(items[0]['vector']) dimension = len(items[0]['vector'])
self._ensure_collection(mt_collection, dimension) self._ensure_collection(mt_collection, dimension)
collection = Collection(mt_collection)
entities = [] entities = []
for item in items: for item in items:
@@ -205,7 +196,7 @@ class MilvusClient(VectorDBBase):
) )
try: try:
collection.insert(entities) self.client.insert(collection_name=mt_collection, data=entities)
except MilvusException as e: except MilvusException as e:
log.error( log.error(
f'Milvus insert failed (collection={mt_collection}, ' 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) mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
_validate_resource_id(resource_id) _validate_resource_id(resource_id)
if not utility.has_collection(mt_collection): if not self.client.has_collection(mt_collection):
return None return None
collection = Collection(mt_collection) self.client.load_collection(mt_collection)
collection.load()
search_params = {'metric_type': MILVUS_METRIC_TYPE, 'params': {}} results = self.client.search(
results = collection.search( collection_name=mt_collection,
data=vectors, data=vectors,
anns_field='vector', anns_field='vector',
param=search_params, search_params={'metric_type': MILVUS_METRIC_TYPE, 'params': {}},
limit=limit, limit=limit,
expr=f"{RESOURCE_ID_FIELD} == '{resource_id}'", filter=f"{RESOURCE_ID_FIELD} == '{resource_id}'",
output_fields=['id', 'text', 'metadata'], output_fields=['id', 'text', 'metadata'],
) )
@@ -245,10 +235,11 @@ class MilvusClient(VectorDBBase):
for hits in results: for hits in results:
batch_ids, batch_docs, batch_metadatas, batch_dists = [], [], [], [] batch_ids, batch_docs, batch_metadatas, batch_dists = [], [], [], []
for hit in hits: for hit in hits:
batch_ids.append(hit.entity.get('id')) entity = hit.get('entity', {})
batch_docs.append(hit.entity.get('text')) batch_ids.append(entity.get('id'))
batch_metadatas.append(hit.entity.get('metadata')) batch_docs.append(entity.get('text'))
batch_dists.append(hit.distance) batch_metadatas.append(entity.get('metadata'))
batch_dists.append(hit.get('distance'))
ids.append(batch_ids) ids.append(batch_ids)
documents.append(batch_docs) documents.append(batch_docs)
metadatas.append(batch_metadatas) metadatas.append(batch_metadatas)
@@ -264,11 +255,9 @@ class MilvusClient(VectorDBBase):
): ):
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
_validate_resource_id(resource_id) _validate_resource_id(resource_id)
if not utility.has_collection(mt_collection): if not self.client.has_collection(mt_collection):
return return
collection = Collection(mt_collection)
expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"] expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"]
if ids: if ids:
# Milvus expects a string list for 'in' operator # Milvus expects a string list for 'in' operator
@@ -280,30 +269,28 @@ class MilvusClient(VectorDBBase):
_validate_metadata_key(key) _validate_metadata_key(key)
expr.append(f"metadata['{key}'] == '{_escape_milvus_string(str(value))}'") 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): def reset(self):
for collection_name in self.shared_collections: for collection_name in self.shared_collections:
if utility.has_collection(collection_name): if self.client.has_collection(collection_name):
utility.drop_collection(collection_name) self.client.drop_collection(collection_name)
def delete_collection(self, collection_name: str): def delete_collection(self, collection_name: str):
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name) mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
_validate_resource_id(resource_id) _validate_resource_id(resource_id)
if not utility.has_collection(mt_collection): if not self.client.has_collection(mt_collection):
return return
collection = Collection(mt_collection) self.client.delete(collection_name=mt_collection, filter=f"{RESOURCE_ID_FIELD} == '{resource_id}'")
collection.delete(f"{RESOURCE_ID_FIELD} == '{resource_id}'")
def query(self, collection_name: str, filter: Dict[str, Any], limit: Optional[int] = None) -> Optional[GetResult]: 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) mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
_validate_resource_id(resource_id) _validate_resource_id(resource_id)
if not utility.has_collection(mt_collection): if not self.client.has_collection(mt_collection):
return None return None
collection = Collection(mt_collection) self.client.load_collection(mt_collection)
collection.load()
expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"] expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"]
if filter: if filter:
@@ -318,8 +305,9 @@ class MilvusClient(VectorDBBase):
else: else:
raise TypeError(f'Unsupported Milvus filter value type for key {key!r}: {type(value).__name__}') raise TypeError(f'Unsupported Milvus filter value type for key {key!r}: {type(value).__name__}')
iterator = collection.query_iterator( iterator = self.client.query_iterator(
expr=' and '.join(expr), collection_name=mt_collection,
filter=' and '.join(expr),
output_fields=['id', 'text', 'metadata'], output_fields=['id', 'text', 'metadata'],
limit=limit if limit else -1, limit=limit if limit else -1,
) )