BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 7d53a6dd527d0bc381f12abe51977679320d99c3

Authored by xiaoxin
1 parent abf5a22c

【优化】AI 知识库: 向量存储补充 knowledgeId

Inline Side-by-side
Showing 2 changed files with 12 additions and 3 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java
  View file @7d53a6d
... ... @@ -15,6 +15,8 @@ import lombok.Data;
15 15 @Data
16 16 public class AiKnowledgeSegmentDO extends BaseDO {
17 17  
  18 + public static final String FIELD_KNOWLEDGE_ID = "knowledgeId";
  19 +
18 20 /**
19 21 * 编号
20 22 */
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
  View file @7d53a6d
... ... @@ -27,6 +27,7 @@ import org.springframework.stereotype.Service;
27 27 import org.springframework.transaction.annotation.Transactional;
28 28  
29 29 import java.util.List;
  30 +import java.util.Map;
30 31  
31 32 /**
32 33 * AI 知识库-文档 Service 实现类
... ... @@ -83,16 +84,22 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
83 84 List<Document> segments = tokenTextSplitter.apply(documents);
84 85 // 2.2 分段内容入库
85 86 List<AiKnowledgeSegmentDO> segmentDOList = CollectionUtils.convertList(segments,
86   - segment -> new AiKnowledgeSegmentDO().setContent(segment.getContent()).setDocumentId(documentId).setKnowledgeId(createReqVO.getKnowledgeId())
  87 + segment -> new AiKnowledgeSegmentDO().setContent(segment.getContent()).setDocumentId(documentId).setKnowledgeId(createReqVO.getKnowledgeId()).setVectorId(segment.getId())
87 88 .setTokens(tokenCountEstimator.estimate(segment.getContent())).setWordCount(segment.getContent().length())
88 89 .setStatus(CommonStatusEnum.ENABLE.getStatus()));
89 90 segmentMapper.insertBatch(segmentDOList);
90 91  
  92 + // 3.1 document 补充源数据
  93 + segments.forEach(segment -> {
  94 + Map<String, Object> metadata = segment.getMetadata();
  95 + metadata.put(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, createReqVO.getKnowledgeId());
  96 + });
  97 +
91 98 AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(createReqVO.getKnowledgeId());
92 99 AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId());
93   - // 3.1 获取向量存储实例
  100 + // 3.2 获取向量存储实例
94 101 VectorStore vectorStore = apiKeyService.getOrCreateVectorStore(model.getKeyId());
95   - // 3.2 向量化并存储
  102 + // 3.3 向量化并存储
96 103 vectorStore.add(segments);
97 104 return documentId;
98 105 }
... ...