fix: 修复milvus无法使用

ruoyi-modules/ruoyi-knowledge/src/main/java/org/ruoyi/knowledge/chain/vectorstore/MilvusVectorStore.java

@@ -2,19 +2,16 @@ package org.ruoyi.knowledge.chain.vectorstore;
 
 import io.milvus.client.MilvusServiceClient;
 import io.milvus.grpc.DataType;
+import io.milvus.grpc.DescribeIndexResponse;
+import io.milvus.grpc.MutationResult;
 import io.milvus.grpc.SearchResults;
-import io.milvus.param.ConnectParam;
-import io.milvus.param.IndexType;
-import io.milvus.param.MetricType;
-import io.milvus.param.R;
-import io.milvus.param.collection.CreateCollectionParam;
-import io.milvus.param.collection.DropCollectionParam;
-import io.milvus.param.collection.FieldType;
-import io.milvus.param.collection.LoadCollectionParam;
+import io.milvus.param.*;
+import io.milvus.param.collection.*;
 import io.milvus.param.dml.DeleteParam;
 import io.milvus.param.dml.InsertParam;
 import io.milvus.param.dml.SearchParam;
 import io.milvus.param.index.CreateIndexParam;
+import io.milvus.param.index.DescribeIndexParam;
 import io.milvus.param.partition.CreatePartitionParam;
 import io.milvus.response.QueryResultsWrapper;
 import io.milvus.response.SearchResultsWrapper;
@@ -32,7 +29,7 @@ import java.util.List;
 
 @Service
 @Slf4j
