95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
|
import asyncio
|
||
|
|
from contextlib import aclosing
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from unittest.mock import AsyncMock
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from src.agent.collect import validated_reply
|
||
|
|
from src.agent.llm import LLMTimeoutError, OpenAILLMService
|
||
|
|
from src.backends.chat import ChatInput, TextDelta
|
||
|
|
from src.backends.langgraph import LangGraphBackend
|
||
|
|
from src.services.app_service import AccidentAppService
|
||
|
|
from src.services.repository import InMemorySessionRepository
|
||
|
|
|
||
|
|
|
||
|
|
class Graph:
|
||
|
|
async def ainvoke(self, data, *, config, context):
|
||
|
|
raw, _ = await validated_reply(context.llm, 'system', 'user', '1002')
|
||
|
|
return {'response_text': raw, 'turn_patch': {}}
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_valid_prefix_reaches_client_before_model_finishes():
|
||
|
|
release = asyncio.Event()
|
||
|
|
finished = asyncio.Event()
|
||
|
|
class LLM:
|
||
|
|
async def stream_tokens(self, *args, **kwargs):
|
||
|
|
yield '<sta'
|
||
|
|
yield 'te>1002</state>'
|
||
|
|
await release.wait()
|
||
|
|
yield 'tail'
|
||
|
|
finished.set()
|
||
|
|
backend = LangGraphBackend(graph=Graph(), llm=LLM())
|
||
|
|
async with aclosing(backend.stream(ChatInput('s', 'hello'))) as events:
|
||
|
|
first = await asyncio.wait_for(anext(events), 1)
|
||
|
|
assert first == TextDelta('<state>1002</state>')
|
||
|
|
assert not finished.is_set()
|
||
|
|
release.set()
|
||
|
|
remaining = [e async for e in events]
|
||
|
|
assert ''.join(e.text for e in remaining) == 'tail'
|
||
|
|
assert finished.is_set()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_closing_application_stream_cancels_model_before_unlocking():
|
||
|
|
cancelled = asyncio.Event()
|
||
|
|
class LLM:
|
||
|
|
async def stream_tokens(self, *args, **kwargs):
|
||
|
|
yield '<state>1002</state>'
|
||
|
|
try:
|
||
|
|
await asyncio.Future()
|
||
|
|
finally:
|
||
|
|
cancelled.set()
|
||
|
|
service = AccidentAppService(LangGraphBackend(graph=Graph(), llm=LLM()), InMemorySessionRepository())
|
||
|
|
async with aclosing(service.stream(ChatInput('s', 'hello'))) as events:
|
||
|
|
await asyncio.wait_for(anext(events), 1)
|
||
|
|
assert cancelled.is_set()
|
||
|
|
assert not service._locks._entries
|
||
|
|
assert await service.set_info('s', 'hphm1', 'test') == ('200', '')
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
@pytest.mark.parametrize('stream', [False, True])
|
||
|
|
async def test_timeout_covers_whole_attempt_and_retries_are_bounded(stream):
|
||
|
|
responses = []
|
||
|
|
class HangingResponse:
|
||
|
|
closed = False
|
||
|
|
def __aiter__(self): return self.generate()
|
||
|
|
async def generate(self):
|
||
|
|
await asyncio.Future()
|
||
|
|
yield None
|
||
|
|
async def close(self): self.closed = True
|
||
|
|
async def create(**kwargs):
|
||
|
|
if stream:
|
||
|
|
response = HangingResponse()
|
||
|
|
responses.append(response)
|
||
|
|
return response
|
||
|
|
await asyncio.Future()
|
||
|
|
llm = OpenAILLMService(model='m', api_key='audit', timeout_seconds=0.01, max_retries=2)
|
||
|
|
call = AsyncMock(side_effect=create)
|
||
|
|
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=call)))
|
||
|
|
llm._retry_delay = AsyncMock()
|
||
|
|
# 先加载 SDK 错误类；外层看门狗不测模块惰性导入耗时。
|
||
|
|
assert llm._retryable(TimeoutError())
|
||
|
|
assert isinstance(llm._error(TimeoutError()), LLMTimeoutError)
|
||
|
|
async with asyncio.timeout(10):
|
||
|
|
with pytest.raises(LLMTimeoutError):
|
||
|
|
if stream:
|
||
|
|
async for _ in llm.stream_tokens('system', 'user'): pass
|
||
|
|
else:
|
||
|
|
await llm.complete('system', 'user')
|
||
|
|
assert call.await_count == 3
|
||
|
|
assert llm._retry_delay.await_count == 2
|
||
|
|
assert all(response.closed for response in responses)
|