大模型输出不再“憋大招”：Axum + SSE 流式响应全攻略

土豆1250
743 阅读10分钟

前言

大家好，我是土豆，欢迎关注我的公众号：土豆学前端

随着大型语言模型（LLM）如 GPT 系列的兴起，用户对于实时、流式获取模型输出的需求日益增强。一次性等待模型生成全部内容再返回，体验往往不佳，尤其对于长文本生成任务。Server-Sent Events (SSE) 作为一种轻量级的、基于 HTTP 的单向流式通信协议，非常适合此类场景。本文将详细介绍 SSE 及其在 Rust Axum 框架中的实现。

目录

  1. 什么是 Server-Sent Events (SSE)？
  2. SSE 协议详解
  3. 为什么选择 SSE 用于大模型输出？
  4. 使用 Rust Axum 实现 SSE
  5. 客户端如何消费 SSE？ (JavaScript 示例)
  6. 高级注意事项
  7. 总结

1. 什么是 Server-Sent Events (SSE)？

Server-Sent Events (SSE) 是一种允许服务器向客户端单向推送更新的技术。它建立在单个持久的 HTTP 连接上，服务器可以通过这个连接持续不断地发送数据给客户端。与 WebSockets 不同，SSE 是单向的（服务器到客户端），并且完全基于 HTTP，这使得它在某些场景下更简单、更轻量。

对于大模型而言，当模型逐个 token 或逐个句子生成文本时，SSE 可以立即将这些部分发送给用户，用户可以实时看到文本逐渐“打印”出来，极大地改善了交互体验。

2. SSE 协议详解

SSE 的核心在于其简洁的协议规范。

HTTP 头部

当服务器响应 SSE 请求时，必须包含以下 HTTP 头部：

  • Content-Type: text/event-stream：告知客户端这是一个事件流。
  • Cache-Control: no-cache：确保客户端不会缓存事件流。
  • Connection: keep-alive：保持连接活跃以持续推送事件。

Axum 的 Sse 响应类型会自动处理这些头部。

事件流格式

事件流是纯文本，由一个或多个“事件”组成，每个事件由一个或多个字段行组成，最后以一个空行（\n\n）结束。

常见的字段有：

  • event: (可选) 事件类型。如果未指定，客户端会触发默认的 message 事件。
    • 示例: event: update
  • data: (必需) 事件的数据。可以有多行 data 字段，它们会被拼接起来作为单个数据负载。
    • 示例: data: This is the first line.
    • 示例: data: This is the second line.
  • id: (可选) 事件的唯一 ID。如果客户端断开连接，它会发送最后一个接收到的 id（通过 Last-Event-ID HTTP 头部），服务器可以据此恢复事件流。
    • 示例: id: 123
  • retry: (可选) 客户端在断开连接后尝试重新连接的等待时间（毫秒）。
    • 示例: retry: 10000 (10 秒)

一个完整的事件示例：

id: event-1
event: token
data: 这是
data: 模型生成的第一个词。

id: event-2
event: token
data: 这是第二个词。

event: done
data: {"status": "completed"}

注意：每条字段行以 field: value 的形式出现，并以 \n 结尾。一个完整的事件以两个连续的 \n（即一个空行）结束。

3. 为什么选择 SSE 用于大模型输出？

  • 实时性强：LLM 生成一个 token 或一个片段，就可以立即推送给用户。
  • 简单高效：基于 HTTP，比 WebSockets 更轻量，实现简单，无需复杂的握手和帧处理。对于 LLM 这种服务器单向推送数据的场景非常适合。
  • 自动重连：浏览器 EventSource API 内置了自动重连机制（可以使用 retry 字段控制重连间隔）。
  • 良好兼容性：主流浏览器都原生支持 EventSource API。
  • 易于调试：事件流是纯文本，可以直接在浏览器开发者工具或 curl 中查看。

4. 使用 Rust Axum 实现 SSE

Axum 提供了 axum::response::sse::{Sse, Event} 类型，可以方便地构建 SSE 服务。

项目设置

首先，创建一个新的 Rust 项目并添加必要的依赖：

cargo new axum_sse_llm_example
cd axum_sse_llm_example

编辑 Cargo.toml 文件：

