app/utils/sse.py

"""SSE 流式工具。"""
from __future__ import annotations

import codecs
import json
from typing import Any, AsyncIterator, Iterable, Optional


def sse_data(obj: Any) -> str:
    return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"


def sse_done() -> str:
    return "data: [DONE]\n\n"


def sse_event(event: str, obj: Any) -> str:
    return f"event: {event}\ndata: {json.dumps(obj, ensure_ascii=False)}\n\n"


async def parse_sse_stream(byte_stream: AsyncIterator[bytes]) -> AsyncIterator[dict]:
    """从字节流解析 SSE，yield 每个 data JSON。

    两个关键点（否则中文/长回复会丢字、换行被破坏）：

    1. 用增量 UTF-8 解码器（codecs.getincrementaldecoder）跨 chunk 组装多字节
       字符。若每个 chunk 单独 decode(errors="replace")，被 TCP 分段切断的中文
       字符的前半字节会被替换成 ``?``，造成乱码/丢字。
    2. 按完整 SSE 帧（``\\n\\n`` 分隔）缓冲后再解析。httpx 的 aiter_bytes()
       字节块边界与 SSE 帧边界无关，若按每个 chunk 独立 split('\\n') 解析，
       跨块切断的 ``data: {...}`` 行会 json.loads 失败被静默丢弃。
    """
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
    buffer = ""
    async for chunk in byte_stream:
        if isinstance(chunk, (bytes, bytearray)):
            text = decoder.decode(bytes(chunk))
        else:
            text = str(chunk)
        buffer += text
        # SSE 以双换行分隔
        while "\n\n" in buffer:
            block, buffer = buffer.split("\n\n", 1)
            data_lines = []
            for line in block.split("\n"):
                if line.startswith("data:"):
                    data_lines.append(line[5:].strip())
            if not data_lines:
                continue
            data_str = "\n".join(data_lines)
            if data_str == "[DONE]":
                yield {"__done__": True}
                return
            try:
                yield json.loads(data_str)
            except json.JSONDecodeError:
                yield {"__raw__": data_str}
    # 收尾：flush 增量解码器残余字节，并尝试解析缓冲区里最后一帧
    tail = decoder.decode(b"", final=True)
    if tail:
        buffer += tail
    if buffer.strip():
        data_lines = []
        for line in buffer.split("\n"):
            if line.startswith("data:"):
                data_lines.append(line[5:].strip())
        if data_lines:
            data_str = "\n".join(data_lines)
            if data_str == "[DONE]":
                yield {"__done__": True}
                return
            try:
                yield json.loads(data_str)
            except json.JSONDecodeError:
                yield {"__raw__": data_str}


def _usage_from_sse_block(block: str) -> Optional[dict]:
    """从一个 SSE 帧文本里提取 usage 字段，没有则返回 None。"""
    data_lines = []
    for line in block.split("\n"):
        if line.startswith("data:"):
            data_lines.append(line[5:].strip())
    if not data_lines:
        return None
    data_str = "\n".join(data_lines)
    if data_str == "[DONE]":
        return None
    try:
        obj = json.loads(data_str)
    except json.JSONDecodeError:
        return None
    if not isinstance(obj, dict):
        return None
    u = obj.get("usage")
    return u if isinstance(u, dict) else None


async def passthrough_with_usage(
    byte_stream: AsyncIterator[bytes], state: dict
) -> AsyncIterator[bytes]:
    """透传 SSE 字节流，同时把最后一个 chunk 的 usage 写入 ``state['usage']``。

    用于 OpenAI 透传流式：把字节原样转发给客户端的同时，旁路解析 SSE 帧，
    提取上游在最后一个 chunk 里返回的 usage。需要请求时带
    ``stream_options.include_usage=true`` 上游才会返回 usage。
    """
    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
    buffer = ""
    async for chunk in byte_stream:
        yield chunk
        if isinstance(chunk, (bytes, bytearray)):
            text = decoder.decode(bytes(chunk))
        else:
            text = str(chunk)
        buffer += text
        while "\n\n" in buffer:
            block, buffer = buffer.split("\n\n", 1)
            u = _usage_from_sse_block(block)
            if u is not None:
                state["usage"] = u
    tail = decoder.decode(b"", final=True)
    if tail:
        buffer += tail
    if buffer.strip():
        u = _usage_from_sse_block(buffer)
        if u is not None:
            state["usage"] = u