| """本地 SSE 响应包装（透传 chunk 并在 result 事件时回调）。""" |
| from __future__ import annotations |
|
|
| import json |
| import time |
| from collections.abc import Callable, Iterator |
| from typing import Any |
|
|
| from flask import Response, stream_with_context |
|
|
|
|
| def _notify_sse_block( |
| block: bytes, |
| on_response: Callable[[Any, float, int], None], |
| started: float, |
| status_code: int, |
| ) -> None: |
| for line in block.decode("utf-8", errors="replace").split("\n"): |
| if not line.startswith("data: "): |
| continue |
| try: |
| payload = json.loads(line[6:].strip()) |
| except json.JSONDecodeError: |
| continue |
| if payload.get("type") == "result": |
| on_response(payload.get("data"), time.perf_counter() - started, status_code) |
| elif payload.get("type") == "error": |
| on_response(payload, time.perf_counter() - started, payload.get("status_code", 500)) |
|
|
|
|
| def _tap_chunk_stream( |
| chunks: Iterator, |
| on_response: Callable[[Any, float, int], None] | None, |
| started: float, |
| status_code: int, |
| ): |
| if on_response is None: |
| yield from chunks |
| return |
| pending = b"" |
| for chunk in chunks: |
| if not chunk: |
| continue |
| yield chunk |
| pending += chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") |
| while b"\n\n" in pending: |
| block, pending = pending.split(b"\n\n", 1) |
| _notify_sse_block(block, on_response, started, status_code) |
|
|
|
|
| def wrap_sse_response( |
| response: Response, |
| on_response: Callable[[Any, float, int], None], |
| started: float, |
| ) -> Response: |
| """本地 SSE 响应：透传 chunk 并在 result 事件时回调。""" |
| status_code = response.status_code |
|
|
| def generating(): |
| yield from _tap_chunk_stream( |
| response.response, on_response, started, status_code |
| ) |
|
|
| return Response( |
| stream_with_context(generating()), |
| status=status_code, |
| headers=dict(response.headers), |
| mimetype=response.mimetype, |
| ) |
|
|
