diff --git a/electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/config/ElePropertyConfig.java b/electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/config/ElePropertyConfig.java index 7bb192c..97fc121 100644 --- a/electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/config/ElePropertyConfig.java +++ b/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 getAiAllowUploadFileTypes() { return List.of(allowUploadFileTypes.replace(" ", "").split(",")); } diff --git a/electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/service/serviceimpl/ChatService.java b/electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/service/serviceimpl/ChatService.java index 8cc09e8..1f9469d 100644 --- a/electrmangnetic/src/main/java/com/electromagnetic/industry/software/manage/service/serviceimpl/ChatService.java +++ b/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.List; -import java.util.Map; -import java.util.Objects; +import java.util.*; @Slf4j @Service @@ -145,10 +143,20 @@ public class ChatService { return ElectromagneticResultUtil.fail("-1", StrFormatter.format("当前格式 {} 不支持", fileType)); } } - documents = new TokenTextSplitter().apply(documents); - documents = documents.subList(0, Math.min(elePropertyConfig.getAiMaxDocs(), documents.size())); - vectorStore.write(documents); + List newDocuments = new ArrayList<>(); + StringBuilder stringBuilder = new StringBuilder(); for (Document document : documents) { + Map 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 documents = vectorStore.similaritySearch(queryDTO.getMsg()); - String augmentedPrompt = "基于以下信息回答:" + documents + "\n\n" + queryDTO.getMsg(); + Set files = new LinkedHashSet<>(); + assert documents != null; + for (Document document : documents) { + files.add(document.getMetadata().get("fileMd5").toString()); + } + List 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 documents = vectorStore.similaritySearch(queryDTO.getMsg()); - String augmentedPrompt = "基于以下信息回答:" + documents + "\n\n" + queryDTO.getMsg(); + + Set files = new LinkedHashSet<>(); + assert documents != null; + for (Document document : documents) { + files.add(document.getMetadata().get("fileMd5").toString()); + } + List 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) diff --git a/electrmangnetic/src/main/resources/application.properties b/electrmangnetic/src/main/resources/application.properties index 8f579d3..0a5890a 100644 --- a/electrmangnetic/src/main/resources/application.properties +++ b/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 diff --git a/electromagnetic-common/src/main/java/com/electromagnetic/industry/software/common/util/AESUtils.java b/electromagnetic-common/src/main/java/com/electromagnetic/industry/software/common/util/AESUtils.java index 476ec90..2167ac2 100644 --- a/electromagnetic-common/src/main/java/com/electromagnetic/industry/software/common/util/AESUtils.java +++ b/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;