From d67cd05ae23013b00751fe0701da5f1d53422ba1 Mon Sep 17 00:00:00 2001 From: ageerle Date: Fri, 14 Mar 2025 11:18:41 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8Dmilvus=E6=97=A0?= =?UTF-8?q?=E6=B3=95=E4=BD=BF=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../chain/vectorstore/MilvusVectorStore.java | 201 +++++++++++++++--- 1 file changed, 175 insertions(+), 26 deletions(-) diff --git a/ruoyi-modules/ruoyi-knowledge/src/main/java/org/ruoyi/knowledge/chain/vectorstore/MilvusVectorStore.java b/ruoyi-modules/ruoyi-knowledge/src/main/java/org/ruoyi/knowledge/chain/vectorstore/MilvusVectorStore.java index c083ef47..eb993f9d 100644 --- a/ruoyi-modules/ruoyi-knowledge/src/main/java/org/ruoyi/knowledge/chain/vectorstore/MilvusVectorStore.java +++ b/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 chunkList, List> vectorList, String kid, String docId, List fidList) { + String fullCollectionName = collectionName + kid; - if (StringUtils.isNotBlank(docId)){ + // 检查集合是否存在 + HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder() + .withCollectionName(fullCollectionName) + .build(); + R booleanR = milvusServiceClient.hasCollection(hasCollectionParam); + + if (booleanR.getStatus() == R.Status.Success.getCode()) { + boolean collectionExists = booleanR.getData().booleanValue(); + if (!collectionExists) { + // 集合不存在，创建集合 + List 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 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 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 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 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 nearest(List queryVector, String kid) { - List search_output_fields = Arrays.asList("content","fv"); + String fullCollectionName = collectionName + kid; + + HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder() + .withCollectionName(fullCollectionName) + .build(); + + R 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 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 loadResponse = milvusServiceClient.loadCollection(loadCollectionParam); +// if (loadResponse.getStatus() != R.Status.Success.getCode()) { +// System.err.println("加载集合 " + fullCollectionName + " 到内存时出错：" + loadResponse.getMessage()); +// return new ArrayList<>(); +// } + + List search_output_fields = Arrays.asList("content", "fv"); List 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 respSearch = milvusServiceClient.search(searchParam); - SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(respSearch.getData().getResults()); - List rowRecords = wrapperSearch.getRowRecords(); + 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 rowRecords = wrapperSearch.getRowRecords(); - List 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); + List 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