【功能新增】AI：增加 MilvusVectorStore 向量库的接入

YunaiV
1 parent 588c9fe3
Showing 5 changed files with 75 additions and 22 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiModelServiceImpl.java
yudao-module-ai/yudao-spring-boot-starter-ai/pom.xml
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/config/YudaoAiAutoConfiguration.java
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactoryImpl.java
yudao-server/src/main/resources/application.yaml
@@ -17,7 +17,7 @@ import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.image.ImageModel;
 import org.springframework.ai.vectorstore.VectorStore;
-import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
+import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
 import org.springframework.stereotype.Service;
 import org.springframework.validation.annotation.Validated;
@@ -114,8 +114,7 @@ public class AiModelServiceImpl implements AiModelService {
     }
     @Override
-    public List<AiModelDO> getModelListByStatusAndType(Integer status, Integer type,
-                                                       String platform) {
+    public List<AiModelDO> getModelListByStatusAndType(Integer status, Integer type, String platform) {
         return modelMapper.selectListByStatusAndType(status, type, platform);
     }
@@ -163,9 +162,10 @@ public class AiModelServiceImpl implements AiModelService {
                 platform, apiKey.getApiKey(), apiKey.getUrl(), model.getModel());
         // 创建或获取 VectorStore 对象
-//        return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel, metadataFields);
-        return modelFactory.getOrCreateVectorStore(QdrantVectorStore.class, embeddingModel, metadataFields);
-//        return modelFactory.getOrCreateVectorStore(RedisVectorStore.class, embeddingModel, metadataFields);
+        // return modelFactory.getOrCreateVectorStore(SimpleVectorStore.class, embeddingModel, metadataFields);
+        // return modelFactory.getOrCreateVectorStore(QdrantVectorStore.class, embeddingModel, metadataFields);
+        // return modelFactory.getOrCreateVectorStore(RedisVectorStore.class, embeddingModel, metadataFields);
+        return modelFactory.getOrCreateVectorStore(MilvusVectorStore.class, embeddingModel, metadataFields);
     }
 }
 \ No newline at end of file
@@ -70,7 +70,6 @@
             <groupId>${spring-ai.groupId}</groupId>
             <artifactId>spring-ai-qdrant-store</artifactId>
             <version>${spring-ai.version}</version>
-<!--            <optional>true</optional>-->
         </dependency>
         <dependency>
@@ -11,8 +11,12 @@ import cn.iocoder.yudao.framework.ai.core.model.siliconflow.SiliconFlowChatModel
 import cn.iocoder.yudao.framework.ai.core.model.suno.api.SunoApi;
 import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatModel;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusServiceClientProperties;
+import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreProperties;
 import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreProperties;
 import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties;
+import org.springframework.ai.embedding.BatchingStrategy;
+import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.openai.OpenAiChatModel;
 import org.springframework.ai.openai.OpenAiChatOptions;
 import org.springframework.ai.openai.api.OpenAiApi;
