feat(streaming): real incremental SSE - model tokens flow through graph to client (assignment 8)

Replaces the 'complete answer sliced into chunks' stopgap with token-
level streaming. Graph nodes are untouched: they still consume the
batched LLMService.complete interface and validate full replies.

- agent/llm: OpenAILLMService.stream_tokens() iterates the provider
  stream with the same message construction, temperature and extra_body
  forwarding as complete(); identical error classification (401 ->
  LLMAuthError, 429 -> LLMRateLimitError, timeout via asyncio.timeout
  around the whole iteration); deliberately no retries (tokens already
  emitted cannot be recalled)
- backends/langgraph: _TokenPassThroughLLM wraps the injected LLM for
  stream turns - dialogue calls push tokens to a queue bridge while
  accumulating the full text for node validation; calls marked with the
  【字段提取】 marker take the silent inner.complete path so extraction
  JSON never leaks into the client stream; LLMs without stream_tokens
  fall back to single-chunk push (fake/test compatibility)
- stream(): graph runs in a background task; the generator yields
  TextDelta per token and FormUpdate after the turn; graph-task errors
  re-raise with the existing ChatBackend error mapping; generator close
  (client disconnect) cancels the graph task and the in-flight LLM call
- fixed-answer turns with no LLM call (welcome) fall back to emitting
  the full response text, preserving previous behavior
- endpoint untouched: its buffering parser already handles arbitrary
  chunk boundaries (state tags split across tokens, stage-first order)

