从0到1,踩坑SSE流式传输

1,172 阅读6分钟

前言

最近项目上接入了大模型,需要实现类ChatGPT聊天框。其中一个技术就是回复内容流式输出,这在市面上不难找到很成熟的方案,那就是SSE(Server-Sent Events)。虽然大方案上很容易确定,但实现细节上还是有很多可以分享的地方。

正文

首先为什么需要实现流式输出这种交互?

基于深度学习的大型语言模型,处理自然语言需要大量的计算资源和时间,响应速度比普通的读数据库要慢的多,普通 http 接口等待时间过长,用户交互感官上很差。

大模型将先计算出的数据“推送”给用户,边计算边返回,避免用户因为等待时间过长关闭页面,而这,可以采用SSE技术

SSE

简单介绍下SSE协议,全称Server-Sent Events,2008年首次出现在HTML5规范中,在2014年随着HTML5被W3C推荐为标准,SSE也登上了舞台。作为HTML5的一部分,旨在提供一种简单的机制,用于服务器向客户端推送实时事件数据

SSE建立在标准的HTTP协议之上,使用普通的HTTP连接,与WebSocket不同的是,SSE是一种单向通信协议,只能是服务器向客户端推送数据,客户端只需要建立连接,而后续的数据推送由服务器单方面完成。

简单示例

服务端

const express = require('express');
const app = express();
const port = 3000;

// 创建一个 SSE 事件源
app.get('/events', (req, res) => {
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');

  // 模拟数据流
  setInterval(() => {
    const data = {
      time: new Date().toISOString(),
      message: 'Hello, client!'
    };

    // 发送数据
    res.write(`data: ${JSON.stringify(data)}\n\n`);

    // 检查客户端是否断开连接
    if (res.writableEnded) {
      clearInterval(intervalId);
    }
  }, 5000); // 每 5 秒发送一次
});

app.listen(port, () => {
  console.log(`Server listening at http://localhost:${port}`);
});

客户端

<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta http-equiv="X-UA-Compatible" content="IE=edge">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>SSE Demo</title>
</head>
<body>
  <h1>SSE Demo</h1>
  <button onclick="connectSSE()">建立 SSE 连接</button>  
  <button onclick="closeSSE()">断开 SSE 连接</button>
  <br />
  <br />
  <div id="message"></div>

  <script>
    const messageElement = document.getElementById('message')

    let eventSource

    // 建立 SSE 连接
    const connectSSE = () => {
      eventSource = new EventSource('/events')

      //监听自定义事件
      eventSource.addEventListener('customEvent', (event) => {
        const data = JSON.parse(event.data)
        messageElement.innerHTML += `${data.id} --- ${data.time} --- params参数:${JSON.stringify(data.params)}` + '<br />'
      })
      // 监听消息事件
      eventSource.onopen = () => {
        messageElement.innerHTML += `SSE 连接成功,状态${eventSource.readyState}<br />`
      }
      
      eventSource.onmessage = (event) => {
       const data = JSON.parse(event.data)
        messageElement.innerHTML += `${data.id} --- ${data.time} --- params参数:${JSON.stringify(data.params)}` + '<br />'
      }

      eventSource.onerror = () => {
        messageElement.innerHTML += `SSE 连接错误,状态${eventSource.readyState}<br />`
      }
    }

    // 断开 SSE 连接
    const closeSSE = () => {
      eventSource.close()
      messageElement.innerHTML += `SSE 连接关闭,状态${eventSource.readyState}<br />`
    }
  </script>
</body>
</html>

以上代码实现了简单的流式传输

不可能,绝对没这么简单。果然,

接口是post方法,需要传一些复杂的数据,而EventSource不支持post,那我们应该怎么办呢?

SSE post实践

思路:SSE (Server-Sent Events) Using A POST Request Without EventSource

办法:用fetch的post

客户端

 const response = await fetch(apiPath, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify(apiParams), // 请求体
});

if (!response.ok) {
          console.error('Error in SSE response', response.statusText);
          return;
}
const reader = response?.body
          ?.pipeThrough(new TextDecoderStream())
          .getReader();
         
if (!reader) {
          console.error('Reader is not defined');
          return;
}

const processStream = async () => {
     while (true) {
         const { done, value: _value } = await reader.read();
         if (done) break;
         // 检查是否是结束标志
         if (_value.data === '"[DONE]"') {
             // 处理结束流的逻辑
             break;
         }else {
             const jsonData = JSON.parse(_value.data);
              // 处理流式报错
              if (jsonData.data === null && jsonData.success === false) {
                  console.log(jsonData.message)
                  break;
              }
              console.log(jsonData.content)
         }
     }
}
processStream();

