Spring AI 落地企业中台：多模型动态路由与 SSE 流式长连接高并发治理

spring:
  application:
    name: enterprise-ai-gateway
  ai:
    openai:
      api-key: ${OPENAI_API_KEY:sk-fake-openai}
      base-url: https://api.openai.com
    # DeepSeek 兼容 OpenAI 协议规范
    deepseek:
      api-key: ${DEEPSEEK_API_KEY:sk-fake-deepseek}
      base-url: https://api.deepseek.com
    # Gemini 原生配置或适配层
    gemini:
      api-key: ${GEMINI_API_KEY:fake-gemini}

# 动态路由权重字典 (支持 Spring Cloud Config / Nacos 运行时刷新)
ai:
  routing:
    rules:
      - provider: deepseek
        weight: 60
        enabled: true
      - provider: openai
        weight: 30
        enabled: true
      - provider: gemini
        weight: 10
        enabled: false # 支持蓝绿与灰度启停

package com.mrwu.ai.gateway.router;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;

@Component
@ConfigurationProperties(prefix = "ai.routing")
public class DynamicChatModelRouter {

    private static final Logger log = LoggerFactory.getLogger(DynamicChatModelRouter.class);

    private final Map<String, ChatModel> registeredModels = new ConcurrentHashMap<>();
    private List<RouteRule> rules = new ArrayList<>();

    // 依赖注入：注入不同 Provider 的 ChatModel Bean
    public DynamicChatModelRouter(Map<String, ChatModel> chatModelBeans) {
        this.registeredModels.putAll(chatModelBeans);
    }

    public static class RouteRule {
        private String provider;
        private int weight;
        private boolean enabled;

        public String getProvider() { return provider; }
        public void setProvider(String provider) { this.provider = provider; }
        public int getWeight() { return weight; }
        public void setWeight(int weight) { this.weight = weight; }
        public boolean isEnabled() { return enabled; }
        public void setEnabled(boolean enabled) { this.enabled = enabled; }
    }

    public void setRules(List<RouteRule> rules) {
        this.rules = rules;
    }

    /**
     * 基于权重随机算法抽取 ChatModel，并支持降级链路
     */
    public ChatModel route() {
        List<RouteRule> activeRules = rules.stream()
                .filter(RouteRule::isEnabled)
                .filter(r -> r.getWeight() > 0)
                .toList();

        if (activeRules.isEmpty()) {
            throw new IllegalStateException("无可用的 LLM 路由策略规则！");
        }

        int totalWeight = activeRules.stream().mapToInt(RouteRule::getWeight).sum();
        int randomVal = ThreadLocalRandom.current().nextInt(totalWeight);
        int currentSum = 0;

        for (RouteRule rule : activeRules) {
            currentSum += rule.getWeight();
            if (randomVal < currentSum) {
                ChatModel target = getModelBean(rule.getProvider());
                if (target != null) {
                    log.debug("命中路由 Provider: [{}], Weight: [{}]", rule.getProvider(), rule.getWeight());
                    return target;
                }
            }
        }

        // 兜底返回第一个可用模型
        return getModelBean(activeRules.get(0).getProvider());
    }

    private ChatModel getModelBean(String provider) {
        // 根据 Bean 名称规范映射，例如 deepseekChatModel, openAiChatModel
        String beanName = provider.toLowerCase() + "ChatModel";
        ChatModel model = registeredModels.get(beanName);
        if (model == null) {
            log.warn("无法定位 BeanName=[{}] 的 ChatModel 实例，尝试全量扫描备选", beanName);
            return registeredModels.values().stream().findFirst().orElse(null);
        }
        return model;
    }
}

package com.mrwu.ai.gateway.controller;

import com.mrwu.ai.gateway.router.DynamicChatModelRouter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.concurrent.atomic.AtomicInteger;

@RestController
public class AiStreamController {

    private static final Logger log = LoggerFactory.getLogger(AiStreamController.class);
    private final DynamicChatModelRouter router;

    public AiStreamController(DynamicChatModelRouter router) {
        this.router = router;
    }

    @GetMapping(value = "/api/v1/ai/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamChat(@RequestParam("prompt") String promptText) {
        ChatModel dynamicModel = router.route();
        Prompt prompt = new Prompt(new UserMessage(promptText));

        AtomicInteger tokenCounter = new AtomicInteger(0);

        // 获取底层 Reactive Stream
        Flux<ChatResponse> responseFlux = dynamicModel.stream(prompt);

        return responseFlux
                .map(chatResponse -> {
                    String content = Optional.ofNullable(chatResponse.getResult())
                            .map(result -> result.getOutput().getContent())
                            .orElse("");
                    tokenCounter.incrementAndGet();
                    return ServerSentEvent.<String>builder()
                            .event("message")
                            .data(content)
                            .build();
                })
                // 1. 应对网络超时：若 20 秒内上游无任何数据下发，触发超时异常
                .timeout(Duration.ofSeconds(20))
                // 2. 核心：客户端主动断开连接治理（关闭浏览器、路由切换等）
                .doOnCancel(() -> {
                    log.warn("检测到客户端主动断开 SSE 链路！立即向上游注销上下文，已输出 Token 片段数: {}", 
                            tokenCounter.get());
                    // 在此处可触发上游 Http 连接物理中断与用户账单回滚
                })
                // 3. 异常拦截并推送前端易读的错误事件
                .onErrorResume(throwable -> {
                    log.error("流式推理过程发生致命异常, 原因: {}", throwable.getMessage());
                    return Flux.just(ServerSentEvent.<String>builder()
                            .event("error")
                            .data("服务暂时不可用: " + throwable.getMessage())
                            .build());
                })
                // 4. 传输完成标志推送
                .concatWith(Flux.just(ServerSentEvent.<String>builder()
                        .event("done")
                        .data("[DONE]")
                        .build()))
                // 5. 补充心跳包，防止代理服务器在 LLM 深度推理（如 DeepSeek R1 思考期）时因空闲切断连接
                .sampleFirst(Duration.ofMillis(50)) // 防止前端消费过载轻微限流
                .log("ai.stream.logger");
    }
}

upstream ai_backend {
    server 127.0.0.1:8080;
    keepalive 32; # 保持上游连接池活跃，降低 TCP 握手开销
}

server {
    listen 80;
    server_name ai.mrwu.com;

    location /api/v1/ai/ {
        proxy_pass http://ai_backend;

        # HTTP/1.1 协议以支持长连接与 Chunked 传输
        proxy_http_version 1.1;
        proxy_set_header Connection "";

        # 【重点】彻底禁用 Nginx Proxy 缓冲，实现真正的实时流式直通
        proxy_buffering off;
        proxy_cache off;

        # 传递 SSE 所需的关键请求头
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_set_header X-Accel-Buffering no; # 告诉上游任何代理均不可缓冲

        # 【重点】LLM 推理长耗时保护，防止 60 秒无数据交互被 Nginx 掐断
        proxy_read_timeout 300s;
        proxy_send_timeout 300s;
    }
}