| """流式辅助：增量释放 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 |
|
|
| 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 |
| |
| 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: |
| 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: |
| break |
|
|
| return _gen() |
|
|