@@ -22,7 +26,6 @@ import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
-import org.springframework.context.annotation.Lazy;
 /**
  * 芋道 AI 自动配置
@@ -33,6 +36,7 @@ import org.springframework.context.annotation.Lazy;
 @EnableConfigurationProperties({YudaoAiProperties.class,
         QdrantVectorStoreProperties.class, // 解析 Qdrant 配置
         RedisVectorStoreProperties.class, // 解析 Redis 配置
+        MilvusVectorStoreProperties.class, MilvusServiceClientProperties.class // 解析 Milvus 配置
 })
 @Slf4j
 public class YudaoAiAutoConfiguration {
@@ -193,18 +197,16 @@ public class YudaoAiAutoConfiguration {
         return new SunoApi(yudaoAiProperties.getSuno().getBaseUrl());
     }
-    // ========== rag 相关 ==========
-    // TODO @xin 免费版本
-//    @Bean
-//    @Lazy // TODO 芋艿：临时注释，避免无法启动」
-//    public TransformersEmbeddingModel transformersEmbeddingClient() {
-//        return new TransformersEmbeddingModel(MetadataMode.EMBED);
-//    }
+    // ========== RAG 相关 ==========
     @Bean
-    @Lazy // TODO 芋艿：临时注释，避免无法启动
     public TokenCountEstimator tokenCountEstimator() {
         return new JTokkitTokenCountEstimator();
     }
+    @Bean
+    public BatchingStrategy batchingStrategy() {
+        return new TokenCountBatchingStrategy();
+    }
+
 }
 \ No newline at end of file
@@ -5,7 +5,6 @@ import cn.hutool.core.lang.Assert;
 import cn.hutool.core.lang.Singleton;
 import cn.hutool.core.lang.func.Func0;
 import cn.hutool.core.util.ArrayUtil;
-import cn.hutool.core.util.ReflectUtil;
 import cn.hutool.core.util.RuntimeUtil;
 import cn.hutool.core.util.StrUtil;
 import cn.hutool.extra.spring.SpringUtil;
@@ -29,6 +28,7 @@ import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions;
 import com.alibaba.cloud.ai.dashscope.image.DashScopeImageModel;
 import com.azure.ai.openai.OpenAIClientBuilder;
 import io.micrometer.observation.ObservationRegistry;
+import io.milvus.client.MilvusServiceClient;
 import io.qdrant.client.QdrantClient;
 import io.qdrant.client.QdrantGrpcClient;
 import lombok.SneakyThrows;
@@ -38,6 +38,10 @@ import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiConnectionPr
 import org.springframework.ai.autoconfigure.ollama.OllamaAutoConfiguration;
 import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
 import org.springframework.ai.autoconfigure.qianfan.QianFanAutoConfiguration;
+import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusServiceClientConnectionDetails;
+import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusServiceClientProperties;
+import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreAutoConfiguration;
+import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreProperties;
 import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreAutoConfiguration;
 import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreProperties;
 import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreAutoConfiguration;
@@ -47,6 +51,7 @@ import org.springframework.ai.autoconfigure.zhipuai.ZhiPuAiConnectionProperties;
 import org.springframework.ai.azure.openai.AzureOpenAiChatModel;
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.document.MetadataMode;
+import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.image.ImageModel;
 import org.springframework.ai.ollama.OllamaChatModel;
@@ -66,6 +71,7 @@ import org.springframework.ai.stabilityai.StabilityAiImageModel;
 import org.springframework.ai.stabilityai.api.StabilityAiApi;
 import org.springframework.ai.vectorstore.SimpleVectorStore;
 import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
 import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
 import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
@@ -247,6 +253,9 @@ public class AiModelFactoryImpl implements AiModelFactory {
             if (type == RedisVectorStore.class) {
                 return buildRedisVectorStore(embeddingModel, metadataFields);
             }
+            if (type == MilvusVectorStore.class) {
+                return buildMilvusVectorStore(embeddingModel);
+            }
             throw new IllegalArgumentException(StrUtil.format("未知类型({})", type));
         });
     }
@@ -482,8 +491,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
         QdrantClient qdrantClient = new QdrantClient(grpcClientBuilder.build());
         // 创建 QdrantVectorStore 对象
         QdrantVectorStore vectorStore = configuration.vectorStore(embeddingModel, properties, qdrantClient,
-                getObservationRegistry(), getCustomObservationConvention(),
-                ReflectUtil.invoke(configuration, "batchingStrategy"));
+                getObservationRegistry(), getCustomObservationConvention(), getBatchingStrategy());
         // 初始化索引
         vectorStore.afterPropertiesSet();
         return vectorStore;
@@ -516,13 +524,48 @@ public class AiModelFactoryImpl implements AiModelFactory {
                 }))
                 .observationRegistry(getObservationRegistry().getObject())
                 .customObservationConvention(getCustomObservationConvention().getObject())
-                .batchingStrategy(ReflectUtil.invoke(configuration, "batchingStrategy"))
+                .batchingStrategy(getBatchingStrategy())
                 .build();
         // 初始化索引
         redisVectorStore.afterPropertiesSet();
         return redisVectorStore;
     }
+    /**
+     * 参考 {@link MilvusVectorStoreAutoConfiguration} 的 vectorStore 方法
+     */
+    @SneakyThrows
+    private MilvusVectorStore buildMilvusVectorStore(EmbeddingModel embeddingModel) {
+        MilvusVectorStoreAutoConfiguration configuration = new MilvusVectorStoreAutoConfiguration();
+        // 获取配置属性
+        MilvusVectorStoreProperties serverProperties = SpringUtil.getBean(MilvusVectorStoreProperties.class);
+        MilvusServiceClientProperties clientProperties = SpringUtil.getBean(MilvusServiceClientProperties.class);
+
+        // 创建 MilvusServiceClient 对象
+        MilvusServiceClient milvusClient = configuration.milvusClient(serverProperties, clientProperties,
+                new MilvusServiceClientConnectionDetails() {
+
+                    @Override
+                    public String getHost() {
+                        return clientProperties.getHost();
+                    }
+
+                    @Override
+                    public int getPort() {
+                        return clientProperties.getPort();
+                    }
+
+                }
+        );
+        // 创建 MilvusVectorStore 对象
+        MilvusVectorStore vectorStore = configuration.vectorStore(milvusClient, embeddingModel, serverProperties,
+                getBatchingStrategy(), getObservationRegistry(), getCustomObservationConvention());
+
+        // 初始化索引
+        vectorStore.afterPropertiesSet();
+        return vectorStore;
+    }
+
     private static ObjectProvider<ObservationRegistry> getObservationRegistry() {
         return new ObjectProvider<>() {
@@ -543,4 +586,8 @@ public class AiModelFactoryImpl implements AiModelFactory {
         };
     }
+    private static BatchingStrategy getBatchingStrategy() {
+        return SpringUtil.getBean(BatchingStrategy.class);
+    }
+
 }
@@ -157,8 +157,13 @@ spring:
         collection-name: knowledge_segment # Qdrant 中向量集合的名称：用于存储向量数据的集合标识符，所有相关的向量操作都会在这个集合中进行
         host: 127.0.0.1
         port: 6334
-        use-tls: false
-        api-key:
+      milvus:
+        initialize-schema: true
+        database-name: default # Milvus 中数据库的名称
+        collection-name: knowledge_segment # Milvus 中集合的名称：用于存储向量数据的集合标识符，所有相关的向量操作都会在这个集合中进行
+        client:
+          host: 127.0.0.1
+          port: 19530
     qianfan: # 文心一言
       api-key: x0cuLZ7XsaTCU08vuJWO87Lg
       secret-key: R9mYF9dl9KASgi5RUq0FQt3wRisSnOcK