package cn.iocoder.yudao.module.ai.framework.ai.core.model;
|
|
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.hutool.extra.spring.SpringUtil;
|
import cn.iocoder.yudao.module.ai.enums.model.AiPlatformEnum;
|
import cn.iocoder.yudao.module.ai.framework.ai.config.AiAutoConfiguration;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.chat.model.ChatModel;
|
import org.springframework.ai.embedding.BatchingStrategy;
|
import org.springframework.ai.embedding.EmbeddingModel;
|
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
|
import org.springframework.ai.openai.OpenAiChatModel;
|
import org.springframework.ai.vectorstore.VectorStore;
|
import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
|
import org.springframework.ai.vectorstore.milvus.autoconfigure.MilvusServiceClientProperties;
|
import org.springframework.ai.vectorstore.milvus.autoconfigure.MilvusVectorStoreProperties;
|
import java.util.Map;
|
|
/**
|
* AI Model 模型工厂实现类
|
*
|
* 使用 OpenAI 兼容接口对接通义千问 + Milvus 向量存储
|
*/
|
@Slf4j
|
public class AiModelFactoryImpl implements AiModelFactory {
|
|
@Override
|
public ChatModel getOrCreateChatModel(AiPlatformEnum platform, String apiKey, String url, String model) {
|
String cacheKey = buildCacheKey(ChatModel.class, platform, apiKey, url, model);
|
return Singleton.get(cacheKey, (Func0<ChatModel>) () -> {
|
if (platform == AiPlatformEnum.TONG_YI) {
|
return AiAutoConfiguration.buildTongYiChatModel(apiKey, model);
|
}
|
throw new IllegalArgumentException(StrUtil.format("不支持的平台({})", platform));
|
});
|
}
|
|
@Override
|
public ChatModel getDefaultChatModel(AiPlatformEnum platform) {
|
if (platform == AiPlatformEnum.TONG_YI) {
|
return SpringUtil.getBean(OpenAiChatModel.class);
|
}
|
throw new IllegalArgumentException(StrUtil.format("不支持的平台({})", platform));
|
}
|
|
@Override
|
public EmbeddingModel getOrCreateEmbeddingModel(AiPlatformEnum platform, String apiKey, String url, String model) {
|
String cacheKey = buildCacheKey(EmbeddingModel.class, platform, apiKey, url, model);
|
return Singleton.get(cacheKey, (Func0<EmbeddingModel>) () -> {
|
if (platform == AiPlatformEnum.TONG_YI) {
|
return AiAutoConfiguration.buildTongYiEmbeddingModel(apiKey, model);
|
}
|
throw new IllegalArgumentException(StrUtil.format("不支持的平台({})", platform));
|
});
|
}
|
|
@Override
|
public VectorStore getOrCreateVectorStore(Class<? extends VectorStore> type,
|
EmbeddingModel embeddingModel,
|
Map<String, Class<?>> metadataFields) {
|
// metadataFields 参与缓存 key，确保不同知识库使用不同配置时不会复用
|
String cacheKey = buildCacheKey(VectorStore.class, embeddingModel, type, metadataFields.hashCode());
|
return Singleton.get(cacheKey, (Func0<VectorStore>) () -> {
|
if (type == MilvusVectorStore.class) {
|
return buildMilvusVectorStore(embeddingModel);
|
}
|
throw new IllegalArgumentException(StrUtil.format("不支持的向量存储类型({})", type));
|
});
|
}
|
|
private MilvusVectorStore buildMilvusVectorStore(EmbeddingModel embeddingModel) {
|
MilvusVectorStoreProperties serverProperties = SpringUtil.getBean(MilvusVectorStoreProperties.class);
|
MilvusServiceClientProperties clientProperties = SpringUtil.getBean(MilvusServiceClientProperties.class);
|
|
var connectParam = io.milvus.param.ConnectParam.newBuilder()
|
.withHost(clientProperties.getHost())
|
.withPort(clientProperties.getPort())
|
.withDatabaseName(serverProperties.getDatabaseName())
|
.build();
|
var milvusClient = new io.milvus.client.MilvusServiceClient(connectParam);
|
|
MilvusVectorStore vectorStore = MilvusVectorStore.builder(milvusClient, embeddingModel)
|
.databaseName(serverProperties.getDatabaseName())
|
.collectionName(serverProperties.getCollectionName())
|
.initializeSchema(serverProperties.isInitializeSchema())
|
.batchingStrategy(new TokenCountBatchingStrategy())
|
.build();
|
try {
|
vectorStore.afterPropertiesSet();
|
} catch (Exception e) {
|
throw new RuntimeException("Milvus 向量存储初始化失败: " + e.getMessage(), e);
|
}
|
return vectorStore;
|
}
|
|
private static String buildCacheKey(Class<?> clazz, Object... params) {
|
if (ArrayUtil.isEmpty(params)) return clazz.getName();
|
return StrUtil.format("{}#{}", clazz.getName(), ArrayUtil.join(params, "_"));
|
}
|
|
}
|
