✨ 支持批量上传RAG资料

Author: twelvet-s

Diff for: twelvet-server/twelvet-server-ai/src/main/java/com/twelvet/server/ai/mapper/AiDocMapper.java

@@ -37,6 +37,13 @@ public interface AiDocMapper {
 	 */
 	int insertAiDoc(AiDoc aiDoc);
 
+	/**
+	 * 批量新增AI知识库文档
+	 * @param aiDocList AI知识库文档
+	 * @return 结果
+	 */
+	int insertAiDocBatch(List<AiDoc> aiDocList);
+
 	/**
 	 * 修改AI知识库文档
 	 * @param aiDoc AI知识库文档

Diff for: twelvet-server/twelvet-server-ai/src/main/java/com/twelvet/server/ai/mapper/AiDocSliceMapper.java

@@ -42,7 +42,7 @@ public interface AiDocSliceMapper {
 	 * @param aiDocSliceList AI知识库文档分片列表
 	 * @return 结果
 	 */
-	void insertAiDocSliceBatch(List<AiDocSlice> aiDocSliceList);
+	int insertAiDocSliceBatch(List<AiDocSlice> aiDocSliceList);
 
 	/**
 	 * 修改AI知识库文档分片

Diff for: twelvet-server/twelvet-server-ai/src/main/java/com/twelvet/server/ai/mq/consumer/RAGMqTopicListener.java

@@ -36,7 +36,11 @@ public RAGMqTopicListener(RAGMqTopicService ragMqTopicService) {
 	public Consumer<Message<AiDocMqDTO>> addRAGDocChannel() {
 		return message -> {
 			log.info("处理添加RAG文档消息: {}", message);
+			long startTime = System.nanoTime();
 			ragMqTopicService.addRAGDocChannel(message);
+			long endTime = System.nanoTime(); // 记录结束时间
+			long duration = endTime - startTime; // 计算持续时间
+			System.out.println("运行时间：" + duration / 1000000);
 		};
 	}
 

Diff for: twelvet-server/twelvet-server-ai/src/main/java/com/twelvet/server/ai/mq/consumer/domain/AiDocMqVO.java

@@ -0,0 +1,72 @@
+package com.twelvet.server.ai.mq.consumer.domain;
+
+import com.twelvet.api.ai.domain.AiDoc;
+import com.twelvet.api.ai.domain.AiDocSlice;
+import io.swagger.v3.oas.annotations.media.Schema;
+import org.springframework.ai.document.Document;
+
+import java.io.Serial;
+import java.io.Serializable;
+import java.util.List;
+
+/**
+ * <p>
+ * 批量处理添加RAG文档消息VO
+ * <p>
+ *
+ * @since 2025/1/10
+ */
+public class AiDocMqVO implements Serializable {
+
+	@Serial
+	private static final long serialVersionUID = 1L;
+
+	/**
+	 * 知识库文档
+	 */
+	@Schema(description = "知识库文档")
+	private AiDoc aiDoc;
+
+	/**
+	 * 知识库文档切片
+	 */
+	@Schema(description = "知识库文档切片")
+	private List<AiDocSlice> aiDocSliceList;
+
+	/**
+	 * 知识库文档切片向量
+	 */
+	@Schema(description = "知识库文档切片向量")
+	private List<Document> documentList;
+
+	public AiDoc getAiDoc() {
+		return aiDoc;
+	}
+
+	public void setAiDoc(AiDoc aiDoc) {
+		this.aiDoc = aiDoc;
+	}
+
+	public List<AiDocSlice> getAiDocSliceList() {
+		return aiDocSliceList;
+	}
+
+	public void setAiDocSliceList(List<AiDocSlice> aiDocSliceList) {
+		this.aiDocSliceList = aiDocSliceList;
+	}
+
+	public List<Document> getDocumentList() {
+		return documentList;
+	}
+
+	public void setDocumentList(List<Document> documentList) {
+		this.documentList = documentList;
+	}
+
+	@Override
+	public String toString() {
+		return "AiDocMqVO{" + "aiDoc=" + aiDoc + ", aiDocSliceList=" + aiDocSliceList + ", documentList=" + documentList
+				+ '}';
+	}
+
+}

Diff for: twelvet-server/twelvet-server-ai/src/main/java/com/twelvet/server/ai/mq/consumer/service/impl/RAGMqTopicServiceImpl.java

@@ -10,6 +10,7 @@
 import com.twelvet.framework.security.utils.SecurityUtils;
 import com.twelvet.server.ai.mapper.AiDocMapper;
 import com.twelvet.server.ai.mapper.AiDocSliceMapper;
+import com.twelvet.server.ai.mq.consumer.domain.AiDocMqVO;
 import com.twelvet.server.ai.mq.consumer.domain.dto.AiDocMqDTO;
 import com.twelvet.server.ai.mq.consumer.service.RAGMqTopicService;
 import org.slf4j.Logger;
@@ -138,8 +139,7 @@ else if (RAGEnums.DocSourceTypeEnums.UPLOAD.equals(sourceType)) { // 处理上
 
 			List<AiDocDTO.FileDTO> fileList = aiDocMqDTO.getFileList();
 
-			List<AiDocSlice> docSliceList = new ArrayList<>();
-			List<Document> docList = new ArrayList<>();
+			List<AiDocMqVO> aiDocMqVOList = new ArrayList<>();
 			// 针对文件进行doc插入
 			for (AiDocDTO.FileDTO file : fileList) {
 				String fileName = file.getFileName();
@@ -148,9 +148,13 @@ else if (RAGEnums.DocSourceTypeEnums.UPLOAD.equals(sourceType)) { // 处理上
 				TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(fileUrl);
 				List<Document> documents = tikaDocumentReader.get();
 				if (CollectionUtil.isEmpty(documents)) {
-					throw new TWTException("上传文件识别为空，空数据");
+					throw new TWTException("上传文件识别为空数据");
 				}
 
+				AiDocMqVO aiDocMqVO = new AiDocMqVO();
+				List<AiDocSlice> docSliceList = new ArrayList<>();
+				List<Document> docList = new ArrayList<>();
+
 				AiDoc aiDoc = new AiDoc();
 				aiDoc.setDocName(fileName);
 				aiDoc.setKnowledgeId(knowledgeId);
@@ -159,10 +163,13 @@ else if (RAGEnums.DocSourceTypeEnums.UPLOAD.equals(sourceType)) { // 处理上
 				aiDoc.setCreateTime(nowDate);
 				aiDoc.setUpdateBy(username);
 				aiDoc.setUpdateTime(nowDate);
-				// TODO 需要优化批量插入
-				aiDocMapper.insertAiDoc(aiDoc);
 
-				Long docId = aiDoc.getDocId();
+				// 设置vo信息
+				aiDocMqVO.setAiDoc(aiDoc);
+				aiDocMqVO.setAiDocSliceList(docSliceList);
+				aiDocMqVO.setDocumentList(docList);
+				// 加入集合
+				aiDocMqVOList.add(aiDocMqVO);
 
 				// 切片文档
 				for (Document document : documents) {
@@ -171,12 +178,10 @@ else if (RAGEnums.DocSourceTypeEnums.UPLOAD.equals(sourceType)) { // 处理上
 					for (Document doc : docs) {
 						Map<String, Object> metadata = doc.getMetadata();
 						metadata.put(RAGEnums.VectorMetadataEnums.KNOWLEDGE_ID.getCode(), knowledgeId);
-						metadata.put(RAGEnums.VectorMetadataEnums.DOC_ID.getCode(), docId);
 
 						AiDocSlice aiDocSlice = new AiDocSlice();
 						aiDocSlice.setKnowledgeId(knowledgeId);
 						aiDocSlice.setVectorId(doc.getId());
-						aiDocSlice.setDocId(docId);
 						aiDocSlice.setSliceName(fileName);
 						aiDocSlice.setContent(doc.getContent());
 						aiDocSlice.setCreateBy(username);
@@ -192,7 +197,32 @@ else if (RAGEnums.DocSourceTypeEnums.UPLOAD.equals(sourceType)) { // 处理上
 				}
 			}
 
-			if (CollectionUtil.isNotEmpty(docSliceList)) {
+			if (CollectionUtil.isNotEmpty(aiDocMqVOList)) {
+
+				List<AiDoc> aiDocList = aiDocMqVOList.stream().map(AiDocMqVO::getAiDoc).toList();
+				// 批量插入文档
+				aiDocMapper.insertAiDocBatch(aiDocList);
+
+				List<AiDocSlice> docSliceList = new ArrayList<>();
+				List<Document> docList = new ArrayList<>();
+
+				for (AiDocMqVO aiDocMqVO : aiDocMqVOList) {
+					AiDoc aiDoc = aiDocMqVO.getAiDoc();
+					Long docId = aiDoc.getDocId();
+
+					List<AiDocSlice> aiDocSliceList = aiDocMqVO.getAiDocSliceList();
+					for (AiDocSlice aiDocSlice : aiDocSliceList) {
+						aiDocSlice.setDocId(docId);
+						docSliceList.add(aiDocSlice);
+					}
+
+					List<Document> documentList = aiDocMqVO.getDocumentList();
+					for (Document document : documentList) {
+						Map<String, Object> metadata = document.getMetadata();
+						metadata.put(RAGEnums.VectorMetadataEnums.DOC_ID.getCode(), docId);
+						docList.add(document);
+					}
+				}
 
 				// 插入切片
 				aiDocSliceMapper.insertAiDocSliceBatch(docSliceList);

Diff for: twelvet-server/twelvet-server-ai/src/main/java/com/twelvet/server/ai/service/impl/AiDocServiceImpl.java

@@ -117,7 +117,8 @@ else if (RAGEnums.DocSourceTypeEnums.UPLOAD.equals(aiDocDTO.getSourceType())) {
 		aiDocMqDTO.setFileList(aiDocDTO.getFileList());
 
 		aiDocMqDTO.setOperatorBy(username);
-		streamBridge.send(RAGChannel.ADD_RAG_DOC, MessageBuilder.withPayload(aiDocMqDTO).build());
+		// streamBridge.send(RAGChannel.ADD_RAG_DOC,
+		// MessageBuilder.withPayload(aiDocMqDTO).build());
 
 		return Boolean.TRUE;
 	}

Diff for: twelvet-server/twelvet-server-ai/src/main/resources/mapper/ai/AiDocMapper.xml

@@ -56,6 +56,15 @@
         </trim>
     </insert>
 
+    <insert id="insertAiDocBatch" parameterType="List" useGeneratedKeys="true" keyProperty="docId">
+        INSERT INTO ai_doc ( knowledge_id, source_type, doc_name, create_by, create_time, update_by, update_time )
+        values
+        <foreach collection="aiDocList" item="aiDoc" separator=",">
+            (#{aiDoc.knowledgeId}, #{aiDoc.sourceType}, #{aiDoc.docName}, #{aiDoc.createBy}, #{aiDoc.createTime},
+            #{aiDoc.updateBy}, #{aiDoc.updateTime})
+        </foreach>
+    </insert>
+
     <update id="updateAiDoc" parameterType="AiDoc">
         update ai_doc
         <trim prefix="SET" suffixOverrides=",">

Diff for: twelvet-server/twelvet-server-dfs/src/main/java/com/twelvet/server/dfs/controller/DFSController.java

@@ -35,7 +35,7 @@ public class DFSController extends TWTController {
 	 * @return JsonResult<List<SysDfs>>
 	 */
 	@Operation(summary = "多文件上传")
-	@Log(service = "多文件上传", businessType = BusinessType.IMPORT)
+	@Log(service = "多文件上传", businessType = BusinessType.OTHER)
 	@PostMapping("/batchUpload")
 	public JsonResult<List<SysDfs>> batchUpload(MultipartFile[] files) {
 		// 上传并返回访问地址
@@ -51,7 +51,7 @@ public JsonResult<List<SysDfs>> batchUpload(MultipartFile[] files) {
 	 */
 	@Operation(summary = "单文件上传")
 	@PostMapping("/commonUpload")
-	@Log(service = "单文件上传", businessType = BusinessType.IMPORT)
+	@Log(service = "单文件上传", businessType = BusinessType.OTHER)
 	public JsonResult<String> commonUpload(MultipartFile file) {
 		// 上传并返回访问地址
 		SysDfs sysDfs = sysFileService.uploadFile(file);

Diff for: twelvet-server/twelvet-server-dfs/src/main/java/com/twelvet/server/dfs/service/impl/DFSServiceImpl.java

@@ -65,7 +65,7 @@ public List<SysDfs> uploadFiles(MultipartFile[] files) {
 				long size = file.getSize();
 
 				sysDfs.setSize(size);
-				sysDfs.setPath("/" + key);
+				sysDfs.setPath(key);
 				sysDfs.setType(FileUtils.getSuffix(originalFilename));
 				sysDfs.setFileName(FileUtils.getName(originalFilename));
 				sysDfs.setOriginalFileName(originalFilename);