code/03-llm-apps/streaming.py

"""流式和非流式各调用一次，比较"第一个字出现的时间"和"全部写完的时间"。"""
import os
import time

from openai import OpenAI

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
)
MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")
QUESTION = [{"role": "user", "content": "用大约 200 字介绍 httpx 和 requests 的主要区别。"}]


def non_streaming(thinking):
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL, messages=QUESTION,
        extra_body={"thinking": {"type": "enabled" if thinking else "disabled"}},
    )
    total = time.time() - start
    # 非流式：要等全部生成完才能拿到结果，所以第一个字和最后一个字是同时出现的
    return total, total, response.usage.completion_tokens


def streaming(thinking, show=False):
    start = time.time()
    first_content = None
    stream = client.chat.completions.create(
        model=MODEL, messages=QUESTION, stream=True,
        stream_options={"include_usage": True},  # 让最后一个数据块带上 usage
        extra_body={"thinking": {"type": "enabled" if thinking else "disabled"}},
    )
    usage = None
    for chunk in stream:
        if chunk.usage:
            usage = chunk.usage
        if not chunk.choices:
            continue
        delta = chunk.choices[0].delta
        if delta.content:
            if first_content is None:
                first_content = time.time() - start
            if show:
                print(delta.content, end="", flush=True)
    if show:
        print()
    return first_content, time.time() - start, usage.completion_tokens


print("流式输出的效果：")
streaming(thinking=False, show=True)
print()

for thinking in [False, True]:
    label = "开思考" if thinking else "不思考"
    first, total, tokens = non_streaming(thinking)
    print(f"{label} 非流式：第一个字 {first:.2f} 秒，全部完成 {total:.2f} 秒，输出 {tokens} 词元")
    first, total, tokens = streaming(thinking)
    print(f"{label} 流式：  第一个字 {first:.2f} 秒，全部完成 {total:.2f} 秒，输出 {tokens} 词元")