Python Generator：一个被低估的性能利器

调用 OpenAI 的 API 时，设置 stream=True ，接着 for chunk in completion: 我们就可以“流式”地获取响应的内容。而非等待远程的模型将所有内容生成完毕，再返回给我们（这通常要等很久）。

本文讨论这背后的 Python Generator 。

从一个经典问题开始

假设我们要处理一个超大的日志文件，需要按行读取并分析。传统的做法是：

def read_log_file(filename):
    result = []
    with open(filename) as f:
        for line in f:
            if "ERROR" in line:
                result.append(line)
    return result

# 使用方式
errors = read_log_file("huge.log")
for error in errors:
    process_error(error)

这段代码有什么问题？它会一次性将所有符合条件的行都读入内存。如果日志文件有 10GB，而符合条件的行有 5GB，那么我们的程序就需要 5GB 的内存。

Generator 版本

我们用 Generator 改写一下：

def read_log_file(filename):
    with open(filename) as f:
        for line in f:
            if "ERROR" in line:
                yield line

# 使用方式
for error in read_log_file("huge.log"):
    process_error(error)

看起来很相似，但运行机制完全不同。这个版本无论日志文件多大，内存占用都很小。

Generator 的工作原理

Generator 的核心特点是"懒加载"（lazy evaluation）。当我们调用一个生成器函数时，它并不会立即执行函数体，而是返回一个生成器对象。只有在实际请求下一个值时，它才会执行到下一个 yield 语句。

来看一个更直观的例子：

def counter():
    print("Starting")
    i = 0
    while True:
        print(f"Generating {
     i}")
        yield i
        i += 1

# 创建生成器对象
c = counter()  # 此时不会打印任何内容
print("Generator created")

# 获取前三个值
print(next(c))  # 打印 "Starting" 和 "Generating 0"，返回 0
print(next(c))  # 打印 "Generating 1"，返回 1
print(next(c))  # 打印 "Generating 2"，返回 2

在流式 API 中的应用

现在我们理解了为什么流式 API 会使用 Generator。以 OpenAI 的流式响应为例：

def stream_completion(prompt):
    # 模拟 API 调用
    response = ["生成", "AI", "回复", "需要", "时间"]
    for token in response:
        yield token

# 使用方式
for chunk in stream_completion("你好"