运行以上代码,发起请求,发现数据确实通过EventStream获取,可以通过控制台查看。

image.png

但是发现服务端的数据并不是流式输出到客户端的,而是等所有数据准备好后一次性返回给了客户端,这不是我想要的。

关闭body压缩功能

经过查阅资料得知,是因为默认情况下,Webpack Dev Server 会自动启用 gzip 压缩,使用 gzip 压缩可以显著减小 HTTP 响应体的大小,从而加快页面加载速度。

请求头:

image.png

响应头:

image.png

但是这影响我们event-stream正常返回,

通过修改配置,可以将其关闭。

exports.default = {
  devServer: {
    // 其他配置...
    compress: false, // 禁用 gzip 压缩
    // 其他配置...
  }
};

修改以上配置,重启项目,

数据流式的出现在了EventSream面板。

这不是最佳方案,最佳方案是让服务端针对针对特定接口禁用gzip压缩。

在控制台中打印数据,打印出来的数据的个数跟EventSream面板的数据不一致且有很多重复的数据格式出现在一个输出字符串里,

数据解析

开始怀疑后端的数据有问题,仔细推敲,已经正确的流式的输出在了EventSream面板,应该不是接口的问题,而是我解析的问题,尝试eventsource-parser

import { EventSourceParserStream } from 'eventsource-parser/stream';


  const reader = response?.body
          ?.pipeThrough(new TextDecoderStream())
          .pipeThrough(new EventSourceParserStream())
          .getReader();

这次终于得到了正确的结果,不容易啊。eventsource-parser/stream作用是将sse接口返回的字符串转为对象且避免了debug断点时接口不间断返回的数据被塞到一个字符串的问题

取消请求

取消请求分为两个阶段

一是发起请求,数据未返回阶段

let abortOperation

const controller = new AbortController();
const { signal } = controller;

abortOperation = () => controller.abort()

const response = await fetch(apiPath, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify(apiParams), // 请求体
          signal,
        });

二是数据正在返回阶段

let abortOperation

const controller = new AbortController();
const { signal } = controller;

//取消请求
abortOperation = () => controller.abort()

const response = await fetch(apiPath, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify(apiParams), // 请求体
          signal,
        });
        
...

const reader = response?.body
          ?.pipeThrough(new TextDecoderStream())
          .pipeThrough(new EventSourceParserStream())
          .getReader();
        if (!reader) {
          console.error('Reader is not defined');
          return;
        }
取消请求
abortOperation = reader.cancel();    


兼容application/json格式返回报错

存在未建立eventStream连接时,就返回错误信息

