05 · 流式输出

data: {"type":"delta","text":"你"}

data: {"type":"delta","text":"好"}

data: {"type":"usage","in":120,"out":8}

data: {"type":"error","message":"upstream timeout","trace_id":"a1b2c3d4e5f6"}

data: {"type":"done"}

# sse.py —— SSE 报文自己编一遍。看清楚字节长什么样，坑就少一半
import json


def sse(payload):
    """把一帧编成 SSE 文本。⭐ 硬规则只有一条：一帧以【空行】结束。
    ⭐ 本板块约定：payload 永远是 dict，帧类型放在它的 "type" 字段里。"""
    body = json.dumps(payload, ensure_ascii=False)          # ⭐ 永远 JSON 包一层
    lines = ["data: %s" % ln for ln in body.split("\n")]    # 兜底：万一真有换行就拆成多行
    return "\n".join(lines) + "\n\n"         # ⭐ 两个 \n：一个收本行，一个收本帧


def comment(text=""):
    return ": %s\n\n" % text                 # ⭐ 冒号开头 = 注释帧，用来做心跳


if __name__ == "__main__":
    wire = "".join([sse({"type": "delta", "text": "退"}),
                    comment("ping"),
                    sse({"type": "delta", "text": "第一行\n第二行"}),   # ⭐ 内容里带换行
                    sse({"type": "usage", "in": 120, "out": 340}),
                    sse({"type": "error", "message": "upstream timeout"}),
                    sse({"type": "done"})])
    print(wire.replace("\n", "\\n\n"))       # 把换行显示出来才看得清帧边界
    # ⭐ 看第三帧：换行已被 json.dumps 转义成 \n 两个字符，所以它仍然只占【一帧】

# probe.py —— 流式探针：每秒吐一帧，再自己量【帧间隔】。只用标准库，可直接跑
import http.server, socketserver, threading, time, urllib.request

N, GAP = 5, 1.0            # 吐 5 帧，每帧间隔 1 秒


class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("X-Accel-Buffering", "no")   # ⭐ 请 Nginx 那一层别缓冲
        self.end_headers()
        for i in range(N):
            self.wfile.write(("data: %d\n\n" % i).encode())
            self.wfile.flush()                        # ⭐ 少了这句，本机这层就先憋住了
            time.sleep(GAP)

    def log_message(self, *a):
        pass                                          # 别把访问日志混进输出


def gaps(url):
    t0, marks = time.perf_counter(), []
    with urllib.request.urlopen(url) as r:
        for line in r:                                # 逐行读，不等整体读完
            if line.strip():
                marks.append(time.perf_counter() - t0)
    return [b - a for a, b in zip(marks, marks[1:])]


if __name__ == "__main__":
    socketserver.TCPServer.allow_reuse_address = True
    srv = socketserver.TCPServer(("127.0.0.1", 8123), Handler)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    g = gaps("http://127.0.0.1:8123/")
    print("帧间隔：", ["%.2f" % x for x in g])
    # ⭐ 间隔接近 GAP = 真流式；全接近 0 而总时长仍是 N*GAP = 中间有人在缓冲
    print("判定：", "流式正常" if g and max(g) > GAP / 2 else "有人在缓冲")
    srv.shutdown()

# 🗓️ 未实跑 —— 需要 FastAPI 和真实模型。要看的是【三件事的位置】。
import json, logging, uuid
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

log, app = logging.getLogger("api"), FastAPI()


@app.post("/api/chat")
async def chat(req: Request, prompt: str = "hi"):
    # ⭐ 鉴权、取历史、检索、输入校验全在这里做完 —— 此刻还能好好地返回 4xx/5xx
    tid = uuid.uuid4().hex[:12]                           # ⭐ 就是 03 章那个 trace_id

    async def gen():
        try:
            async for piece in call_model(prompt):        # 你的调用层（04 章）
                if await req.is_disconnected():           # ⭐ 客户端走了就别再烧钱
                    break
                yield "data: %s\n\n" % json.dumps(
                    {"type": "delta", "text": piece}, ensure_ascii=False)
        except Exception:                                 # ⭐ 状态码早是 200 了，错误只能顺流走
            log.exception("stream failed trace=%s", tid)  # ⚠️ 细节只进日志，不进帧
            yield "data: %s\n\n" % json.dumps(            # ⚠️ message 是给人看的一句话
                {"type": "error", "message": "上游调用失败，请重试",
                 "trace_id": tid}, ensure_ascii=False)    # ⭐ 用户能报给你的那个号
        yield 'data: {"type":"done"}\n\n'                 # ⭐ 无论如何都要有结束帧

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache",
                                      "X-Accel-Buffering": "no"})

# cancel.py —— 客户端断开后，上游那条流会不会跟着停？（可直接跑，对照看得见）
import asyncio


class Upstream:
    """假的上游模型流：打开后就在后台自己生成，只有 close() 能让它停。
    真实 SDK 就是这个形状 —— 你这边不读了，对面并不知道。"""

    def __init__(self):
        self.emitted, self._task = 0, None

    def open(self):
        async def produce():
            while True:
                await asyncio.sleep(0.05)
                self.emitted += 1          # ⭐ 每 +1 都是真金白银的 output token
        self._task = asyncio.create_task(produce())
        return self

    def close(self):
        self._task.cancel()

    async def next_piece(self):            # 后端每次从上游取一段往下发
        await asyncio.sleep(0.05)
        return "块%d" % self.emitted


async def bad(up):
    """❌ 没有 finally：任务被取消了，上游还在后台继续生成、继续计费。"""
    while True:
        await up.next_piece()


async def good(up):
    """✅ finally 里关上游：取消、异常、正常结束三条路都会走到。"""
    try:
        while True:
            await up.next_piece()
    finally:
        up.close()                         # ⭐ 这一行就是「别再烧钱了」


async def main():
    for name, handler in (("❌ 没有 finally", bad), ("✅ finally 里 close", good)):
        up = Upstream().open()
        task = asyncio.create_task(handler(up))
        await asyncio.sleep(0.3)
        task.cancel()                      # ⭐ 模拟：用户关了页面，连接断开
        try:
            await task
        except asyncio.CancelledError:
            pass
        n0 = up.emitted
        await asyncio.sleep(0.5)           # 断开后再等半秒，看它还在不在生成
        print("%-18s 断开时已产出 %d 段，半秒后 %d 段" % (name, n0, up.emitted))
        up.close()


asyncio.run(main())