spring-ai/llm-gateway/src/main/java/com/ankurm/gateway/GatewayConfig.java

package com.ankurm.gateway;

import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import io.github.bucket4j.TimeMeter;
import io.github.resilience4j.circuitbreaker.CircuitBreakerConfig;
import io.github.resilience4j.circuitbreaker.CircuitBreakerRegistry;
import org.springframework.ai.anthropic.AnthropicChatModel;
import org.springframework.ai.anthropic.AnthropicChatOptions;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
@EnableConfigurationProperties(GatewayProperties.class)
public class GatewayConfig {

    @Bean
    Routes routes(GatewayProperties p) {
        Map<String, Target> byName = new LinkedHashMap<>();
        for (GatewayProperties.Provider c : p.providers()) {
            byName.put(c.name(), new Target(c.name(), c.model(), chatModel(c), Price.of(c.inputPrice(), c.outputPrice())));
        }
        Map<String, List<Target>> table = new LinkedHashMap<>();
        p.routes().forEach((hint, names) -> table.put(hint, names.stream().map(byName::get).toList()));
        return new Routes(table);
    }

    /** No retries inside the SDK: the gateway's own failover is the retry policy. */
    public static ChatModel chatModel(GatewayProperties.Provider c) {
        return switch (c.type()) {
            case "openai" -> OpenAiChatModel.builder().options(OpenAiChatOptions.builder()
                    .baseUrl(c.baseUrl()).apiKey(c.apiKey()).model(c.model()).maxRetries(0).build()).build();
            case "anthropic" -> AnthropicChatModel.builder().options(AnthropicChatOptions.builder()
                    .baseUrl(c.baseUrl()).apiKey(c.apiKey()).model(c.model()).maxTokens(1024).maxRetries(0).build())
                    .build();
            default -> throw new IllegalArgumentException("Unknown provider type " + c.type());
        };
    }

    @Bean
    CircuitBreakerRegistry breakers(GatewayProperties p) {
        return registry(p.breaker());
    }

    /** One breaker per provider name, all built from the same settings. */
    public static CircuitBreakerRegistry registry(GatewayProperties.Breaker b) {
        return CircuitBreakerRegistry.of(CircuitBreakerConfig.custom()
                .slidingWindowSize(b.slidingWindowSize())
                .minimumNumberOfCalls(b.minimumCalls())
                .failureRateThreshold(b.failureRateThreshold())
                .waitDurationInOpenState(Duration.ofMillis(b.waitOpenMillis()))
                .permittedNumberOfCallsInHalfOpenState(1)
                .build());
    }

    @Bean
    Budgets budgets(GatewayProperties p) {
        return new Budgets(t -> p.tenant(t).budgetMicros());
    }

    @Bean
    TokenLimiter tokenLimiter(GatewayProperties p) {
        return new TokenLimiter(t -> p.tenant(t).tokensPerHour(), TimeMeter.SYSTEM_MILLISECONDS);
    }

    @Bean
    Gateway gateway(Routes routes, CircuitBreakerRegistry breakers, Budgets budgets, TokenLimiter limiter,
            ObjectProvider<EmbeddingModel> embeddings) {
        EmbeddingModel em = embeddings.getIfAvailable();
        return new Gateway(routes, targets -> new Failover(targets, breakers, budgets), limiter,
                em == null ? null : new SemanticCache(em, 0.92), Map.of("refund_order", new OrderTools()));
    }
}