173 lines
7.0 KiB
Python
173 lines
7.0 KiB
Python
import asyncio
|
|
import json
|
|
from contextlib import aclosing
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import httpx
|
|
import pytest
|
|
from langchain_core.messages import AIMessageChunk
|
|
|
|
from src.agent.chat_model import ManagedChatOpenAI
|
|
|
|
from src.agent.llm import LLMAPIError, LLMRateLimitError, OpenAILLMService
|
|
from src.agent.prompts import SINGLE_CONFIRM, render_prompt
|
|
from src.backends.chat import ChatInput, FormUpdate, TextDelta
|
|
from src.backends.langgraph import LangGraphBackend
|
|
from src.services.app_service import AccidentAppService
|
|
from src.services.repository import InMemorySessionRepository
|
|
|
|
|
|
class PatchGraph:
|
|
async def ainvoke(self, *args, **kwargs):
|
|
return {'response_text': '<state>1002</state>ok', 'turn_patch': {'has_injury': False}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_person_synonyms_share_storage_and_prompt_values():
|
|
repo = InMemorySessionRepository()
|
|
service = AccidentAppService(LangGraphBackend(graph=PatchGraph()), repo)
|
|
for index in (1, 2):
|
|
for api, native in ((f'sfzmhm{index}', f'driver{index}ID'), (f'sjhm{index}', f'driver{index}Phonenumber')):
|
|
await service.set_info('s', native, 'old')
|
|
await service.set_info('s', api, 'new')
|
|
assert json.loads(await service.get_info('s', api)) == 'new'
|
|
assert json.loads(await service.get_info('s', native)) == 'new'
|
|
assert api not in await repo.load('s')
|
|
fields = (await service._injected_fields('s'))['person_fields']
|
|
prompt = render_prompt(SINGLE_CONFIRM, {'person_fields': fields})
|
|
assert 'old' not in prompt
|
|
assert 'new' in prompt
|
|
await service.set_info('s', 'driver1ID', 'latest')
|
|
group = json.loads(await service.get_info('s', 'acdhuman1'))
|
|
assert group['driver1ID'] == group['sfzmhm1'] == 'latest'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('stream', [False, True])
|
|
@pytest.mark.parametrize('visible', [False, True])
|
|
async def test_form_visibility_does_not_control_internal_commit(stream, visible):
|
|
backend = LangGraphBackend(graph=PatchGraph())
|
|
service = AccidentAppService(backend, InMemorySessionRepository())
|
|
request = ChatInput('s', 'hello', need_form_update=visible)
|
|
if stream:
|
|
events = [e async for e in service.stream(request)]
|
|
assert any(isinstance(e, FormUpdate) for e in events) is visible
|
|
else:
|
|
result = await service.chat(request)
|
|
assert bool(result.form_update) is visible
|
|
assert json.loads(await service.get_info('s', 'ywrysw')) == '0'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_commit_is_not_lost_when_consumer_closes_after_graph_completes():
|
|
backend = LangGraphBackend(graph=PatchGraph())
|
|
repo = InMemorySessionRepository()
|
|
service = AccidentAppService(backend, repo)
|
|
async with aclosing(service.stream(ChatInput('s', 'hello'))) as events:
|
|
assert isinstance(await anext(events), TextDelta)
|
|
assert json.loads(await service.get_info('s', 'ywrysw')) == '0'
|
|
assert not service._locks._entries
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_factory_disables_sdk_retries():
|
|
llm = OpenAILLMService(model='m', api_key='audit')
|
|
try:
|
|
model = llm._get_client()
|
|
assert model.max_retries == 0
|
|
assert model.root_async_client.max_retries == 0
|
|
assert model.root_client.max_retries == 0
|
|
finally:
|
|
await llm.aclose()
|
|
|
|
|
|
def completion(stream):
|
|
if not stream:
|
|
return httpx.Response(200, json={
|
|
'id': 'r', 'object': 'chat.completion', 'created': 0, 'model': 'm',
|
|
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
|
|
})
|
|
data = {'id': 'r', 'object': 'chat.completion.chunk', 'created': 0, 'model': 'm',
|
|
'choices': [{'index': 0, 'delta': {'content': 'ok'}, 'finish_reason': None}]}
|
|
return httpx.Response(200, headers={'content-type': 'text/event-stream'},
|
|
content='data: '+json.dumps(data)+'\n\ndata: [DONE]\n\n')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('stream', [False, True])
|
|
@pytest.mark.parametrize('status', [429, 503])
|
|
async def test_actual_http_attempt_count_is_not_multiplied_by_sdk(stream, status):
|
|
calls = []
|
|
def handler(request):
|
|
calls.append(json.loads(request.content))
|
|
if len(calls) <= 3:
|
|
return httpx.Response(status, json={'error': {'message': 'retry'}})
|
|
return completion(stream)
|
|
llm = OpenAILLMService(model='m', api_key='audit', max_retries=3,
|
|
extra_body={'reasoning_effort': 'none'})
|
|
llm._client = ManagedChatOpenAI(
|
|
model='m', api_key='audit', base_url='http://audit/v1', max_retries=0,
|
|
use_responses_api=False, stream_usage=False, temperature=None,
|
|
extra_body={'reasoning_effort': 'none'},
|
|
http_client=httpx.Client(trust_env=False),
|
|
http_async_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
|
)
|
|
llm._retry_delay = AsyncMock()
|
|
try:
|
|
result = ''.join([t async for t in llm.stream_tokens('system', 'user')]) if stream else await llm.complete('system', 'user')
|
|
assert result == 'ok'
|
|
assert len(calls) == 4
|
|
assert llm._retry_delay.await_count == 3
|
|
assert all(c['reasoning_effort'] == 'none' for c in calls)
|
|
finally:
|
|
await llm.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_failure_after_first_token_closes_response_without_retry():
|
|
class Response:
|
|
def __init__(self): self.closed = False
|
|
def __aiter__(self): return self.generate()
|
|
async def generate(self):
|
|
yield AIMessageChunk(content='partial')
|
|
raise OSError('broken socket')
|
|
async def aclose(self): self.closed = True
|
|
response = Response()
|
|
create = Mock(return_value=response)
|
|
llm = OpenAILLMService(model='m', api_key='audit')
|
|
llm._client = SimpleNamespace(astream=create)
|
|
tokens = []
|
|
with pytest.raises(LLMAPIError):
|
|
async for token in llm.stream_tokens('system', 'user'):
|
|
tokens.append(token)
|
|
assert tokens == ['partial']
|
|
assert create.call_count == 1
|
|
assert response.closed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('stream', [False, True])
|
|
async def test_retry_budget_exhaustion_is_exact(stream):
|
|
calls = []
|
|
def handler(request):
|
|
calls.append(request)
|
|
return httpx.Response(429, json={'error': {'message': 'retry'}})
|
|
llm = OpenAILLMService(model='m', api_key='audit', max_retries=2)
|
|
llm._retry_delay = AsyncMock()
|
|
llm._client = ManagedChatOpenAI(
|
|
model='m', api_key='audit', base_url='http://audit/v1', max_retries=0,
|
|
use_responses_api=False, stream_usage=False, temperature=None,
|
|
http_client=httpx.Client(trust_env=False),
|
|
http_async_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
|
)
|
|
try:
|
|
with pytest.raises(LLMRateLimitError):
|
|
if stream:
|
|
async for _ in llm.stream_tokens('system', 'user'): pass
|
|
else:
|
|
await llm.complete('system', 'user')
|
|
assert len(calls) == 3
|
|
finally:
|
|
await llm.aclose()
|