feat: Enhance knowledge base management with collection renaming, improved chunk retrieval, and error handling

Dallas98 · Dallas98 · commit 72a6522892c9 · 2025-11-20T14:33:20.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
@@ -16,8 +16,14 @@
 import com.datamate.rag.indexer.infrastructure.event.DataInsertedEvent;
 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 lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
 import org.springframework.beans.BeanUtils;
@@ -26,7 +32,9 @@
 import org.springframework.transaction.annotation.Transactional;
 import org.springframework.util.StringUtils;
 
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -44,6 +52,8 @@ public class KnowledgeBaseService {
     private final ModelConfigRepository modelConfigRepository;
     private final MilvusService milvusService;
 
+    private JsonObject jsonObject = new JsonObject();
+
     /**
      * 创建知识库
      *
@@ -63,10 +73,15 @@ public String create(KnowledgeBaseCreateReq request) {
      * @param knowledgeBaseId 知识库 ID
      * @param request         知识库更新请求
      */
+    @Transactional(rollbackFor = Exception.class)
     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())
+                    .build());
             knowledgeBase.setName(request.getName());
         }
         if (StringUtils.hasText(request.getDescription())) {
@@ -75,7 +90,13 @@ public void update(String knowledgeBaseId, KnowledgeBaseUpdateReq request) {
         knowledgeBaseRepository.updateById(knowledgeBase);
     }
 
-    @Transactional
+
+    /**
+     * 删除知识库
+     *
+     * @param knowledgeBaseId 知识库 ID
+     */
+    @Transactional(rollbackFor = Exception.class)
     public void delete(String knowledgeBaseId) {
         KnowledgeBase knowledgeBase = Optional.ofNullable(knowledgeBaseRepository.getById(knowledgeBaseId))
                 .orElseThrow(() -> BusinessException.of(KnowledgeBaseErrorCode.KNOWLEDGE_BASE_NOT_FOUND));
@@ -154,7 +175,34 @@ public void deleteFiles(String knowledgeBaseId, DeleteFilesReq request) {
     }
 
     public PagedResponse<RagChunk> getChunks(String knowledgeBaseId, String ragFileId, PagingQuery pagingQuery) {
-        IPage<RagChunk> page = new Page<>(pagingQuery.getPage(), pagingQuery.getSize());
-        return PagedResponse.of(page.getRecords(), page.getCurrent(), page.getTotal(), page.getPages());
+        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())
+                .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();
+
+        // 获取总数
+        R<QueryResults> countResults = milvusService.getMilvusClient().query(QueryParam.newBuilder()
+                .withCollectionName(knowledgeBase.getName())
+                .withExpr("metadata[\"rag_file_id\"] == \"" + ragFileId + "\"")
+                .withOutFields(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());
+
+        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/domain/model/RagChunk.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/domain/model/RagChunk.java
@@ -6,5 +6,10 @@
  * @author dallas
  * @since 2025-10-29
  */
-public class RagChunk {
-}
+
+public record RagChunk(
+        String id,
+        String text,
+        String metadata
+) {
+}
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/domain/model/RagFile.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/domain/model/RagFile.java
@@ -44,4 +44,6 @@ public class RagFile extends BaseEntity<String> {
     private Map<String, Object> metadata;
 
     private FileStatus status;
+
+    private String errMsg;
 }
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
@@ -9,7 +9,7 @@
 import com.datamate.rag.indexer.domain.model.RagFile;
 import com.datamate.rag.indexer.domain.repository.RagFileRepository;
 import com.datamate.rag.indexer.infrastructure.milvus.MilvusService;
-import com.datamate.rag.indexer.interfaces.dto.ProcessType;
+import com.datamate.rag.indexer.interfaces.dto.AddFilesReq;
 import com.google.common.collect.Lists;
 import dev.langchain4j.data.document.Document;
 import dev.langchain4j.data.document.DocumentParser;
@@ -20,10 +20,7 @@
 import dev.langchain4j.data.document.parser.apache.poi.ApachePoiDocumentParser;
 import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
 import dev.langchain4j.data.document.parser.markdown.MarkdownDocumentParser;
-import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
-import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
-import dev.langchain4j.data.document.splitter.DocumentBySentenceSplitter;
-import dev.langchain4j.data.document.splitter.DocumentByWordSplitter;
+import dev.langchain4j.data.document.splitter.*;
 import dev.langchain4j.data.document.transformer.jsoup.HtmlToTextDocumentTransformer;
 import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
@@ -85,6 +82,7 @@ public void processAfterCommit(DataInsertedEvent event) {
                                 // 处理异常
                                 log.error("Error processing RAG file: {}", ragFile.getFileId(), e);
                                 ragFile.setStatus(FileStatus.PROCESS_FAILED);
+                                ragFile.setErrMsg(e.getMessage());
                                 ragFileRepository.updateById(ragFile);
                             } finally {
                                 SEMAPHORE.release();
@@ -109,7 +107,7 @@ private void processRagFile(RagFile ragFile, DataInsertedEvent event) {
         }
         document.metadata().put("rag_file_id", ragFile.getId());
         // 使用文档分块器对文档进行分块
-        DocumentSplitter splitter = documentSplitter(event.addFilesReq().getProcessType());
+        DocumentSplitter splitter = documentSplitter(event.addFilesReq());
         List<TextSegment> split = splitter.split(document);
 
         // 更新分块数量
@@ -145,13 +143,14 @@ public DocumentParser documentParser(String fileType) {
         };
     }
 
-    public DocumentSplitter documentSplitter(ProcessType processType) {
-        return switch (processType) {
-            case PARAGRAPH_CHUNK -> new DocumentByParagraphSplitter(1000, 100);
-            case CHAPTER_CHUNK -> new DocumentByLineSplitter(1000, 100);
-            case CUSTOM_SEPARATOR_CHUNK -> new DocumentBySentenceSplitter(1000, 100);
-            case LENGTH_CHUNK -> new DocumentByWordSplitter(1000, 100);
-            case DEFAULT_CHUNK -> new DocumentByLineSplitter(1000, 100);
+    public DocumentSplitter documentSplitter(AddFilesReq req) {
+        return switch (req.getProcessType()) {
+            case PARAGRAPH_CHUNK -> new DocumentByParagraphSplitter(req.getChunkSize(), req.getOverlapSize());
+            case SENTENCE_CHUNK -> new DocumentBySentenceSplitter(req.getChunkSize(), req.getOverlapSize());
+            case LENGTH_CHUNK -> new DocumentByCharacterSplitter(req.getChunkSize(), req.getOverlapSize());
+            case DEFAULT_CHUNK -> new DocumentByWordSplitter(req.getChunkSize(), req.getOverlapSize());
+            case CUSTOM_SEPARATOR_CHUNK ->
+                    new DocumentByRegexSplitter(req.getDelimiter(), "", req.getChunkSize(), req.getOverlapSize());
         };
     }
-}
+}
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
@@ -36,6 +36,11 @@ public EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel,
                 .build();
     }
 
+    /**
+     * 单例模式获取 Milvus 客户端，不依赖 Spring 容器
+     *
+     * @return MilvusClient
+     */
     public MilvusClient getMilvusClient() {
         if (milvusClient == null) {
             synchronized (this) {
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
@@ -11,7 +11,6 @@
 import org.springframework.web.bind.annotation.*;
 
 
-
 /**
  * 知识库控制器
  *
@@ -124,8 +123,20 @@ public void deleteFile(@PathVariable("knowledgeBaseId") String knowledgeBaseId,
      */
     @GetMapping("/{knowledgeBaseId}/files/{ragFileId}")
     public PagedResponse<RagChunk> getChunks(@PathVariable("knowledgeBaseId") String knowledgeBaseId,
-                                             @PathVariable("ragFileId") String ragFileId,
-                                             PagingQuery pagingQuery) {
+                                                  @PathVariable("ragFileId") String ragFileId,
+                                                  PagingQuery pagingQuery) {
         return knowledgeBaseService.getChunks(knowledgeBaseId, ragFileId, pagingQuery);
     }
+
+    /**
+     * 检索知识库内容
+     *
+     * @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);
+    }
 }
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/interfaces/dto/AddFilesReq.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/interfaces/dto/AddFilesReq.java
@@ -18,7 +18,7 @@ public class AddFilesReq {
     private ProcessType processType;
     private Integer chunkSize;
     private Integer overlapSize;
-    private String customSeparator;
+    private String delimiter;
     private List<FileInfo> files;
 
     public record FileInfo(String id, String name) {
diff --git a/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/interfaces/dto/ProcessType.java b/backend/services/rag-indexer-service/src/main/java/com/datamate/rag/indexer/interfaces/dto/ProcessType.java
@@ -7,27 +7,28 @@
  * @since 2025-10-29
  */
 public enum ProcessType {
-    /**
-     * 章节分块
-     */
-    CHAPTER_CHUNK,
     /**
      * 段落分块
      */
     PARAGRAPH_CHUNK,
 
     /**
-     * 按长度分块
+     * 按句子分块
      */
-    LENGTH_CHUNK,
+    SENTENCE_CHUNK,
 
     /**
-     * 自定义分割符分块
+     * 按长度分块，字符串分块
      */
-    CUSTOM_SEPARATOR_CHUNK,
+    LENGTH_CHUNK,
 
     /**
-     * 默认分块
+     * 默认分块，按单词分块
      */
     DEFAULT_CHUNK,
+
+    /**
+     * 自定义分割符分块
+     */
+    CUSTOM_SEPARATOR_CHUNK,
 }
diff --git a/backend/shared/domain-common/src/main/java/com/datamate/common/infrastructure/exception/KnowledgeBaseErrorCode.java b/backend/shared/domain-common/src/main/java/com/datamate/common/infrastructure/exception/KnowledgeBaseErrorCode.java
@@ -16,7 +16,12 @@ public enum KnowledgeBaseErrorCode implements ErrorCode {
     /**
      * 知识库不存在
      */
-    KNOWLEDGE_BASE_NOT_FOUND("knowledge.0001", "知识库不存在");
+    KNOWLEDGE_BASE_NOT_FOUND("knowledge.0001", "知识库不存在"),
+
+    /**
+     * 文件不存在
+     */
+    RAG_FILE_NOT_FOUND("knowledge.0002", "文件不存在");
 
     private final String code;
     private final String message;
diff --git a/scripts/db/rag-management-init.sql b/scripts/db/rag-management-init.sql
@@ -22,6 +22,7 @@ create table if not exists t_rag_file
   chunk_count       INT COMMENT '切片数',
   metadata          JSON COMMENT '元数据',
   status            VARCHAR(50) COMMENT '文件状态',
+  err_msg           text NULL COMMENT '错误信息',
   created_at        TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
   updated_at        TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
   created_by        VARCHAR(255) COMMENT '创建者',

Original file line number	Diff line number	Diff line change
`@@ -44,4 +44,6 @@ public class RagFile extends BaseEntity<String> {`
`44`	`44`	`private Map<String, Object> metadata;`
`45`	`45`
`46`	`46`	`private FileStatus status;`
	`47`	`+`
	`48`	`+ private String errMsg;`
`47`	`49`	`}`