[dependencies]
axum = { version = "0.7", features = ["macros"] }
tokio = { version = "1", features = ["full"] }
futures-util = { version = "0.3", default-features = false, features = ["alloc", "stream"] }
serde = { version = "1.0", features = ["derive"] } # 可选，如果数据是复杂结构
serde_json = "1.0" # 可选，用于序列化数据为 JSON

创建 SSE Handler

Handler 函数需要返回一个实现了 Stream<Item = Result<Event, Infallible>> 的类型，并用 Sse 包装器包裹。Infallible 表示这个流本身不会产生错误（如果产生，通常意味着连接中断）。

// src/main.rs
use axum::{
    response::sse::{Event, Sse, KeepAlive}, // KeepAlive 用于发送心跳包
    routing::get,
    Router,
    extract::State, // 用于共享状态，比如 LLM 客户端
};
use futures_util::stream::{self, Stream, StreamExt}; // StreamExt 用于 map 等操作
use std::{convert::Infallible, time::Duration, sync::Arc};
use tokio::sync::mpsc; // 用于在任务间传递消息

// 模拟 LLM 服务配置或客户端
struct AppState {
    // llm_client: YourLlmClient, // 实际项目中可能是 LLM 客户端
    greeting: String,
}

模拟 LLM 流式输出

我们将创建一个函数，它模拟 LLM 逐块生成文本，并通过 mpsc channel 发送。

// src/main.rs (续)

async fn sse_handler(
    State(state): State<Arc<AppState>>, // 注入共享状态
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    println!("Client connected to SSE endpoint. Greeting: {}", state.greeting);

    // 创建一个 mpsc channel 用于从生成任务传递数据到 SSE 流
    let (tx, mut rx) = mpsc::channel::<String>(10); // Buffer size 10

    // 异步任务模拟 LLM 生成内容
    tokio::spawn(async move {
        // 模拟 LLM 输出分块
        let simulated_llm_output = vec![
            "这是", "一个", "来自", "大语言模型", "的", "实时", "流式", "输出", "示例。",
            "\n\n", "它", "逐块", "发送", "文本，", "就像", "打字机", "一样。",
            "\n", "希望", "你", "喜欢", "这个", "效果！"
        ];

        for (i, chunk) in simulated_llm_output.iter().enumerate() {
            if tx.send(chunk.to_string()).await.is_err() {
                // 接收端已关闭，停止发送
                eprintln!("SSE client disconnected, stopping generation.");
                return;
            }
            // 模拟 LLM 生成每个块的时间间隔
            tokio::time::sleep(Duration::from_millis(200)).await;

            // 可以在这里打印日志，确认发送
            // println!("Sent chunk {}: {}", i + 1, chunk);
        }
        // （可选）发送一个特殊的完成信号
        // tx.send("[DONE]".to_string()).await.ok();
        // 当 tx 被 drop 时，rx.recv() 会返回 None，流自然结束
        println!("LLM simulation finished sending data.");
    });

    // 将 mpsc::Receiver 转换为 Stream
    // 这里我们使用 stream::unfold 来创建流，直到 rx.recv() 返回 None
    let stream = stream::unfold(rx, |mut rx_receiver| async move {
        match rx_receiver.recv().await {
            Some(data) => {
                // if data == "[DONE]" { // 如果使用特殊完成信号
                //     // 发送一个特殊的 "done" 事件
                //     let event = Event::default().event("done").data("Stream completed.");
                //     Some((Ok(event), rx_receiver)) 
                //     // 之后可以返回 None 来结束流，或者让rx_receiver在drop时自动结束
                // } else {
                    let event = Event::default().data(data);
                    Some((Ok(event), rx_receiver))
                // }
            }
            None => {
                // Channel 关闭，流结束
                None
            }
        }
    });

    Sse::new(stream).keep_alive(
        KeepAlive::new()
            .interval(Duration::from_secs(10)) // 每10秒
            .text("keep-alive-text"), // 发送一个注释行作为心跳包
    )
}

