llm学习——流式传输

技术揭秘：用 Python yield 实现自己的流式 API

在与大语言模型（如ChatGPT）交互时，我们都对它逐字吐出的“打字机效果”印象深刻。这背后的核心技术就是流式传输 (Streaming)。本文将带你了解流式传输的定义，剖析其实现关键 yield，并手把手教你用 Flask 构建一个与 OpenAI 类似效果的流式 API。

1. 流式传输的定义

服务器不等待完整响应生成完毕，而是边生成边将数据块（chunk）发送给客户端的技术。

• 优点：极大降低首字节时间（TTFB），提升用户体验，尤其适合耗时较长的任务。
• 缺点：相比一次性返回，状态管理更复杂。

2. 业界标杆：OpenAI 的 ChatCompletionChunk

ChatCompletionChunk 是 OpenAI API 在流式响应模式下返回的数据单元。它不是一次性返回完整回答，而是将回答拆分成一个个“数据块”持续发送。

• 核心属性：每个 chunk 对象中，最重要的部分是 delta，它代表了“增量内容”，即本次新生成的文字或信息。客户端通过不断拼接 delta 的内容，最终还原出完整回答。

3. 实现核心：Python 的 yield 函数

yield 是理解和构建流式 API 的关键。一个包含 yield 的函数不再是一个普通函数，而是一个生成器 (Generator)。

执行步骤分析

步骤1： 调用包含yield的函数时，不会立即执行，而是返回一个生成器对象。

步骤2： 通过迭代（如 for 循环）或调用 next() 时，函数开始执行，直到遇到第一个 yield。

步骤3： yield 产生一个值并发送出去，然后函数在当前位置**暂停**，并保持所有局部变量的状态。

步骤4： 当再次需要值时，函数从上次暂停处继续执行，直到遇到下一个 yield。

步骤5： 重复步骤3-4，直到函数执行结束。

yield 的“暂停与继续”特性，完美契合了流式传输“生成一块，发送一块”的需求。

---

深入一层：yield 与 next() 的关系

你可能会问：是什么在“驱动”生成器一步步地执行呢？答案就是 next() 函数。

yield 和 next() 是一对紧密的 生产者-消费者 关系：

• yield：是生产者。它在生成器内部生产数据，并设置“暂停点”。
• next()：是消费者。它从外部“拉取”数据，命令生成器从上次暂停的地方开始，执行到下一个 yield 为止。

让我们手动消费一个生成器，来直观地感受一下：

def simple_generator():
    print("准备生产第1个数据...")
    yield "数据1"
    print("准备生产第2个数据...")
    yield "数据2"
    print("生产结束。")

# 得到生成器对象，此时函数内部代码还未执行
gen = simple_generator()

# 调用 next()，驱动生成器执行到第一个 yield
chunk1 = next(gen)  # 输出 "准备生产第1个数据..."
print(f"消费了: {chunk1}") # 输出 "消费了: 数据1"

# 再次调用 next()，从上次暂停处继续，直到第二个 yield
chunk2 = next(gen)  # 输出 "准备生产第2个数据..."
print(f"消费了: {chunk2}") # 输出 "消费了: 数据2"

# 当生成器耗尽，再次调用 next() 会抛出 StopIteration 异常
# next(gen) # 输出 "生产结束。" 然后抛出 StopIteration

那么，在 Flask 中，是谁在调用 next() 呢？
答案是 Web 服务器（WSGI）。当你把生成器交给 Response 对象，Flask 底层的服务器会自动为你处理迭代，它会不断地调用 next()，拿到 yield 出来的数据块，然后立即发送给客户端，直到捕获 StopIteration 异常，便知流已结束。

---

4. 实践：用 Flask 和 yield 构建流式 API

现在，我们将以上概念结合，用 Flask 创建一个模拟打字机效果的流式接口。

import time
from flask import Flask, Response

app = Flask(__name__)

def generate_text_stream():
    """一个模拟生成文本的生成器函数"""
    text = "你好，我是由 Flask 和 yield 构建的流式AI助手。我能逐字为你生成回答。"
    for char in text:
        # yield 将每个字符作为数据块发送
        yield char
        # 模拟生成每个字所需的时间
        time.sleep(0.1)

@app.route('/stream')
def stream():
    """流式响应路由"""
    # Response 对象接收一个生成器函数作为参数
    # mimetype='text/event-stream' 是实现服务器发送事件（SSE）的关键
    return Response(generate_text_stream(), mimetype='text/event-stream')

if __name__ == '__main__':
    app.run(debug=True)

代码解读：

1. generate_text_stream 是一个生成器，每次循环 yield 一个字符。
2. /stream 路由直接将这个生成器作为 Response 的内容返回。
3. Flask (及其底层的服务器) 检测到响应内容是生成器时，便会自动在后台循环调用 next()，将每次 yield 的结果作为数据块发送出去。

5. 前端如何消费流式数据

前端不能用 await response.json() 来一次性接收，而是需要使用 Fetch API 的 ReadableStream 来处理持续到来的数据。

const outputElement = document.getElementById('assistant-output');

fetch('/stream')
  .then(response => {
    // 获取响应的读取器
    const reader = response.body.getReader();
    
    function readStream() {
      reader.read().then(({ done, value }) => {
        if (done) {
          // 流结束
          console.log('Stream complete');
          return;
        }
        
        // 将接收到的 Uint8Array 数据块解码为字符串
        const chunk = new TextDecoder().decode(value);
        
        // 将新字符追加到页面元素上
        outputElement.textContent += chunk;
        
        // 继续读取下一块数据
        readStream();
      });
    }
    
    readStream();
  })
  .catch(error => {
    console.error('Error fetching stream:', error);
  });

总结

通过本文，我们发现构建一个看似复杂的流式 API，其核心原理非常清晰：

后端利用 yield 将任务拆分成多个步骤，创建一个作为“生产者”的生成器；Web 框架则在后台默默扮演“消费者”，通过 next() 不断拉取数据并包装成流式响应；前端则通过 ReadableStream 消费这些数据块，实现局部和渐进式更新。

掌握了 yield 和它背后的 next() 驱动机制，你就掌握了构建高性能、体验友好型 API 的一个强大武器。