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

import httpx
import pytest
from openai import AsyncOpenAI

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 repo.load('s')
    fields = 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:
        assert llm._get_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 = AsyncOpenAI(api_key='audit', base_url='http://audit/v1', max_retries=0,
                             http_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 SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content='partial'))])
            raise OSError('broken socket')
        async def close(self): self.closed = True
    response = Response()
    create = AsyncMock(return_value=response)
    llm = OpenAILLMService(model='m', api_key='audit')
    llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
    tokens = []
    with pytest.raises(LLMAPIError):
        async for token in llm.stream_tokens('system', 'user'):
            tokens.append(token)
    assert tokens == ['partial']
    assert create.await_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 = AsyncOpenAI(api_key='audit', base_url='http://audit/v1', max_retries=0,
                             http_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()