解释:

  1. sse_handler: 这是我们的 Axum handler。
  2. mpsc::channel: 创建了一个多生产者单消费者通道。tokio::spawn 中的任务是生产者，SSE 流是消费者。
  3. tokio::spawn: 启动一个独立的异步任务来模拟 LLM 的行为。这个任务会定期发送文本块到 tx。
    • 如果 tx.send() 失败（通常因为客户端断开，rx 被丢弃），任务会优雅退出。
  4. stream::unfold: 这是一个方便的工具，用于从一个初始状态（这里是 rx）逐步生成流项。
    • rx_receiver.recv().await 会等待从通道接收数据。
    • 如果接收到 Some(data)，就创建一个 Event::default().data(data) 并将其作为流的一部分返回，同时也将 rx_receiver 传回以供下次迭代。
    • 如果接收到 None（意味着 tx 端已被丢弃，所有数据都发送完毕），则返回 None，表示流结束。
  5. Sse::new(stream): 将生成的流包装成 SSE 响应。
  6. .keep_alive(...): (可选但推荐) 配置 SSE 的心跳机制。这会定期发送注释行（例如 : keep-alive-text\n\n）到客户端，以防止某些代理或防火墙因长时间无数据活动而关闭连接。

配置路由并运行服务

// src/main.rs (续)

