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

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

View File

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

View File

@ -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<Document> segments = tokenTextSplitter.apply(documents);
// 2.2 分段内容入库
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())
.setStatus(CommonStatusEnum.ENABLE.getStatus()));
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());
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;
}