ZNJJ-api-server/test/services/test_field_commit_and_retry.py

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()