SpringBoot中承接SSE流式输出的核心代码

常用作对接AI大模型，AI大模型为了保证推理速度和用户体验，会进行sse流式返回。

一、引入依赖

<!-- okhttp依赖，okhttp对SSE流式输出的稳定性和速度最优 -->
		<dependency>
			<groupId>com.squareup.okhttp</groupId>
			<artifactId>okhttp</artifactId>
			<version>2.7.5</version>
		</dependency>

二、编写大模型的连接器

比如我要连接多个大模型，文心一言、智谱AI、以及公司内部的自有模型，就需要将连接器抽离出来进行单独编写。下面就是一个连接器的例子：

@Component
public class AppLLMChainClient {

    @Value("${ai.app.KnowledgeBase.url}")
    private String COMPLETION_ENDPOINT_ONE;
    private static Logger logger = LoggerFactory.getLogger(AppLLMChainClient.class);

    OkHttpClient client = new OkHttpClient();
    MediaType mediaType;
    Request.Builder requestBuilder;

    /**
     * 初始化连接器
     */
    @PostConstruct
    private void init() {
        //client.setProxy(new Proxy(Proxy.Type.HTTP, new InetSocketAddress("xxx.xx.xx.x0", 9999)));
        client.setConnectTimeout(60, TimeUnit.SECONDS);
        client.setReadTimeout(60, TimeUnit.SECONDS);
        mediaType = MediaType.parse("application/json; charset=utf-8");
        requestBuilder = new Request.Builder()
                .url(COMPLETION_ENDPOINT_ONE + "chat/chat")
                .header("Content-Type", "application/json; charset=utf-8")
                .header("Accept", "text/event-stream");
    }



    /**
     * 与LLM模型对话，通过LLMChain功能
     */
    public Response chatChat(AppLLMChainRequestBody query) throws ChatException {
        RequestBody bodyOk = RequestBody.create(mediaType, JSONObject.toJSONString(query));
        Request requestOk = requestBuilder.post(bodyOk).build();
        Call call = client.newCall(requestOk);
        Response response;
        try {
            response = call.execute();
        } catch (IOException e) {
            throw new ChatException("LLMChain 请求时IO异常: " + e.getMessage());
        }
        if (response.isSuccessful()) {
            return response;
        }
        try(ResponseBody body = response.body()) {
            throw new ChatException("LLMChain 请求异常, code: " + response.code() + "body: " + body.string());
        } catch (IOException e) {
            throw new ChatException("LLMChain 请求后IO异常: " + e.getMessage());
        }
    }
}

三、应用大模型LLM的核心处理器，SSE流式返回

public class AppLLMConverseHandleWrapper {

    public static AppLLMChainClient appLLMChainClient;

    private static final ExecutorService appLLMChainES = Executors.newFixedThreadPool(10);

    private final static Pattern contentPattern = Pattern.compile(AppConsts.SSE_RESULT_PREFIX2);

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

    // 用于数据传输的 SseEmitter
    private final SseEmitter emitter = new SseEmitter(0L);

    // 对话上下文
    private List<AppLLMChainQAHistoryEntity> history;

    // 当前用户问题内容
    private String query;

    /**
     * 对话处理
     * @return SseEmitter
     */
    public SseEmitter handle() {
        if (!messageListCheck()) {
            return emitter;
        }
        doConverse();
        return emitter;
    }

    /**
     * 流式对话，异步的，在新的线程的
     */
    public SseEmitter doConverse() {
        appLLMChainES.execute(this::run);
        return emitter;
    }

    /**
     * 对话上下文检查
     * @return 是否通过
     */
    private boolean messageListCheck() {
        //限制会话次数
        if (history.size() > AppConsts.MSG_LIST_MAX) {
            sendData2Client(AppConsts.EVENT_ERROR, "上下文对话次数超限啦！");
            return false;
        }
        return true;
    }


    /**
     * 向客户端发送数据
     * @param event 事件类型
     * @param data 数据
     */
    private boolean sendData2Client(String event, String data) {
        try {
            emitter.send(SseEmitter.event().name(event).data(JSON.toJSONString(data)));
            return true;
        } catch (IOException e) {
            log.error("向客户端发送消息时出现异常");
            e.printStackTrace();
        }
        return false;
    }


    private void run() {
        AppLLMChainRequestBody chatRequestBody = new AppLLMChainRequestBody();
        chatRequestBody.setHistory(history);
        chatRequestBody.setQuery(query);
        Response chatResponse;
        try {
            chatResponse = appLLMChainClient.chatChat(chatRequestBody);
        } catch (ChatException e) {
            sendData2Client(AppConsts.EVENT_ERROR, "网络连接异常，请稍后重试");
            emitter.complete();
            e.printStackTrace();
            return;
        }
        try (ResponseBody responseBody = chatResponse.body();
             InputStream inputStream = responseBody.byteStream();
             BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            String content = null;
            //循环发送主体内容
            while ((line = bufferedReader.readLine()) != null) {
                if (StringUtils.hasLength(line)) {
                    /*log.info("data:{}", line);
                    Matcher matcher = contentPattern.matcher(line);
                    if (matcher.find()) {
                        log.info("进入data:{}", line);
                    }*/
                    if (!sendData2Client(AppConsts.EVENT_DATA, line)) {
                        break;
                    }
                }
            }
        } catch (IOException e) {
            log.error("ResponseBody读取错误");
            e.printStackTrace();
        } finally {
            emitter.complete();
        }
    }

    public AppLLMConverseHandleWrapper(){}
    public AppLLMConverseHandleWrapper(List<AppLLMChainQAHistoryEntity> history, String query) {
        this.history = history;
        this.query = query;
    }

}

四、将上方编写的各个大模型的连接器抽象成配置类并注册，方便拓展和统一维护

@Configuration
public class ConverseHandleConfig {

    @Autowired
    WYKnowledgeBase01Client wyKnowledgeBase01Client;
    @Autowired
    WYLLMClient wyllmClient;
    @Autowired
    AppKnowledgeBaseClient appKnowledgeBaseClient;
    @Autowired
    AppLLMChainClient appLLMChainClient;

    /**
     * 注册各个连接器
     */
    @PostConstruct
    private void init() {
        WYKnowledgeOneConverseHandleWrapper.wyKnowledgeBase01Client = this.wyKnowledgeBase01Client; //文心一言知识库01的连接器
        WYLLMConverseHandleWrapper.wyllmClient = this.wyllmClient; //文心一言LLM的连接器
        AppKnowledgeOneConverseHandleWrapper.appKnowledgeBaseClient = this.appKnowledgeBaseClient; //应用知识库连接器
        AppLLMConverseHandleWrapper.appLLMChainClient = this.appLLMChainClient;//应用大模型LLMChain连接器
    }


}

调用的话就像这样调用就可以啦：