package com.ruoyi.approve.service.impl;
|
|
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
|
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
|
import com.ruoyi.approve.dto.KnowledgeBaseVectorVO;
|
import com.ruoyi.approve.mapper.KnowledgeBaseVectorMapper;
|
import com.ruoyi.approve.pojo.KnowledgeBase;
|
import com.ruoyi.approve.pojo.KnowledgeBaseVector;
|
import com.ruoyi.approve.service.KnowledgeBaseService;
|
import com.ruoyi.approve.service.KnowledgeBaseVectorService;
|
import com.ruoyi.ai.service.KnowledgeRagService;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.context.annotation.Lazy;
|
import org.springframework.stereotype.Service;
|
import org.springframework.transaction.annotation.Transactional;
|
|
import java.util.List;
|
|
/**
|
* 知识库文件向量记录 Service实现
|
*/
|
@Slf4j
|
@Service
|
public class KnowledgeBaseVectorServiceImpl extends ServiceImpl<KnowledgeBaseVectorMapper, KnowledgeBaseVector>
|
implements KnowledgeBaseVectorService {
|
|
private final KnowledgeBaseService knowledgeBaseService;
|
private final KnowledgeRagService knowledgeRagService;
|
|
public KnowledgeBaseVectorServiceImpl(
|
KnowledgeBaseService knowledgeBaseService,
|
@Lazy KnowledgeRagService knowledgeRagService) {
|
this.knowledgeBaseService = knowledgeBaseService;
|
this.knowledgeRagService = knowledgeRagService;
|
}
|
|
@Override
|
public List<KnowledgeBaseVectorVO> getVectorStatusByKnowledgeBaseId(Long knowledgeBaseId) {
|
return baseMapper.selectByKnowledgeBaseId(knowledgeBaseId);
|
}
|
|
@Override
|
public KnowledgeBaseVector createVectorRecord(Long knowledgeBaseId, Long storageBlobId,
|
String fileName, String fileType) {
|
KnowledgeBase knowledgeBase = knowledgeBaseService.getById(knowledgeBaseId);
|
if (knowledgeBase == null) {
|
throw new RuntimeException("知识库不存在: " + knowledgeBaseId);
|
}
|
|
KnowledgeBaseVector vector = new KnowledgeBaseVector();
|
vector.setKnowledgeBaseId(knowledgeBaseId);
|
vector.setStorageBlobId(storageBlobId);
|
vector.setFileName(fileName);
|
vector.setFileType(fileType);
|
vector.setVectorStatus(KnowledgeBaseVector.STATUS_PENDING);
|
vector.setNamespace("kb-" + knowledgeBaseId);
|
vector.setChunkCount(0);
|
save(vector);
|
|
// 异步触发向量化处理
|
knowledgeRagService.processVectorAsync(vector.getId());
|
|
return vector;
|
}
|
|
@Override
|
public void updateVectorStatus(Long id, Integer status, Integer chunkCount, String error) {
|
KnowledgeBaseVector vector = getById(id);
|
if (vector == null) {
|
return;
|
}
|
vector.setVectorStatus(status);
|
if (chunkCount != null) {
|
vector.setChunkCount(chunkCount);
|
}
|
if (error != null) {
|
vector.setVectorError(error);
|
}
|
updateById(vector);
|
|
// 如果完成，更新知识库统计
|
if (status == KnowledgeBaseVector.STATUS_COMPLETED) {
|
updateKnowledgeBaseStats(vector.getKnowledgeBaseId());
|
}
|
}
|
|
@Override
|
@Transactional(rollbackFor = Exception.class)
|
public void reprocessVector(Long id) {
|
KnowledgeBaseVector vector = getById(id);
|
if (vector == null) {
|
throw new RuntimeException("向量记录不存在: " + id);
|
}
|
vector.setVectorStatus(KnowledgeBaseVector.STATUS_PENDING);
|
vector.setVectorError(null);
|
vector.setChunkCount(0);
|
updateById(vector);
|
|
// 异步重新处理
|
knowledgeRagService.processVectorAsync(id);
|
}
|
|
@Override
|
@Transactional(rollbackFor = Exception.class)
|
public void deleteVector(Long id) {
|
KnowledgeBaseVector vector = getById(id);
|
if (vector == null) {
|
return;
|
}
|
|
// 删除向量库中的数据
|
try {
|
knowledgeRagService.deleteEmbeddings(vector.getNamespace(), vector.getStorageBlobId());
|
} catch (Exception e) {
|
log.error("删除向量库数据失败", e);
|
}
|
|
// 删除记录
|
removeById(id);
|
|
// 更新知识库统计
|
updateKnowledgeBaseStats(vector.getKnowledgeBaseId());
|
}
|
|
@Override
|
@Transactional(rollbackFor = Exception.class)
|
public void deleteVectors(List<Long> ids) {
|
for (Long id : ids) {
|
deleteVector(id);
|
}
|
}
|
|
private void updateKnowledgeBaseStats(Long knowledgeBaseId) {
|
KnowledgeBase knowledgeBase = knowledgeBaseService.getById(knowledgeBaseId);
|
if (knowledgeBase == null) {
|
return;
|
}
|
int fileCount = baseMapper.countByKnowledgeBaseId(knowledgeBaseId);
|
int totalChunkCount = baseMapper.sumChunkCountByKnowledgeBaseId(knowledgeBaseId);
|
knowledgeBase.setFileCount(fileCount);
|
knowledgeBase.setTotalChunkCount(totalChunkCount);
|
knowledgeBaseService.updateById(knowledgeBase);
|
}
|
}
|
