【功能新增】AI：知识库文档的模型变化时，触发重索引（异步）

YunaiV
1 parent 32e1ef4d
Showing 5 changed files with 60 additions and 6 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentService.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentServiceImpl.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeServiceImpl.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiModelServiceImpl.java
@@ -42,6 +42,11 @@ public interface AiKnowledgeSegmentMapper extends BaseMapperX&lt;AiKnowledgeSegment
                 .orderByDesc(AiKnowledgeSegmentDO::getId));
     }
  
+    default List<AiKnowledgeSegmentDO> selectListByKnowledgeIdAndStatus(Long knowledgeId, Integer status) {
+        return selectList(AiKnowledgeSegmentDO::getKnowledgeId, knowledgeId,
+                AiKnowledgeSegmentDO::getStatus, status);
+    }
+
     default List<AiKnowledgeSegmentProcessRespVO> selectProcessList(Collection<Long> documentIds) {
         MPJLambdaWrapper<AiKnowledgeSegmentDO> wrapper = new MPJLambdaWrapperX<AiKnowledgeSegmentDO>()
                 .selectAs(AiKnowledgeSegmentDO::getDocumentId, AiKnowledgeSegmentProcessRespVO::getDocumentId)
@@ -54,7 +59,7 @@ public interface AiKnowledgeSegmentMapper extends BaseMapperX&lt;AiKnowledgeSegment
     }
  
     default void updateRetrievalCountIncrByIds(List<Long> ids) {
-        update( new LambdaUpdateWrapper<AiKnowledgeSegmentDO>()
+        update(new LambdaUpdateWrapper<AiKnowledgeSegmentDO>()
                 .setSql(" retrieval_count = retrieval_count + 1")
                 .in(AiKnowledgeSegmentDO::getId, ids));
     }
@@ -99,6 +99,23 @@ public interface AiKnowledgeSegmentService {
     void updateKnowledgeSegmentStatus(AiKnowledgeSegmentUpdateStatusReqVO reqVO);
  
     /**
+     * 重新索引知识库下的所有文档段落
+     *
+     * @param knowledgeId 知识库编号
+     */
+    void reindexKnowledgeSegmentByKnowledgeId(Long knowledgeId);
+
+    /**
+     * 【异步】重新索引知识库下的所有文档段落
+     *
+     * @param knowledgeId 知识库编号
+     */
+    @Async
+    default void reindexByKnowledgeIdAsync(Long knowledgeId) {
+        reindexKnowledgeSegmentByKnowledgeId(knowledgeId);
+    }
+
+    /**
      * 根据文档编号删除段落
      *
      * @param documentId 文档编号
@@ -158,6 +158,30 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService 
         }
     }
  
+    @Override
+    public void reindexKnowledgeSegmentByKnowledgeId(Long knowledgeId) {
+        // 1.1 校验知识库存在
+        AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(knowledgeId);
+        // 1.2 获取知识库向量实例
+        VectorStore vectorStore = getVectorStoreById(knowledge);
+
+        // 2.1 查询知识库下的所有启用状态的段落
+        List<AiKnowledgeSegmentDO> segments = segmentMapper.selectListByKnowledgeIdAndStatus(
+                knowledgeId, CommonStatusEnum.ENABLE.getStatus());
+        if (CollUtil.isEmpty(segments)) {
+            return;
+        }
+        // 2.2 遍历所有段落，重新索引
+        for (AiKnowledgeSegmentDO segment : segments) {
+            // 删除旧的向量
+            deleteVectorStore(vectorStore, segment);
+            // 重新创建向量
+            writeVectorStore(vectorStore, segment, new Document(segment.getContent()));
+        }
+        log.info("[reindexKnowledgeSegmentByKnowledgeId][知识库({}) 重新索引完成，共处理 {} 个段落]",
+                knowledgeId, segments.size());
+    }
+
     private void writeVectorStore(VectorStore vectorStore, AiKnowledgeSegmentDO segmentDO, Document segment) {
         // 1. 向量存储
         // 为什么要 toString 呢？因为部分 VectorStore 实现，不支持 Long 类型，例如说 QdrantVectorStore
 package cn.iocoder.yudao.module.ai.service.knowledge;
  
+import cn.hutool.core.util.ObjUtil;
+import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
 import cn.iocoder.yudao.framework.common.pojo.PageResult;
 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgePageReqVO;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgeSaveReqVO;
 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeMapper;
 import cn.iocoder.yudao.module.ai.service.model.AiModelService;
@@ -31,6 +34,8 @@ public class AiKnowledgeServiceImpl implements AiKnowledgeService {
  
     @Resource
     private AiModelService modelService;
+    @Resource
+    private AiKnowledgeSegmentService knowledgeSegmentService;
  
     @Override
     public Long createKnowledge(AiKnowledgeSaveReqVO createReqVO) {
@@ -47,7 +52,7 @@ public class AiKnowledgeServiceImpl implements AiKnowledgeService {
     @Override
     public void updateKnowledge(AiKnowledgeSaveReqVO updateReqVO) {
         // 1.1 校验知识库存在
-        validateKnowledgeExists(updateReqVO.getId());
+        AiKnowledgeDO oldKnowledge = validateKnowledgeExists(updateReqVO.getId());
         // 1.2 校验模型配置
         AiModelDO model = modelService.validateModel(updateReqVO.getEmbeddingModelId());
  
@@ -56,7 +61,10 @@ public class AiKnowledgeServiceImpl implements AiKnowledgeService {
                 .setEmbeddingModel(model.getModel());
         knowledgeMapper.updateById(updateObj);
  
-        // TODO @芋艿：如果模型变化，需要 reindex 所有的文档
+        // 3. 如果模型变化，需要 reindex 所有的文档
+        if (ObjUtil.notEqual(oldKnowledge.getEmbeddingModelId(), updateReqVO.getEmbeddingModelId())) {
+            knowledgeSegmentService.reindexByKnowledgeIdAsync(updateReqVO.getId());
+        }
     }
  
     @Override
@@ -16,8 +16,8 @@ import jakarta.annotation.Resource;
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.image.ImageModel;
+import org.springframework.ai.vectorstore.SimpleVectorStore;
 import org.springframework.ai.vectorstore.VectorStore;
-import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
 import org.springframework.stereotype.Service;
 import org.springframework.validation.annotation.Validated;
  
@@ -162,10 +162,10 @@ public class AiModelServiceImpl implements AiModelService {
                 platform, apiKey.getApiKey(), apiKey.getUrl(), model.getModel());
  
         // 创建或获取 VectorStore 对象
-        // return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel, metadataFields);
+         return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel, metadataFields);
         // return modelFactory.getOrCreateVectorStore(QdrantVectorStore.class, embeddingModel, metadataFields);
         // return modelFactory.getOrCreateVectorStore(RedisVectorStore.class, embeddingModel, metadataFields);
-        return modelFactory.getOrCreateVectorStore(MilvusVectorStore.class, embeddingModel, metadataFields);
+//        return modelFactory.getOrCreateVectorStore(MilvusVectorStore.class, embeddingModel, metadataFields);
     }
  
 }
 \ No newline at end of file