feat: 知识库优化

This commit is contained in:
moyangzhan 2024-05-16 18:34:31 +08:00
parent be9aecc45f
commit f4fa8d026f
11 changed files with 116 additions and 21 deletions

View File

@ -49,6 +49,7 @@ mybatis-plus:
mapper-locations: classpath*:/mapper/*.xml
configuration:
log-impl: org.apache.ibatis.logging.stdout.StdOutImpl
map-underscore-to-camel-case: true
logging:
file: /data/logs

View File

@ -1,21 +1,28 @@
package com.moyz.adi.chat.controller;
import com.moyz.adi.common.entity.AdiFile;
import com.moyz.adi.common.exception.BaseException;
import com.moyz.adi.common.service.FileService;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.hibernate.validator.constraints.Length;
import org.springframework.http.MediaType;
import org.springframework.core.io.InputStreamResource;
import org.springframework.http.*;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import static com.moyz.adi.common.enums.ErrorEnum.A_FILE_NOT_EXIST;
@Slf4j
@RestController
@Validated
@ -35,6 +42,25 @@ public class FileController {
}
}
@GetMapping(value = "/file/{uuid}")
public ResponseEntity<org.springframework.core.io.Resource> file(@Length(min = 32, max = 32) @PathVariable String uuid) {
AdiFile adiFile = fileService.getByUuid(uuid);
if (null == adiFile) {
throw new BaseException(A_FILE_NOT_EXIST);
}
byte[] bytes = fileService.readBytes(adiFile);
InputStreamResource inputStreamResource = new InputStreamResource(new ByteArrayInputStream(bytes));
String fileName = adiFile.getName();
if (StringUtils.isBlank(fileName)) {
fileName = adiFile.getUuid() + "." + adiFile.getExt();
}
HttpHeaders headers = new HttpHeaders();
headers.setContentDisposition(ContentDisposition.attachment().filename(fileName).build());
return new ResponseEntity<>(inputStreamResource, headers, HttpStatus.OK);
}
@PostMapping(path = "/file/upload", headers = "content-type=multipart/form-data", produces = MediaType.APPLICATION_JSON_VALUE)
public Map<String, String> upload(@RequestPart(value = "file") MultipartFile file) {
Map<String, String> result = new HashMap<>();

View File

@ -1,6 +1,7 @@
package com.moyz.adi.chat.controller;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.moyz.adi.common.dto.KbItemDto;
import com.moyz.adi.common.dto.KbItemEditReq;
import com.moyz.adi.common.dto.KbItemEmbeddingBatchReq;
import com.moyz.adi.common.entity.KnowledgeBaseItem;
@ -11,8 +12,6 @@ import jakarta.validation.constraints.NotNull;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import java.util.List;
@RestController
@RequestMapping("/knowledge-base-item")
@Validated
@ -27,7 +26,7 @@ public class KnowledgeBaseItemController {
}
@GetMapping("/search")
public Page<KnowledgeBaseItem> search(String kbUuid, String keyword, @NotNull @Min(1) Integer currentPage, @NotNull @Min(10) Integer pageSize) {
public Page<KbItemDto> search(String kbUuid, String keyword, @NotNull @Min(1) Integer currentPage, @NotNull @Min(10) Integer pageSize) {
return knowledgeBaseItemService.search(kbUuid, keyword, currentPage, pageSize);
}

View File

@ -27,6 +27,8 @@ public class ResponseWrapper implements ResponseBodyAdvice<Object> {
return result;
} else if (result instanceof String) {
return JsonUtil.toJson(new BaseResponse(true, result));
} else if (result instanceof org.springframework.core.io.Resource) {
return result;
}
log.info("result:" + result);
return new BaseResponse(true, result);

View File

@ -0,0 +1,34 @@
package com.moyz.adi.common.dto;
import com.baomidou.mybatisplus.annotation.TableField;
import lombok.Data;
import java.time.LocalDateTime;
@Data
public class KbItemDto {
private Long kbId;
private String kbUuid;
private Long sourceFileId;
private String uuid;
private String title;
private String brief;
private String remark;
private Boolean isEmbedded;
private String sourceFileName;
private String sourceFileUuid;
private LocalDateTime createTime;
private LocalDateTime updateTime;
}

View File

@ -1,9 +1,16 @@
package com.moyz.adi.common.mapper;
import com.baomidou.mybatisplus.annotation.InterceptorIgnore;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.moyz.adi.common.dto.KbItemDto;
import com.moyz.adi.common.entity.KnowledgeBaseItem;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;
@Mapper
public interface KnowledgeBaseItemMapper extends BaseMapper<KnowledgeBaseItem> {
@InterceptorIgnore(tenantLine = "true")
Page<KbItemDto> searchByKb(Page<KbItemDto> page, @Param("kbUuid") String kbUuid, @Param("keyword") String keyword);
}

View File

@ -21,7 +21,6 @@ import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.annotation.Async;
@ -49,9 +48,6 @@ public class AiImageService extends ServiceImpl<AiImageMapper, AiImage> {
@Resource
private QuotaHelper quotaHelper;
@Value("${local.images}")
private String localImagesPath;
@Resource
private RateLimitHelper rateLimitHelper;
@ -180,7 +176,7 @@ public class AiImageService extends ServiceImpl<AiImageMapper, AiImage> {
}
List<String> imageUuids = new ArrayList();
images.forEach(imageUrl -> {
String imageUuid = fileService.saveToLocal(user, imageUrl);
String imageUuid = fileService.saveImageToLocal(user, imageUrl);
imageUuids.add(imageUuid);
});
String imageUuidsJoin = imageUuids.stream().collect(Collectors.joining(","));

View File

@ -50,6 +50,7 @@ public class FileService extends ServiceImpl<FileMapper, AdiFile> {
String uuid = UUID.randomUUID().toString().replace("-", "");
Pair<String, String> originalFile = FileUtil.saveToLocal(file, imagePath, uuid);
AdiFile adiFile = new AdiFile();
adiFile.setName(file.getOriginalFilename());
adiFile.setUuid(uuid);
adiFile.setMd5(md5);
adiFile.setPath(originalFile.getLeft());
@ -59,7 +60,7 @@ public class FileService extends ServiceImpl<FileMapper, AdiFile> {
return adiFile;
}
public String saveToLocal(User user, String sourceImageUrl) {
public String saveImageToLocal(User user, String sourceImageUrl) {
String uuid = UUID.randomUUID().toString().replace("-", "");
String localPath = imagePath + uuid + ".png";
File target = new File(localPath);
@ -71,10 +72,12 @@ public class FileService extends ServiceImpl<FileMapper, AdiFile> {
throw new BaseException(B_SAVE_IMAGE_ERROR);
}
AdiFile adiFile = new AdiFile();
adiFile.setName(target.getName());
adiFile.setUuid(uuid);
adiFile.setMd5(MD5Utils.calculateMD5(localPath));
adiFile.setPath(localPath);
adiFile.setUserId(user.getId());
adiFile.setExt("png");
this.getBaseMapper().insert(adiFile);
return uuid;
}
@ -103,6 +106,20 @@ public class FileService extends ServiceImpl<FileMapper, AdiFile> {
return this.softDel(uuid);
}
public AdiFile getByUuid(String uuid) {
return this.lambdaQuery()
.eq(AdiFile::getUuid, uuid)
.eq(AdiFile::getUserId, ThreadContext.getCurrentUserId())
.oneOpt().orElse(null);
}
public byte[] readBytes(AdiFile adiFile) {
if (null == adiFile) {
throw new BaseException(A_FILE_NOT_EXIST);
}
return FileUtil.readBytes(adiFile.getPath());
}
public byte[] readBytes(String uuid) {
AdiFile adiFile = this.lambdaQuery()
.eq(AdiFile::getUuid, uuid)

View File

@ -1,11 +1,11 @@
package com.moyz.adi.common.service;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.baomidou.mybatisplus.extension.toolkit.ChainWrappers;
import com.moyz.adi.common.base.ThreadContext;
import com.moyz.adi.common.cosntant.AdiConstant;
import com.moyz.adi.common.dto.KbItemDto;
import com.moyz.adi.common.dto.KbItemEditReq;
import com.moyz.adi.common.entity.KnowledgeBase;
import com.moyz.adi.common.entity.KnowledgeBaseItem;
@ -25,7 +25,6 @@ import org.springframework.transaction.annotation.Transactional;
import java.util.Optional;
import java.util.UUID;
import static com.moyz.adi.common.cosntant.AdiConstant.RAG_TYPE_KB;
import static com.moyz.adi.common.enums.ErrorEnum.*;
@Slf4j
@ -74,15 +73,8 @@ public class KnowledgeBaseItemService extends ServiceImpl<KnowledgeBaseItemMappe
.one();
}
public Page<KnowledgeBaseItem> search(String kbUuid, String keyword, Integer currentPage, Integer pageSize) {
LambdaQueryChainWrapper<KnowledgeBaseItem> wrapper = ChainWrappers.lambdaQueryChain(baseMapper);
wrapper.select(KnowledgeBaseItem::getId, KnowledgeBaseItem::getUuid, KnowledgeBaseItem::getTitle, KnowledgeBaseItem::getBrief, KnowledgeBaseItem::getKbUuid, KnowledgeBaseItem::getIsEmbedded, KnowledgeBaseItem::getCreateTime, KnowledgeBaseItem::getUpdateTime);
wrapper.eq(KnowledgeBaseItem::getIsDeleted, false);
wrapper.eq(KnowledgeBaseItem::getKbUuid, kbUuid);
if (StringUtils.isNotBlank(keyword)) {
wrapper.eq(KnowledgeBaseItem::getTitle, keyword);
}
return wrapper.page(new Page<>(currentPage, pageSize));
public Page<KbItemDto> search(String kbUuid, String keyword, Integer currentPage, Integer pageSize) {
return baseMapper.searchByKb(new Page<>(currentPage, pageSize), kbUuid, keyword);
}
public boolean checkAndEmbedding(String[] uuids) {

View File

@ -10,6 +10,7 @@ import org.springframework.web.multipart.MultipartFile;
import java.io.File;
import java.io.IOException;
import java.nio.file.NoSuchFileException;
import static com.moyz.adi.common.enums.ErrorEnum.A_FILE_NOT_EXIST;
@ -37,6 +38,9 @@ public class FileUtil {
public static byte[] readBytes(String localPath) {
try {
return FileUtils.readFileToByteArray(new File(localPath));
} catch (NoSuchFileException noSuchFileException) {
log.error("file not found", noSuchFileException);
throw new BaseException(A_FILE_NOT_EXIST);
} catch (IOException e) {
throw new RuntimeException(e);
}

View File

@ -0,0 +1,17 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
<mapper namespace="com.moyz.adi.common.mapper.KnowledgeBaseItemMapper">
<select id="searchByKb" parameterType="Map" resultType="com.moyz.adi.common.dto.KbItemDto">
select a.*, b.name source_file_name, b.uuid source_file_uuid
from adi_knowledge_base_item a
left join adi_file b on a.source_file_id = b.id
where a.is_deleted = false
and a.kb_uuid = #{kbUuid}
<if test="keyword!='' and keyword!=null">
and a.title like CONCAT('%', #{keyword}, '%')
</if>
order by create_time desc
</select>
</mapper>