fix: 修复milvus无法使用

ageerle

2025-03-14 d67cd05ae23013b00751fe0701da5f1d53422ba1

fix: 修复milvus无法使用

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

@@ -2,19 +2,16 @@

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;
@@ -137,11 +134,103 @@

    @Override
    public void storeEmbeddings(List<String> chunkList, List<List<Double>> vectorList, String kid, String docId, List<String> fidList) {
        String fullCollectionName = collectionName + kid;

        // 检查集合是否存在
        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 @@
        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,6 +305,38 @@

    @Override
    public List<String> nearest(List<Double> queryVector, String kid) {
        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++) {
@@ -219,22 +354,36 @@
                .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());
        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 (resultList!=null && resultList.size() > 0){
            for (int i = 0; i < rowRecords.size(); i++) {
                String content = rowRecords.get(i).get("content").toString();
                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 new ArrayList<>();

    }

    /**
     * milvus 不支持通过文本检索相似性
     *
     * @param query
     * @param kid
     * @return