Commit 5f5e77a3923bc39b3ef160195e00df5f3baee1aa
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
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
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<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<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 | } | ... | ... |