ewilliams28/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-23 22:14:58 -06:00
Code Issues 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
2026-07-27 00:12:08 +02:00
committed by GitHub
parent f2ff310b2a
commit a15e44a5ff
2 changed files with 67 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/vector/dbs/milvus.py
+5 -7
View File
@@ -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',
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,
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,
)