123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182 |
- import json
- from datetime import timedelta
- from django.http import JsonResponse
- from minio import Minio
- from pymilvus import Collection, FieldSchema, CollectionSchema, DataType, utility, connections
- from DCbackend import settings
- from DCbackend.settings import MILVUS_HOST, MILVUS_PORT, MILVUS_USER, MILVUS_PASSWORD
- from DCbackend.utils.common import success, fail, pageData
- from backend.models import TaskSublist, DocumentKbm, Knowledgebase
- from base import logger
- minio_client = Minio(
- settings.MINIO_ENDPOINT,
- access_key=settings.MINIO_ACCESS_KEY,
- secret_key=settings.MINIO_SECRET_KEY,
- secure=settings.MINIO_SECURE
- )
- class MilvusService:
- @staticmethod
- def list_all_collections(request):
- """
- 列出 Milvus 中的所有集合
- """
- try:
- # 连接到 Milvus
- connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
- # 获取所有集合名称
- collection_names = utility.list_collections()
- # 获取每个集合的详细信息
- collections_info = []
- for name in collection_names:
- collection = Collection(name)
- schema = collection.schema
- schema_dict = {
- "fields": [
- {
- "name": field.name,
- "dtype": str(field.dtype),
- "is_primary": field.is_primary,
- "auto_id": field.auto_id,
- "description": field.description,
- "max_length": field.max_length if hasattr(field, 'max_length') else None,
- "dim": field.dim if field.dtype == DataType.FLOAT_VECTOR else None
- } for field in schema.fields
- ],
- "description": schema.description
- }
- info = {
- "name": name,
- "entities": collection.num_entities,
- "schema": schema_dict
- }
- collections_info.append(info)
- return success(collections_info) # 假设 success 函数可以处理字典列表
- except Exception as e:
- print(f"列出集合时发生错误: {str(e)}")
- raise
- finally:
- # 断开 Milvus 连接
- connections.disconnect("default")
- @staticmethod
- def delete_collection(request):
- """
- 删除 Milvus 中的指定集合
- :param request: HTTP 请求对象
- :return: JsonResponse 对象
- """
- collection_name = request.POST.get("collection_name")
- if not collection_name:
- return fail("集合名称未提供")
- try:
- # 连接到 Milvus
- connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
- # 检查集合是否存在
- if not utility.has_collection(collection_name):
- return fail(f"集合 '{collection_name}' 不存在")
- # 删除集合
- utility.drop_collection(collection_name)
- return success( f"集合 '{collection_name}' 已成功删除")
- except Exception as e:
- return fail(f"删除集合时发生错误: {str(e)}")
- finally:
- # 断开 Milvus 连接
- connections.disconnect("default")
- @staticmethod
- def delete_milvus_data(request):
- """
- 从 Milvus 中删除指定 ID 的数据
- :param request: HTTP 请求对象
- :return: JsonResponse 对象
- """
- collection_name = request.POST.get("collection_name")
- id_to_delete = request.POST.get("id")
- if not collection_name:
- return fail("集合名称未提供")
- if not id_to_delete:
- return fail("要删除的 ID 未提供")
- try:
- # 连接到 Milvus
- connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
- # 检查集合是否存在
- if not utility.has_collection(collection_name):
- return fail(f"集合 '{collection_name}' 不存在")
- # 获取集合对象
- collection = Collection(collection_name)
- # 执行删除操作
- expr = f'id in [{id_to_delete}]'
- delete_result = collection.delete(expr)
- if delete_result.delete_count > 0:
- return success(f"ID 为 {id_to_delete} 的数据已成功从集合 '{collection_name}' 中删除")
- else:
- return fail(f"未找到 ID 为 {id_to_delete} 的数据")
- except Exception as e:
- return fail(f"删除数据时发生错误: {str(e)}")
- finally:
- # 断开 Milvus 连接
- connections.disconnect("default")
- @staticmethod
- def getMinioURl(request):
- data = json.loads(request.body)
- id = data.get("id")
- logger.info(f"request:{request}")
- if not id:
- return fail("id为空")
- task = TaskSublist.objects.filter(milvus_id=id).first()
- if not task:
- return fail("无此数据")
- document = DocumentKbm.objects.filter(id = task.doc_id).first()
- if not document:
- return fail("未找到文件")
- kmb = Knowledgebase.objects.filter(id=document.kb_id).first()
- object_name = document.location
- bucket_name = kmb.location
- if not object_name:
- return fail('Object name is required')
- url = minio_client.presigned_get_object(
- bucket_name,
- object_name,
- expires=timedelta(days=1)
- )
- bucket_info = {
- 'object_name': object_name,
- 'url': url,
- }
- return success(bucket_info)
|