Eric Wang
952bffe6f0
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)
231 lines
8.6 KiB
Python
231 lines
8.6 KiB
Python
"""作业 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
|