BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 9b8136ef303568302131a1fd3914b47a19268d4f

Authored by xiaoxin
1 parent 92d32b65

【新增】AI 知识库: 配置自定义、段落启禁用

Inline Side-by-side
Showing 9 changed files with 149 additions and 27 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/vo/knowledge/AiKnowledgeCreateMyReqVO.java
  View file @9b8136e
... ... @@ -25,4 +25,11 @@ public class AiKnowledgeCreateMyReqVO {
25 25 @NotNull(message = "嵌入模型不能为空")
26 26 private Long modelId;
27 27  
  28 + @Schema(description = "相似性阈值", requiredMode = Schema.RequiredMode.REQUIRED, example = "0.5")
  29 + @NotNull(message = "相似性阈值不能为空")
  30 + private Double similarityThreshold;
  31 +
  32 + @Schema(description = "topK", requiredMode = Schema.RequiredMode.REQUIRED, example = "3")
  33 + @NotNull(message = "topK 不能为空")
  34 + private Integer topK;
28 35 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/vo/knowledge/AiKnowledgeDocumentCreateReqVO.java
  View file @9b8136e
... ... @@ -23,4 +23,23 @@ public class AiKnowledgeDocumentCreateReqVO {
23 23 @URL(message = "文档 URL 格式不正确")
24 24 private String url;
25 25  
  26 + @Schema(description = "每个文本块的目标 token 数", requiredMode = Schema.RequiredMode.REQUIRED, example = "800")
  27 + @NotNull(message = "每个文本块的目标 token 数不能为空")
  28 + private Integer defaultChunkSize;
  29 +
  30 + @Schema(description = "每个文本块的最小字符数", requiredMode = Schema.RequiredMode.REQUIRED, example = "350")
  31 + @NotNull(message = "每个文本块的最小字符数不能为空")
  32 + private Integer minChunkSizeChars;
  33 +
  34 + @Schema(description = "丢弃阈值", requiredMode = Schema.RequiredMode.REQUIRED, example = "5")
  35 + @NotNull(message = "丢弃阈值不能为空")
  36 + private Integer minChunkLengthToEmbed;
  37 +
  38 + @Schema(description = "最大块数", requiredMode = Schema.RequiredMode.REQUIRED, example = "10000")
  39 + @NotNull(message = "最大块数不能为空")
  40 + private Integer maxNumChunks;
  41 +
  42 + @Schema(description = "分块是否保留分隔符", requiredMode = Schema.RequiredMode.REQUIRED, example = "true")
  43 + @NotNull(message = "分块是否保留分隔符不能为空")
  44 + private Boolean keepSeparator;
26 45 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeDO.java
  View file @9b8136e
... ... @@ -52,6 +52,18 @@ public class AiKnowledgeDO extends BaseDO {
52 52 * 模型标识
53 53 */
54 54 private String model;
  55 +
  56 + /**
  57 + * topK
  58 + */
  59 + private Integer topK;
  60 +
  61 + /**
  62 + * 相似度阈值
  63 + */
  64 + private Double similarityThreshold;
  65 +
  66 +
55 67 /**
56 68 * 状态
57 69 * <p>
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeDocumentDO.java
  View file @9b8136e
... ... @@ -23,7 +23,7 @@ public class AiKnowledgeDocumentDO extends BaseDO {
23 23 private Long id;
24 24 /**
25 25 * 知识库编号
26   - *
  26 + * <p>
27 27 * 关联 {@link AiKnowledgeDO#getId()}
28 28 */
29 29 private Long knowledgeId;
... ... @@ -48,6 +48,26 @@ public class AiKnowledgeDocumentDO extends BaseDO {
48 48 */
49 49 private Integer wordCount;
50 50 /**
  51 + * 每个文本块的目标 token 数
  52 + */
  53 + private Integer defaultChunkSize;
  54 + /**
  55 + * 每个文本块的最小字符数
  56 + */
  57 + private Integer minChunkSizeChars;
  58 + /**
  59 + * 低于此值的块会被丢弃
  60 + */
  61 + private Integer minChunkLengthToEmbed;
  62 + /**
  63 + * 最大块数
  64 + */
  65 + private Integer maxNumChunks;
  66 + /**
  67 + * 分块是否保留分隔符
  68 + */
  69 + private Boolean keepSeparator;
  70 + /**
51 71 * 切片状态
52 72 * <p>
53 73 * 枚举 {@link AiKnowledgeDocumentStatusEnum}
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java
  View file @9b8136e
... ... @@ -2,6 +2,8 @@ package cn.iocoder.yudao.module.ai.dal.dataobject.knowledge;
2 2  
3 3 import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
4 4 import cn.iocoder.yudao.framework.mybatis.core.dataobject.BaseDO;
  5 +import com.baomidou.mybatisplus.annotation.FieldStrategy;
  6 +import com.baomidou.mybatisplus.annotation.TableField;
5 7 import com.baomidou.mybatisplus.annotation.TableId;
6 8 import com.baomidou.mybatisplus.annotation.TableName;
7 9 import lombok.Data;
... ... @@ -25,16 +27,17 @@ public class AiKnowledgeSegmentDO extends BaseDO {
25 27 /**
26 28 * 向量库的编号
27 29 */
  30 + @TableField(updateStrategy = FieldStrategy.ALWAYS)
28 31 private String vectorId;
29 32 /**
30 33 * 知识库编号
31   - *
  34 + * <p>
32 35 * 关联 {@link AiKnowledgeDO#getId()}
33 36 */
34 37 private Long knowledgeId;
35 38 /**
36 39 * 文档编号
37   - *
  40 + * <p>
38 41 * 关联 {@link AiKnowledgeDocumentDO#getId()}
39 42 */
40 43 private Long documentId;
... ... @@ -52,7 +55,7 @@ public class AiKnowledgeSegmentDO extends BaseDO {
52 55 private Integer tokens;
53 56 /**
54 57 * 状态
55   - *
  58 + * <p>
56 59 * 枚举 {@link CommonStatusEnum}
57 60 */
58 61 private Integer status;
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
  View file @9b8136e
... ... @@ -9,15 +9,11 @@ import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
9 9 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentPageReqVO;
10 10 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentUpdateReqVO;
11 11 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgeDocumentCreateReqVO;
12   -import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
13 12 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
14 13 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
15   -import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
16 14 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeDocumentMapper;
17 15 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeSegmentMapper;
18 16 import cn.iocoder.yudao.module.ai.enums.knowledge.AiKnowledgeDocumentStatusEnum;
19   -import cn.iocoder.yudao.module.ai.service.model.AiApiKeyService;
20   -import cn.iocoder.yudao.module.ai.service.model.AiChatModelService;
21 17 import jakarta.annotation.Resource;
22 18 import lombok.extern.slf4j.Slf4j;
23 19 import org.springframework.ai.document.Document;
... ... @@ -49,23 +45,15 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
49 45 private AiKnowledgeSegmentMapper segmentMapper;
50 46  
51 47 @Resource
52   - private TokenTextSplitter tokenTextSplitter;
53   - @Resource
54 48 private TokenCountEstimator tokenCountEstimator;
55   -
56   - @Resource
57   - private AiApiKeyService apiKeyService;
58 49 @Resource
59 50 private AiKnowledgeService knowledgeService;
60   - @Resource
61   - private AiChatModelService chatModelService;
62 51  
63 52 @Override
64 53 @Transactional(rollbackFor = Exception.class)
65 54 public Long createKnowledgeDocument(AiKnowledgeDocumentCreateReqVO createReqVO) {
66   - // 0. 校验
67   - AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(createReqVO.getKnowledgeId());
68   - AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId());
  55 + // 0. 校验并获取向量存储实例
  56 + VectorStore vectorStore = knowledgeService.getVectorStoreById(createReqVO.getKnowledgeId());
69 57  
70 58 // 1.1 下载文档
71 59 TikaDocumentReader loader = new TikaDocumentReader(downloadFile(createReqVO.getUrl()));
... ... @@ -82,6 +70,9 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
82 70 return documentId;
83 71 }
84 72  
  73 + // 2 构造文本分段器
  74 + TokenTextSplitter tokenTextSplitter = new TokenTextSplitter(createReqVO.getDefaultChunkSize(), createReqVO.getMinChunkSizeChars(), createReqVO.getMinChunkLengthToEmbed(),
  75 + createReqVO.getMaxNumChunks(), createReqVO.getKeepSeparator());
85 76 // 2.1 文档分段
86 77 List<Document> segments = tokenTextSplitter.apply(documents);
87 78 // 2.2 分段内容入库
... ... @@ -92,8 +83,6 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
92 83 .setStatus(CommonStatusEnum.ENABLE.getStatus()));
93 84 segmentMapper.insertBatch(segmentDOList);
94 85  
95   - // 3.1 获取向量存储实例
96   - VectorStore vectorStore = apiKeyService.getOrCreateVectorStore(model.getKeyId());
97 86 // 3.2 向量化并存储
98 87 segments.forEach(segment -> segment.getMetadata().put(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, createReqVO.getKnowledgeId()));
99 88 vectorStore.add(segments);
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentServiceImpl.java
  View file @9b8136e
... ... @@ -2,6 +2,7 @@ package cn.iocoder.yudao.module.ai.service.knowledge;
2 2  
3 3 import cn.hutool.core.collection.CollUtil;
4 4 import cn.hutool.core.collection.ListUtil;
  5 +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
5 6 import cn.iocoder.yudao.framework.common.pojo.PageResult;
6 7 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
7 8 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentPageReqVO;
... ... @@ -23,6 +24,10 @@ import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
23 24 import org.springframework.stereotype.Service;
24 25  
25 26 import java.util.List;
  27 +import java.util.Objects;
  28 +
  29 +import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
  30 +import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.KNOWLEDGE_SEGMENT_NOT_EXISTS;
26 31  
27 32 /**
28 33 * AI 知识库分片 Service 实现类
... ... @@ -50,14 +55,45 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
50 55  
51 56 @Override
52 57 public void updateKnowledgeSegment(AiKnowledgeSegmentUpdateReqVO reqVO) {
53   - segmentMapper.updateById(BeanUtils.toBean(reqVO, AiKnowledgeSegmentDO.class));
54   - // TODO @xin 重新向量化
  58 + // 0 校验
  59 + AiKnowledgeSegmentDO oldKnowledgeSegment = validateKnowledgeSegmentExists(reqVO.getId());
  60 + // 2.1 获取知识库向量实例
  61 + VectorStore vectorStore = knowledgeService.getVectorStoreById(oldKnowledgeSegment.getKnowledgeId());
  62 + // 2.2 删除原向量
  63 + vectorStore.delete(List.of(oldKnowledgeSegment.getVectorId()));
  64 +
  65 + // 2.3 重新向量化
  66 + Document document = new Document(reqVO.getContent());
  67 + document.getMetadata().put(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, oldKnowledgeSegment.getKnowledgeId());
  68 + vectorStore.add(List.of(document));
  69 +
  70 + // 2.1 更新段落内容
  71 + AiKnowledgeSegmentDO knowledgeSegment = BeanUtils.toBean(reqVO, AiKnowledgeSegmentDO.class);
  72 + knowledgeSegment.setVectorId(document.getId());
  73 + segmentMapper.updateById(knowledgeSegment);
55 74 }
56 75  
57 76 @Override
58 77 public void updateKnowledgeSegmentStatus(AiKnowledgeSegmentUpdateStatusReqVO reqVO) {
59   - segmentMapper.updateById(BeanUtils.toBean(reqVO, AiKnowledgeSegmentDO.class));
60   - // TODO @xin 1.禁用删除向量 2.启用重新向量化
  78 + // 0 校验
  79 + AiKnowledgeSegmentDO oldKnowledgeSegment = validateKnowledgeSegmentExists(reqVO.getId());
  80 + // 1 获取知识库向量实例
  81 + VectorStore vectorStore = knowledgeService.getVectorStoreById(oldKnowledgeSegment.getKnowledgeId());
  82 + AiKnowledgeSegmentDO knowledgeSegment = BeanUtils.toBean(reqVO, AiKnowledgeSegmentDO.class);
  83 +
  84 + if (Objects.equals(reqVO.getStatus(), CommonStatusEnum.ENABLE.getStatus())) {
  85 + // 2.1 启用重新向量化
  86 + Document document = new Document(oldKnowledgeSegment.getContent());
  87 + document.getMetadata().put(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, oldKnowledgeSegment.getKnowledgeId());
  88 + vectorStore.add(List.of(document));
  89 + knowledgeSegment.setVectorId(document.getId());
  90 + } else {
  91 + // 2.2 禁用删除向量
  92 + vectorStore.delete(List.of(oldKnowledgeSegment.getVectorId()));
  93 + knowledgeSegment.setVectorId(null);
  94 + }
  95 + // 3 更新段落状态
  96 + segmentMapper.updateById(knowledgeSegment);
61 97 }
62 98  
63 99 @Override
... ... @@ -71,9 +107,8 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
71 107  
72 108 // 1.2 向量检索
73 109 List<Document> documentList = vectorStore.similaritySearch(SearchRequest.query(reqVO.getContent())
74   - //TODO @xin 配置提取
75   - .withTopK(5)
76   - .withSimilarityThreshold(0.5d)
  110 + .withTopK(knowledge.getTopK())
  111 + .withSimilarityThreshold(knowledge.getSimilarityThreshold())
77 112 .withFilterExpression(new FilterExpressionBuilder().eq(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, reqVO.getKnowledgeId()).build()));
78 113 if (CollUtil.isEmpty(documentList)) {
79 114 return ListUtil.empty();
... ... @@ -81,4 +116,19 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
81 116 // 2.1 段落召回
82 117 return segmentMapper.selectList(CollUtil.getFieldValues(documentList, "id", String.class));
83 118 }
  119 +
  120 +
  121 + /**
  122 + * 校验段落是否存在
  123 + *
  124 + * @param id 文档编号
  125 + * @return 段落信息
  126 + */
  127 + private AiKnowledgeSegmentDO validateKnowledgeSegmentExists(Long id) {
  128 + AiKnowledgeSegmentDO knowledgeSegment = segmentMapper.selectById(id);
  129 + if (knowledgeSegment == null) {
  130 + throw exception(KNOWLEDGE_SEGMENT_NOT_EXISTS);
  131 + }
  132 + return knowledgeSegment;
  133 + }
84 134 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeService.java
  View file @9b8136e
... ... @@ -5,6 +5,7 @@ import cn.iocoder.yudao.framework.common.pojo.PageResult;
5 5 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgeCreateMyReqVO;
6 6 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgeUpdateMyReqVO;
7 7 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
  8 +import org.springframework.ai.vectorstore.VectorStore;
8 9  
9 10 /**
10 11 * AI 知识库-基础信息 Service 接口
... ... @@ -47,4 +48,12 @@ public interface AiKnowledgeService {
47 48 * @return 知识库分页
48 49 */
49 50 PageResult<AiKnowledgeDO> getKnowledgePageMy(Long userId, PageParam pageReqVO);
  51 +
  52 + /**
  53 + * 根据知识库编号获取向量存储实例
  54 + *
  55 + * @param knowledgeId 知识库编号
  56 + * @return 向量存储实例
  57 + */
  58 + VectorStore getVectorStoreById(Long knowledgeId);
50 59 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeServiceImpl.java
  View file @9b8136e
... ... @@ -10,9 +10,11 @@ import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnow
10 10 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
11 11 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
12 12 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeMapper;
  13 +import cn.iocoder.yudao.module.ai.service.model.AiApiKeyService;
13 14 import cn.iocoder.yudao.module.ai.service.model.AiChatModelService;
14 15 import jakarta.annotation.Resource;
15 16 import lombok.extern.slf4j.Slf4j;
  17 +import org.springframework.ai.vectorstore.VectorStore;
16 18 import org.springframework.stereotype.Service;
17 19  
18 20 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
... ... @@ -32,6 +34,10 @@ public class AiKnowledgeServiceImpl implements AiKnowledgeService {
32 34  
33 35 @Resource
34 36 private AiKnowledgeMapper knowledgeMapper;
  37 + @Resource
  38 + private AiChatModelService chatModelService;
  39 + @Resource
  40 + private AiApiKeyService apiKeyService;
35 41  
36 42 @Override
37 43 public Long createKnowledgeMy(AiKnowledgeCreateMyReqVO createReqVO, Long userId) {
... ... @@ -75,4 +81,11 @@ public class AiKnowledgeServiceImpl implements AiKnowledgeService {
75 81 return knowledgeMapper.selectPageByMy(userId, pageReqVO);
76 82 }
77 83  
  84 + @Override
  85 + public VectorStore getVectorStoreById(Long knowledgeId) {
  86 + AiKnowledgeDO knowledge = validateKnowledgeExists(knowledgeId);
  87 + AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId());
  88 + return apiKeyService.getOrCreateVectorStore(model.getKeyId());
  89 + }
  90 +
78 91 }
... ...