Java实现ChatGpt根据SseEmitter流式输出

背景

在工作中基于使用大模型需要进行流式输出效果，公司客服助手需要接入大模型，并传参prompt给大模型参考知识，由大模型根据理解进行答案润色，润色后的结果进行流式输出。

为什么需要使用SSE

Server-Sent Events (SSE) 使用 SseEmitter:

• SSE是基于HTTP协议的，它允许服务器端推送数据到客户端。
• SseEmitter 是Spring框架中的一个类，用于在Spring MVC应用程序中实现SSE。
• SSE仅支持单向通信，即服务器到客户端的单向数据流。
• SSE通常用于简单的事件通知和消息推送，如股票价格更新、新闻feed等。
• SSE在客户端使用简单的EventSource API进行实现，不需要复杂的协议支持。
• SseEmitter是非阻塞的，并且可以很好地和Spring的异步特性结合使用。

WebSocket:

• WebSocket是一种独立的、创建在TCP上的协议，允许全双工通信。
• WebSocket可以支持服务器到客户端和客户端到服务器的双向实时通信。
• WebSocket适用于需要复杂交互和高频率通信的应用，如在线游戏、聊天应用等。
• WebSocket在客户端和服务端都需要特定的协议支持。
• WebSocket连接一旦建立，就会保持活动状态，直到被客户端或服务器关闭。

SSE工作原理

1. 建立连接：客户端通过发送HTTP请求与服务器建立连接。在请求中，客户端指定了接收事件的终点（Endpoint）。
2. 保持连接：服务器接收到连接请求后，保持连接打开，并定期发送事件数据给客户端。
3. 事件流：服务器使用 “Content-Type: text/event-stream”
   头部标识SSE连接，并使用特定格式的数据（事件流）发送给客户端。
4. 客户端处理事件：客户端通过JavaScript的 EventSource 接口监听SSE连接，一旦接收到事件，就可以处理数据并更新页面。
5. SSE协议默认是以\n\n来作为消息的结束标志，所以需要注意消息中不能包含两个换行符，如果必须要包含的话，可以选择对换行符进行转移

SseEmitter的优势:

• 简单性: SSE在客户端的实现比WebSocket简单，因为它是基于标准的HTTP协议。
• 兼容性: 对于只需要服务器到客户端单向通信的应用，SSE更容易兼容老旧浏览器。
• 效率: 对于不需要频繁从客户端到服务器的通信的应用，SSE可能更高效，因为它没有WebSocket的握手和帧控制开销。
• HTTP友好: SSE易于通过现有的HTTP基础设施进行传输，比如可以利用HTTP缓存机制。

服务端实现SSE

1、添加pom.xml依赖 --因为Springweb中就存在SseEmitter，所以只需要引入Spring即可。

SpringBoot应用

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>

若非SpringBoot应用，则可以引入以下依赖

<!-- Spring Web MVC -->
<dependency>
	<groupId>org.springframework</groupId>
	<artifactId>spring-webmvc</artifactId>
	<version>4.3.30.RELEASE</version>
</dependency>

2、创建一个Controller来处理SSE连接请求，并使用SseEmitter作为响应。

     * 建立sse连接，流式返回test
     *
     * @return
     */
    @GetMapping(value = "/streamTest")
    public SseEmitter stream() {
        log.info("=======================streamTest start");
        // 用于创建一个 SSE 连接对象
        SseEmitter emitter = new SseEmitter();
        try {
            // 在后台线程中模拟实时数据
            new Thread(() -> {
                log.info("=======================streamTest new Thread start");
                try {
                    for (int i = 0; i < 10; i++) {
                        // emitter.send() 方法向客户端发送消息
                        // 使用SseEmitter.event()创建一个事件对象，设置事件名称和数据
                        emitter.send(SseEmitter.event().name("message").data("===========>[" + new Date() + "] Data #" + i));
                        log.info("=======================streamTest push data==>" + i);
                        Thread.sleep(1000);
                    }
                    // 数据发送完成后，关闭连接
                    emitter.complete();
                } catch (IOException | InterruptedException e) {
                    // 发生错误时，关闭连接并报错
                    emitter.completeWithError(e);
                    log.error("=======================streamTest push data error==>" + e);
                }
            }).start();
            log.info("=======================streamTest end");
        } catch (Exception e) {
            log.error("streamTest error!", e);
        }
        return emitter;
    }

这个是示例代码，可以创建一个工具类对SseEmitter 进行管理，好处是线程安全支持多客户端，可以自动管理连接等。
工具类示例：

import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @ClassName SseEmitterStoreUtils
 * @Description Sse存储Utils
 * @Version 1.0
 **/