-public class MilvusVectorStore implements VectorStore{
+public class MilvusVectorStore implements VectorStore {
 
     private volatile Integer dimension;
     private volatile String collectionName;
@@ -48,7 +45,7 @@ public class MilvusVectorStore implements VectorStore{
     }
 
     @PostConstruct
-    public void init(){
+    public void init() {
         String milvusHost = configService.getConfigValue("milvus", "host");
         String milvausPort = configService.getConfigValue("milvus", "port");
         milvusServiceClient = new MilvusServiceClient(
@@ -137,11 +134,103 @@ public class MilvusVectorStore implements VectorStore{
 
     @Override
     public void storeEmbeddings(List<String> chunkList, List<List<Double>> vectorList, String kid, String docId, List<String> fidList) {
+        String fullCollectionName = collectionName + kid;
 
-        if (StringUtils.isNotBlank(docId)){
+        // 检查集合是否存在
+        HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder()
+                .withCollectionName(fullCollectionName)
+                .build();
+        R<Boolean> booleanR = milvusServiceClient.hasCollection(hasCollectionParam);
+
+        if (booleanR.getStatus() == R.Status.Success.getCode()) {
+            boolean collectionExists = booleanR.getData().booleanValue();
+            if (!collectionExists) {
+                // 集合不存在，创建集合
+                List<FieldType> fieldTypes = new ArrayList<>();
+                // 假设这里定义 id 字段，根据实际情况修改
+                FieldType idField = FieldType.newBuilder()
+                        .withName("id")
+                        .withDataType(DataType.Int64)
+                        .withPrimaryKey(true)
+                        .withAutoID(true)
+                        .build();
+                fieldTypes.add(idField);
+
+                // 定义向量字段
+                FieldType vectorField = FieldType.newBuilder()
+                        .withName("fv")
+                        .withDataType(DataType.FloatVector)
+                        .withDimension(vectorList.get(0).size())
+                        .build();
+                fieldTypes.add(vectorField);
+
+                // 定义其他字段
+                FieldType contentField = FieldType.newBuilder()
+                        .withName("content")
+                        .withDataType(DataType.VarChar)
+                        .withMaxLength(chunkList.size() * 1024) // 根据实际情况修改
+                        .build();
+                fieldTypes.add(contentField);
+
+                FieldType kidField = FieldType.newBuilder()
+                        .withName("kid")
+                        .withDataType(DataType.VarChar)
+                        .withMaxLength(256) // 根据实际情况修改
+                        .build();
+                fieldTypes.add(kidField);
+
+                FieldType docIdField = FieldType.newBuilder()
+                        .withName("docId")
+                        .withDataType(DataType.VarChar)
+                        .withMaxLength(256) // 根据实际情况修改
+                        .build();
+                fieldTypes.add(docIdField);
+
+                FieldType fidField = FieldType.newBuilder()
+                        .withName("fid")
+                        .withDataType(DataType.VarChar)
+                        .withMaxLength(256) // 根据实际情况修改
+                        .build();
+                fieldTypes.add(fidField);
+
+                CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
+                        .withCollectionName(fullCollectionName)
+                        .withFieldTypes(fieldTypes)
+                        .build();
+
+                R<RpcStatus> collection = milvusServiceClient.createCollection(createCollectionParam);
+                if (collection.getStatus() == R.Status.Success.getCode()) {
+                    System.out.println("集合 " + fullCollectionName + " 创建成功");
+
+                    // 创建索引
+                    CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
+                            .withCollectionName(fullCollectionName)
+                            .withFieldName("fv") // 向量字段名
+                            .withIndexType(IndexType.IVF_FLAT) // 索引类型
+                            .withMetricType(MetricType.IP)
+                            .withExtraParam("{\"nlist\":1024}") // 索引参数
+                            .build();
+                    R<RpcStatus> indexResponse = milvusServiceClient.createIndex(createIndexParam);
+                    if (indexResponse.getStatus() == R.Status.Success.getCode()) {
+                        System.out.println("索引创建成功");
+                    } else {
+                        System.err.println("索引创建失败: " + indexResponse.getMessage());
+                        return;
+                    }
+                } else {
+                    System.err.println("集合创建失败: " + collection.getMessage());
+                    return;
+                }
+            }
+        } else {
+            System.err.println("检查集合是否存在时出错: " + booleanR.getMessage());
+            return;
+        }
+
+        if (StringUtils.isNotBlank(docId)) {
             milvusServiceClient.createPartition(
                     CreatePartitionParam.newBuilder()
-                            .withCollectionName(collectionName + kid)
+                            .withCollectionName(fullCollectionName)
                             .withPartitionName(docId)
                             .build()
             );
@@ -169,15 +258,29 @@ public class MilvusVectorStore implements VectorStore{
         fields.add(new InsertParam.Field("fv", vectorFloatList));
 
         InsertParam insertParam = InsertParam.newBuilder()
-                .withCollectionName(collectionName + kid)
+                .withCollectionName(fullCollectionName)
                 .withPartitionName(docId)
                 .withFields(fields)
                 .build();
-        milvusServiceClient.insert(insertParam);
-        // milvus在将数据装载到内存后才能进行向量计算
-        milvusServiceClient.loadCollection(LoadCollectionParam.newBuilder().withCollectionName(collectionName + kid).build());
-    }
+        System.out.println("=========================");
 
+        R<MutationResult> insert = milvusServiceClient.insert(insertParam);
+        if (insert.getStatus() == R.Status.Success.getCode()) {
+            System.out.println("插入成功，插入的行数: " + insert.getData().getInsertCnt());
+        } else {
+            System.err.println("插入失败: " + insert.getMessage());
+        }
+        System.out.println("=========================");
+        // milvus在将数据装载到内存后才能进行向量计算.
+        LoadCollectionParam loadCollectionParam = LoadCollectionParam.newBuilder()
+                .withCollectionName(fullCollectionName)
+                .build();
+        R<RpcStatus> loadResponse = milvusServiceClient.loadCollection(loadCollectionParam);
+        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
+            System.err.println("加载集合 " + fullCollectionName + " 到内存时出错：" + loadResponse.getMessage());
+        }
+//        milvusServiceClient.loadCollection(LoadCollectionParam.newBuilder().withCollectionName(fullCollectionName).build());
+    }
 
 
     @Override
@@ -202,7 +305,39 @@ public class MilvusVectorStore implements VectorStore{
 
     @Override
     public List<String> nearest(List<Double> queryVector, String kid) {
-        List<String> search_output_fields = Arrays.asList("content","fv");
+        String fullCollectionName = collectionName + kid;
+
+        HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder()
+                .withCollectionName(fullCollectionName)
+                .build();
+
+        R<Boolean> booleanR = milvusServiceClient.hasCollection(hasCollectionParam);
+        if (booleanR.getStatus() != R.Status.Success.getCode() || !booleanR.getData().booleanValue()) {
+            System.err.println("集合 " + fullCollectionName + " 不存在或检查集合存在性时出错。");
+            return new ArrayList<>();
+        }
+
+        DescribeIndexParam describeIndexParam = DescribeIndexParam.newBuilder().withCollectionName(fullCollectionName).build();
+
+        R<DescribeIndexResponse> describeIndexResponseR = milvusServiceClient.describeIndex(describeIndexParam);
+
+        if (describeIndexResponseR.getStatus() == R.Status.Success.getCode()) {
+            System.out.println("索引信息: " + describeIndexResponseR.getData().getIndexDescriptionsCount());
+        } else {
+            System.err.println("获取索引失败: " + describeIndexResponseR.getMessage());
+        }
+
+//        // 加载集合到内存
+//        LoadCollectionParam loadCollectionParam = LoadCollectionParam.newBuilder()
+//                .withCollectionName(fullCollectionName)
+//                .build();
+//        R<RpcStatus> loadResponse = milvusServiceClient.loadCollection(loadCollectionParam);
+//        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
+//            System.err.println("加载集合 " + fullCollectionName + " 到内存时出错：" + loadResponse.getMessage());
+//            return new ArrayList<>();
+//        }
+
+        List<String> search_output_fields = Arrays.asList("content", "fv");
         List<Float> fv = new ArrayList<>();
         for (int i = 0; i < queryVector.size(); i++) {
             fv.add(queryVector.get(i).floatValue());
@@ -219,22 +354,36 @@ public class MilvusVectorStore implements VectorStore{
                 .withVectorFieldName("fv")
                 .withParams(search_param)
                 .build();
+        System.out.println("SearchParam: " + searchParam.toString());
         R<SearchResults> respSearch = milvusServiceClient.search(searchParam);
-        SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(respSearch.getData().getResults());
-        List<QueryResultsWrapper.RowRecord> rowRecords = wrapperSearch.getRowRecords();
-
-        List<String> resultList = new ArrayList<>();
-        if (resultList!=null && resultList.size() > 0){
-            for (int i = 0; i < rowRecords.size(); i++) {
-                String content = rowRecords.get(i).get("content").toString();
-                resultList.add(content);
+        if (respSearch.getStatus() == R.Status.Success.getCode()) {
+            SearchResults searchResults = respSearch.getData();
+            if (searchResults != null) {
+                System.out.println(searchResults.getResults());
+                SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(searchResults.getResults());
+                List<QueryResultsWrapper.RowRecord> rowRecords = wrapperSearch.getRowRecords();
+
+                List<String> resultList = new ArrayList<>();
+                if (rowRecords != null && !rowRecords.isEmpty()) {
+                    for (QueryResultsWrapper.RowRecord rowRecord : rowRecords) {
+                        String content = rowRecord.get("content").toString();
+                        resultList.add(content);
+                    }
+                }
+                return resultList;
+            } else {
+                System.err.println("搜索结果为空");
             }
+        } else {
+            System.err.println("搜索操作失败: " + respSearch.getMessage());
         }
-        return resultList;
+        return new ArrayList<>();
+
     }
 
     /**
      * milvus 不支持通过文本检索相似性
+     *
      * @param query
      * @param kid
      * @return