Commit 5f5e77a3923bc39b3ef160195e00df5f3baee1aa

Authored by YunaiV
1 parent ebd93514

【功能新增】AI:新增 document 向量的进度查询

Showing 20 changed files with 385 additions and 61 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/AiKnowledgeController.http
... ... @@ -9,7 +9,8 @@ tenant-id: {{adminTenantId}}
9 9 "description": "测试描述",
10 10 "embeddingModelId": 30,
11 11 "topK": 3,
12   - "similarityThreshold": 0.5
  12 + "similarityThreshold": 0.5,
  13 + "status": 0
13 14 }
14 15  
15 16 ### 更新知识库
... ... @@ -24,7 +25,8 @@ tenant-id: {{adminTenantId}}
24 25 "description": "测试描述",
25 26 "embeddingModelId": 30,
26 27 "topK": 5,
27   - "similarityThreshold": 0.6
  28 + "similarityThreshold": 0.6,
  29 + "status": 0
28 30 }
29 31  
30 32 ### 获取知识库分页
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/AiKnowledgeDocumentController.http
... ... @@ -5,7 +5,7 @@ Authorization: Bearer {{token}}
5 5 tenant-id: {{adminTenantId}}
6 6  
7 7 {
8   - "knowledgeId": 1,
  8 + "knowledgeId": 2,
9 9 "name": "测试文档",
10 10 "url": "https://static.iocoder.cn/README.md",
11 11 "segmentMaxTokens": 800
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/AiKnowledgeSegmentController.http
... ... @@ -4,9 +4,14 @@ Content-Type: application/json
4 4 Authorization: Bearer {{token}}
5 5 tenant-id: {{adminTenantId}}
6 6  
7   -### 获取文档处理列表
8   -GET {{baseUrl}}/ai/knowledge/segment/get-process-list?documentIds=1,2,3
  7 +### 搜索段落内容
  8 +GET {{baseUrl}}/ai/knowledge/segment/search?knowledgeId=2&content=如何使用这个产品&topK=5&similarityThreshold=0.1
9 9 Content-Type: application/json
10 10 Authorization: Bearer {{token}}
11 11 tenant-id: {{adminTenantId}}
12 12  
  13 +### 获取文档处理列表
  14 +GET {{baseUrl}}/ai/knowledge/segment/get-process-list?documentIds=1,2,3
  15 +Content-Type: application/json
  16 +Authorization: Bearer {{token}}
  17 +tenant-id: {{adminTenantId}}
13 18 \ No newline at end of file
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/AiKnowledgeSegmentController.java
1 1 package cn.iocoder.yudao.module.ai.controller.admin.knowledge;
2 2  
  3 +import cn.hutool.core.collection.CollUtil;
3 4 import cn.iocoder.yudao.framework.common.pojo.CommonResult;
4 5 import cn.iocoder.yudao.framework.common.pojo.PageResult;
  6 +import cn.iocoder.yudao.framework.common.util.collection.MapUtils;
5 7 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
6   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentPageReqVO;
7   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentRespVO;
8   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateReqVO;
9   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateStatusReqVO;
10   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentProcessRespVO;
  8 +import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.*;
  9 +import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
11 10 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
  11 +import cn.iocoder.yudao.module.ai.service.knowledge.AiKnowledgeDocumentService;
12 12 import cn.iocoder.yudao.module.ai.service.knowledge.AiKnowledgeSegmentService;
  13 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchReqBO;
  14 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchRespBO;
13 15 import io.swagger.v3.oas.annotations.Operation;
14 16 import io.swagger.v3.oas.annotations.Parameter;
15 17 import io.swagger.v3.oas.annotations.Parameters;
... ... @@ -20,9 +22,12 @@ import org.hibernate.validator.constraints.URL;
20 22 import org.springframework.validation.annotation.Validated;
21 23 import org.springframework.web.bind.annotation.*;
22 24  
  25 +import java.util.Collections;
23 26 import java.util.List;
  27 +import java.util.Map;
24 28  
25 29 import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
  30 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertSet;
26 31  
27 32 // TODO @芋艿:增加权限标识
28 33 @Tag(name = "管理后台 - AI 知识库段落")
... ... @@ -34,6 +39,9 @@ public class AiKnowledgeSegmentController {
34 39 @Resource
35 40 private AiKnowledgeSegmentService segmentService;
36 41  
  42 + @Resource
  43 + private AiKnowledgeDocumentService documentService;
  44 +
37 45 @GetMapping("/page")
38 46 @Operation(summary = "获取段落分页")
39 47 public CommonResult<PageResult<AiKnowledgeSegmentRespVO>> getKnowledgeSegmentPage(
... ... @@ -79,4 +87,23 @@ public class AiKnowledgeSegmentController {
79 87 return success(list);
80 88 }
81 89  
  90 + @GetMapping("/search")
  91 + @Operation(summary = "搜索段落内容")
  92 + public CommonResult<List<AiKnowledgeSegmentSearchRespVO>> searchKnowledgeSegment(
  93 + @Valid AiKnowledgeSegmentSearchReqVO reqVO) {
  94 + // 1. 搜索段落
  95 + List<AiKnowledgeSegmentSearchRespBO> segments = segmentService
  96 + .searchKnowledgeSegment(BeanUtils.toBean(reqVO, AiKnowledgeSegmentSearchReqBO.class));
  97 + if (CollUtil.isEmpty(segments)) {
  98 + return success(Collections.emptyList());
  99 + }
  100 +
  101 + // 2. 拼接 VO
  102 + Map<Long, AiKnowledgeDocumentDO> documentMap = documentService.getKnowledgeDocumentMap(convertSet(
  103 + segments, AiKnowledgeSegmentSearchRespBO::getDocumentId));
  104 + return success(BeanUtils.toBean(segments, AiKnowledgeSegmentSearchRespVO.class,
  105 + segment -> MapUtils.findAndThen(documentMap, segment.getDocumentId(),
  106 + document -> segment.setDocumentName(document.getName()))));
  107 + }
  108 +
82 109 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/vo/document/AiKnowledgeDocumentUpdateStatusReqVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document;
  2 +
  3 +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
  4 +import cn.iocoder.yudao.framework.common.validation.InEnum;
  5 +import io.swagger.v3.oas.annotations.media.Schema;
  6 +import jakarta.validation.constraints.NotNull;
  7 +import lombok.Data;
  8 +
  9 +@Schema(description = "管理后台 - AI 知识库文档更新状态 Request VO")
  10 +@Data
  11 +public class AiKnowledgeDocumentUpdateStatusReqVO {
  12 +
  13 + @Schema(description = "编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "15583")
  14 + @NotNull(message = "编号不能为空")
  15 + private Long id;
  16 +
  17 + @Schema(description = "状态", requiredMode = Schema.RequiredMode.REQUIRED, example = "0")
  18 + @NotNull(message = "状态不能为空")
  19 + @InEnum(CommonStatusEnum.class)
  20 + private Integer status;
  21 +
  22 +}
0 23 \ No newline at end of file
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/vo/segment/AiKnowledgeSegmentSearchReqVO.java
... ... @@ -3,15 +3,25 @@ package cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment;
3 3 import io.swagger.v3.oas.annotations.media.Schema;
4 4 import lombok.Data;
5 5  
  6 +import jakarta.validation.constraints.NotEmpty;
  7 +import jakarta.validation.constraints.NotNull;
6 8  
7   -@Schema(description = "管理后台 - AI 知识库段落召回 Request VO")
  9 +@Schema(description = "管理后台 - AI 知识库段落搜索 Request VO")
8 10 @Data
9 11 public class AiKnowledgeSegmentSearchReqVO {
10 12  
11   - @Schema(description = "知识库编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "24790")
  13 + @Schema(description = "知识库编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024")
  14 + @NotNull(message = "知识库编号不能为空")
12 15 private Long knowledgeId;
13 16  
14   - @Schema(description = "内容", requiredMode = Schema.RequiredMode.REQUIRED, example = "Java 学习路线")
  17 + @Schema(description = "内容", requiredMode = Schema.RequiredMode.REQUIRED, example = "如何使用这个产品")
  18 + @NotEmpty(message = "内容不能为空")
15 19 private String content;
16 20  
  21 + @Schema(description = "最大返回数量", example = "5")
  22 + private Integer topK;
  23 +
  24 + @Schema(description = "相似度阈值", example = "0.7")
  25 + private Double similarityThreshold;
  26 +
17 27 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/vo/segment/AiKnowledgeSegmentSearchRespVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment;
  2 +
  3 +import io.swagger.v3.oas.annotations.media.Schema;
  4 +import lombok.Data;
  5 +
  6 +@Schema(description = "管理后台 - AI 知识库段落搜索 Response VO")
  7 +@Data
  8 +public class AiKnowledgeSegmentSearchRespVO extends AiKnowledgeSegmentRespVO {
  9 +
  10 + @Schema(description = "文档名称", requiredMode = Schema.RequiredMode.REQUIRED, example = "产品使用手册")
  11 + private String documentName;
  12 +
  13 + @Schema(description = "相似度分数", requiredMode = Schema.RequiredMode.REQUIRED, example = "0.95")
  14 + private Double score;
  15 +
  16 +}
0 17 \ No newline at end of file
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/chat/AiChatConversationDO.java
... ... @@ -65,6 +65,7 @@ public class AiChatConversationDO extends BaseDO {
65 65 */
66 66 private Long roleId;
67 67  
  68 + // TODO @芋艿:可优化,绑定多个知识库。前提,spring ai 支持 RerankModel 的封装
68 69 /**
69 70 * 知识库编号
70 71 * <p>
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeDocumentMapper.java
... ... @@ -5,8 +5,11 @@ import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
5 5 import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
6 6 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentPageReqVO;
7 7 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
  8 +import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
8 9 import org.apache.ibatis.annotations.Mapper;
9 10  
  11 +import java.util.Collection;
  12 +
10 13 /**
11 14 * AI 知识库文档 Mapper
12 15 *
... ... @@ -22,4 +25,10 @@ public interface AiKnowledgeDocumentMapper extends BaseMapperX&lt;AiKnowledgeDocume
22 25 .orderByDesc(AiKnowledgeDocumentDO::getId));
23 26 }
24 27  
  28 + default void updateRetrievalCountIncr(Collection<Long> ids) {
  29 + update( new LambdaUpdateWrapper<AiKnowledgeDocumentDO>()
  30 + .setSql(" retrieval_count = retrieval_count + 1")
  31 + .in(AiKnowledgeDocumentDO::getId, ids));
  32 + }
  33 +
25 34 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java
... ... @@ -7,11 +7,12 @@ import cn.iocoder.yudao.framework.mybatis.core.query.MPJLambdaWrapperX;
7 7 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentPageReqVO;
8 8 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentProcessRespVO;
9 9 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
  10 +import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
10 11 import com.github.yulichang.wrapper.MPJLambdaWrapper;
11 12 import org.apache.ibatis.annotations.Mapper;
12 13  
13   -import java.util.List;
14 14 import java.util.Collection;
  15 +import java.util.List;
15 16  
16 17 /**
17 18 * AI 知识库分片 Mapper
... ... @@ -52,4 +53,10 @@ public interface AiKnowledgeSegmentMapper extends BaseMapperX&lt;AiKnowledgeSegment
52 53 return selectJoinList(AiKnowledgeSegmentProcessRespVO.class, wrapper);
53 54 }
54 55  
  56 + default void updateRetrievalCountIncrByIds(List<Long> ids) {
  57 + update( new LambdaUpdateWrapper<AiKnowledgeSegmentDO>()
  58 + .setSql(" retrieval_count = retrieval_count + 1")
  59 + .in(AiKnowledgeSegmentDO::getId, ids));
  60 + }
  61 +
55 62 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
... ... @@ -12,7 +12,6 @@ import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils;
12 12 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessagePageReqVO;
13 13 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendReqVO;
14 14 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendRespVO;
15   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentSearchReqVO;
16 15 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatConversationDO;
17 16 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatMessageDO;
18 17 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
... ... @@ -133,7 +132,6 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
133 132 Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
134 133  
135 134 // 3.4 流式返回
136   - // TODO 注意:Schedulers.immediate() 目的是,避免默认 Schedulers.parallel() 并发消费 chunk 导致 SSE 响应前端会乱序问题
137 135 StringBuffer contentBuffer = new StringBuffer();
138 136 return streamResponse.map(chunk -> {
139 137 String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getText() : null;
... ... @@ -159,7 +157,8 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
159 157 if (Objects.isNull(knowledgeId)) {
160 158 return Collections.emptyList();
161 159 }
162   - return knowledgeSegmentService.similaritySearch(new AiKnowledgeSegmentSearchReqVO().setKnowledgeId(knowledgeId).setContent(content));
  160 +// return knowledgeSegmentService.similaritySearch(new AiKnowledgeSegmentSearchReqVO().setKnowledgeId(knowledgeId).setContent(content));
  161 + return null;
163 162 }
164 163  
165 164 private Prompt buildPrompt(AiChatConversationDO conversation, List<AiChatMessageDO> messages,List<AiKnowledgeSegmentDO> segmentList,
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentService.java
1 1 package cn.iocoder.yudao.module.ai.service.knowledge;
2 2  
3 3 import cn.iocoder.yudao.framework.common.pojo.PageResult;
  4 +import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentCreateListReqVO;
4 5 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentPageReqVO;
5 6 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentUpdateReqVO;
6 7 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentUpdateStatusReqVO;
7 8 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgeDocumentCreateReqVO;
8   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentCreateListReqVO;
9 9 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
10 10  
  11 +import java.util.Collection;
11 12 import java.util.List;
  13 +import java.util.Map;
  14 +
  15 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertMap;
12 16  
13 17 /**
14 18 * AI 知识库文档 Service 接口
... ... @@ -64,6 +68,13 @@ public interface AiKnowledgeDocumentService {
64 68 void updateKnowledgeDocumentStatus(AiKnowledgeDocumentUpdateStatusReqVO reqVO);
65 69  
66 70 /**
  71 + * 更新文档检索次数(增加 +1)
  72 + *
  73 + * @param ids 文档编号列表
  74 + */
  75 + void updateKnowledgeDocumentRetrievalCountIncr(Collection<Long> ids);
  76 +
  77 + /**
67 78 * 校验文档是否存在
68 79 *
69 80 * @param id 文档编号
... ... @@ -79,4 +90,22 @@ public interface AiKnowledgeDocumentService {
79 90 */
80 91 String readUrl(String url);
81 92  
  93 + /**
  94 + * 获取文档列表
  95 + *
  96 + * @param ids 文档编号列表
  97 + * @return 文档列表
  98 + */
  99 + List<AiKnowledgeDocumentDO> getKnowledgeDocumentList(Collection<Long> ids);
  100 +
  101 + /**
  102 + * 获取文档 Map
  103 + *
  104 + * @param ids 文档编号列表
  105 + * @return 文档 Map
  106 + */
  107 + default Map<Long, AiKnowledgeDocumentDO> getKnowledgeDocumentMap(Collection<Long> ids) {
  108 + return convertMap(getKnowledgeDocumentList(ids), AiKnowledgeDocumentDO::getId);
  109 + }
  110 +
82 111 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
... ... @@ -7,10 +7,10 @@ import cn.hutool.http.HttpUtil;
7 7 import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
8 8 import cn.iocoder.yudao.framework.common.pojo.PageResult;
9 9 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
  10 +import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentCreateListReqVO;
10 11 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentPageReqVO;
11 12 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentUpdateReqVO;
12 13 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentUpdateStatusReqVO;
13   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.document.AiKnowledgeDocumentCreateListReqVO;
14 14 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.knowledge.AiKnowledgeDocumentCreateReqVO;
15 15 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
16 16 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeDocumentMapper;
... ... @@ -25,6 +25,7 @@ import org.springframework.stereotype.Service;
25 25 import org.springframework.transaction.annotation.Transactional;
26 26  
27 27 import java.util.ArrayList;
  28 +import java.util.Collection;
28 29 import java.util.List;
29 30  
30 31 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
... ... @@ -149,6 +150,14 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
149 150 }
150 151  
151 152 @Override
  153 + public void updateKnowledgeDocumentRetrievalCountIncr(Collection<Long> ids) {
  154 + if (CollUtil.isEmpty(ids)) {
  155 + return;
  156 + }
  157 + knowledgeDocumentMapper.updateRetrievalCountIncr(ids);
  158 + }
  159 +
  160 + @Override
152 161 public AiKnowledgeDocumentDO validateKnowledgeDocumentExists(Long id) {
153 162 AiKnowledgeDocumentDO knowledgeDocument = knowledgeDocumentMapper.selectById(id);
154 163 if (knowledgeDocument == null) {
... ... @@ -182,4 +191,12 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
182 191 return document.getText();
183 192 }
184 193  
  194 + @Override
  195 + public List<AiKnowledgeDocumentDO> getKnowledgeDocumentList(Collection<Long> ids) {
  196 + if (CollUtil.isEmpty(ids)) {
  197 + return new ArrayList<>();
  198 + }
  199 + return knowledgeDocumentMapper.selectByIds(ids);
  200 + }
  201 +
185 202 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentService.java
... ... @@ -2,11 +2,12 @@ package cn.iocoder.yudao.module.ai.service.knowledge;
2 2  
3 3 import cn.iocoder.yudao.framework.common.pojo.PageResult;
4 4 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentPageReqVO;
5   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentSearchReqVO;
  5 +import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentProcessRespVO;
6 6 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateReqVO;
7 7 import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentUpdateStatusReqVO;
8   -import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.AiKnowledgeSegmentProcessRespVO;
9 8 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
  9 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchReqBO;
  10 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchRespBO;
10 11 import org.springframework.scheduling.annotation.Async;
11 12  
12 13 import java.util.List;
... ... @@ -67,12 +68,12 @@ public interface AiKnowledgeSegmentService {
67 68 void deleteKnowledgeSegmentByDocumentId(Long documentId);
68 69  
69 70 /**
70   - * 召回段落
  71 + * 搜索知识库段落,并返回结果
71 72 *
72   - * @param reqVO 召回请求信息
73   - * @return 召回的段落
  73 + * @param reqBO 搜索请求信息
  74 + * @return 搜索结果段落列表
74 75 */
75   - List<AiKnowledgeSegmentDO> similaritySearch(AiKnowledgeSegmentSearchReqVO reqVO);
  76 + List<AiKnowledgeSegmentSearchRespBO> searchKnowledgeSegment(AiKnowledgeSegmentSearchReqBO reqBO);
76 77  
77 78 /**
78 79 * 根据 URL 内容,切片创建多个段落
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentServiceImpl.java
... ... @@ -2,6 +2,7 @@ package cn.iocoder.yudao.module.ai.service.knowledge;
2 2  
3 3 import cn.hutool.core.collection.CollUtil;
4 4 import cn.hutool.core.collection.ListUtil;
  5 +import cn.hutool.core.util.ObjUtil;
5 6 import cn.hutool.core.util.StrUtil;
6 7 import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
7 8 import cn.iocoder.yudao.framework.common.pojo.PageResult;
... ... @@ -12,6 +13,8 @@ import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO
12 13 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
13 14 import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeSegmentMapper;
14 15 import cn.iocoder.yudao.module.ai.service.model.AiApiKeyService;
  16 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchReqBO;
  17 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchRespBO;
15 18 import jakarta.annotation.Resource;
16 19 import lombok.extern.slf4j.Slf4j;
17 20 import org.springframework.ai.document.Document;
... ... @@ -171,25 +174,45 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
171 174 }
172 175  
173 176 @Override
174   - public List<AiKnowledgeSegmentDO> similaritySearch(AiKnowledgeSegmentSearchReqVO reqVO) {
  177 + public List<AiKnowledgeSegmentSearchRespBO> searchKnowledgeSegment(AiKnowledgeSegmentSearchReqBO reqBO) {
175 178 // 1. 校验
176   - AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(reqVO.getKnowledgeId());
  179 + AiKnowledgeDO knowledge = knowledgeService.validateKnowledgeExists(reqBO.getKnowledgeId());
177 180  
178   - // 2. 获取向量存储实例
  181 + // 2.1 向量检索
179 182 VectorStore vectorStore = apiKeyService.getOrCreateVectorStoreByModelId(knowledge.getEmbeddingModelId());
180   -
181   - // 3.1 向量检索
182 183 List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder()
183   - .query(reqVO.getContent())
184   - .topK(knowledge.getTopK()).similarityThreshold(knowledge.getSimilarityThreshold())
  184 + .query(reqBO.getContent())
  185 + .topK(ObjUtil.defaultIfNull(reqBO.getTopK(), knowledge.getTopK()))
  186 + .similarityThreshold(ObjUtil.defaultIfNull(reqBO.getSimilarityThreshold(), knowledge.getSimilarityThreshold()))
185 187 .filterExpression(new FilterExpressionBuilder()
186   - .eq(VECTOR_STORE_METADATA_KNOWLEDGE_ID, reqVO.getKnowledgeId()).build())
  188 + .eq(VECTOR_STORE_METADATA_KNOWLEDGE_ID, reqBO.getKnowledgeId()).build())
187 189 .build());
188 190 if (CollUtil.isEmpty(documents)) {
189 191 return ListUtil.empty();
190 192 }
191   - // 3.2 段落召回
192   - return segmentMapper.selectListByVectorIds(convertList(documents, Document::getId));
  193 + // 2.2 段落召回
  194 + List<AiKnowledgeSegmentDO> segments = segmentMapper
  195 + .selectListByVectorIds(convertList(documents, Document::getId));
  196 + if (CollUtil.isEmpty(segments)) {
  197 + return ListUtil.empty();
  198 + }
  199 +
  200 + // 3. 增加召回次数
  201 + segmentMapper.updateRetrievalCountIncrByIds(convertList(segments, AiKnowledgeSegmentDO::getId));
  202 +
  203 + // 4. 构建结果
  204 + List<AiKnowledgeSegmentSearchRespBO> result = convertList(segments, segment -> {
  205 + Document document = CollUtil.findOne(documents, // 找到对应的文档
  206 + doc -> Objects.equals(doc.getId(), segment.getVectorId()));
  207 + if (document == null) {
  208 + return null;
  209 + }
  210 + return BeanUtils.toBean(segment, AiKnowledgeSegmentSearchRespBO.class)
  211 + .setScore(document.getScore());
  212 + });
  213 + result.sort((o1, o2)
  214 + -> Double.compare(o2.getScore(), o1.getScore())); // 按照分数降序排序
  215 + return result;
193 216 }
194 217  
195 218 @Override
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/bo/AiKnowledgeSegmentSearchReqBO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.service.knowledge.bo;
  2 +
  3 +import lombok.Data;
  4 +
  5 +import javax.validation.constraints.NotNull;
  6 +
  7 +import jakarta.validation.constraints.NotEmpty;
  8 +
  9 +/**
  10 + * AI 知识库段落搜索 Request BO
  11 + *
  12 + * @author 芋道源码
  13 + */
  14 +@Data
  15 +public class AiKnowledgeSegmentSearchReqBO {
  16 +
  17 + /**
  18 + * 知识库编号
  19 + */
  20 + @NotNull(message = "知识库编号不能为空")
  21 + private Long knowledgeId;
  22 +
  23 + /**
  24 + * 内容
  25 + */
  26 + @NotEmpty(message = "内容不能为空")
  27 + private String content;
  28 +
  29 + /**
  30 + * 最大返回数量
  31 + */
  32 + private Integer topK;
  33 +
  34 + /**
  35 + * 相似度阈值
  36 + */
  37 + private Double similarityThreshold;
  38 +
  39 +}
0 40 \ No newline at end of file
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/bo/AiKnowledgeSegmentSearchRespBO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai.service.knowledge.bo;
  2 +
  3 +import lombok.Data;
  4 +
  5 +/**
  6 + * AI 知识库段落搜索 Response BO
  7 + *
  8 + * @author 芋道源码
  9 + */
  10 +@Data
  11 +public class AiKnowledgeSegmentSearchRespBO {
  12 +
  13 + /**
  14 + * 段落编号
  15 + */
  16 + private Long id;
  17 + /**
  18 + * 文档编号
  19 + */
  20 + private Long documentId;
  21 + /**
  22 + * 知识库编号
  23 + */
  24 + private Long knowledgeId;
  25 +
  26 + /**
  27 + * 内容
  28 + */
  29 + private String content;
  30 + /**
  31 + * 内容长度
  32 + */
  33 + private Integer contentLength;
  34 +
  35 + /**
  36 + * Token 数量
  37 + */
  38 + private Integer tokens;
  39 +
  40 + /**
  41 + * 相似度分数
  42 + */
  43 + private Double score;
  44 +
  45 +}
0 46 \ No newline at end of file
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiApiKeyServiceImpl.java
... ... @@ -16,6 +16,7 @@ import jakarta.annotation.Resource;
16 16 import org.springframework.ai.chat.model.ChatModel;
17 17 import org.springframework.ai.embedding.EmbeddingModel;
18 18 import org.springframework.ai.image.ImageModel;
  19 +import org.springframework.ai.vectorstore.SimpleVectorStore;
19 20 import org.springframework.ai.vectorstore.VectorStore;
20 21 import org.springframework.context.annotation.Lazy;
21 22 import org.springframework.stereotype.Service;
... ... @@ -154,7 +155,7 @@ public class AiApiKeyServiceImpl implements AiApiKeyService {
154 155 apiKey.getUrl(), chatModel.getModel());
155 156  
156 157 // 创建或获取 VectorStore 对象
157   - return modelFactory.getOrCreateVectorStore(embeddingModel);
  158 + return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel);
158 159 }
159 160  
160 161 }
161 162 \ No newline at end of file
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactory.java
... ... @@ -99,9 +99,10 @@ public interface AiModelFactory {
99 99 * <p>
100 100 * 如果不存在,则进行创建
101 101 *
  102 + * @param type 向量存储类型
102 103 * @param embeddingModel 向量模型
103 104 * @return VectorStore 对象
104 105 */
105   - VectorStore getOrCreateVectorStore(EmbeddingModel embeddingModel);
  106 + VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type, EmbeddingModel embeddingModel);
106 107  
107 108 }
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactoryImpl.java
1 1 package cn.iocoder.yudao.framework.ai.core.factory;
2 2  
  3 +import cn.hutool.core.io.FileUtil;
3 4 import cn.hutool.core.lang.Assert;
4 5 import cn.hutool.core.lang.Singleton;
5 6 import cn.hutool.core.lang.func.Func0;
6 7 import cn.hutool.core.util.ArrayUtil;
  8 +import cn.hutool.core.util.RuntimeUtil;
7 9 import cn.hutool.core.util.StrUtil;
8 10 import cn.hutool.extra.spring.SpringUtil;
9 11 import cn.iocoder.yudao.framework.ai.config.YudaoAiAutoConfiguration;
... ... @@ -24,6 +26,7 @@ import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
24 26 import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions;
25 27 import com.alibaba.cloud.ai.dashscope.image.DashScopeImageModel;
26 28 import com.azure.ai.openai.OpenAIClientBuilder;
  29 +import lombok.SneakyThrows;
27 30 import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration;
28 31 import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiChatProperties;
29 32 import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiConnectionProperties;
... ... @@ -60,7 +63,11 @@ import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
60 63 import org.springframework.ai.zhipuai.api.ZhiPuAiImageApi;
61 64 import org.springframework.web.client.RestClient;
62 65  
  66 +import java.io.File;
  67 +import java.time.Duration;
63 68 import java.util.List;
  69 +import java.util.Timer;
  70 +import java.util.TimerTask;
64 71  
65 72 /**
66 73 * AI Model 模型工厂的实现类
... ... @@ -73,7 +80,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
73 80 public ChatModel getOrCreateChatModel(AiPlatformEnum platform, String apiKey, String url) {
74 81 String cacheKey = buildClientCacheKey(ChatModel.class, platform, apiKey, url);
75 82 return Singleton.get(cacheKey, (Func0<ChatModel>) () -> {
76   - //noinspection EnhancedSwitchMigration
  83 + // noinspection EnhancedSwitchMigration
77 84 switch (platform) {
78 85 case TONG_YI:
79 86 return buildTongYiChatModel(apiKey);
... ... @@ -105,7 +112,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
105 112  
106 113 @Override
107 114 public ChatModel getDefaultChatModel(AiPlatformEnum platform) {
108   - //noinspection EnhancedSwitchMigration
  115 + // noinspection EnhancedSwitchMigration
109 116 switch (platform) {
110 117 case TONG_YI:
111 118 return SpringUtil.getBean(DashScopeChatModel.class);
... ... @@ -136,7 +143,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
136 143  
137 144 @Override
138 145 public ImageModel getDefaultImageModel(AiPlatformEnum platform) {
139   - //noinspection EnhancedSwitchMigration
  146 + // noinspection EnhancedSwitchMigration
140 147 switch (platform) {
141 148 case TONG_YI:
142 149 return SpringUtil.getBean(DashScopeImageModel.class);
... ... @@ -155,7 +162,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
155 162  
156 163 @Override
157 164 public ImageModel getOrCreateImageModel(AiPlatformEnum platform, String apiKey, String url) {
158   - //noinspection EnhancedSwitchMigration
  165 + // noinspection EnhancedSwitchMigration
159 166 switch (platform) {
160 167 case TONG_YI:
161 168 return buildTongYiImagesModel(apiKey);
... ... @@ -174,9 +181,11 @@ public class AiModelFactoryImpl implements AiModelFactory {
174 181  
175 182 @Override
176 183 public MidjourneyApi getOrCreateMidjourneyApi(String apiKey, String url) {
177   - String cacheKey = buildClientCacheKey(MidjourneyApi.class, AiPlatformEnum.MIDJOURNEY.getPlatform(), apiKey, url);
  184 + String cacheKey = buildClientCacheKey(MidjourneyApi.class, AiPlatformEnum.MIDJOURNEY.getPlatform(), apiKey,
  185 + url);
178 186 return Singleton.get(cacheKey, (Func0<MidjourneyApi>) () -> {
179   - YudaoAiProperties.MidjourneyProperties properties = SpringUtil.getBean(YudaoAiProperties.class).getMidjourney();
  187 + YudaoAiProperties.MidjourneyProperties properties = SpringUtil.getBean(YudaoAiProperties.class)
  188 + .getMidjourney();
180 189 return new MidjourneyApi(url, apiKey, properties.getNotifyUrl());
181 190 });
182 191 }
... ... @@ -204,25 +213,31 @@ public class AiModelFactoryImpl implements AiModelFactory {
204 213 }
205 214  
206 215 @Override
207   - public VectorStore getOrCreateVectorStore(EmbeddingModel embeddingModel) {
208   -// String cacheKey = buildClientCacheKey(VectorStore.class, platform, apiKey, url);
209   - String cacheKey = buildClientCacheKey(VectorStore.class, embeddingModel);
  216 + public VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type, EmbeddingModel embeddingModel) {
  217 + // String cacheKey = buildClientCacheKey(VectorStore.class, platform, apiKey,
  218 + // url);
  219 + String cacheKey = buildClientCacheKey(VectorStore.class, embeddingModel, type);
210 220 return Singleton.get(cacheKey, (Func0<VectorStore>) () -> {
  221 + if (type == SimpleVectorStore.class) {
  222 + return buildSimpleVectorStore(embeddingModel);
  223 + }
  224 + throw new IllegalArgumentException(StrUtil.format("未知类型({})", type));
211 225 // TODO @芋艿:先临时使用 store
212   - return SimpleVectorStore.builder(embeddingModel).build();
213 226 // TODO @芋艿:@xin:后续看看,是不是切到阿里云之类的
214   -// String prefix = StrUtil.format("{}#{}:", platform.getPlatform(), apiKey);
215   -// var config = RedisVectorStore.RedisVectorStoreConfig.builder()
216   -// .withIndexName(cacheKey)
217   -// .withPrefix(prefix)
218   -// .withMetadataFields(new RedisVectorStore.MetadataField("knowledgeId", Schema.FieldType.NUMERIC))
219   -// .build();
220   -// RedisProperties redisProperties = SpringUtils.getBean(RedisProperties.class);
221   -// RedisVectorStore redisVectorStore = new RedisVectorStore(config, embeddingModel,
222   -// new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
223   -// true);
224   -// redisVectorStore.afterPropertiesSet();
225   -// return redisVectorStore;
  227 + // String prefix = StrUtil.format("{}#{}:", platform.getPlatform(), apiKey);
  228 + // var config = RedisVectorStore.RedisVectorStoreConfig.builder()
  229 + // .withIndexName(cacheKey)
  230 + // .withPrefix(prefix)
  231 + // .withMetadataFields(new RedisVectorStore.MetadataField("knowledgeId",
  232 + // Schema.FieldType.NUMERIC))
  233 + // .build();
  234 + // RedisProperties redisProperties = SpringUtils.getBean(RedisProperties.class);
  235 + // RedisVectorStore redisVectorStore = new RedisVectorStore(config,
  236 + // embeddingModel,
  237 + // new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
  238 + // true);
  239 + // redisVectorStore.afterPropertiesSet();
  240 + // return redisVectorStore;
226 241 });
227 242 }
228 243  
... ... @@ -307,7 +322,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
307 322 */
308 323 private ChatModel buildSiliconFlowChatModel(String apiKey) {
309 324 YudaoAiProperties.SiliconFlowProperties properties = new YudaoAiProperties.SiliconFlowProperties()
310   - .setApiKey(apiKey);
  325 + .setApiKey(apiKey);
311 326 return new YudaoAiAutoConfiguration().buildSiliconFlowChatClient(properties);
312 327 }
313 328  
... ... @@ -397,7 +412,8 @@ public class AiModelFactoryImpl implements AiModelFactory {
397 412 */
398 413 private DashScopeEmbeddingModel buildTongYiEmbeddingModel(String apiKey, String model) {
399 414 DashScopeApi dashScopeApi = new DashScopeApi(apiKey);
400   - DashScopeEmbeddingOptions dashScopeEmbeddingOptions = DashScopeEmbeddingOptions.builder().withModel(model).build();
  415 + DashScopeEmbeddingOptions dashScopeEmbeddingOptions = DashScopeEmbeddingOptions.builder().withModel(model)
  416 + .build();
401 417 return new DashScopeEmbeddingModel(dashScopeApi, MetadataMode.EMBED, dashScopeEmbeddingOptions);
402 418 }
403 419  
... ... @@ -407,4 +423,58 @@ public class AiModelFactoryImpl implements AiModelFactory {
407 423 return OllamaEmbeddingModel.builder().ollamaApi(ollamaApi).defaultOptions(ollamaOptions).build();
408 424 }
409 425  
  426 + // ========== 各种创建 VectorStore 的方法 ==========
  427 +
  428 + /**
  429 + * 注意:仅适合本地测试使用,生产建议还是使用 Qdrant、Milvus 等
  430 + */
  431 + @SneakyThrows
  432 + @SuppressWarnings("ResultOfMethodCallIgnored")
  433 + private SimpleVectorStore buildSimpleVectorStore(EmbeddingModel embeddingModel) {
  434 + SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
  435 + // 启动加载
  436 + File file = new File(StrUtil.format("{}/vector_store/simple_{}.json",
  437 + FileUtil.getUserHomePath(), embeddingModel.getClass().getSimpleName()));
  438 + if (!file.exists()) {
  439 + FileUtil.mkParentDirs(file);
  440 + file.createNewFile();
  441 + } else if (file.length() > 0) {
  442 + vectorStore.load(file);
  443 + }
  444 + // 定时持久化,每分钟一次
  445 + Timer timer = new Timer("SimpleVectorStoreTimer-" + file.getAbsolutePath());
  446 + timer.scheduleAtFixedRate(new TimerTask() {
  447 +
  448 + @Override
  449 + public void run() {
  450 + vectorStore.save(file);
  451 + }
  452 +
  453 + }, Duration.ofMinutes(1).toMillis(), Duration.ofMinutes(1).toMillis());
  454 + // 关闭时,进行持久化
  455 + RuntimeUtil.addShutdownHook(() -> vectorStore.save(file));
  456 + return vectorStore;
  457 + }
  458 +
  459 + /**
  460 + * 创建向量存储文件
  461 + *
  462 + * @param embeddingModel 嵌入模型
  463 + * @return 向量存储文件
  464 + */
  465 + private File createVectorStoreFile(EmbeddingModel embeddingModel) {
  466 + // 获取简单类名
  467 + String simpleClassName = embeddingModel.getClass().getSimpleName();
  468 + // 获取用户主目录
  469 + String userHome = FileUtil.getUserHomePath();
  470 + // 创建vector_store目录
  471 + File vectorStoreDir = new File(userHome, "vector_store");
  472 + if (!vectorStoreDir.exists()) {
  473 + vectorStoreDir.mkdirs();
  474 + }
  475 +
  476 + // 创建文件
  477 + return new File(vectorStoreDir, "simple_" + simpleClassName + ".json");
  478 + }
  479 +
410 480 }
... ...