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

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/vector/dbs/milvus_multitenancy.py
+55 -73
View File
@@ -33,26 +33,26 @@ from pymilvus import (
log = logging.getLogger(__name__)
RESOURCE_ID_FIELD = "resource_id"
RESOURCE_ID_FIELD = 'resource_id'
class MilvusClient(VectorDBBase):
def __init__(self):
# 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(
alias="default",
alias='default',
uri=MILVUS_URI,
token=MILVUS_TOKEN,
db_name=MILVUS_DB,
)
# Main collection types for multi-tenancy
self.MEMORY_COLLECTION = f"{self.collection_prefix}_memories"
self.KNOWLEDGE_COLLECTION = f"{self.collection_prefix}_knowledge"
self.FILE_COLLECTION = f"{self.collection_prefix}_files"
self.WEB_SEARCH_COLLECTION = f"{self.collection_prefix}_web_search"
self.HASH_BASED_COLLECTION = f"{self.collection_prefix}_hash_based"
self.MEMORY_COLLECTION = f'{self.collection_prefix}_memories'
self.KNOWLEDGE_COLLECTION = f'{self.collection_prefix}_knowledge'
self.FILE_COLLECTION = f'{self.collection_prefix}_files'
self.WEB_SEARCH_COLLECTION = f'{self.collection_prefix}_web_search'
self.HASH_BASED_COLLECTION = f'{self.collection_prefix}_hash_based'
self.shared_collections = [
self.MEMORY_COLLECTION,
self.KNOWLEDGE_COLLECTION,
@@ -74,15 +74,13 @@ class MilvusClient(VectorDBBase):
"""
resource_id = collection_name
if collection_name.startswith("user-memory-"):
if collection_name.startswith('user-memory-'):
return self.MEMORY_COLLECTION, resource_id
elif collection_name.startswith("file-"):
elif collection_name.startswith('file-'):
return self.FILE_COLLECTION, resource_id
elif collection_name.startswith("web-search-"):
elif collection_name.startswith('web-search-'):
return self.WEB_SEARCH_COLLECTION, resource_id
elif len(collection_name) == 63 and all(
c in "0123456789abcdef" for c in collection_name
):
elif len(collection_name) == 63 and all(c in '0123456789abcdef' for c in collection_name):
return self.HASH_BASED_COLLECTION, resource_id
else:
return self.KNOWLEDGE_COLLECTION, resource_id
@@ -90,36 +88,36 @@ class MilvusClient(VectorDBBase):
def _create_shared_collection(self, mt_collection_name: str, dimension: int):
fields = [
FieldSchema(
name="id",
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='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")
schema = CollectionSchema(fields, 'Shared collection for multi-tenancy')
collection = Collection(mt_collection_name, schema)
index_params = {
"metric_type": MILVUS_METRIC_TYPE,
"index_type": MILVUS_INDEX_TYPE,
"params": {},
'metric_type': MILVUS_METRIC_TYPE,
'index_type': MILVUS_INDEX_TYPE,
'params': {},
}
if MILVUS_INDEX_TYPE == "HNSW":
index_params["params"] = {
"M": MILVUS_HNSW_M,
"efConstruction": MILVUS_HNSW_EFCONSTRUCTION,
if MILVUS_INDEX_TYPE == 'HNSW':
index_params['params'] = {
'M': MILVUS_HNSW_M,
'efConstruction': MILVUS_HNSW_EFCONSTRUCTION,
}
elif MILVUS_INDEX_TYPE == "IVF_FLAT":
index_params["params"] = {"nlist": MILVUS_IVF_FLAT_NLIST}
elif MILVUS_INDEX_TYPE == 'IVF_FLAT':
index_params['params'] = {'nlist': MILVUS_IVF_FLAT_NLIST}
collection.create_index("vector", index_params)
collection.create_index('vector', index_params)
collection.create_index(RESOURCE_ID_FIELD)
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):
@@ -127,9 +125,7 @@ class MilvusClient(VectorDBBase):
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
)
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
if not utility.has_collection(mt_collection):
return False
@@ -141,19 +137,17 @@ class MilvusClient(VectorDBBase):
def upsert(self, collection_name: str, items: List[VectorItem]):
if not items:
return
mt_collection, resource_id = self._get_collection_and_resource_id(
collection_name
)
dimension = len(items[0]["vector"])
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
dimension = len(items[0]['vector'])
self._ensure_collection(mt_collection, dimension)
collection = Collection(mt_collection)
entities = [
{
"id": item["id"],
"vector": item["vector"],
"text": item["text"],
"metadata": item["metadata"],
'id': item['id'],
'vector': item['vector'],
'text': item['text'],
'metadata': item['metadata'],
RESOURCE_ID_FIELD: resource_id,
}
for item in items
@@ -170,41 +164,37 @@ class MilvusClient(VectorDBBase):
if not vectors:
return None
mt_collection, resource_id = self._get_collection_and_resource_id(
collection_name
)
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
if not utility.has_collection(mt_collection):
return None
collection = Collection(mt_collection)
collection.load()
search_params = {"metric_type": MILVUS_METRIC_TYPE, "params": {}}
search_params = {'metric_type': MILVUS_METRIC_TYPE, 'params': {}}
results = collection.search(
data=vectors,
anns_field="vector",
anns_field='vector',
param=search_params,
limit=limit,
expr=f"{RESOURCE_ID_FIELD} == '{resource_id}'",
output_fields=["id", "text", "metadata"],
output_fields=['id', 'text', 'metadata'],
)
ids, documents, metadatas, distances = [], [], [], []
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_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)
ids.append(batch_ids)
documents.append(batch_docs)
metadatas.append(batch_metadatas)
distances.append(batch_dists)
return SearchResult(
ids=ids, documents=documents, metadatas=metadatas, distances=distances
)
return SearchResult(ids=ids, documents=documents, metadatas=metadatas, distances=distances)
def delete(
self,
@@ -212,9 +202,7 @@ class MilvusClient(VectorDBBase):
ids: Optional[List[str]] = None,
filter: Optional[Dict[str, Any]] = None,
):
mt_collection, resource_id = self._get_collection_and_resource_id(
collection_name
)
mt_collection, resource_id = self._get_collection_and_resource_id(collection_name)
if not utility.has_collection(mt_collection):
return
@@ -224,14 +212,14 @@ class MilvusClient(VectorDBBase):
expr = [f"{RESOURCE_ID_FIELD} == '{resource_id}'"]
if ids:
# Milvus expects a string list for 'in' operator
id_list_str = ", ".join([f"'{id_val}'" for id_val in ids])
expr.append(f"id in [{id_list_str}]")
id_list_str = ', '.join([f"'{id_val}'" for id_val in ids])
expr.append(f'id in [{id_list_str}]')
if filter:
for key, value in filter.items():
expr.append(f"metadata['{key}'] == '{value}'")
collection.delete(" and ".join(expr))
collection.delete(' and '.join(expr))
def reset(self):
for collection_name in self.shared_collections:
@@ -239,21 +227,15 @@ class MilvusClient(VectorDBBase):
utility.drop_collection(collection_name)
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)
if not utility.has_collection(mt_collection):
return
collection = Collection(mt_collection)
collection.delete(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
)
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)
if not utility.has_collection(mt_collection):
return None
@@ -269,8 +251,8 @@ class MilvusClient(VectorDBBase):
expr.append(f"metadata['{key}'] == {value}")
iterator = collection.query_iterator(
expr=" and ".join(expr),
output_fields=["id", "text", "metadata"],
expr=' and '.join(expr),
output_fields=['id', 'text', 'metadata'],
limit=limit if limit else -1,
)
@@ -282,9 +264,9 @@ class MilvusClient(VectorDBBase):
break
all_results.extend(batch)
ids = [res["id"] for res in all_results]
documents = [res["text"] for res in all_results]
metadatas = [res["metadata"] for res in all_results]
ids = [res['id'] for res in all_results]
documents = [res['text'] for res in all_results]
metadatas = [res['metadata'] for res in all_results]
return GetResult(ids=[ids], documents=[documents], metadatas=[metadatas])