diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java index 6d5da06ea..fc758ce31 100644 --- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java +++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java @@ -15,6 +15,8 @@ import lombok.Data; @Data public class AiKnowledgeSegmentDO extends BaseDO { + public static final String FIELD_KNOWLEDGE_ID = "knowledgeId"; + /** * 编号 */ diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java index 9be3dd1af..bcfb64c55 100644 --- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java +++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java @@ -27,6 +27,7 @@ import org.springframework.stereotype.Service; import org.springframework.transaction.annotation.Transactional; import java.util.List; +import java.util.Map; /** * AI 知识库-文档 Service 实现类 @@ -83,16 +84,22 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic List segments = tokenTextSplitter.apply(documents); // 2.2 分段内容入库 List segmentDOList = CollectionUtils.convertList(segments, - segment -> new AiKnowledgeSegmentDO().setContent(segment.getContent()).setDocumentId(documentId).setKnowledgeId(createReqVO.getKnowledgeId()) + segment -> new AiKnowledgeSegmentDO().setContent(segment.getContent()).setDocumentId(documentId).setKnowledgeId(createReqVO.getKnowledgeId()).setVectorId(segment.getId()) .setTokens(tokenCountEstimator.estimate(segment.getContent())).setWordCount(segment.getContent().length()) .setStatus(CommonStatusEnum.ENABLE.getStatus())); segmentMapper.insertBatch(segmentDOList); + // 3.1 document 补充源数据 + segments.forEach(segment -> { + Map metadata = segment.getMetadata(); + metadata.put(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, createReqVO.getKnowledgeId()); + }); + AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(createReqVO.getKnowledgeId()); AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId()); - // 3.1 获取向量存储实例 + // 3.2 获取向量存储实例 VectorStore vectorStore = apiKeyService.getOrCreateVectorStore(model.getKeyId()); - // 3.2 向量化并存储 + // 3.3 向量化并存储 vectorStore.add(segments); return documentId; }