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

package com.ankurm.gateway;

import java.util.List;
import java.util.Map;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.ToolCallingAdvisor;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;

/**
 * The one method every calling service uses: rate limit, cache, route, record. Each request gets
 * its own {@link Failover} over the targets for its hint, wrapped in a {@code ChatClient} whose
 * tool loop is the standard {@link ToolCallingAdvisor}.
 */
public final class Gateway {

    private final Routes routes;

    private final java.util.function.Function<List<Target>, Failover> failovers;

    private final TokenLimiter limiter;

    private final SemanticCache cache;

    private final Map<String, ToolCallback[]> toolbox;

    public Gateway(Routes routes, java.util.function.Function<List<Target>, Failover> failovers,
            TokenLimiter limiter, SemanticCache cache, Map<String, Object> toolObjects) {
        this.routes = routes;
        this.failovers = failovers;
        this.limiter = limiter;
        this.cache = cache;
        this.toolbox = new java.util.LinkedHashMap<>();
        toolObjects.forEach((name, obj) -> toolbox.put(name, ToolCallbacks.from(obj)));
    }

    public GatewayResponse complete(String tenant, GatewayRequest req) {
        List<Target> targets = routes.get(req.modelHint());
        List<ToolCallback> callbacks = resolveTools(req.tools());

        long estimate = (req.systemPrompt().length() + req.userMessage().length() + 3) / 4 + req.maxTokens();
        limiter.take(tenant, estimate);
        CallContext.Meter meter = new CallContext.Meter();
        try {
            var cached = cache == null ? java.util.Optional.<String>empty()
                    : cache.lookup(tenant, req.featureTag(), req.userMessage());
            if (cached.isPresent()) {
                limiter.settle(tenant, estimate, 0);
                return new GatewayResponse(cached.get(), "cache", "cache", 0, 0, 0, 0, true, List.of("cache: hit"));
            }
            ChatClient client = ChatClient.builder(failovers.apply(targets))
                    .defaultAdvisors(ToolCallingAdvisor.builder().build())
                    .build();
            CallContext ctx = new CallContext(tenant, req.featureTag(), meter);
            ChatResponse response = ScopedValue.where(CallContext.CURRENT, ctx).call(() -> {
                var spec = client.prompt().user(req.userMessage())
                        .options(org.springframework.ai.model.tool.ToolCallingChatOptions.builder()
                                .maxTokens(req.maxTokens()));
                if (!req.systemPrompt().isBlank()) {
                    spec = spec.system(req.systemPrompt());
                }
                if (!callbacks.isEmpty()) {
                    spec = spec.toolCallbacks(callbacks);
                }
                return spec.call().chatResponse();
            });
            String content = response.getResult().getOutput().getText();
            limiter.settle(tenant, estimate, meter.promptTokens() + meter.completionTokens());
            if (cache != null) {
                cache.put(tenant, req.featureTag(), req.userMessage(), content);
            }
            return new GatewayResponse(content, meter.provider(), meter.model(), meter.promptTokens(),
                    meter.completionTokens(), meter.costMicros(), meter.modelCalls(), false, meter.trail());
        }
        catch (RuntimeException e) {
            limiter.settle(tenant, estimate, meter.promptTokens() + meter.completionTokens());
            throw e;
        }
        catch (Exception e) {
            throw new IllegalStateException(e);
        }
    }

    private List<ToolCallback> resolveTools(List<String> names) {
        java.util.ArrayList<ToolCallback> out = new java.util.ArrayList<>();
        for (String n : names) {
            ToolCallback[] cb = toolbox.get(n);
            if (cb == null) {
                throw new IllegalArgumentException("Tool '" + n + "' is not on the gateway allow-list "
                        + toolbox.keySet());
            }
            out.addAll(List.of(cb));
        }
        return out;
    }
}