springai优雅的实现对多个key进行轮转
分析springai源码得知 springai 默认是取配置文件的key放到请求头里的,因此我们只需要重写加载key与放置key的方法即可:
import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
import org.springframework.ai.model.SpringAIModelProperties;
import org.springframework.ai.model.SpringAIModels;
import org.springframework.ai.model.openai.autoconfigure.*;
import org.springframework.ai.model.tool.DefaultToolExecutionEligibilityPredicate;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionEligibilityPredicate;
import org.springframework.ai.model.tool.autoconfigure.ToolCallingAutoConfiguration;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.retry.autoconfigure.SpringAiRetryAutoConfiguration;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.ImportAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
import org.springframework.boot.autoconfigure.web.reactive.function.client.WebClientAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
import static org.springframework.ai.model.openai.autoconfigure.OpenAIAutoConfigurationUtil.resolveConnectionProperties;
/**
* @author Jingway
* @date 2025/7/14 16:03
* @description: 重写 OpenAiApi 的配置,使用随机的 appKey
*/
@Primary // 确保这个配置在其他 OpenAiApi 配置之前加载
@AutoConfiguration(after = { RestClientAutoConfiguration.class, WebClientAutoConfiguration.class,
SpringAiRetryAutoConfiguration.class, ToolCallingAutoConfiguration.class })
@ConditionalOnClass(OpenAiApi.class)
@EnableConfigurationProperties({ OpenAiConnectionProperties.class, OpenAiChatProperties.class })
@ConditionalOnProperty(name = SpringAIModelProperties.CHAT_MODEL, havingValue = SpringAIModels.OPENAI,
matchIfMissing = true)
@ImportAutoConfiguration(classes = { SpringAiRetryAutoConfiguration.class, RestClientAutoConfiguration.class,
WebClientAutoConfiguration.class, ToolCallingAutoConfiguration.class })
public class OpenAiAppKeyAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public OpenAiChatModel openAiChatModel(OpenAiConnectionProperties commonProperties,
OpenAiChatProperties chatProperties, ObjectProvider<RestClient.Builder> restClientBuilderProvider,
ObjectProvider<WebClient.Builder> webClientBuilderProvider, ToolCallingManager toolCallingManager,
RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler,
ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<ChatModelObservationConvention> observationConvention,
ObjectProvider<ToolExecutionEligibilityPredicate> openAiToolExecutionEligibilityPredicate) {
var openAiApi = openAiApi(chatProperties, commonProperties, responseErrorHandler, "chat");
var chatModel = OpenAiChatModel.builder()
.openAiApi(openAiApi)
.defaultOptions(chatProperties.getOptions())
.toolCallingManager(toolCallingManager)
.toolExecutionEligibilityPredicate(openAiToolExecutionEligibilityPredicate.getIfUnique(DefaultToolExecutionEligibilityPredicate::new))
.retryTemplate(retryTemplate)
.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
.build();
observationConvention.ifAvailable(chatModel::setObservationConvention);
return chatModel;
}
public OpenAiApi openAiApi(OpenAiChatProperties chatProperties, OpenAiConnectionProperties commonProperties,
ResponseErrorHandler responseErrorHandler, String modelType) {
OpenAIAutoConfigurationUtil.ResolvedConnectionProperties resolved = resolveConnectionProperties(
commonProperties, chatProperties, modelType);
var rotatingKey = new RotatingApiKey(resolved.apiKey());
var restClient = RotatingApiKey.buildRestClient(resolved.baseUrl(), rotatingKey, resolved.headers(), responseErrorHandler);
var webClient = RotatingApiKey.buildWebClient(resolved.baseUrl(), rotatingKey, resolved.headers());
return OpenAiApi.builder()
.baseUrl(resolved.baseUrl())
.apiKey(rotatingKey)
.headers(resolved.headers())
.completionsPath(chatProperties.getCompletionsPath())
.embeddingsPath(OpenAiEmbeddingProperties.DEFAULT_EMBEDDINGS_PATH)
.restClientBuilder(restClient.mutate())
.webClientBuilder(webClient.mutate())
.responseErrorHandler(responseErrorHandler)
.build();
}
}
import org.springframework.ai.model.ApiKey;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;
/**
* @author Jingway
* @date 2025/7/14 15:17
* @description: 通过代理随机取apikey
*/
public record RotatingApiKey(String appKeys) implements ApiKey {
@Override
public String getValue() {
var apiKeys = List.of(appKeys.split(","));
return apiKeys.get(ThreadLocalRandom.current().nextInt(apiKeys.size()));
}
@Override
public String toString() {
return "RotatingApiKey{value='***'}";
}
public static RestClient buildRestClient(String baseUrl, ApiKey apiKey, MultiValueMap<String, String> headers,
ResponseErrorHandler errorHandler) {
return RestClient.builder()
.baseUrl(baseUrl)
.defaultStatusHandler(errorHandler)
.requestInterceptor((request, body, execution) -> {
HttpHeaders httpHeaders = request.getHeaders();
httpHeaders.setBearerAuth(apiKey.getValue()); // 每次调用动态 key
httpHeaders.setContentType(MediaType.APPLICATION_JSON);
httpHeaders.addAll(headers);
return execution.execute(request, body);
})
.build();
}
public static WebClient buildWebClient(String baseUrl, ApiKey apiKey, MultiValueMap<String, String> headers) {
return WebClient.builder()
.baseUrl(baseUrl)
.defaultHeaders(h -> {
h.setBearerAuth(apiKey.getValue()); // 每次调用动态 key
h.setContentType(MediaType.APPLICATION_JSON);
h.addAll(headers);
})
.build();
}
}
将以上两个类copy到我们的springboot程序即可。无需其他配置。