From 7d53a6dd527d0bc381f12abe51977679320d99c3 Mon Sep 17 00:00:00 2001 From: xiaoxin <718949661@qq.com> Date: Thu, 29 Aug 2024 15:56:26 +0800 Subject: [PATCH] =?UTF-8?q?=E3=80=90=E4=BC=98=E5=8C=96=E3=80=91AI=20?= =?UTF-8?q?=E7=9F=A5=E8=AF=86=E5=BA=93:=20=E5=90=91=E9=87=8F=E5=AD=98?= =?UTF-8?q?=E5=82=A8=E8=A1=A5=E5=85=85=20knowledgeId?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../dataobject/knowledge/AiKnowledgeSegmentDO.java | 2 ++ .../knowledge/AiKnowledgeDocumentServiceImpl.java | 13 ++++++++++--- 2 files changed, 12 insertions(+), 3 deletions(-) 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; }