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程序即可。无需其他配置。

posted @ 2025-07-16 14:20  景伟·郭  阅读(63)  评论(0)    收藏  举报