MilvusService.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182
  1. import json
  2. from datetime import timedelta
  3. from django.http import JsonResponse
  4. from minio import Minio
  5. from pymilvus import Collection, FieldSchema, CollectionSchema, DataType, utility, connections
  6. from DCbackend import settings
  7. from DCbackend.settings import MILVUS_HOST, MILVUS_PORT, MILVUS_USER, MILVUS_PASSWORD
  8. from DCbackend.utils.common import success, fail, pageData
  9. from backend.models import TaskSublist, DocumentKbm, Knowledgebase
  10. from base import logger
  11. minio_client = Minio(
  12. settings.MINIO_ENDPOINT,
  13. access_key=settings.MINIO_ACCESS_KEY,
  14. secret_key=settings.MINIO_SECRET_KEY,
  15. secure=settings.MINIO_SECURE
  16. )
  17. class MilvusService:
  18. @staticmethod
  19. def list_all_collections(request):
  20. """
  21. 列出 Milvus 中的所有集合
  22. """
  23. try:
  24. # 连接到 Milvus
  25. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
  26. # 获取所有集合名称
  27. collection_names = utility.list_collections()
  28. # 获取每个集合的详细信息
  29. collections_info = []
  30. for name in collection_names:
  31. collection = Collection(name)
  32. schema = collection.schema
  33. schema_dict = {
  34. "fields": [
  35. {
  36. "name": field.name,
  37. "dtype": str(field.dtype),
  38. "is_primary": field.is_primary,
  39. "auto_id": field.auto_id,
  40. "description": field.description,
  41. "max_length": field.max_length if hasattr(field, 'max_length') else None,
  42. "dim": field.dim if field.dtype == DataType.FLOAT_VECTOR else None
  43. } for field in schema.fields
  44. ],
  45. "description": schema.description
  46. }
  47. info = {
  48. "name": name,
  49. "entities": collection.num_entities,
  50. "schema": schema_dict
  51. }
  52. collections_info.append(info)
  53. return success(collections_info) # 假设 success 函数可以处理字典列表
  54. except Exception as e:
  55. print(f"列出集合时发生错误: {str(e)}")
  56. raise
  57. finally:
  58. # 断开 Milvus 连接
  59. connections.disconnect("default")
  60. @staticmethod
  61. def delete_collection(request):
  62. """
  63. 删除 Milvus 中的指定集合
  64. :param request: HTTP 请求对象
  65. :return: JsonResponse 对象
  66. """
  67. collection_name = request.POST.get("collection_name")
  68. if not collection_name:
  69. return fail("集合名称未提供")
  70. try:
  71. # 连接到 Milvus
  72. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
  73. # 检查集合是否存在
  74. if not utility.has_collection(collection_name):
  75. return fail(f"集合 '{collection_name}' 不存在")
  76. # 删除集合
  77. utility.drop_collection(collection_name)
  78. return success( f"集合 '{collection_name}' 已成功删除")
  79. except Exception as e:
  80. return fail(f"删除集合时发生错误: {str(e)}")
  81. finally:
  82. # 断开 Milvus 连接
  83. connections.disconnect("default")
  84. @staticmethod
  85. def delete_milvus_data(request):
  86. """
  87. 从 Milvus 中删除指定 ID 的数据
  88. :param request: HTTP 请求对象
  89. :return: JsonResponse 对象
  90. """
  91. collection_name = request.POST.get("collection_name")
  92. id_to_delete = request.POST.get("id")
  93. if not collection_name:
  94. return fail("集合名称未提供")
  95. if not id_to_delete:
  96. return fail("要删除的 ID 未提供")
  97. try:
  98. # 连接到 Milvus
  99. connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT,user=MILVUS_USER,password=MILVUS_PASSWORD)
  100. # 检查集合是否存在
  101. if not utility.has_collection(collection_name):
  102. return fail(f"集合 '{collection_name}' 不存在")
  103. # 获取集合对象
  104. collection = Collection(collection_name)
  105. # 执行删除操作
  106. expr = f'id in [{id_to_delete}]'
  107. delete_result = collection.delete(expr)
  108. if delete_result.delete_count > 0:
  109. return success(f"ID 为 {id_to_delete} 的数据已成功从集合 '{collection_name}' 中删除")
  110. else:
  111. return fail(f"未找到 ID 为 {id_to_delete} 的数据")
  112. except Exception as e:
  113. return fail(f"删除数据时发生错误: {str(e)}")
  114. finally:
  115. # 断开 Milvus 连接
  116. connections.disconnect("default")
  117. @staticmethod
  118. def getMinioURl(request):
  119. data = json.loads(request.body)
  120. id = data.get("id")
  121. logger.info(f"request:{request}")
  122. if not id:
  123. return fail("id为空")
  124. task = TaskSublist.objects.filter(milvus_id=id).first()
  125. if not task:
  126. return fail("无此数据")
  127. document = DocumentKbm.objects.filter(id = task.doc_id).first()
  128. if not document:
  129. return fail("未找到文件")
  130. kmb = Knowledgebase.objects.filter(id=document.kb_id).first()
  131. object_name = document.location
  132. bucket_name = kmb.location
  133. if not object_name:
  134. return fail('Object name is required')
  135. url = minio_client.presigned_get_object(
  136. bucket_name,
  137. object_name,
  138. expires=timedelta(days=1)
  139. )
  140. bucket_info = {
  141. 'object_name': object_name,
  142. 'url': url,
  143. }
  144. return success(bucket_info)