【新增】AI 知识库: 段落召回

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/vo/segment/AiKnowledgeSegmentSearchReqVO.java

@@ -0,0 +1,17 @@
+package cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment;
+
+import io.swagger.v3.oas.annotations.media.Schema;
+import lombok.Data;
+
+
+@Schema(description = "管理后台 - AI 知识库段落召回 Request VO")
+@Data
+public class AiKnowledgeSegmentSearchReqVO {
+
+    @Schema(description = "知识库编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "24790")
+    private Long knowledgeId;
+
+    @Schema(description = "内容", requiredMode = Schema.RequiredMode.REQUIRED, example = "Java 学习路线")
+    private String content;
+
+}

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java

@@ -7,6 +7,8 @@ import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowle
 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
 import org.apache.ibatis.annotations.Mapper;
 
+import java.util.List;
+
 /**
  * AI 知识库-分片 Mapper
  *
@@ -22,4 +24,10 @@ public interface AiKnowledgeSegmentMapper extends BaseMapperX<AiKnowledgeSegment
                 .likeIfPresent(AiKnowledgeSegmentDO::getContent, reqVO.getKeyword())
                 .orderByDesc(AiKnowledgeSegmentDO::getId));
     }
+
+    default List<AiKnowledgeSegmentDO> selectList(List<String> vectorIdList) {
+        return selectList(new LambdaQueryWrapperX<AiKnowledgeSegmentDO>()
+                .in(AiKnowledgeSegmentDO::getVectorId, vectorIdList)
+                .orderByDesc(AiKnowledgeSegmentDO::getId));
+    }
 }

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentService.java

@@ -2,10 +2,13 @@ package cn.iocoder.yudao.module.ai.service.knowledge;
 
 import cn.iocoder.yudao.framework.common.pojo.PageResult;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentPageReqVO;
+import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentSearchReqVO;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateReqVO;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateStatusReqVO;
 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
 
+import java.util.List;
+
 /**
  * AI 知识库段落 Service 接口
  *
@@ -35,4 +38,13 @@ public interface AiKnowledgeSegmentService {
      */
     void updateKnowledgeSegmentStatus(AiKnowledgeSegmentUpdateStatusReqVO reqVO);
 
+
+    /**
+     * 段落召回
+     *
+     * @param reqVO 召回请求信息
+     * @return 召回的段落
+     */
+    List<AiKnowledgeSegmentDO> similaritySearch(AiKnowledgeSegmentSearchReqVO reqVO);
+
 }

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentServiceImpl.java

@@ -1,16 +1,29 @@
 package cn.iocoder.yudao.module.ai.service.knowledge;
 
+import cn.hutool.core.collection.CollUtil;
+import cn.hutool.core.collection.ListUtil;
 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.segment.AiKnowledgeSegmentPageReqVO;
+import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentSearchReqVO;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateReqVO;
 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateStatusReqVO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeSegmentMapper;
+import cn.iocoder.yudao.module.ai.service.model.AiApiKeyService;
+import cn.iocoder.yudao.module.ai.service.model.AiChatModelService;
 import jakarta.annotation.Resource;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
 import org.springframework.stereotype.Service;
 
+import java.util.List;
+
 /**
  * AI 知识库分片 Service 实现类
  *
@@ -23,6 +36,13 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
     @Resource
     private AiKnowledgeSegmentMapper segmentMapper;
 
+    @Resource
+    private AiKnowledgeService knowledgeService;
+    @Resource
+    private AiChatModelService chatModelService;
+    @Resource
+    private AiApiKeyService apiKeyService;
+
     @Override
     public PageResult<AiKnowledgeSegmentDO> getKnowledgeSegmentPage(AiKnowledgeSegmentPageReqVO pageReqVO) {
         return segmentMapper.selectPage(pageReqVO);
@@ -39,4 +59,26 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
         segmentMapper.updateById(BeanUtils.toBean(reqVO, AiKnowledgeSegmentDO.class));
         // TODO @xin 1.禁用删除向量 2.启用重新向量化
     }
+
+    @Override
+    public List<AiKnowledgeSegmentDO> similaritySearch(AiKnowledgeSegmentSearchReqVO reqVO) {
+        // 0. 校验
+        AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(reqVO.getKnowledgeId());
+        AiChatModelDO model = chatModelService.validateChatModel(knowledge.getModelId());
+
+        // 1.1 获取向量存储实例
+        VectorStore vectorStore = apiKeyService.getOrCreateVectorStore(model.getKeyId());
+
+        // 1.2 向量检索
+        List<Document> documentList = vectorStore.similaritySearch(SearchRequest.query(reqVO.getContent())
+                //TODO  @xin 配置提取
+                .withTopK(5)
+                .withSimilarityThreshold(0.5d)
+                .withFilterExpression(new FilterExpressionBuilder().eq(AiKnowledgeSegmentDO.FIELD_KNOWLEDGE_ID, reqVO.getKnowledgeId()).build()));
+        if (CollUtil.isEmpty(documentList)) {
+            return ListUtil.empty();
+        }
+        // 2.1 段落召回
+        return segmentMapper.selectList(CollUtil.getFieldValues(documentList, "id", String.class));
+    }
 }