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
@@ -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 保留
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user