@Slf4j
public class SseEmitterStoreUtils {
    /** data */
    private static final String DATA = "data";
    /** type */
    private static final String TYPE = "type";
    /** end */
    private static final String TYPE_END = "end";
    /** sensitive */
    private static final String TYPE_SEN = "sensitive";

    /**
     * messageId的 SseEmitter对象映射集
     */
    private static final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    public static int getEmitterSize() {
        return sseEmitterMap.size();
    }

    public static SseEmitter getEmitter(String messageId) {
        return sseEmitterMap.get(messageId);
    }

    public static void addEmitter(String messageId, SseEmitter emitter) {
        sseEmitterMap.put(messageId, emitter);
        log.info("当前messageId={},sseEmitterMap.size()={}",messageId,sseEmitterMap.size());
    }


    /**
     * 给指定messageId发消息
     *
     * @param messageId - 消息id（唯一）
     * @param message   - 消息文本
     */
    public static boolean sendMessage(String messageId, String message) {
        SseEmitter emitter = sseEmitterMap.get(messageId);
        if (emitter != null) {
            try {
                // 构造并发送自定义事件格式的数据
                SseEmitter.SseEventBuilder event = SseEmitter.event()
                        .name(DATA)
                        .data(message, MediaType.TEXT_PLAIN);
                emitter.send(event);
                log.info("当前messageId={},发送消息={}", messageId, message);
                return true;
            } catch (Exception e) {
                log.error("发送消息异常 ==> messageId={}, 异常信息：{}", messageId, e.getMessage());
                // 关闭连接
                emitter.complete();
            }
        } else {
            throw new RuntimeException("连接不存在或者超时， messageId=" + messageId);
        }
        return false;
    }


    /**
     * 主动断开连接
     * @param messageId
     */
    public static void sensitiveCloseSseEmitter(String messageId){
        closeSseEmitter(messageId, TYPE_SEN);
    }

    /**
     * 主动断开连接
     * @param messageId
     */
    public static void closeSseEmitter(String messageId){
        closeSseEmitter(messageId, TYPE_END);
    }

    /**
     * 主动断开连接 关闭sse
     * @param messageId
     */
    public static void closeSseEmitter(String messageId, String type) {
        if (sseEmitterMap.containsKey(messageId)) {
            SseEmitter emitter = sseEmitterMap.get(messageId);
            if (emitter != null) {
                try {
                    // 构造并发送自定义事件格式的数据
                    SseEmitter.SseEventBuilder event = SseEmitter.event()
                            .name(type)
                            .data("", MediaType.TEXT_PLAIN);
                    emitter.send(event);
                } catch (IOException e) {
                    log.error("关闭链接异常 ==> messageId={}", messageId, e);
                    // 关闭连接
                    emitter.complete();
                }
                log.info("当前messageId={},主动连接断开",messageId);
                // 关闭连接
                emitter.complete();
                //移除
                sseEmitterMap.remove(messageId);
            }
        }
    }

    /**
     * 主动断开所有连接 关闭sse
     * @param
     */
    public static void closeAllSseEmitter() {
        for (Map.Entry<String, SseEmitter> entry : sseEmitterMap.entrySet()) {
            SseEmitter emitter = entry.getValue();
            if (emitter != null) {
                log.info("当前messageId={},主动连接断开",entry.getKey());
                // 关闭连接
                emitter.complete();
            }
        }
        sseEmitterMap.clear();
    }
}

示例HTML代码调用：

<!doctype html>
<html lang="en">
<head>
    <title>Sse测试文档</title>
</head>
<body>
<div>sse测试</div>
<div id="result"></div>

<script>
    var source = new EventSource('http:xxxx/sse/streamTest');

    source.addEventListener("message", function (event) {
        var resultDiv = document.getElementById('result');
        // 使用innerText可能会导致换行符被忽略，改用textContent
        resultDiv.textContent += event.data + '\n';
    });

    // 添加一个开启回调
    source.onopen = function (event) {
        console.log('连接已开启', event);
    };

    // 添加一个错误回调
    source.onerror = function (event) {
        if (source.readyState === EventSource.CLOSED) {
            console.error('连接被关闭');
        } else {
            console.error('发生错误', event);
        }
    };
</script>

</body>
</html>

这样直接打开这个html文件，就可以查看到流式效果。

当然也可以通过postMan进行调用，使用postman调用推荐使用post调用方式。
请求头设置Connection=keep-alive

