feat: Update Milvus integration with new API, enhance collection management, and improve embedding storage

Dallas98 · Dallas98 · commit 4ca7dc01aa9b · 2025-11-20T20:13:27.000+08:00
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/application/KnowledgeBaseService.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/application/KnowledgeBaseService.java
@@ -17,13 +17,11 @@
 import com.datamate.rag.indexer.infrastructure.milvus.MilvusService;
 import com.datamate.rag.indexer.interfaces.dto.*;
 import com.google.gson.JsonObject;
-import io.milvus.grpc.QueryResults;
-import io.milvus.param.R;
-import io.milvus.param.collection.DropCollectionParam;
-import io.milvus.param.collection.RenameCollectionParam;
-import io.milvus.param.dml.DeleteParam;
-import io.milvus.param.dml.QueryParam;
-import io.milvus.response.QueryResultsWrapper;
+import io.milvus.v2.service.collection.request.DropCollectionReq;
+import io.milvus.v2.service.collection.request.RenameCollectionReq;
+import io.milvus.v2.service.vector.request.DeleteReq;
+import io.milvus.v2.service.vector.request.QueryReq;
+import io.milvus.v2.service.vector.response.QueryResp;
 import lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
 import org.springframework.beans.BeanUtils;
@@ -34,7 +32,6 @@
 
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -78,9 +75,9 @@ public void update(String knowledgeBaseId, KnowledgeBaseUpdateReq request) {
         KnowledgeBase knowledgeBase = Optional.ofNullable(knowledgeBaseRepository.getById(knowledgeBaseId))
                 .orElseThrow(() -> BusinessException.of(KnowledgeBaseErrorCode.KNOWLEDGE_BASE_NOT_FOUND));
         if (StringUtils.hasText(request.getName())) {
-            milvusService.getMilvusClient().renameCollection(RenameCollectionParam.newBuilder()
-                    .withOldCollectionName(knowledgeBase.getName())
-                    .withNewCollectionName(request.getName())
+            milvusService.getMilvusClient().renameCollection(RenameCollectionReq.builder()
+                    .collectionName(knowledgeBase.getName())
+                    .newCollectionName(request.getName())
                     .build());
             knowledgeBase.setName(request.getName());
         }
@@ -102,7 +99,7 @@ public void delete(String knowledgeBaseId) {
                 .orElseThrow(() -> BusinessException.of(KnowledgeBaseErrorCode.KNOWLEDGE_BASE_NOT_FOUND));
         knowledgeBaseRepository.removeById(knowledgeBaseId);
         ragFileRepository.removeByKnowledgeBaseId(knowledgeBaseId);
-        milvusService.getMilvusClient().dropCollection(DropCollectionParam.newBuilder().withCollectionName(knowledgeBase.getName()).build());
+        milvusService.getMilvusClient().dropCollection(DropCollectionReq.builder().collectionName(knowledgeBase.getName()).build());
     }
 
     public KnowledgeBaseResp getById(String knowledgeBaseId) {
@@ -168,41 +165,39 @@ public void deleteFiles(String knowledgeBaseId, DeleteFilesReq request) {
         KnowledgeBase knowledgeBase = Optional.ofNullable(knowledgeBaseRepository.getById(knowledgeBaseId))
                 .orElseThrow(() -> BusinessException.of(KnowledgeBaseErrorCode.KNOWLEDGE_BASE_NOT_FOUND));
         ragFileRepository.removeByIds(request.getIds());
-        milvusService.getMilvusClient().delete(DeleteParam.newBuilder()
-                .withCollectionName(knowledgeBase.getName())
-                .withExpr("metadata[\"rag_file_id\"] in [" + org.apache.commons.lang3.StringUtils.join(request.getIds().stream().map(id -> "\"" + id + "\"").toArray(), ",") + "]")
+        milvusService.getMilvusClient().delete(DeleteReq.builder()
+                .collectionName(knowledgeBase.getName())
+                .filter("metadata[\"rag_file_id\"] in [" + org.apache.commons.lang3.StringUtils.join(request.getIds().stream().map(id -> "\"" + id + "\"").toArray(), ",") + "]")
                 .build());
     }
 
     public PagedResponse<RagChunk> getChunks(String knowledgeBaseId, String ragFileId, PagingQuery pagingQuery) {
         KnowledgeBase knowledgeBase = Optional.ofNullable(knowledgeBaseRepository.getById(knowledgeBaseId))
                 .orElseThrow(() -> BusinessException.of(KnowledgeBaseErrorCode.KNOWLEDGE_BASE_NOT_FOUND));
-        R<QueryResults> results = milvusService.getMilvusClient().query(QueryParam.newBuilder()
-                .withCollectionName(knowledgeBase.getName())
-                .withExpr("metadata[\"rag_file_id\"] == \"" + ragFileId + "\"")
-                .withOutFields(Collections.singletonList("*"))
-                .withLimit(Long.valueOf(pagingQuery.getSize()))
-                .withOffset((long) (pagingQuery.getPage() - 1) * pagingQuery.getSize())
+        QueryResp results = milvusService.getMilvusClient().query(QueryReq.builder()
+                .collectionName(knowledgeBase.getName())
+                .filter("metadata[\"rag_file_id\"] == \"" + ragFileId + "\"")
+                .outputFields(Collections.singletonList("*"))
+                .limit(Long.valueOf(pagingQuery.getSize()))
+                .offset((long) (pagingQuery.getPage() - 1) * pagingQuery.getSize())
                 .build());
-        QueryResultsWrapper wrapper = new QueryResultsWrapper(results.getData());
-        List<Map<String, Object>> list = wrapper.getRowRecords().stream().map(QueryResultsWrapper.RowRecord::getFieldValues).toList();
-
-        List<RagChunk> ragChunks = list.stream().map(item -> new RagChunk(
-                item.get("id").toString(),
-                item.get("text").toString(),
-                item.get("metadata").toString()
-        )).toList();
+        List<QueryResp.QueryResult> queryResults = results.getQueryResults();
+        List<RagChunk> ragChunks = queryResults.stream()
+                .map(QueryResp.QueryResult::getEntity)
+                .map(item -> new RagChunk(
+                        item.get("id").toString(),
+                        item.get("text").toString(),
+                        item.get("metadata").toString()
+                )).toList();
 
         // 获取总数
-        R<QueryResults> countResults = milvusService.getMilvusClient().query(QueryParam.newBuilder()
-                .withCollectionName(knowledgeBase.getName())
-                .withExpr("metadata[\"rag_file_id\"] == \"" + ragFileId + "\"")
-                .withOutFields(Collections.singletonList("count(*)"))
+        QueryResp countResults = milvusService.getMilvusClient().query(QueryReq.builder()
+                .collectionName(knowledgeBase.getName())
+                .filter("metadata[\"rag_file_id\"] == \"" + ragFileId + "\"")
+                .outputFields(Collections.singletonList("count(*)"))
                 .build());
-        QueryResultsWrapper countWrapper = new QueryResultsWrapper(countResults.getData());
-        List<Map<String, Object>> countList = countWrapper.getRowRecords().stream().map(QueryResultsWrapper.RowRecord::getFieldValues).toList();
-        long totalCount = Long.parseLong(countList.getFirst().get("count(*)").toString());
 
+        long totalCount = Long.parseLong(countResults.getQueryResults().getFirst().getEntity().get("count(*)").toString());
         return PagedResponse.of(ragChunks, pagingQuery.getPage(), totalCount, (int) Math.ceil((double) totalCount / pagingQuery.getSize()));
     }
 }
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/infrastructure/event/RagEtlService.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/infrastructure/event/RagEtlService.java
@@ -119,16 +119,19 @@ private void processRagFile(RagFile ragFile, DataInsertedEvent event) {
         EmbeddingModel embeddingModel = ModelClient.invokeEmbeddingModel(model);
         // 调用嵌入模型获取嵌入向量
 
+        if (!milvusService.hasCollection(event.knowledgeBase().getName())) {
+            milvusService.createCollection(event.knowledgeBase().getName(), embeddingModel.dimension());
+        }
+
         Lists.partition(split, 20).forEach(partition -> {
-            List<Embedding> content = embeddingModel.embedAll(partition).content();
-            // 存储嵌入向量到 Milvus
-            milvusService.embeddingStore(embeddingModel, event.knowledgeBase().getName()).addAll(content, partition);
+            List<Embedding> embeddings = embeddingModel.embedAll(partition).content();
+            milvusService.addAll(event.knowledgeBase().getName(),partition, embeddings);
         });
     }
 
     /**
      * 根据文件类型返回对应的文档解析器
-     *
+     *x
      * @param fileType 文件类型
      * @return 文档解析器
      */
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/infrastructure/milvus/MilvusService.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/infrastructure/milvus/MilvusService.java
@@ -1,16 +1,28 @@
 package com.datamate.rag.indexer.infrastructure.milvus;
 
+import com.google.gson.*;
+import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 import dev.langchain4j.store.embedding.EmbeddingStore;
 import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
-import io.milvus.client.MilvusClient;
-import io.milvus.client.MilvusServiceClient;
-import io.milvus.param.ConnectParam;
+import io.milvus.common.clientenum.FunctionType;
+import io.milvus.v2.client.ConnectConfig;
+import io.milvus.v2.client.MilvusClientV2;
+import io.milvus.v2.common.DataType;
+import io.milvus.v2.common.IndexParam;
+import io.milvus.v2.service.collection.request.AddFieldReq;
+import io.milvus.v2.service.collection.request.CreateCollectionReq;
+import io.milvus.v2.service.collection.request.HasCollectionReq;
+import io.milvus.v2.service.vector.request.InsertReq;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 
+import java.util.*;
+
+import static dev.langchain4j.internal.Utils.randomUUID;
+
 /**
  * Milvus 服务类
  *
@@ -24,13 +36,19 @@ public class MilvusService {
     private String milvusHost;
     @Value("${datamate.rag.milvus-port:19530}")
     private int milvusPort;
+    @Value("${datamate.rag.milvus-uri:http://milvus-standalone:19530}")
+    private String milvusUri;
+    private static final Gson GSON;
+
+    static {
+        GSON = (new GsonBuilder()).setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE).create();
+    }
 
-    private volatile MilvusClient milvusClient;
+    private volatile MilvusClientV2 milvusClient;
 
     public EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel, String knowledgeBaseName) {
         return MilvusEmbeddingStore.builder()
-                .host(milvusHost)
-                .port(milvusPort)
+                .uri(milvusUri)
                 .collectionName(knowledgeBaseName)
                 .dimension(embeddingModel.dimension())
                 .build();
@@ -41,16 +59,15 @@ public EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel,
      *
      * @return MilvusClient
      */
-    public MilvusClient getMilvusClient() {
+    public MilvusClientV2 getMilvusClient() {
         if (milvusClient == null) {
             synchronized (this) {
                 if (milvusClient == null) {
                     try {
-                        ConnectParam connectParam = ConnectParam.newBuilder()
-                                .withHost(milvusHost)
-                                .withPort(milvusPort)
+                        ConnectConfig connectConfig = ConnectConfig.builder()
+                                .uri(milvusUri)
                                 .build();
-                        milvusClient = new MilvusServiceClient(connectParam);
+                        milvusClient = new MilvusClientV2(connectConfig);
                         log.info("Milvus client connected successfully");
                     } catch (Exception e) {
                         log.error("Milvus client connection failed: {}", e.getMessage());
@@ -61,4 +78,107 @@ public MilvusClient getMilvusClient() {
         }
         return milvusClient;
     }
+
+
+    public boolean hasCollection(String collectionName) {
+        HasCollectionReq request = HasCollectionReq.builder().collectionName(collectionName).build();
+        return getMilvusClient().hasCollection(request);
+    }
+
+    public void createCollection(String collectionName, int dimension) {
+        CreateCollectionReq.CollectionSchema schema = CreateCollectionReq.CollectionSchema.builder()
+                .build();
+        schema.addField(AddFieldReq.builder()
+                .fieldName("id")
+                .dataType(DataType.VarChar)
+                .maxLength(36)
+                .isPrimaryKey(true)
+                .autoID(false)
+                .build());
+        schema.addField(AddFieldReq.builder()
+                .fieldName("text")
+                .dataType(DataType.VarChar)
+                .maxLength(65535)
+                .enableAnalyzer(true)
+                .build());
+        schema.addField(AddFieldReq.builder()
+                .fieldName("metadata")
+                .dataType(DataType.JSON)
+                .build());
+        schema.addField(AddFieldReq.builder()
+                .fieldName("vector")
+                .dataType(DataType.FloatVector)
+                .dimension(dimension)
+                .build());
+        schema.addField(AddFieldReq.builder()
+                .fieldName("sparse")
+                .dataType(DataType.SparseFloatVector)
+                .build());
+        schema.addFunction(CreateCollectionReq.Function.builder()
+                .functionType(FunctionType.BM25)
+                .name("text_bm25_emb")
+                .inputFieldNames(Collections.singletonList("text"))
+                .outputFieldNames(Collections.singletonList("sparse"))
+                .build());
+
+        Map<String, Object> params = new HashMap<>();
+        params.put("inverted_index_algo", "DAAT_MAXSCORE");
+        params.put("bm25_k1", 1.2);
+        params.put("bm25_b", 0.75);
+
+        List<IndexParam> indexes = new ArrayList<>();
+        indexes.add(IndexParam.builder()
+                .fieldName("sparse")
+                .indexType(IndexParam.IndexType.SPARSE_INVERTED_INDEX)
+                .metricType(IndexParam.MetricType.BM25)
+                .extraParams(params)
+                .build());
+        indexes.add(IndexParam.builder()
+                .fieldName("vector")
+                .indexType(IndexParam.IndexType.FLAT)
+                .metricType(IndexParam.MetricType.COSINE)
+                .extraParams(Map.of())
+                .build());
+
+        CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
+                .collectionName(collectionName)
+                .collectionSchema(schema)
+                .indexParams(indexes)
+                .build();
+        this.getMilvusClient().createCollection(createCollectionReq);
+    }
+
+    public void addAll(String collectionName, List<TextSegment> textSegments, List<Embedding> embeddings) {
+        List<JsonObject> data = convertToJsonObjects(textSegments, embeddings);
+        InsertReq insertReq = InsertReq.builder()
+                .collectionName(collectionName)
+                .data(data)
+                .build();
+        this.getMilvusClient().insert(insertReq);
+    }
+
+    public List<JsonObject> convertToJsonObjects(List<TextSegment> textSegments, List<Embedding> embeddings) {
+        List<JsonObject> data = new ArrayList<>();
+        for (int i = 0; i < textSegments.size(); i++) {
+            JsonObject jsonObject = new JsonObject();
+            jsonObject.addProperty("id", randomUUID());
+            jsonObject.addProperty("text", textSegments.get(i).text());
+            jsonObject.add("metadata", GSON.toJsonTree(textSegments.get(i).metadata().toMap()).getAsJsonObject());
+            JsonArray vectorArray = new JsonArray();
+            for (float f : embeddings.get(i).vector()) {
+                vectorArray.add(f);
+            }
+            jsonObject.add("vector", vectorArray);
+            data.add(jsonObject);
+        }
+        return data;
+    }
+
+    List<String> generateIds(int n) {
+        List<String> ids = new ArrayList<>();
+        for (int i = 0; i < n; i++) {
+            ids.add(randomUUID());
+        }
+        return ids;
+    }
 }
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/interfaces/KnowledgeBaseController.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/interfaces/KnowledgeBaseController.java
@@ -133,10 +133,10 @@ public PagedResponse<RagChunk> getChunks(@PathVariable("knowledgeBaseId") String
      *
      * @param knowledgeBaseId 知识库 ID
      */
-    @PostMapping("/{knowledgeBaseId}/retrieve")
-    public PagedResponse<RagChunk> retrieve(@PathVariable("knowledgeBaseId") String knowledgeBaseId,
-                                            @RequestBody @Valid RetrieveReq request,
-                                            PagingQuery pagingQuery) {
-        return knowledgeBaseService.retrieve(knowledgeBaseId, request, pagingQuery);
-    }
+//    @PostMapping("/{knowledgeBaseId}/retrieve")
+//    public PagedResponse<RagChunk> retrieve(@PathVariable("knowledgeBaseId") String knowledgeBaseId,
+//                                            @RequestBody @Valid RetrieveReq request,
+//                                            PagingQuery pagingQuery) {
+//        return knowledgeBaseService.retrieve(knowledgeBaseId, request, pagingQuery);
+//    }
 }
