【新增】AI 知识库: AiVectorFactory 负责管理不同 EmbeddingModel 对应的 VectorStore

2025-10-31 10:18:42 +08:00 · 2024-08-29 14:13:37 +08:00
parent 024109dac9
commit f97fb0a8fe
13 changed files with 239 additions and 52 deletions
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/config/YudaoAiAutoConfiguration.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/config/YudaoAiAutoConfiguration.java
@@ -2,6 +2,8 @@ package cn.iocoder.yudao.framework.ai.config;

 import cn.iocoder.yudao.framework.ai.core.factory.AiModelFactory;
 import cn.iocoder.yudao.framework.ai.core.factory.AiModelFactoryImpl;
+import cn.iocoder.yudao.framework.ai.core.factory.AiVectorFactory;
+import cn.iocoder.yudao.framework.ai.core.factory.AiVectorFactoryImpl;
 import cn.iocoder.yudao.framework.ai.core.model.deepseek.DeepSeekChatModel;
 import cn.iocoder.yudao.framework.ai.core.model.deepseek.DeepSeekChatOptions;
 import cn.iocoder.yudao.framework.ai.core.model.midjourney.api.MidjourneyApi;
@@ -10,22 +12,15 @@ import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatModel;
 import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatOptions;
 import com.alibaba.cloud.ai.tongyi.TongYiAutoConfiguration;
 import lombok.extern.slf4j.Slf4j;
-import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties;
-import org.springframework.ai.document.MetadataMode;
-import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
 import org.springframework.ai.tokenizer.TokenCountEstimator;
 import org.springframework.ai.transformer.splitter.TokenTextSplitter;
-import org.springframework.ai.transformers.TransformersEmbeddingModel;
-import org.springframework.ai.vectorstore.RedisVectorStore;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
-import org.springframework.boot.autoconfigure.data.redis.RedisProperties;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Import;
 import org.springframework.context.annotation.Lazy;
-import redis.clients.jedis.JedisPooled;

 /**
 * 芋道 AI 自动配置
@@ -43,6 +38,12 @@ public class YudaoAiAutoConfiguration {
        return new AiModelFactoryImpl();
    }

+    @Bean
+    public AiVectorFactory aiVectorFactory() {
+        return new AiVectorFactoryImpl();
+    }
+
+
    // ========== 各种 AI Client 创建 ==========

    @Bean
@@ -85,30 +86,31 @@ public class YudaoAiAutoConfiguration {
    }

    // ========== rag 相关 ==========
-    @Bean
-    @Lazy // TODO 芋艿：临时注释，避免无法启动
-    public EmbeddingModel transformersEmbeddingClient() {
-        return new TransformersEmbeddingModel(MetadataMode.EMBED);
-    }
+    // TODO @xin 免费版本
+//    @Bean
+//    @Lazy // TODO 芋艿：临时注释，避免无法启动」
+//    public EmbeddingModel transformersEmbeddingClient() {
+//        return new TransformersEmbeddingModel(MetadataMode.EMBED);
+//    }

    /**
-     * TODO @xin 抽离出去，根据具体模型走
+     * TODO @xin 默认版本先不弄，目前都先取对应的 EmbeddingModel
     */
-    @Bean
-    @Lazy // TODO 芋艿：临时注释，避免无法启动
-    public RedisVectorStore vectorStore(TransformersEmbeddingModel transformersEmbeddingModel, RedisVectorStoreProperties properties,
-                                        RedisProperties redisProperties) {
-        var config = RedisVectorStore.RedisVectorStoreConfig.builder()
-                .withIndexName(properties.getIndex())
-                .withPrefix(properties.getPrefix())
-                .build();
-
-        RedisVectorStore redisVectorStore = new RedisVectorStore(config, transformersEmbeddingModel,
-                new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
-                properties.isInitializeSchema());
-        redisVectorStore.afterPropertiesSet();
-        return redisVectorStore;
-    }
+//    @Bean
+//    @Lazy // TODO 芋艿：临时注释，避免无法启动
+//    public RedisVectorStore vectorStore(TongYiTextEmbeddingModel tongYiTextEmbeddingModel, RedisVectorStoreProperties properties,
+//                                        RedisProperties redisProperties) {
+//        var config = RedisVectorStore.RedisVectorStoreConfig.builder()
+//                .withIndexName(properties.getIndex())
+//                .withPrefix(properties.getPrefix())
+//                .build();
+//
+//        RedisVectorStore redisVectorStore = new RedisVectorStore(config, tongYiTextEmbeddingModel,
+//                new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
+//                properties.isInitializeSchema());
+//        redisVectorStore.afterPropertiesSet();
+//        return redisVectorStore;
+//    }

    @Bean
    @Lazy // TODO 芋艿：临时注释，避免无法启动
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactory.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactory.java
@@ -4,6 +4,7 @@ import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
 import cn.iocoder.yudao.framework.ai.core.model.midjourney.api.MidjourneyApi;
 import cn.iocoder.yudao.framework.ai.core.model.suno.api.SunoApi;
 import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.image.ImageModel;

 /**
@@ -25,6 +26,18 @@ public interface AiModelFactory {
     */
    ChatModel getOrCreateChatModel(AiPlatformEnum platform, String apiKey, String url);

