ZNJJ-api-server/test/backends/test_langgraph_streaming.py

"""作业 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