From d67cd05ae23013b00751fe0701da5f1d53422ba1 Mon Sep 17 00:00:00 2001
From: ageerle <ageerle@163.com>
Date: 星期五, 14 三月 2025 11:18:41 +0800
Subject: [PATCH] fix: 修复milvus无法使用
---
ruoyi-modules/ruoyi-knowledge/src/main/java/org/ruoyi/knowledge/chain/vectorstore/MilvusVectorStore.java | 201 +++++++++++++++++++++++++++++++++++++++++++------
1 files 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 c083ef4..eb993f9 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 @@
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 @@
@Service
@Slf4j
-public class MilvusVectorStore implements VectorStore{
+public class MilvusVectorStore implements VectorStore {
private volatile Integer dimension;
private volatile String collectionName;
@@ -48,7 +45,7 @@
}
@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 @@
@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 @@
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 @@
@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 @@
.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();
+ 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();
- resultList.add(content);
+ 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
--
Gitblit v1.9.3