Commit 588c9fe323e2f48f1c3a6fa7159755279049cc09

Authored by YunaiV
1 parent 44bcc947

【功能新增】AI:增加 RedisVectorStore 向量库的接入

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentServiceImpl.java
... ... @@ -32,6 +32,7 @@ import org.springframework.stereotype.Service;
32 32  
33 33 import java.util.Collections;
34 34 import java.util.List;
  35 +import java.util.Map;
35 36 import java.util.Objects;
36 37  
37 38 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
... ... @@ -48,9 +49,14 @@ import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.KNOWLEDGE_SEGM
48 49 @Slf4j
49 50 public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService {
50 51  
51   - public static final String VECTOR_STORE_METADATA_KNOWLEDGE_ID = "knowledgeId";
52   - public static final String VECTOR_STORE_METADATA_DOCUMENT_ID = "documentId";
53   - public static final String VECTOR_STORE_METADATA_SEGMENT_ID = "segmentId";
  52 + private static final String VECTOR_STORE_METADATA_KNOWLEDGE_ID = "knowledgeId";
  53 + private static final String VECTOR_STORE_METADATA_DOCUMENT_ID = "documentId";
  54 + private static final String VECTOR_STORE_METADATA_SEGMENT_ID = "segmentId";
  55 +
  56 + private static final Map<String, Class<?>> VECTOR_STORE_METADATA_TYPES = Map.of(
  57 + VECTOR_STORE_METADATA_KNOWLEDGE_ID, String.class,
  58 + VECTOR_STORE_METADATA_DOCUMENT_ID, String.class,
  59 + VECTOR_STORE_METADATA_SEGMENT_ID, String.class);
54 60  
55 61 @Resource
56 62 private AiKnowledgeSegmentMapper segmentMapper;
... ... @@ -257,7 +263,7 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
257 263 }
258 264  
259 265 private VectorStore getVectorStoreById(AiKnowledgeDO knowledge) {
260   - return modelService.getOrCreateVectorStore(knowledge.getEmbeddingModelId());
  266 + return modelService.getOrCreateVectorStore(knowledge.getEmbeddingModelId(), VECTOR_STORE_METADATA_TYPES);
261 267 }
262 268  
263 269 private VectorStore getVectorStoreById(Long knowledgeId) {
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiModelService.java
... ... @@ -13,6 +13,7 @@ import org.springframework.ai.vectorstore.VectorStore;
13 13  
14 14 import javax.annotation.Nullable;
15 15 import java.util.List;
  16 +import java.util.Map;
16 17  
17 18 /**
18 19 * AI 模型 Service 接口
... ... @@ -125,8 +126,9 @@ public interface AiModelService {
125 126 * 获得 VectorStore 对象
126 127 *
127 128 * @param id 编号
  129 + * @param metadataFields 元数据的定义
128 130 * @return VectorStore 对象
129 131 */
130   - VectorStore getOrCreateVectorStore(Long id);
  132 + VectorStore getOrCreateVectorStore(Long id, Map<String, Class<?>> metadataFields);
131 133  
132 134 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiModelServiceImpl.java
... ... @@ -22,6 +22,7 @@ import org.springframework.stereotype.Service;
22 22 import org.springframework.validation.annotation.Validated;
23 23  
24 24 import java.util.List;
  25 +import java.util.Map;
25 26  
26 27 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
27 28 import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.*;
... ... @@ -151,7 +152,7 @@ public class AiModelServiceImpl implements AiModelService {
151 152 }
152 153  
153 154 @Override
154   - public VectorStore getOrCreateVectorStore(Long id) {
  155 + public VectorStore getOrCreateVectorStore(Long id, Map<String, Class<?>> metadataFields) {
155 156 // 获取模型 + 密钥
156 157 AiModelDO model = validateModel(id);
157 158 AiApiKeyDO apiKey = apiKeyService.validateApiKey(model.getKeyId());
... ... @@ -162,8 +163,9 @@ public class AiModelServiceImpl implements AiModelService {
162 163 platform, apiKey.getApiKey(), apiKey.getUrl(), model.getModel());
163 164  
164 165 // 创建或获取 VectorStore 对象
165   -// return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel);
166   - return modelFactory.getOrCreateVectorStore(QdrantVectorStore.class, embeddingModel);
  166 +// return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel, metadataFields);
  167 + return modelFactory.getOrCreateVectorStore(QdrantVectorStore.class, embeddingModel, metadataFields);
  168 +// return modelFactory.getOrCreateVectorStore(RedisVectorStore.class, embeddingModel, metadataFields);
167 169 }
168 170  
169 171 }
170 172 \ No newline at end of file
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/config/YudaoAiAutoConfiguration.java
... ... @@ -12,12 +12,12 @@ import cn.iocoder.yudao.framework.ai.core.model.suno.api.SunoApi;
12 12 import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatModel;
13 13 import lombok.extern.slf4j.Slf4j;
14 14 import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreProperties;
  15 +import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties;
15 16 import org.springframework.ai.openai.OpenAiChatModel;
16 17 import org.springframework.ai.openai.OpenAiChatOptions;
17 18 import org.springframework.ai.openai.api.OpenAiApi;
18 19 import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
19 20 import org.springframework.ai.tokenizer.TokenCountEstimator;
20   -import org.springframework.ai.transformer.splitter.TokenTextSplitter;
21 21 import org.springframework.boot.autoconfigure.AutoConfiguration;
22 22 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
23 23 import org.springframework.boot.context.properties.EnableConfigurationProperties;
... ... @@ -31,7 +31,8 @@ import org.springframework.context.annotation.Lazy;
31 31 */
32 32 @AutoConfiguration
33 33 @EnableConfigurationProperties({YudaoAiProperties.class,
34   - QdrantVectorStoreProperties.class // 解析 Qdrant 配置
  34 + QdrantVectorStoreProperties.class, // 解析 Qdrant 配置
  35 + RedisVectorStoreProperties.class, // 解析 Redis 配置
35 36 })
36 37 @Slf4j
37 38 public class YudaoAiAutoConfiguration {
... ... @@ -200,32 +201,6 @@ public class YudaoAiAutoConfiguration {
200 201 // return new TransformersEmbeddingModel(MetadataMode.EMBED);
201 202 // }
202 203  
203   - /**
204   - * TODO @xin 默认版本先不弄,目前都先取对应的 EmbeddingModel
205   - */
206   -// @Bean
207   -// @Lazy // TODO 芋艿:临时注释,避免无法启动
208   -// public RedisVectorStore vectorStore(TransformersEmbeddingModel embeddingModel, RedisVectorStoreProperties properties,
209   -// RedisProperties redisProperties) {
210   -// var config = RedisVectorStore.RedisVectorStoreConfig.builder()
211   -// .withIndexName(properties.getIndex())
212   -// .withPrefix(properties.getPrefix())
213   -// .withMetadataFields(new RedisVectorStore.MetadataField("knowledgeId", Schema.FieldType.NUMERIC))
214   -// .build();
215   -//
216   -// RedisVectorStore redisVectorStore = new RedisVectorStore(config, embeddingModel,
217   -// new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
218   -// properties.isInitializeSchema());
219   -// redisVectorStore.afterPropertiesSet();
220   -// return redisVectorStore;
221   -// }
222   - @Bean
223   - @Lazy // TODO 芋艿:临时注释,避免无法启动
224   - public TokenTextSplitter tokenTextSplitter() {
225   - //TODO @xin 配置提取
226   - return new TokenTextSplitter(500, 100, 5, 10000, true);
227   - }
228   -
229 204 @Bean
230 205 @Lazy // TODO 芋艿:临时注释,避免无法启动
231 206 public TokenCountEstimator tokenCountEstimator() {
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactory.java
... ... @@ -8,6 +8,8 @@ import org.springframework.ai.embedding.EmbeddingModel;
8 8 import org.springframework.ai.image.ImageModel;
9 9 import org.springframework.ai.vectorstore.VectorStore;
10 10  
  11 +import java.util.Map;
  12 +
11 13 /**
12 14 * AI Model 模型工厂的接口类
13 15 *
... ... @@ -96,13 +98,16 @@ public interface AiModelFactory {
96 98  
97 99 /**
98 100 * 基于指定配置,获得 VectorStore 对象
99   - * <p>
  101 + *
100 102 * 如果不存在,则进行创建
101 103 *
102 104 * @param type 向量存储类型
103 105 * @param embeddingModel 向量模型
  106 + * @param metadataFields 元数据字段
104 107 * @return VectorStore 对象
105 108 */
106   - VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type, EmbeddingModel embeddingModel);
  109 + VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type,
  110 + EmbeddingModel embeddingModel,
  111 + Map<String, Class<?>> metadataFields);
107 112  
108 113 }
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactoryImpl.java
... ... @@ -19,6 +19,7 @@ import cn.iocoder.yudao.framework.ai.core.model.midjourney.api.MidjourneyApi;
19 19 import cn.iocoder.yudao.framework.ai.core.model.siliconflow.SiliconFlowChatModel;
20 20 import cn.iocoder.yudao.framework.ai.core.model.suno.api.SunoApi;
21 21 import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatModel;
  22 +import cn.iocoder.yudao.framework.common.util.spring.SpringUtils;
22 23 import com.alibaba.cloud.ai.autoconfigure.dashscope.DashScopeAutoConfiguration;
23 24 import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
24 25 import com.alibaba.cloud.ai.dashscope.api.DashScopeImageApi;
... ... @@ -39,12 +40,13 @@ import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
39 40 import org.springframework.ai.autoconfigure.qianfan.QianFanAutoConfiguration;
40 41 import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreAutoConfiguration;
41 42 import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreProperties;
  43 +import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreAutoConfiguration;
  44 +import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties;
42 45 import org.springframework.ai.autoconfigure.zhipuai.ZhiPuAiAutoConfiguration;
43 46 import org.springframework.ai.autoconfigure.zhipuai.ZhiPuAiConnectionProperties;
44 47 import org.springframework.ai.azure.openai.AzureOpenAiChatModel;
45 48 import org.springframework.ai.chat.model.ChatModel;
46 49 import org.springframework.ai.document.MetadataMode;
47   -import org.springframework.ai.embedding.BatchingStrategy;
48 50 import org.springframework.ai.embedding.EmbeddingModel;
49 51 import org.springframework.ai.image.ImageModel;
50 52 import org.springframework.ai.ollama.OllamaChatModel;
... ... @@ -67,20 +69,26 @@ import org.springframework.ai.vectorstore.VectorStore;
67 69 import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
68 70 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
69 71 import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
  72 +import org.springframework.ai.vectorstore.redis.RedisVectorStore;
70 73 import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
71 74 import org.springframework.ai.zhipuai.ZhiPuAiImageModel;
72 75 import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
73 76 import org.springframework.ai.zhipuai.api.ZhiPuAiImageApi;
74 77 import org.springframework.beans.BeansException;
75 78 import org.springframework.beans.factory.ObjectProvider;
  79 +import org.springframework.boot.autoconfigure.data.redis.RedisProperties;
76 80 import org.springframework.web.client.RestClient;
  81 +import redis.clients.jedis.JedisPooled;
77 82  
78 83 import java.io.File;
79 84 import java.time.Duration;
80 85 import java.util.List;
  86 +import java.util.Map;
81 87 import java.util.Timer;
82 88 import java.util.TimerTask;
83 89  
  90 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList;
  91 +
84 92 /**
85 93 * AI Model 模型工厂的实现类
86 94 *
... ... @@ -225,7 +233,9 @@ public class AiModelFactoryImpl implements AiModelFactory {
225 233 }
226 234  
227 235 @Override
228   - public VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type, EmbeddingModel embeddingModel) {
  236 + public VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type,
  237 + EmbeddingModel embeddingModel,
  238 + Map<String, Class<?>> metadataFields) {
229 239 String cacheKey = buildClientCacheKey(VectorStore.class, embeddingModel, type);
230 240 return Singleton.get(cacheKey, (Func0<VectorStore>) () -> {
231 241 if (type == SimpleVectorStore.class) {
... ... @@ -234,23 +244,10 @@ public class AiModelFactoryImpl implements AiModelFactory {
234 244 if (type == QdrantVectorStore.class) {
235 245 return buildQdrantVectorStore(embeddingModel);
236 246 }
  247 + if (type == RedisVectorStore.class) {
  248 + return buildRedisVectorStore(embeddingModel, metadataFields);
  249 + }
237 250 throw new IllegalArgumentException(StrUtil.format("未知类型({})", type));
238   - // TODO @芋艿:先临时使用 store
239   - // TODO @芋艿:@xin:后续看看,是不是切到阿里云之类的
240   - // String prefix = StrUtil.format("{}#{}:", platform.getPlatform(), apiKey);
241   - // var config = RedisVectorStore.RedisVectorStoreConfig.builder()
242   - // .withIndexName(cacheKey)
243   - // .withPrefix(prefix)
244   - // .withMetadataFields(new RedisVectorStore.MetadataField("knowledgeId",
245   - // Schema.FieldType.NUMERIC))
246   - // .build();
247   - // RedisProperties redisProperties = SpringUtils.getBean(RedisProperties.class);
248   - // RedisVectorStore redisVectorStore = new RedisVectorStore(config,
249   - // embeddingModel,
250   - // new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
251   - // true);
252   - // redisVectorStore.afterPropertiesSet();
253   - // return redisVectorStore;
254 251 });
255 252 }
256 253  
... ... @@ -469,21 +466,65 @@ public class AiModelFactoryImpl implements AiModelFactory {
469 466 return vectorStore;
470 467 }
471 468  
  469 + /**
  470 + * 参考 {@link QdrantVectorStoreAutoConfiguration} 的 vectorStore 方法
  471 + */
  472 + @SneakyThrows
472 473 private QdrantVectorStore buildQdrantVectorStore(EmbeddingModel embeddingModel) {
473 474 QdrantVectorStoreAutoConfiguration configuration = new QdrantVectorStoreAutoConfiguration();
474   - QdrantVectorStoreProperties vectorStoreProperties = SpringUtil.getBean(QdrantVectorStoreProperties.class);
  475 + QdrantVectorStoreProperties properties = SpringUtil.getBean(QdrantVectorStoreProperties.class);
475 476 // 参考 QdrantVectorStoreAutoConfiguration 实现,创建 QdrantClient 对象
476 477 QdrantGrpcClient.Builder grpcClientBuilder = QdrantGrpcClient.newBuilder(
477   - vectorStoreProperties.getHost(), vectorStoreProperties.getPort(), vectorStoreProperties.isUseTls());
478   - if (StrUtil.isNotEmpty(vectorStoreProperties.getApiKey())) {
479   - grpcClientBuilder.withApiKey(vectorStoreProperties.getApiKey());
  478 + properties.getHost(), properties.getPort(), properties.isUseTls());
  479 + if (StrUtil.isNotEmpty(properties.getApiKey())) {
  480 + grpcClientBuilder.withApiKey(properties.getApiKey());
480 481 }
481 482 QdrantClient qdrantClient = new QdrantClient(grpcClientBuilder.build());
482   - // 参考 QdrantVectorStoreAutoConfiguration 实现,实现 batchingStrategy
483   - BatchingStrategy batchingStrategy = ReflectUtil.invoke(configuration, "batchingStrategy");
484   -
485 483 // 创建 QdrantVectorStore 对象
486   - ObjectProvider<ObservationRegistry> observationRegistry = new ObjectProvider<>() {
  484 + QdrantVectorStore vectorStore = configuration.vectorStore(embeddingModel, properties, qdrantClient,
  485 + getObservationRegistry(), getCustomObservationConvention(),
  486 + ReflectUtil.invoke(configuration, "batchingStrategy"));
  487 + // 初始化索引
  488 + vectorStore.afterPropertiesSet();
  489 + return vectorStore;
  490 + }
  491 +
  492 + /**
  493 + * 参考 {@link RedisVectorStoreAutoConfiguration} 的 vectorStore 方法
  494 + */
  495 + private RedisVectorStore buildRedisVectorStore(EmbeddingModel embeddingModel,
  496 + Map<String, Class<?>> metadataFields) {
  497 + // 创建 JedisPooled 对象
  498 + RedisProperties redisProperties = SpringUtils.getBean(RedisProperties.class);
  499 + JedisPooled jedisPooled = new JedisPooled(redisProperties.getHost(), redisProperties.getPort());
  500 + // 创建 RedisVectorStoreProperties 对象
  501 + RedisVectorStoreAutoConfiguration configuration = new RedisVectorStoreAutoConfiguration();
  502 + RedisVectorStoreProperties properties = SpringUtil.getBean(RedisVectorStoreProperties.class);
  503 + RedisVectorStore redisVectorStore = RedisVectorStore.builder(jedisPooled, embeddingModel)
  504 + .indexName(properties.getIndex()).prefix(properties.getPrefix())
  505 + .initializeSchema(properties.isInitializeSchema())
  506 + .metadataFields(convertList(metadataFields.entrySet(), entry -> {
  507 + String fieldName = entry.getKey();
  508 + Class<?> fieldType = entry.getValue();
  509 + if (Number.class.isAssignableFrom(fieldType)) {
  510 + return RedisVectorStore.MetadataField.numeric(fieldName);
  511 + }
  512 + if (Boolean.class.isAssignableFrom(fieldType)) {
  513 + return RedisVectorStore.MetadataField.tag(fieldName);
  514 + }
  515 + return RedisVectorStore.MetadataField.text(fieldName);
  516 + }))
  517 + .observationRegistry(getObservationRegistry().getObject())
  518 + .customObservationConvention(getCustomObservationConvention().getObject())
  519 + .batchingStrategy(ReflectUtil.invoke(configuration, "batchingStrategy"))
  520 + .build();
  521 + // 初始化索引
  522 + redisVectorStore.afterPropertiesSet();
  523 + return redisVectorStore;
  524 + }
  525 +
  526 + private static ObjectProvider<ObservationRegistry> getObservationRegistry() {
  527 + return new ObjectProvider<>() {
487 528  
488 529 @Override
489 530 public ObservationRegistry getObject() throws BeansException {
... ... @@ -491,16 +532,15 @@ public class AiModelFactoryImpl implements AiModelFactory {
491 532 }
492 533  
493 534 };
494   - ObjectProvider <VectorStoreObservationConvention> customObservationConvention = new ObjectProvider<>() {
  535 + }
495 536  
  537 + private static ObjectProvider<VectorStoreObservationConvention> getCustomObservationConvention() {
  538 + return new ObjectProvider<>() {
496 539 @Override
497 540 public VectorStoreObservationConvention getObject() throws BeansException {
498 541 return new DefaultVectorStoreObservationConvention();
499 542 }
500   -
501 543 };
502   - return configuration.vectorStore(embeddingModel, vectorStoreProperties, qdrantClient,
503   - observationRegistry, customObservationConvention, batchingStrategy);
504 544 }
505 545  
506 546 }
... ...
yudao-server/src/main/resources/application.yaml
... ... @@ -149,10 +149,12 @@ spring:
149 149 ai:
150 150 vectorstore: # 向量存储
151 151 redis:
152   - index: default-index
153   - prefix: "default:"
  152 + initialize-schema: true
  153 + index: knowledge_index # Redis 中向量索引的名称:用于存储和检索向量数据的索引标识符,所有相关的向量搜索操作都会基于这个索引进行
  154 + prefix: "knowledge_segment:" # Redis 中存储向量数据的键名前缀:这个前缀会添加到每个存储在 Redis 中的向量数据键名前,每个 document 都是一个 hash 结构
154 155 qdrant:
155   - collection-name: knowledge_segment
  156 + initialize-schema: true
  157 + collection-name: knowledge_segment # Qdrant 中向量集合的名称:用于存储向量数据的集合标识符,所有相关的向量操作都会在这个集合中进行
156 158 host: 127.0.0.1
157 159 port: 6334
158 160 use-tls: false
... ...