const contentType = response.headers.get('Content-Type');
        if (contentType && contentType.includes('application/json')) {
          const data = await response.json();
          if (!data.success) {
            console.log(data.message)
          }

完整代码

包含状态处理、数据处理

listenToSSE() {
      try {
        const controller = new AbortController();
        const { signal } = controller;

        handleSSEAbort.current = () => {
          stopTalk({ conversationId: conversationInfo.conversationId }).then((res => {
            if (res.success && res.data) {
              setTaskStatus && setTaskStatus(AI_TASK_STATUS.处理成功);
              setChatMessage(preChatMessage => {
                const _preChatMessage = [...preChatMessage];
                _preChatMessage[_preChatMessage.length - 1] = {
                  ..._preChatMessage[_preChatMessage.length - 1],
                  content: '用户取消',
                  id: res.data as number,
                  isLoading: false,
                  error: true,
                };
                return _preChatMessage;
              });
              controller.abort();
            } else {
              message.error(res.msg || '取消失败');
            }
          }));
        };

        const response = await fetch(apiPath, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify(apiParams), // 请求体
          signal,
        });

        if (!response.ok) {
          console.error('Error in SSE response', response.statusText);
          setChatMessage((prevMessages) => {
            const _prevMessages = [...prevMessages];
            _prevMessages[_prevMessages.length - 1] = {
              ..._prevMessages[_prevMessages.length - 1],
              content: response.statusText,
              isLoading: false,
              error: true,
            };
            return _prevMessages;
          });
          setTaskStatus && setTaskStatus(AI_TASK_STATUS.处理成功);
          return;
        }

        // 首先检查响应头中的 'Content-Type',以确定响应类型
        // 兼容会话被删除的情况
        const contentType = response.headers.get('Content-Type');
        if (contentType && contentType.includes('application/json')) {
          const data = await response.json();
          if (!data.success) {
            if (data.code === '-100') {
              message.warning('会话已删除');
              createConversation();
            } else {
              setChatMessage((prevMessages) => {
                const _prevMessages = [...prevMessages];
                _prevMessages[_prevMessages.length - 1] = {
                  ..._prevMessages[_prevMessages.length - 1],
                  isLoading: false,
                  content: data.message || '系统异常',
                  id: Math.random(),
                  error: true,
                };
                return _prevMessages;
              });
            }
          }
          setTaskStatus(AI_TASK_STATUS.处理成功);
          return;
        }

        const reader = response?.body
          ?.pipeThrough(new TextDecoderStream())
          .pipeThrough(new EventSourceParserStream())
          .getReader();
        if (!reader) {
          console.error('Reader is not defined');
          return;
        }

        // 用户取消
        handleSSEAbort.current = () => {
          stopTalk({ conversationId: conversationInfo.conversationId, messageRecordId: chatMessage[chatMessage.length - 1].id }).then((res) => {
            if (res.success && res.data) {
              setTaskStatus && setTaskStatus(AI_TASK_STATUS.处理成功);
              setChatMessage(preChatMessage => {
                const _preChatMessage = [...preChatMessage];
                _preChatMessage[_preChatMessage.length - 1] = {
                  ..._preChatMessage[_preChatMessage.length - 1],
                  isLoading: false,
                };
                return _preChatMessage;
              });
              reader.cancel();
            } else {
              message.error(res.msg || '取消失败');
            }
          });
        };
        const processStream = async () => {
          let first = true;
          try {
            while (!signal.aborted) {
              const { done, value: _value } = await reader.read();
              if (done) break;

              // 检查是否是结束标志
              if (_value.data === '"[DONE]"') {
                // 处理结束流的逻辑
                setChatMessage((prevMessages) => {
                  const _prevMessages = [...prevMessages];
                  _prevMessages[_prevMessages.length - 1] = {
                    ..._prevMessages[_prevMessages.length - 1],
                    isLoading: false,
                  };
                  return _prevMessages;
                });
                setTaskStatus && setTaskStatus(AI_TASK_STATUS.处理成功);
                break;
              } else {
                let jsonData;
                try {
                  jsonData = JSON.parse(_value.data);
                } catch (parseError) {
                  console.error('JSON 解析错误:', parseError);
                  console.log('原始数据:', _value.data);
                  continue; // 跳过这个无效的数据块
                }

                // 处理流式报错
                if (jsonData.data === null && jsonData.success === false) {
                  setChatMessage((prevMessages) => {
                    const _prevMessages = [...prevMessages];
                    _prevMessages[_prevMessages.length - 1] = {
                      ..._prevMessages[_prevMessages.length - 1],
                      isLoading: false,
                      content: jsonData.message,
                      id: jsonData.id,
                      error: true,
                    };
                    return _prevMessages;
                  });
                  setTaskStatus && setTaskStatus(AI_TASK_STATUS.处理成功);
                  break;
                }

                if (first) {
                  setChatMessage((prevMessages) => {
                    const _prevMessages = [...prevMessages];
                    _prevMessages[_prevMessages.length - 1] = {
                      userInput: false,
                      content: (jsonData.choices && jsonData.choices[0]?.delta?.content) || '',
                      sourceSize: jsonData.ext?.sourceSize,
                      contentType: 'text',
                      id: jsonData.id,
                      isLoading: true,
                    };
                    return _prevMessages;
                  });
                  first = false;
                } else {
                  setChatMessage((prevMessages) => {
                    const updatedMessages = prevMessages.map(
                      (msg: ChatMessage) => {
                        if (msg.id === jsonData.id) {
                          // 更新该元素的 content
                          return {
                            ...msg,
                            content:
                                msg.content +
                                (jsonData.choices[0]?.delta?.content || ''),
                          };
                        }
                        return msg;
                      },
                    );
                    return updatedMessages;
                  });
                }
              }
            }
          } catch (err) {
            console.error('处理流时发生错误:', err);
            setTaskStatus && setTaskStatus(AI_TASK_STATUS.处理成功);
            setChatMessage((prevMessages) => {
              const _prevMessages = [...prevMessages];
              _prevMessages[_prevMessages.length - 1] = {
                ..._prevMessages[_prevMessages.length - 1],
                isLoading: false,
              };
              return _prevMessages;
            });
          }
        };

        processStream();
        // 清理工作,例如关闭控制器
      } catch (err: any) {
        containerScrollRef.current?.scrollTo(0, containerScrollRef.current?.scrollHeight);
        if (err.name === 'AbortError') {
          console.log('Request was aborted');
        } else {
          console.error('Error in SSE request', err);
        }
      }
    }

感谢您的阅读,有不足之处请为我指出!