huangge1199/vue-pro
1
0
Fork 0
mirror of https://gitee.com/huangge1199_admin/vue-pro.git synced 2024-11-26 17:21:53 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
xiaoxin 2024-08-29 15:56:26 +08:00
parent abf5a22cd0
commit 7d53a6dd52
2 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/knowledge/AiKnowledgeSegmentDO.java
View File

@ -15,6 +15,8 @@ import lombok.Data;
@Data @Data
public class AiKnowledgeSegmentDO extends BaseDO { public class AiKnowledgeSegmentDO extends BaseDO {
public static final String FIELD_KNOWLEDGE_ID = "knowledgeId";
/** /**
* 编号 * 编号
*/ */

13
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
View File

@ -27,6 +27,7 @@ import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional; import org.springframework.transaction.annotation.Transactional;
import java.util.List; import java.util.List;
import java.util.Map;
/** /**
* AI 知识库-文档 Service 实现类 * AI 知识库-文档 Service 实现类
@ -83,16 +84,22 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
List<Document> segments = tokenTextSplitter.apply(documents); List<Document> segments = tokenTextSplitter.apply(documents);
// 2.2 分段内容入库 // 2.2 分段内容入库
List<AiKnowledgeSegmentDO> segmentDOList = CollectionUtils.convertList(segments, List<AiKnowledgeSegmentDO> 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()) .setTokens(tokenCountEstimator.estimate(segment.getContent())).setWordCount(segment.getContent().length())
.setStatus(CommonStatusEnum.ENABLE.getStatus())); .setStatus(CommonStatusEnum.ENABLE.getStatus()));
segmentMapper.insertBatch(segmentDOList); segmentMapper.insertBatch(segmentDOList);
// 3.1 document 补充源数据
segments.forEach(segment -> {
Map<String, Object> metadata = segment.getMetadata();
metadata.put(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, createReqVO.getKnowledgeId());
});
AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(createReqVO.getKnowledgeId()); AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(createReqVO.getKnowledgeId());
AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId()); AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId());
// 3.1 获取向量存储实例 // 3.2 获取向量存储实例
VectorStore vectorStore = apiKeyService.getOrCreateVectorStore(model.getKeyId()); VectorStore vectorStore = apiKeyService.getOrCreateVectorStore(model.getKeyId());
// 3.2 向量化并存储 // 3.3 向量化并存储
vectorStore.add(segments); vectorStore.add(segments);
return documentId; return documentId;
} }