#[tokio::main]
async fn main() {
    let shared_state = Arc::new(AppState {
        greeting: "Hello from Axum SSE LLM!".to_string(),
    });

    let app = Router::new()
        .route("/sse-llm", get(sse_handler))
        .with_state(shared_state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
        .await
        .unwrap();
    println!("Listening on http://127.0.0.1:3000");
    axum::serve(listener, app).await.unwrap();
}

完整示例代码 (src/main.rs)

use axum::{
    response::sse::{Event, Sse, KeepAlive},
    routing::get,
    Router,
    extract::State,
};
use futures_util::stream::{self, Stream}; // StreamExt is not strictly needed for this simple example
use std::{convert::Infallible, time::Duration, sync::Arc};
use tokio::sync::mpsc;

// 模拟 LLM 服务配置或客户端
struct AppState {
    greeting: String,
}

async fn sse_handler(
    State(state): State<Arc<AppState>>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    println!("Client connected to SSE endpoint. Greeting: {}", state.greeting);

    let (tx, mut rx) = mpsc::channel::<String>(10);

    tokio::spawn(async move {
        let simulated_llm_output = vec![
            "这是", "一个", "来自", "大语言模型", "的", "实时", "流式", "输出", "示例。",
            "\n\n", "它", "逐块", "发送", "文本，", "就像", "打字机", "一样。",
            "\n", "希望", "你", "喜欢", "这个", "效果！"
        ];
        let mut event_id_counter = 1;

        for chunk in simulated_llm_output.iter() {
            // 准备发送的数据，可以是纯文本，也可以是 JSON 字符串
            let data_to_send = chunk.to_string();
            // let data_to_send = serde_json::json!({ "token": chunk, "sequence_id": event_id_counter }).to_string();


            if tx.send(data_to_send).await.is_err() {
                eprintln!("SSE client disconnected, stopping generation.");
                return;
            }
            event_id_counter += 1;
            tokio::time::sleep(Duration::from_millis(300)).await;
        }
        // 当 tx 被 drop 时 (函数结束时自动发生), rx.recv() 会返回 None，流自然结束。
        println!("LLM simulation finished sending data. Sender (tx) is being dropped.");
    });

    let stream = stream::unfold(rx, move |mut rx_receiver| async move {
        match rx_receiver.recv().await {
            Some(data) => {
                // 创建一个 Event 对象
                // 你可以设置 event name, id 等
                // let event = Event::default().data(data).event("llm_token").id(id_counter.to_string());
                let event = Event::default().data(data);
                Some((Ok(event), rx_receiver))
            }
            None => {
                // Channel 关闭，流结束
                None
            }
        }
    });

    Sse::new(stream).keep_alive(
        KeepAlive::new()
            .interval(Duration::from_secs(15)) // 每15秒发送一次心跳
            .text("ping"), // 心跳内容将是: `:ping\n\n`
    )
}

#[tokio::main]
async fn main() {
    let shared_state = Arc::new(AppState {
        greeting: "Hello from Axum SSE LLM!".to_string(),
    });

    let app = Router::new()
        .route("/sse-llm", get(sse_handler))
        .with_state(shared_state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
        .await
        .unwrap();
    println!("Server listening on http://127.0.0.1:3000/sse-llm");
    axum::serve(listener, app).await.unwrap();
}

运行项目： cargo run

5. 客户端如何消费 SSE？ (JavaScript 示例)

在客户端，通常使用 EventSource API 来接收 SSE。

创建一个简单的 index.html 文件：

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>SSE LLM Demo</title>
    <style>
        body { font-family: sans-serif; margin: 20px; }
        #output { 
            white-space: pre-wrap; /* Preserve whitespace and newlines */
            border: 1px solid #ccc; 
            padding: 10px; 
            min-height: 100px; 
            background-color: #f9f9f9;
        }
    </style>
</head>
<body>
    <h1>LLM Output (via SSE)</h1>
    <div id="output"></div>

    <script>
        const outputDiv = document.getElementById('output');
        // 注意这里的URL要和你Axum服务配置的路由一致
        const eventSource = new EventSource('http://127.0.0.1:3000/sse-llm');

        eventSource.onopen = function() {
            console.log("Connection to SSE server opened.");
            outputDiv.innerHTML += "连接已建立...\n";
        };

        // 监听默认的 'message' 事件
        // 如果服务器端 Event::default().event("custom_event_name") 设置了事件名
        // 则客户端需要 eventSource.addEventListener("custom_event_name", function(event) { ... });
        eventSource.onmessage = function(event) {
            // event.data 是服务器发送的数据
            console.log("Message received:", event.data);
            // 如果服务器发送的是JSON字符串，需要解析
            // const parsedData = JSON.parse(event.data);
            // outputDiv.innerHTML += parsedData.token;
            outputDiv.innerHTML += event.data;
        };

        eventSource.onerror = function(err) {
            console.error("EventSource failed:", err);
            outputDiv.innerHTML += "\n连接错误或已关闭。\n";
            eventSource.close(); // 出错时通常需要关闭
        };

        // 你也可以监听特定的事件名，如果服务器端发送时指定了 event 字段
        // eventSource.addEventListener('llm_token', function(event) {
        //     const data = JSON.parse(event.data);
        //     console.log("LLM Token:", data);
        //     outputDiv.innerHTML += data.token;
        // });

        // eventSource.addEventListener('done', function(event) {
        //     console.log("Stream finished:", event.data);
        //     outputDiv.innerHTML += "\n--- 流结束 ---";
        //     eventSource.close(); // 完成后关闭连接
        // });

    </script>
</body>
</html>

用浏览器打开此 HTML 文件，你将看到文本逐块显示在页面上。

6. 高级注意事项

  • 错误处理：虽然我们的流声明为 Result<Event, Infallible>，这意味着流本身不应该产生可恢复的错误。真正的 I/O 错误或 panic 会导致连接中断。客户端 EventSource 的 onerror 回调会处理这些。
  • 背压 (Backpressure)：mpsc::channel 有一个缓冲区。如果 LLM 生成速度远快于网络传输或客户端处理速度，缓冲区可能会满，导致 tx.send() 阻塞或报错（取决于通道类型）。在真实应用中，需要考虑如何处理这种情况。
  • 认证与授权：SSE 请求也是 HTTP 请求，可以使用标准的 HTTP 认证机制（如 Cookies, Authorization Header）来保护 SSE 端点。Axum 的中间件系统可以很好地处理这些。
  • 状态管理：如果每个 SSE 连接都需要特定的上下文（例如，基于用户会话的 LLM 对话），你需要一种方式来管理和传递这些状态到 sse_handler 中，State extractor 可以用于此。
  • 优雅关闭 (Graceful Shutdown)：确保在服务器关闭时，正在进行的 SSE 连接能够得到妥善处理，例如发送一个结束信号或允许它们自然完成。Axum 支持优雅关闭。
  • 事件 ID (id 字段)：对于需要断线重连并从上次中断处继续的场景，服务器应为每个事件发送 id。客户端重连时，浏览器会自动在请求头中包含 Last-Event-ID。服务器可以根据这个 ID 决定从哪里继续发送事件。

7. 总结

Server-Sent Events (SSE) 为大语言模型等需要流式输出的应用提供了一种简单、高效且用户友好的解决方案。Rust 配合 Axum 框架，可以轻松构建健壮、高性能的 SSE 服务端。通过 mpsc 通道和异步任务，可以清晰地分离数据生成逻辑和网络推送逻辑，使得代码结构更清晰，易于维护。希望本文能帮助你理解并应用 SSE 技术。

avatar
文章
阅读
粉丝