+    /**
+     * 基于指定配置，获得 EmbeddingModel 对象
+     * <p>
+     * 如果不存在，则进行创建
+     *
+     * @param platform 平台
+     * @param apiKey   API KEY
+     * @param url      API URL
+     * @return ChatModel 对象
+     */
+    EmbeddingModel getOrCreateEmbeddingModel(AiPlatformEnum platform, String apiKey, String url);
+
    /**
     * 基于默认配置，获得 ChatModel 对象
     *
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactoryImpl.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactoryImpl.java
@@ -21,6 +21,7 @@ import com.alibaba.cloud.ai.tongyi.image.TongYiImagesModel;
 import com.alibaba.cloud.ai.tongyi.image.TongYiImagesProperties;
 import com.alibaba.dashscope.aigc.generation.Generation;
 import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesis;
+import com.alibaba.dashscope.embeddings.TextEmbedding;
 import com.azure.ai.openai.OpenAIClient;
 import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration;
 import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiChatProperties;
@@ -37,6 +38,7 @@ import org.springframework.ai.autoconfigure.zhipuai.ZhiPuAiConnectionProperties;
 import org.springframework.ai.autoconfigure.zhipuai.ZhiPuAiImageProperties;
 import org.springframework.ai.azure.openai.AzureOpenAiChatModel;
 import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.image.ImageModel;
 import org.springframework.ai.model.function.FunctionCallbackContext;
 import org.springframework.ai.ollama.OllamaChatModel;
@@ -97,6 +99,21 @@ public class AiModelFactoryImpl implements AiModelFactory {
        });
    }

+    @Override
+    public EmbeddingModel getOrCreateEmbeddingModel(AiPlatformEnum platform, String apiKey, String url) {
+        String cacheKey = buildClientCacheKey(EmbeddingModel.class, platform, apiKey, url);
+        return Singleton.get(cacheKey, (Func0<EmbeddingModel>) () -> {
+            // TODO @xin 先测试一个
+            switch (platform) {
+                case TONG_YI:
+                    return buildTongYiEmbeddingModel(apiKey);
+                default:
+                    throw new IllegalArgumentException(StrUtil.format("未知平台({})", platform));
+            }
+        });
+    }
+
+
    @Override
    public ChatModel getDefaultChatModel(AiPlatformEnum platform) {
        //noinspection EnhancedSwitchMigration
@@ -239,7 +256,7 @@ public class AiModelFactoryImpl implements AiModelFactory {

    /**
     * 可参考 {@link ZhiPuAiAutoConfiguration#zhiPuAiChatModel(
-     * ZhiPuAiConnectionProperties, ZhiPuAiChatProperties, RestClient.Builder, List, FunctionCallbackContext, RetryTemplate, ResponseErrorHandler)}
+     *ZhiPuAiConnectionProperties, ZhiPuAiChatProperties, RestClient.Builder, List, FunctionCallbackContext, RetryTemplate, ResponseErrorHandler)}
     */
    private ZhiPuAiChatModel buildZhiPuChatModel(String apiKey, String url) {
        url = StrUtil.blankToDefault(url, ZhiPuAiConnectionProperties.DEFAULT_BASE_URL);
@@ -249,7 +266,7 @@ public class AiModelFactoryImpl implements AiModelFactory {

    /**
     * 可参考 {@link ZhiPuAiAutoConfiguration#zhiPuAiImageModel(
-     * ZhiPuAiConnectionProperties, ZhiPuAiImageProperties, RestClient.Builder, RetryTemplate, ResponseErrorHandler)}
+     *ZhiPuAiConnectionProperties, ZhiPuAiImageProperties, RestClient.Builder, RetryTemplate, ResponseErrorHandler)}
     */
    private ZhiPuAiImageModel buildZhiPuAiImageModel(String apiKey, String url) {
        url = StrUtil.blankToDefault(url, ZhiPuAiConnectionProperties.DEFAULT_BASE_URL);
@@ -315,4 +332,15 @@ public class AiModelFactoryImpl implements AiModelFactory {
        return new StabilityAiImageModel(stabilityAiApi);
    }

+    // ========== 各种创建 EmbeddingModel 的方法 ==========
+
+    /**
+     * 可参考 {@link TongYiAutoConfiguration#tongYiTextEmbeddingClient(TextEmbedding, TongYiConnectionProperties)}
+     */
+    private EmbeddingModel buildTongYiEmbeddingModel(String apiKey) {
+        TongYiConnectionProperties connectionProperties = new TongYiConnectionProperties();
+        connectionProperties.setApiKey(apiKey);
+        return new TongYiAutoConfiguration().tongYiTextEmbeddingClient(SpringUtil.getBean(TextEmbedding.class), connectionProperties);
+    }
+
 }
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiVectorFactory.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiVectorFactory.java
@@ -0,0 +1,27 @@
+package cn.iocoder.yudao.framework.ai.core.factory;
+
+import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.vectorstore.VectorStore;
+
+/**
+ * AI Vector 模型工厂的接口类
+ * @author xiaoxin
+ */
+public interface AiVectorFactory {
+
+
+    /**
+     * 基于指定配置，获得 VectorStore 对象
+     * <p>
+     * 如果不存在，则进行创建
+     *
+     * @param embeddingModel 嵌入模型
+     * @param platform       平台
+     * @param apiKey         API KEY
+     * @param url            API URL
+     * @return VectorStore 对象
+     */
+    VectorStore getOrCreateVectorStore(EmbeddingModel embeddingModel, AiPlatformEnum platform, String apiKey, String url);
+
+}
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiVectorFactoryImpl.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiVectorFactoryImpl.java
@@ -0,0 +1,51 @@
+package cn.iocoder.yudao.framework.ai.core.factory;
+
+import cn.hutool.core.lang.Singleton;
+import cn.hutool.core.lang.func.Func0;
+import cn.hutool.core.util.ArrayUtil;
+import cn.hutool.core.util.StrUtil;
+import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
+import cn.iocoder.yudao.framework.common.util.spring.SpringUtils;
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.vectorstore.RedisVectorStore;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.boot.autoconfigure.data.redis.RedisProperties;
+import redis.clients.jedis.JedisPooled;
+
+/**
+ * AI Vector 模型工厂的实现类
+ * 使用 redisVectorStore 实现 VectorStore
+ *
+ * @author xiaoxin
+ */
+public class AiVectorFactoryImpl implements AiVectorFactory {
+
+    @Override
+    public VectorStore getOrCreateVectorStore(EmbeddingModel embeddingModel, AiPlatformEnum platform, String apiKey, String url) {
+        String cacheKey = buildClientCacheKey(VectorStore.class, platform, apiKey, url);
+        return Singleton.get(cacheKey, (Func0<VectorStore>) () -> {
+            // TODO 芋艿 @xin 这两个配置取哪好呢
+            // TODO 不同模型的向量维度可能会不一样，目前看貌似是以 index 来做区分的，维度不一样存不到一个 index 上
+            String index = "default-index";
+            String prefix = "default:";
+            var config = RedisVectorStore.RedisVectorStoreConfig.builder()
+                    .withIndexName(index)
+                    .withPrefix(prefix)
+                    .build();
+            RedisProperties redisProperties = SpringUtils.getBean(RedisProperties.class);
+            RedisVectorStore redisVectorStore = new RedisVectorStore(config, embeddingModel,
+                    new JedisPooled(redisProperties.getHost(), redisProperties.getPort()),
+                    true);
+            redisVectorStore.afterPropertiesSet();
+            return redisVectorStore;
+        });
+    }
+
+
+    private static String buildClientCacheKey(Class<?> clazz, Object... params) {
+        if (ArrayUtil.isEmpty(params)) {
+            return clazz.getName();
+        }
+        return StrUtil.format("{}#{}", clazz.getName(), ArrayUtil.join(params, "_"));
+    }
+}
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreAutoConfiguration.java
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreAutoConfiguration.java
@@ -19,6 +19,7 @@ import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.vectorstore.RedisVectorStore;
 import org.springframework.ai.vectorstore.RedisVectorStore.RedisVectorStoreConfig;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.data.redis.RedisAutoConfiguration;
@@ -38,7 +39,7 @@ import redis.clients.jedis.JedisPooled;
 */
@AutoConfiguration(after = RedisAutoConfiguration.class)
@ConditionalOnClass({JedisPooled.class, JedisConnectionFactory.class, RedisVectorStore.class, EmbeddingModel.class})
-//@ConditionalOnBean(JedisConnectionFactory.class)
+@ConditionalOnBean(JedisConnectionFactory.class)
@EnableConfigurationProperties(RedisVectorStoreProperties.class)
 public class RedisVectorStoreAutoConfiguration {