Verified: 241 unit tests (11 new streaming tests: token incrementality,
extraction bypass, FormUpdate ordering, error mapping on both paths,
disconnect cancellation, fallback chunking, SSE protocol end-to-end
with per-token text_delta events)
This commit is contained in:
Eric Wang
2026-10-23 00:20:03 +08:00
parent 20e151b097
commit 952bffe6f0
4 changed files with 522 additions and 24 deletions
+63 -1
View File
@@ -9,7 +9,7 @@ fake 与真实 provider 实现同一 ``LLMService`` Protocol,节点不感知
选择实现,见 core.config。
"""
from collections.abc import Awaitable, Callable
from collections.abc import AsyncIterator, Awaitable, Callable
from typing import Protocol
Responder = Callable[[str, str], str | Awaitable[str]]
@@ -196,6 +196,68 @@ class OpenAILLMService:
raise LLMAPIError("LLM 返回了空内容")
return content
async def stream_tokens(
self, system_prompt: str, user_message: str, *,
history: list[dict[str, str]] | None = None,
) -> AsyncIterator[str]:
"""逐 token 产出对话文本(作业 8 真实增量的模型侧)。
语义约束(与 complete 对齐的部分):
- 同样的消息构造(system + role-preserving history + user)、
temperature 与 extra_body 透传;
- 同样的错误分类:401 -> LLMAuthError、429 -> LLMRateLimitError、
超时 -> LLMTimeoutError、其余 -> LLMAPIError。
有意差异:
- **不做重试**——首个 token 发出后重试没有意义(已流出的内容
无法收回);首包前的连接/鉴权错误照样按上述分类抛出,
调用方(透传包装器)决定如何呈现;
- 超时用 asyncio.timeout 包整个迭代(末 token 迟到同样算超时)。
"""
import asyncio
from openai import (
APIConnectionError,
AuthenticationError,
RateLimitError,
)
client = self._get_client()
messages = [
{"role": "system", "content": system_prompt},
*[dict(message) for message in (history or [])],
{"role": "user", "content": user_message},
]
kwargs: dict = {"model": self._model, "messages": messages, "stream": True}
if self._temperature is not None:
kwargs["temperature"] = self._temperature
kwargs.update(self._extra_body)
try:
async with asyncio.timeout(self._timeout_seconds):
response = await client.chat.completions.create(**kwargs)
async for event in response:
choices = getattr(event, "choices", None)
if not choices:
continue
delta = getattr(choices[0], "delta", None)
token = getattr(delta, "content", None)
if token:
yield token
except AuthenticationError as exc:
raise LLMAuthError("LLM provider 认证失败") from exc
except RateLimitError as exc:
raise LLMRateLimitError("LLM provider 限流") from exc
except TimeoutError as exc:
raise LLMTimeoutError(
f"LLM 流式调用超时({self._timeout_seconds}s)"
) from exc
except APIConnectionError as exc:
raise LLMAPIError(f"LLM provider 连接失败:{type(exc).__name__}") from exc
except LLMError:
raise
except Exception as exc:
raise LLMAPIError(f"LLM 流式调用失败:{type(exc).__name__}") from exc
def _default_fake_responder(system_prompt: str, user_message: str) -> str:
"""离线兜底假回复:回显输入,保持服务在无配置时可启动。"""
+111 -23
View File
@@ -3,19 +3,23 @@
实现 ``ChatBackend`` Protocol,把图执行结果翻译为公共契约:
- ``complete()`` 返回 ``ChatResult``,``content`` 保留
``<state>XXXX</state>`` 前缀(API 层负责解析);
- ``stream()`` 产生 ``TextDelta`` / ``FormUpdate``;初期允许把固定完整答案
切成若干片段,真实增量输出在作业 8 接入。
- ``stream()`` 产生 ``TextDelta`` / ``FormUpdate``。
作业 8 起为真实增量:图内对话调用的模型 token 逐条透出,
字段提取等辅助调用经标记旁路(见 _TokenPassThroughLLM)。
通过 ``build_graph()`` 编译的图构造;会话用
``configurable.thread_id = chat_input.session_id`` 绑定。
对应测试:test/backends/test_langgraph_backend.py
对应测试:test/backends/test_langgraph_backend.py、
test/backends/test_langgraph_streaming.py
"""
from collections.abc import AsyncIterator
import asyncio
from collections.abc import AsyncIterator, Callable
from ..agent.collect import _EXTRACTION_MARKER
from ..agent.graph import GraphDeps, build_graph, workflow_fake_llm
from ..agent.llm import LLMAuthError, LLMRateLimitError
from ..agent.llm import LLMAuthError, LLMRateLimitError, LLMService
from .chat import (
ChatBackend,
ChatBackendAuthenticationError,
@@ -27,6 +31,48 @@ from .chat import (
TextDelta,
)
# 流式桥接的队列哨兵:图任务完成后入队,消费端据此收尾
_DONE = object()
class _TokenPassThroughLLM:
"""LLMService 包装器:对话调用逐 token 旁路给 sink,提取调用静默。
图节点仍按批处理接口 ``complete()`` 使用 LLM(校验/解析依赖完整
文本,语义不变);本包装器在收集 token 的同时推给 SSE 桥:
- 系统提示含【字段提取】标记的辅助调用走 inner.complete,
不推流(提取 JSON 不应出现在客户端流里);
- 其余调用若 inner 提供 stream_tokens 则逐 token 推送;
否则退化为整段推送(旧行为,fake/测试 LLM 兼容)。
"""
def __init__(self, inner: LLMService, sink: Callable[[str], None]):
self._inner = inner
self._sink = sink
async def complete(
self, system_prompt: str, user_message: str, *,
history: list[dict[str, str]] | None = None,
) -> str:
if _EXTRACTION_MARKER in system_prompt:
return await self._inner.complete(
system_prompt, user_message, history=history
)
stream_tokens = getattr(self._inner, "stream_tokens", None)
if stream_tokens is None:
text = await self._inner.complete(
system_prompt, user_message, history=history
)
self._sink(text)
return text
chunks: list[str] = []
async for token in stream_tokens(
system_prompt, user_message, history=history
):
chunks.append(token)
self._sink(token)
return "".join(chunks)
class LangGraphBackend(ChatBackend):
"""Adapt a compiled LangGraph to the neutral chat contract."""
@@ -59,9 +105,14 @@ class LangGraphBackend(ChatBackend):
form_update = result.get("turn_patch") or {}
)
async def _complete_turn(self, chat_input: ChatInput) -> dict:
async def _complete_turn(
self, chat_input: ChatInput, llm: LLMService | None = None,
) -> dict:
"""执行一轮图,把 LLM 层错误翻译为后端中立错误。
llm:运行时 LLM 依赖;缺省用构造注入的实例。stream 路径传入
token 透传包装器(_TokenPassThroughLLM)。
错误码等价(对照 FastGPT 后端):
LLMAuthError -> ChatBackendAuthenticationError(API 层 401);
LLMRateLimitError -> ChatBackendRateLimitError(API 层 429)。
@@ -79,7 +130,7 @@ class LangGraphBackend(ChatBackend):
return await self._graph.ainvoke(
graph_input,
config={"configurable": {"thread_id": chat_input.session_id}},
context=GraphDeps(llm=self._llm),
context=GraphDeps(llm=llm or self._llm),
)
except LLMAuthError as exc:
raise ChatBackendAuthenticationError(str(exc)) from exc
@@ -87,21 +138,58 @@ class LangGraphBackend(ChatBackend):
raise ChatBackendRateLimitError(str(exc)) from exc
async def stream(self, chat_input: ChatInput) -> AsyncIterator[ChatStreamEvent]:
"""流式:产生中立流事件。
"""流式(作业 8 真实增量):模型 token 逐条透出。
TODO(作业 2): 实现该方法。最低要求:
- 先产出含完整 ``<state>...</state>`` 前缀的 ``TextDelta``,
再产出正文片段(保证 API 层的 stage 先于正文约束);
- 本阶段允许“完整答案切片段”,提交说明须标注
“尚未接入模型增量输出”(作业 8 改为真实增量)。
结构:后台任务执行一轮图(LLM 经 _TokenPassThroughLLM 把
对话 token 推入队列),本生成器从队列消费并产出
``TextDelta``;图完成后产出 ``FormUpdate``(差量)。
- API 层的缓冲解析(state 标签跨片段拼接、stage 事件优先)
本就支持任意分片,无需改动;
- 生成器被关闭(客户端断开)时取消图任务,停止 LLM 调用;
- 图任务异常在消费端重新抛出(错误映射沿用 _complete_turn)。
"""
result = await self.complete(chat_input)
prefix, sep, body = result.content.partition("</state>")
if not sep:
yield TextDelta(result.content)
return
yield TextDelta(prefix + "</state>")
if result.form_update:
yield FormUpdate(result.form_update)
if body:
yield TextDelta(body)
queue: asyncio.Queue = asyncio.Queue()
token_count = {"n": 0} # sink 同步回调不可直接闭包递增,用盒
def sink(token: str) -> None:
# LLM 在图任务的事件循环线程内同步回调
token_count["n"] += 1
queue.put_nowait(token)
async def run_turn() -> dict:
try:
result = await self._complete_turn(
chat_input, llm=_TokenPassThroughLLM(self._llm, sink)
)
except BaseException as exc:
queue.put_nowait(exc)
return {}
queue.put_nowait(_DONE)
return result
task = asyncio.ensure_future(run_turn())
try:
while True:
item = await queue.get()
if item is _DONE:
break
if isinstance(item, BaseException):
raise item
yield TextDelta(item)
result = await task
if token_count["n"] == 0:
# 本回合无 LLM 调用(welcome 等固定应答轮):
# sink 无 token 可推,回退为完整文本切片(旧行为)
text = result.get("response_text") or ""
if text:
yield TextDelta(text)
form_update = result.get("turn_patch") or {}
if form_update:
yield FormUpdate(form_update)
finally:
if not task.done():
task.cancel()
try:
await task
except (asyncio.CancelledError, Exception):
pass
+118
View File
@@ -0,0 +1,118 @@
"""OpenAILLMService.stream_tokens 测试(模型侧流式)。
mock 客户端返回 token 事件流;验证消息构造与 complete 一致、
错误分类(401/429/超时/连接)与不重试语义。
"""
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
import httpx
import pytest
from openai import AuthenticationError, RateLimitError
from src.agent.llm import (
LLMAPIError,
LLMAuthError,
LLMRateLimitError,
LLMTimeoutError,
OpenAILLMService,
)
def _event(content):
return SimpleNamespace(
choices=[SimpleNamespace(delta=SimpleNamespace(content=content))]
)
def _service_with_stream(events):
llm = OpenAILLMService(
model="test-model", api_key="test", base_url="http://llm.local/v1",
timeout_seconds=5, temperature=0, extra_body={"reasoning_effort": "none"},
)
class _Stream:
def __init__(self, items):
self._items = list(items)
def __aiter__(self):
return self._aiter()
async def _aiter(self):
for item in self._items:
yield item
create = AsyncMock(return_value=_Stream(events))
llm._client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
)
return llm, create
@pytest.mark.asyncio
async def test_stream_tokens_yields_incrementally_and_forwards_kwargs():
llm, create = _service_with_stream([_event("你"), _event(""), _event("好")])
tokens = [
token async for token in
llm.stream_tokens("系统", "输入", history=[{"role": "assistant", "content": "h"}])
]
assert tokens == ["你", "好"] # 空 delta 被跳过
kwargs = create.call_args.kwargs
assert kwargs["model"] == "test-model"
assert kwargs["stream"] is True
assert kwargs["temperature"] == 0
assert kwargs["reasoning_effort"] == "none" # extra_body 透传
assert kwargs["messages"][0] == {"role": "system", "content": "系统"}
assert kwargs["messages"][-1] == {"role": "user", "content": "输入"}
@pytest.mark.asyncio
async def test_stream_tokens_maps_authentication_error():
request = httpx.Request("POST", "http://llm.local/v1/chat/completions")
exc = AuthenticationError("bad key", response=httpx.Response(401, request=request), body=None)
llm, create = _service_with_stream([])
create.side_effect = exc
with pytest.raises(LLMAuthError):
async for _ in llm.stream_tokens("系统", "输入"):
pass
@pytest.mark.asyncio
async def test_stream_tokens_maps_rate_limit_error():
request = httpx.Request("POST", "http://llm.local/v1/chat/completions")
exc = RateLimitError("slow down", response=httpx.Response(429, request=request), body=None)
llm, create = _service_with_stream([])
create.side_effect = exc
with pytest.raises(LLMRateLimitError):
async for _ in llm.stream_tokens("系统", "输入"):
pass
@pytest.mark.asyncio
async def test_stream_tokens_times_out_on_stale_stream():
"""末 token 迟到同样算超时(asyncio.timeout 包整个迭代)。"""
class _HangingStream:
def __aiter__(self):
return self._gen()
async def _gen(self):
yield _event("你")
await asyncio.sleep(30) # 超过 timeout_seconds=5
yield _event("好")
llm = OpenAILLMService(model="m", api_key="k", timeout_seconds=0.2)
create = AsyncMock(return_value=_HangingStream())
llm._client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
)
tokens = []
with pytest.raises(LLMTimeoutError):
async for token in llm.stream_tokens("系统", "输入"):
tokens.append(token)
assert tokens == ["你"] # 超时前已流出的 token 保留
+230
View File
@@ -0,0 +1,230 @@
"""作业 8 真实增量流式测试:模型 token 逐条透出到 SSE 事件。
验证点:
- 对话调用的 token 逐条到达(不是整段切片);
- 字段提取辅助调用(【字段提取】标记)不推流;
- 图完成后产出 FormUpdate(差量);
- 图任务异常(LLM 鉴权/限流)在流路径正确映射;
- 客户端断开时图任务被取消(停止 LLM 调用);
- 无 stream_tokens 的 LLM(fake 兼容路径)退化为整段推送。
"""
import asyncio
from types import SimpleNamespace
import pytest
from src.agent.llm import (
LLMAuthError,
LLMRateLimitError,
LLMService,
)
from src.backends.chat import (
ChatBackendAuthenticationError,
ChatBackendRateLimitError,
ChatInput,
FormUpdate,
TextDelta,
)
from src.backends.langgraph import LangGraphBackend
from src.agent.graph import build_graph
class _Event:
def __init__(self, content):
self.choices = [SimpleNamespace(delta=SimpleNamespace(content=content))]
class StreamingFakeLLM:
"""带 stream_tokens 的测试 LLM:对话逐 token,提取返回 JSON。"""
def __init__(self, reply_tokens=None, extraction_json="{}"):
self._reply_tokens = reply_tokens or ["<state>1002</state>", "你", "好"]
self._extraction_json = extraction_json
self.complete_calls = [] # inner.complete 实际收到的调用
self.streamed_sessions = [] # 推过流的 system_prompt
async def complete(self, system_prompt, user_message, *, history=None):
self.complete_calls.append(system_prompt)
if "【字段提取】" in system_prompt:
return self._extraction_json
return "".join(self._reply_tokens)
async def stream_tokens(self, system_prompt, user_message, *, history=None):
self.streamed_sessions.append(system_prompt)
for token in self._reply_tokens:
await asyncio.sleep(0) # 让出调度,模拟真实流式
yield token
def make_input(session="stream-s1", text="继续办理"):
return ChatInput(session_id=session, text=text, client_mode="direct")
async def _second_turn_events(backend, session="stream-s1"):
"""先完成 welcome 轮(固定应答、无 LLM 调用),再收集第二轮的流事件。"""
await backend.complete(make_input(session, "新对话"))
return [
event async for event in backend.stream(make_input(session, "继续办理"))
]
@pytest.mark.asyncio
async def test_stream_emits_tokens_incrementally():
"""对话 token 逐条到达(多片段、非整段),且 state 前缀在流中。
首轮 welcome 是固定应答(无 LLM 调用),走整段回退;
第二轮进入采集、LLM 被调用,才是真实增量路径。
"""
backend = LangGraphBackend(llm=StreamingFakeLLM())
events = await _second_turn_events(backend)
deltas = [event for event in events if isinstance(event, TextDelta)]
# 逐 token:至少 3 个片段(旧实现只有 2 个:前缀+正文)
assert len(deltas) >= 3
joined = "".join(event.text for event in deltas)
assert joined == "<state>1002</state>你好"
assert deltas[0].text.startswith("<state>")
@pytest.mark.asyncio
async def test_extraction_call_does_not_stream():
"""字段提取调用(【字段提取】标记)走 inner.complete,不进客户端流。"""
llm = StreamingFakeLLM(extraction_json='{"has_injury": false}')
backend = LangGraphBackend(llm=llm)
events = await _second_turn_events(backend)
streamed = "".join(
event.text for event in events if isinstance(event, TextDelta)
)
assert "has_injury" not in streamed # 提取 JSON 不漏出
assert "【字段提取】" not in streamed
assert any("【字段提取】" in call for call in llm.complete_calls)
# 提取标记的调用没有走 stream_tokens
assert all(
"【字段提取】" not in prompt for prompt in llm.streamed_sessions
)
@pytest.mark.asyncio
async def test_form_update_emitted_after_stream():
"""差量(turn_patch)在文本流之后产出。"""
llm = StreamingFakeLLM(extraction_json='{"has_injury": false, "vehicle_count": 2}')
backend = LangGraphBackend(llm=llm)
events = await _second_turn_events(backend)
kinds = [type(event).__name__ for event in events]
assert kinds[-1] == "FormUpdate"
update = events[-1].data
assert update.get("has_injury") is False
@pytest.mark.asyncio
async def test_stream_maps_llm_errors():
class _Failing(LLMService):
async def complete(self, *a, **k):
raise AssertionError("不应走批处理")
async def stream_tokens(self, *a, **k):
raise LLMAuthError("bad key")
yield # pragma: no cover - 保持生成器形态
backend = LangGraphBackend(llm=_Failing())
await backend.complete(make_input("err-s1", "新对话")) # welcome 轮
with pytest.raises(ChatBackendAuthenticationError):
async for _ in backend.stream(make_input("err-s1", "继续办理")):
pass
class _Failing429(LLMService):
async def complete(self, *a, **k):
raise AssertionError("不应走批处理")
async def stream_tokens(self, *a, **k):
raise LLMRateLimitError("限流")
yield
backend = LangGraphBackend(llm=_Failing429())
await backend.complete(make_input("err-s2", "新对话"))
with pytest.raises(ChatBackendRateLimitError):
async for _ in backend.stream(make_input("err-s2", "继续办理")):
pass
@pytest.mark.asyncio
async def test_client_disconnect_cancels_graph_task():
"""生成器被关闭后,图任务(LLM 调用)随之取消。
首轮 welcome 是固定应答、不调 LLM;先完成一轮,再对第二轮
(进入采集、会调 LLM)断开。
"""
started = asyncio.Event()
cancelled = asyncio.Event()
class _SlowLLM:
async def complete(self, system_prompt, user_message, *, history=None):
return "<state>1002</state>你好"
async def stream_tokens(self, system_prompt, user_message, *, history=None):
started.set()
try:
await asyncio.sleep(60)
yield "不应到达"
except asyncio.CancelledError:
cancelled.set()
raise
backend = LangGraphBackend(llm=_SlowLLM())
await backend.complete(make_input("disc-s1", "新对话")) # welcome 轮
generator = backend.stream(make_input("disc-s1", "继续办理"))
# 启动消费(endpoint 开始读取 SSE),等 LLM 流式开始
pending = asyncio.ensure_future(generator.__anext__())
await asyncio.wait_for(started.wait(), timeout=5)
# 模拟客户端断开:取消挂起的读取 -> 生成器 finally 取消图任务
pending.cancel()
await asyncio.wait_for(cancelled.wait(), timeout=5)
@pytest.mark.asyncio
async def test_llm_without_stream_tokens_falls_back_to_single_chunk():
"""无 stream_tokens 的 LLM(fake 兼容)退化为整段推送,语义不炸。"""
from src.agent.llm import FakeLLMService
backend = LangGraphBackend(
llm=FakeLLMService(lambda s, u: "<state>1002</state>你好")
)
events = await _second_turn_events(backend, "fb-s1")
deltas = [event for event in events if isinstance(event, TextDelta)]
assert "".join(event.text for event in deltas) == "<state>1002</state>你好"
@pytest.mark.asyncio
async def test_streaming_backend_through_endpoint_sse_protocol():
"""真实流式 backend 过 endpoint SSE 协议层:token 级分片被正确重组,
stage_code 事件先于正文(API 层缓冲解析对任意分片的兼容性验证)。
"""
from src.api.endpoints import chat
from src.schemas.models import ProcessRequest_chat
backend = LangGraphBackend(llm=StreamingFakeLLM())
await backend.complete(make_input("e2e-s1", "新对话")) # welcome 轮
response = await chat(
ProcessRequest_chat(
sessionId="e2e-s1", timeStamp="20261022120000",
text="继续办理", clientMode="direct",
),
stream=True, backend=backend,
)
chunks = []
async for chunk in response.body_iterator:
chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk)
body = "".join(chunks)
assert body.index("event: stage_code") < body.index("event: text_delta")
assert '"nextStageCode": "1002"' in body
# token 级流式:"你"、"好" 作为独立 text_delta 事件逐条到达
assert body.count("event: text_delta") == 2
assert '{"text": "你"}' in body
assert '{"text": "好"}' in body
assert "event: done" in body