更新ai功能

electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/config/ElePropertyConfig.java

@@ -1,11 +1,13 @@
 package com.electromagnetic.industry.software.manage.config;
 
+import cn.hutool.core.io.FileUtil;
 import com.electromagnetic.industry.software.common.enums.DataOwnEnum;
 import com.electromagnetic.industry.software.common.util.EleCommonUtil;
 import lombok.Getter;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 
+import java.io.File;
 import java.util.List;
 
 @Component
@@ -78,6 +80,16 @@ public class ElePropertyConfig {
     @Value("${allow.upload.file.types}")
     private String allowUploadFileTypes;
 
+    @Value("${ai.file.upload.store.dir}")
+    private String aiFileUploadStoreDir;
+
+    public String getAiFileUploadStoreDir() {
+        if (EleCommonUtil.isWinOs()) {
+            return FileUtil.normalize(winPrefix + File.separator + aiFileUploadStoreDir);
+        }
+        return FileUtil.normalize(aiFileUploadStoreDir);
+    }
+
     public List<String> getAiAllowUploadFileTypes() {
         return List.of(allowUploadFileTypes.replace(" ", "").split(","));
     }

electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/service/serviceimpl/ChatService.java

@@ -39,12 +39,10 @@ import org.springframework.web.multipart.MultipartFile;
 import reactor.core.publisher.Flux;
 
 import java.io.File;
+import java.nio.charset.Charset;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.ArrayList;
+import java.util.*;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
 
 @Slf4j
 @Service
@@ -145,10 +143,20 @@ public class ChatService {
                 return ElectromagneticResultUtil.fail("-1", StrFormatter.format("当前格式 {} 不支持", fileType));
             }
         }
-        documents = new TokenTextSplitter().apply(documents);
+        List<Document> newDocuments = new ArrayList<>();
-        documents = documents.subList(0, Math.min(elePropertyConfig.getAiMaxDocs(), documents.size()));
+        StringBuilder stringBuilder = new StringBuilder();
-        vectorStore.write(documents);
         for (Document document : documents) {
+            Map<String, Object> metadata = new HashMap<>();
+            metadata.put("fileName", file.getOriginalFilename());
+            metadata.put("fileMd5", fileMd5);
+            newDocuments.add(new Document(Objects.requireNonNull(document.getText()), metadata));
+            stringBuilder.append(document.getText());
+        }
+        FileUtil.writeString(stringBuilder.toString(), elePropertyConfig.getAiFileUploadStoreDir() + File.separator + fileMd5, Charset.defaultCharset());
+        newDocuments = new TokenTextSplitter().apply(newDocuments);
+        newDocuments = newDocuments.subList(0, Math.min(elePropertyConfig.getAiMaxDocs(), newDocuments.size()));
+        vectorStore.write(newDocuments);
+        for (Document document : newDocuments) {
             aiFileUploadRecordMapper.insert(new AiFileUploadRecord()
                     .setId(IdWorker.getSnowFlakeIdString())
                     .setVectorId(document.getId())
@@ -169,7 +177,17 @@ public class ChatService {
         aiQuestionRecordMapper.insert(record);
         UserThreadLocal.remove();
         List<Document> documents = vectorStore.similaritySearch(queryDTO.getMsg());
-        String augmentedPrompt = "基于以下信息回答：" + documents + "\n\n" + queryDTO.getMsg();
+        Set<String> files = new LinkedHashSet<>();
+        assert documents != null;
+        for (Document document : documents) {
+            files.add(document.getMetadata().get("fileMd5").toString());
+        }
+        List<Document> newDocuments = new ArrayList<>();
+        for (String fileMd5 : files) {
+            String s = FileUtil.readString(elePropertyConfig.getAiFileUploadStoreDir() + File.separator + fileMd5, Charset.defaultCharset());
+            newDocuments.add(new Document(s));
+        }
+        String augmentedPrompt = "基于以下信息回答：" + newDocuments + "\n\n" + queryDTO.getMsg();
         return ChatClient.builder(model)
                 .defaultSystem("必须用中文回答")
                 .defaultAdvisors(messageChatMemoryAdvisor)
@@ -191,7 +209,19 @@ public class ChatService {
         aiQuestionRecordMapper.insert(record);
         UserThreadLocal.remove();
         List<Document> documents = vectorStore.similaritySearch(queryDTO.getMsg());
-        String augmentedPrompt = "基于以下信息回答：" + documents + "\n\n" + queryDTO.getMsg();
+
+        Set<String> files = new LinkedHashSet<>();
+        assert documents != null;
+        for (Document document : documents) {
+            files.add(document.getMetadata().get("fileMd5").toString());
+        }
+        List<Document> newDocuments = new ArrayList<>();
+        for (String fileMd5 : files) {
+            String s = FileUtil.readString(elePropertyConfig.getAiFileUploadStoreDir() + File.separator + fileMd5, Charset.defaultCharset());
+            newDocuments.add(new Document(s));
+        }
+
+        String augmentedPrompt = "基于以下信息回答：" + newDocuments + "\n\n" + queryDTO.getMsg();
         return ChatClient.builder(model)
                 .defaultSystem("必须用中文回答")
                 .defaultAdvisors(messageChatMemoryAdvisor)

electrmangnetic/src/main/resources/application.properties

@@ -21,6 +21,7 @@ spring.ai.vectorstore.mariadb.distance-type=COSINE
 spring.ai.vectorstore.mariadb.dimensions=1536
 spring.ai.vectorstore.mariadb.table-name=rag
 allow.upload.file.types=doc, docx, txt, csv, xls, xlsx, pdf, ppt, pptx
+ai.file.upload.store.dir=/szsd/data/ele/aifile/
 # mybatis
 mybatis-plus.mapper-locations=classpath:sqlmapper/*.xml
 mybatis-plus.configuration.log-impl=org.apache.ibatis.logging.stdout.StdOutImpl

electromagnetic-common/src/main/java/com/electromagnetic/industry/software/common/util/AESUtils.java

@@ -4,7 +4,10 @@ package com.electromagnetic.industry.software.common.util;
 import cn.hutool.core.codec.Base64;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
 
-import javax.crypto.*;
+import javax.crypto.BadPaddingException;
+import javax.crypto.Cipher;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
 import javax.crypto.spec.SecretKeySpec;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;