ZNJJ-api-server/test/agent/test_llm_stream_tokens.py

"""OpenAILLMService.stream_tokens 测试（模型侧流式）。

mock 客户端返回 token 事件流；验证消息构造与 complete 一致、
错误分类（401/429/超时/连接）与不重试语义。
"""

import asyncio
from types import SimpleNamespace
from unittest.mock import Mock

from langchain_core.messages import AIMessageChunk

import httpx
import pytest
from openai import AuthenticationError, RateLimitError

from src.agent.llm import (
    LLMAPIError,
    LLMAuthError,
    LLMRateLimitError,
    LLMTimeoutError,
    OpenAILLMService,
)


def _event(content):
    return AIMessageChunk(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 = Mock(return_value=_Stream(events))
    llm._client = SimpleNamespace(astream=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 被跳过

    messages = create.call_args.args[0]
    assert messages[0] == {"role": "system", "content": "系统"}
    assert messages[1] == {"role": "assistant", "content": "h"}
    assert messages[-1] == {"role": "user", "content": "输入"}
    assert create.call_args.kwargs["config"]["run_name"] == "llm.stream"


@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 = Mock(return_value=_HangingStream())
    llm._client = SimpleNamespace(astream=create)
    tokens = []
    with pytest.raises(LLMTimeoutError):
        async for token in llm.stream_tokens("系统", "输入"):
            tokens.append(token)
    assert tokens == ["你"]   # 超时前已流出的 token 保留