流式响应

演示如何处理 LLM 的流式（Server-Sent Events）响应：逐 token 产出、聚合、 以及解析单条 SSE 数据块。流式输出能显著降低首 token 延迟，改善用户体验。 对应教程：第3卷《OpenAI API》第2章 流式响应。

文件路径: examples/volume-03-openai-api\streaming.py

---

源代码

"""流式响应示例。

演示如何处理 LLM 的流式（Server-Sent Events）响应：逐 token 产出、聚合、
以及解析单条 SSE 数据块。流式输出能显著降低首 token 延迟，改善用户体验。

对应教程：第3卷《OpenAI API》第2章 流式响应。
"""

from __future__ import annotations

from collections.abc import Iterator
from typing import Any, Protocol


class StreamClient(Protocol):
    """支持流式的聊天补全客户端协议。"""

    @property
    def chat(self) -> Any: ...


def parse_sse_data(line: str) -> str | None:
    """解析单行 SSE 数据。

    Args:
        line: 一行 SSE 文本，形如 ``data: {"...": "..."}`` 或 ``data: [DONE]``。

    Returns:
        提取到的 JSON 字符串载荷；若为 ``[DONE]`` 返回 ``None``；非 data 行返回 ``None``。
    """
    if not line or not line.startswith("data:"):
        return None
    payload = line[len("data:") :].strip()
    if payload == "[DONE]" or payload == "":
        return None
    return payload


def extract_delta_text(chunk: Any) -> str:
    """从流式 chunk 中提取增量文本。

    兼容两种结构：
        - ``chunk.choices[0].delta.content``
        - ``chunk.choices[0].text``（旧版 completions）

    Args:
        chunk: 流式响应的单个 chunk。

    Returns:
        增量文本，若不存在则返回空字符串。
    """
    choices = getattr(chunk, "choices", None)
    if not choices:
        return ""
    delta = getattr(choices[0], "delta", None)
    if delta is not None:
        return getattr(delta, "content", "") or ""
    text = getattr(choices[0], "text", None)
    return text or ""


def simulate_stream(tokens: list[str]) -> Iterator[str]:
    """模拟流式输出：逐个产出 token。

    Args:
        tokens: token 列表。

    Yields:
        单个 token 字符串。
    """
    yield from tokens


def aggregate_stream(chunks: Iterator[Any]) -> str:
    """聚合流式 chunk 为完整文本。

    Args:
        chunks: chunk 迭代器。

    Returns:
        拼接后的完整文本。
    """
    return "".join(extract_delta_text(c) for c in chunks)


def stream_chat(
    client: StreamClient,
    prompt: str,
    model: str = "gpt-4o-mini",
) -> str:
    """以流式方式调用聊天补全并返回完整文本。

    Args:
        client: 符合 ``StreamClient`` 协议的客户端。
        prompt: 用户输入。
        model: 模型名称。

    Returns:
        聚合后的完整回复文本。
    """
    stream = client.chat.completions.create(  # type: ignore[union-attr]
        model=model,
        messages=[{"role": "user", "content": prompt}],
        stream=True,
    )
    return aggregate_stream(iter(stream))


def main() -> None:
    """演示纯模拟的流式聚合（不调用真实 API）。"""
    # 模拟服务端按 token 推送
    fake_tokens = ["你", "好", "，", "我", "是", "AI", "助手", "。"]
    print("模拟流式输出：")
    full = ""
    for token in simulate_stream(fake_tokens):
        full += token
        print(f"  收到: {token!r}  累积: {full}")

    # 演示 SSE 解析
    print("\nSSE 解析：")
    for line in ['data: {"id": "1"}', "data: [DONE]", ": comment", "data: "]:
        print(f"  {line!r} -> {parse_sse_data(line)!r}")


if __name__ == "__main__":
    main()