代码实现：

     * 建立sse连接，流式返回,以messageId维度
     *
     * @return
     */
    @PostMapping(value = "/connection",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatGptConnection(HttpServletResponse response, HttpServletRequest request) throws IOException {
        // 获得请求体(这里requestBody如果需要使用就用这种方式获取)
        String requestBody = getRequestBody(request);
        log.info("=======================chatGptConnection requestBody==>" + requestBody);
        // 解析请求体
        if(StringUtils.isBlank(requestBody)){
            log.error("=======================chatGptConnection requestBody为空");
            return null;
        }
        ObjectMapper mapper = new ObjectMapper();
        ChatGptDto chatGptDto = null;
        try {
            chatGptDto = mapper.readValue(requestBody, ChatGptDto.class);
        } catch (Exception e) {
            log.error("=======================chatGptConnection requestBody解析失败",e);
        }
        // 创建一个 SSE 连接对象 超时时间设置为1小时
        long timeout = 3600_000L;
        SseEmitter emitter = new SseEmitter(timeout);
         // 在线程池中执行模拟实时数据
         threadPool.execute(() -> simulateRealtimeData(emitter));
        return emitter;
    }

	/**
     * 获得请求体信息
     */
    private static String getRequestBody(HttpServletRequest request) throws IOException {
        if (request == null) {
            return null;
        }
        StringBuilder requestBody = new StringBuilder();
        BufferedReader reader = request.getReader();
        String line;
        while ((line = reader.readLine()) != null) {
            requestBody.append(line);
        }
        return requestBody.toString();
    }

	/**
     * 模拟实时数据
     *
     * @param emitter
     */
    private void simulateRealtimeData(SseEmitter emitter) {
        try {
            //获取配置Mock内容
            String introduction = "大家好！欢迎来到我们的聊天世界。我是小助手。我随时准备着回答你的问题，与你分享知识，或者仅仅是为了愉快的对话。无论你有任何问题，都可以随时向我提出，我将尽我所能为你提供帮助。我们可以聊天、探讨各种话题，甚至是编程问题或者日常生活中遇到的小困难。现在，告诉我，我怎样才能协助你呢？";
            int limit = 50;
            for (int i = 0; i < introduction.length(); i++) {
                // emitter.send() 方法向客户端发送消息 
                // 使用SseEmitter.event()创建一个事件对象，设置事件名称和数据
                emitter.send(SseEmitter.event().name("message").data(introduction.substring(i, i + 1)));
                // 这里也可以改为使用SseEmitterStoreUtils.sendMessage(messageId, introduction.substring(i, i + 1));
                Thread.sleep(limit);
            }
            // 数据发送完成后，关闭连接
            emitter.complete();
        } catch (IOException | InterruptedException e) {
            // 发生错误时，关闭连接并报错
            emitter.completeWithError(e);
        }
    }

效果：

3、分享几点SSE的踩坑
报错：Async support must be enabled on a servlet and for all filters involved in async request processing
原因：异步请求必须对servlet和所有涉及到的filter都需要开启异步支持
解决方法：

一、Spring web应用

<servlet>
        <servlet-name>springmvc</servlet-name>
        <servlet-class>org.springframework.web.servlet.DispatcherServlet</servlet-class>
        <init-param>
            <param-name>contextConfigLocation</param-name>
            <param-value>classpath:spring.xml</param-value>
        </init-param>
        <load-on-startup>1</load-on-startup>
        <async-supported>true</async-supported>---主要是添加这个
    </servlet>

这里注意true是web.xml 3.0的新特性，web.xml需要配置的是 web-app_3_0.xsd

<web-app xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
 xmlns="http://java.sun.com/xml/ns/javaee" 
 xsi:schemaLocation="http://java.sun.com/xml/ns/javaee http://java.sun.com/xml/ns/javaee/web-app_3_0.xsd" id="WebApp_ID" version="3.0">

二、SpringBoot应用
Spring Boot自动配置了DispatcherServlet，并默认启用了异步支持。如果需要显式地配置异步支持，可以通过继承WebMvcConfigurerAdapter类或实现WebMvcConfigurer接口，并重写其中的方法来配置

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.web.servlet.config.annotation.AsyncSupportConfigurer;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

@SpringBootApplication
public class MyApp {
    public static void main(String[] args) {
        SpringApplication.run(MyApp.class, args);
    }
    @Bean
    public WebMvcConfigurer webMvcConfigurer() {
        return new WebMvcConfigurer() {
            @Override
            public void configureAsyncSupport(AsyncSupportConfigurer configurer) {
            		// 设置异步请求的超时时间
                configurer.setDefaultTimeout(5000); 
                // 还可以配置其他异步相关的选项，比如TaskExecutor
            }
        };
    }
}