File size: 3,856 Bytes
fa1140b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """流式辅助：增量释放 warmup/guard 双缓冲 + safe_sse_stream 安全包装。
**warmup/guard 双缓冲**（参考 cursor2api streaming-text）：先缓冲预热文本，确认不是应拦截
的前缀（如拒绝文本）再开始释放；释放时永远保留尾部 guard 窗口，给跨 chunk 的清洗规则留
上下文；超过 guard 强制放行。避免「拒绝文本泄漏给客户端」。
**safe_sse_stream**（参考 grok2api）：把 SSE 流中途异常转译成 error chunk，避免 ASGI 在
响应已开始后崩溃（连接已关则吞掉二次错误）。
"""
from __future__ import annotations
from collections.abc import AsyncIterator, Callable
DEFAULT_WARMUP = 96
DEFAULT_GUARD = 256
class IncrementalStreamer:
"""增量释放器：warmup 预热 + guard 尾部缓冲。
- :meth:`push` 返回当前可安全释放给客户端的文本。
- :meth:`finish` 返回剩余全部（含 guard 窗口）。
- ``is_blocked`` 回调在 warmup 阶段判断文本是否为应拦截的前缀（如拒绝）。
"""
def __init__(
self,
*,
warmup: int = DEFAULT_WARMUP,
guard: int = DEFAULT_GUARD,
is_blocked: Callable[[str], bool] | None = None,
) -> None:
self._warmup = warmup
self._guard = guard
self._is_blocked = is_blocked or (lambda _: False)
self._buf: list[str] = []
self._buflen = 0
self._unlocked = False # warmup 通过后置 True
def push(self, chunk: str) -> str:
"""喂入一段文本，返回可释放的部分（可能为空，因 warmup/guard 缓冲）。"""
if not chunk:
return ""
self._buf.append(chunk)
self._buflen += len(chunk)
if not self._unlocked:
if self._buflen < self._warmup:
return ""
head = "".join(self._buf)
if self._is_blocked(head):
# 命中拦截前缀（拒绝）：丢弃已缓冲，不再释放
self._buf.clear()
self._buflen = 0
return ""
self._unlocked = True
# 释放时保留尾部 guard 窗口
full = "".join(self._buf)
if len(full) <= self._guard:
return ""
release_len = len(full) - self._guard
self._buf = [full[release_len:]]
self._buflen = len(self._buf[0])
return full[:release_len]
def finish(self) -> str:
"""收尾：返回剩余全部文本（若全程未过 warmup，最后再判一次拦截前缀）。"""
if not self._unlocked:
head = "".join(self._buf)
if head and self._is_blocked(head):
self._buf.clear()
self._buflen = 0
return ""
self._unlocked = True
out = "".join(self._buf)
self._buf.clear()
self._buflen = 0
return out
def safe_sse_stream(
stream: AsyncIterator[str],
*,
on_error: Callable[[BaseException], list[str]] | None = None,
) -> AsyncIterator[str]:
"""包装 SSE 流：正常 yield；中途异常 → on_error 产出的 chunk（默认通用 OpenAI 风格 error + [DONE]）。
各 adapter 可传入自定义 ``on_error`` 以产出符合自家协议的 error 帧。
"""
async def _gen() -> AsyncIterator[str]:
try:
async for chunk in stream:
yield chunk
except Exception as e: # noqa: BLE001
chunks = on_error(e) if on_error else [
f'data: {{"error":{{"message":"{str(e)[:200]}"}}}}\n\n',
"data: [DONE]\n\n",
]
for c in chunks:
try:
yield c
except Exception: # noqa: BLE001 连接已断，吞掉二次错误
break
return _gen()
|