feat(agent): use native parallel nodes and early reply streaming
This commit is contained in:
Eric Wang
@@ -0,0 +1,298 @@
|
||||
"""并行节点和原生流事件的时序、提交与取消契约。"""
|
||||
|
||||
import asyncio
|
||||
from contextlib import aclosing
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from src.agent.collect import CollectReplyError
|
||||
from src.agent.graph import build_graph
|
||||
from src.agent.llm import OpenAILLMService
|
||||
from src.api import endpoints
|
||||
from src.backends.chat import ChatInput, FormUpdate, ReplyFinished, TextDelta
|
||||
from src.backends.langgraph import LangGraphBackend
|
||||
from src.core.config import Config
|
||||
from src.schemas.models import ProcessRequest_chat
|
||||
from src.services.app_service import AccidentAppService
|
||||
from src.storage.runtime import StorageSettings, open_storage
|
||||
from src.services.repository import InMemorySessionRepository
|
||||
|
||||
|
||||
class ParallelLLM:
|
||||
def __init__(self, patch=None):
|
||||
self.patch = patch if patch is not None else {'has_injury': False}
|
||||
self.extraction_started = asyncio.Event()
|
||||
self.reply_started = asyncio.Event()
|
||||
self.release_extraction = asyncio.Event()
|
||||
self.extraction_cancelled = asyncio.Event()
|
||||
self.reply_done = asyncio.Event()
|
||||
|
||||
async def complete(self, system, user, **kwargs):
|
||||
if '【字段提取】' in system:
|
||||
self.extraction_started.set()
|
||||
try:
|
||||
await self.release_extraction.wait()
|
||||
except asyncio.CancelledError:
|
||||
self.extraction_cancelled.set()
|
||||
raise
|
||||
return json.dumps(self.patch)
|
||||
self.reply_started.set()
|
||||
# 若是串行实现，会卡在这里；两个原生图节点必须同时运行。
|
||||
await self.extraction_started.wait()
|
||||
self.reply_done.set()
|
||||
return '<state>1002</state>请问您还在事故现场吗？'
|
||||
|
||||
async def stream_tokens(self, system, user, **kwargs):
|
||||
self.reply_started.set()
|
||||
await self.extraction_started.wait()
|
||||
yield '<state>1002</state>'
|
||||
yield '请问您还在事故现场吗？'
|
||||
self.reply_done.set()
|
||||
|
||||
|
||||
async def service_for(llm):
|
||||
repository = InMemorySessionRepository()
|
||||
backend = LangGraphBackend(llm=llm)
|
||||
service = AccidentAppService(backend, repository)
|
||||
await service.chat(ChatInput('session', 'start'))
|
||||
return service, repository
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_stream_nodes_fan_out_before_join_and_commit():
|
||||
llm = ParallelLLM()
|
||||
service, repo = await service_for(llm)
|
||||
task = asyncio.create_task(service.chat(ChatInput('session', '继续')))
|
||||
try:
|
||||
await asyncio.wait_for(llm.reply_done.wait(), 3)
|
||||
assert llm.extraction_started.is_set()
|
||||
assert not task.done()
|
||||
assert await repo.load('session') == {}
|
||||
llm.release_extraction.set()
|
||||
await asyncio.wait_for(task, 3)
|
||||
assert (await repo.load('session'))['ywrysw'] is False
|
||||
finally:
|
||||
llm.release_extraction.set()
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reply_finished_precedes_slow_extraction_and_commit():
|
||||
llm = ParallelLLM()
|
||||
service, repo = await service_for(llm)
|
||||
events = []
|
||||
async with aclosing(service.stream(ChatInput('session', '继续', need_form_update=True))) as source:
|
||||
while True:
|
||||
event = await asyncio.wait_for(anext(source), 3)
|
||||
events.append(event)
|
||||
if isinstance(event, ReplyFinished):
|
||||
break
|
||||
assert llm.reply_done.is_set()
|
||||
assert not llm.release_extraction.is_set()
|
||||
assert all(not isinstance(event, FormUpdate) for event in events)
|
||||
assert await repo.load('session') == {}
|
||||
llm.release_extraction.set()
|
||||
events.extend([event async for event in source])
|
||||
assert isinstance(events[-1], FormUpdate)
|
||||
assert (await repo.load('session'))['ywrysw'] is False
|
||||
assert not service._locks._entries
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('patch', [{}, {'has_injury': False}])
|
||||
@pytest.mark.parametrize('visible', [False, True])
|
||||
async def test_sse_flushes_final_sentence_before_extraction_with_unchanged_public_events(monkeypatch, patch, visible):
|
||||
llm = ParallelLLM(patch)
|
||||
service, repo = await service_for(llm)
|
||||
monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph')
|
||||
monkeypatch.setattr(endpoints, 'get_app_service', lambda: service)
|
||||
response = await endpoints.chat(
|
||||
ProcessRequest_chat(sessionId='session', timeStamp='probe', text='继续',
|
||||
needFormUpdate=visible, useTextChunk=True),
|
||||
stream=True, backend=service._backend, idempotency_key=None,
|
||||
)
|
||||
received = []
|
||||
async with aclosing(response.body_iterator) as source:
|
||||
while True:
|
||||
event = await asyncio.wait_for(anext(source), 3)
|
||||
received.append(event)
|
||||
if 'event: text_delta' in event:
|
||||
break
|
||||
assert '请问您还在事故现场吗？' in received[-1]
|
||||
assert llm.extraction_started.is_set()
|
||||
assert not llm.release_extraction.is_set()
|
||||
assert not any('event: done' in event for event in received)
|
||||
assert await repo.load('session') == {}
|
||||
llm.release_extraction.set()
|
||||
received.extend([event async for event in source])
|
||||
body = ''.join(received)
|
||||
assert body.count('event: text_delta') == 1
|
||||
assert body.count('event: done') == 1
|
||||
assert ('event: formUpdate' in body) is (visible and bool(patch))
|
||||
assert 'reply_finished' not in body
|
||||
if not visible:
|
||||
assert 'has_injury' not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closing_after_reply_finished_cancels_extraction_before_session_unlock():
|
||||
llm = ParallelLLM()
|
||||
service, repo = await service_for(llm)
|
||||
async with aclosing(service.stream(ChatInput('session', '继续'))) as source:
|
||||
while not isinstance(await asyncio.wait_for(anext(source), 3), ReplyFinished):
|
||||
pass
|
||||
assert llm.extraction_cancelled.is_set()
|
||||
assert not service._locks._entries
|
||||
assert await repo.load('session') == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_full_reply_never_emits_finished_or_commits_and_cancels_parallel_extraction():
|
||||
class Invalid(ParallelLLM):
|
||||
async def stream_tokens(self, *args, **kwargs):
|
||||
await self.extraction_started.wait()
|
||||
yield '<state>1002</state>partial'
|
||||
yield '<state>9999</state>'
|
||||
llm = Invalid()
|
||||
service, repo = await service_for(llm)
|
||||
events = []
|
||||
with pytest.raises(CollectReplyError):
|
||||
async for event in service.stream(ChatInput('session', '继续')):
|
||||
events.append(event)
|
||||
assert any(isinstance(event, TextDelta) for event in events)
|
||||
assert not any(isinstance(event, ReplyFinished) for event in events)
|
||||
assert llm.extraction_cancelled.is_set()
|
||||
assert await repo.load('session') == {}
|
||||
assert not service._locks._entries
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_graph_nodes_are_visible_and_join_once():
|
||||
llm = ParallelLLM()
|
||||
llm.release_extraction.set()
|
||||
graph = build_graph()
|
||||
from src.agent.graph import GraphDeps
|
||||
names = []
|
||||
async for event in graph.astream_events(
|
||||
{'session_id': 'native', 'input_text': '继续', 'stage_code': '1002'},
|
||||
config={'configurable': {'thread_id': 'native'}}, context=GraphDeps(llm=llm), version='v2',
|
||||
):
|
||||
if event['event'] == 'on_chain_end':
|
||||
names.append(event['name'])
|
||||
assert names.count('collect') == names.count('extract_accident_fields') == names.count('merge_collect') == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_confirmation_reply_finishes_before_its_dependent_extraction():
|
||||
class LLM:
|
||||
def __init__(self):
|
||||
self.release = asyncio.Event()
|
||||
async def complete(self, *args, **kwargs):
|
||||
await self.release.wait()
|
||||
return '{}'
|
||||
async def stream_tokens(self, *args, **kwargs):
|
||||
yield '<state>3002</state>请问您的姓名是什么？'
|
||||
llm = LLM()
|
||||
backend = LangGraphBackend(llm=llm)
|
||||
async with aclosing(backend.stream(ChatInput(
|
||||
'confirm', '是的', injected_fields={'stage_code': '3002'},
|
||||
))) as source:
|
||||
received = []
|
||||
while True:
|
||||
event = await asyncio.wait_for(anext(source), 3)
|
||||
received.append(event)
|
||||
if isinstance(event, ReplyFinished):
|
||||
break
|
||||
assert ''.join(event.text for event in received if isinstance(event, TextDelta)) == '<state>3002</state>请问您的姓名是什么？'
|
||||
assert not llm.release.is_set()
|
||||
llm.release.set()
|
||||
assert [event async for event in source] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sqlite_disconnect_after_reply_finishes_preserves_committed_snapshot(tmp_path):
|
||||
llm = ParallelLLM()
|
||||
settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'parallel.sqlite3'), auto_migrate=True)
|
||||
async with open_storage(settings) as storage:
|
||||
backend = LangGraphBackend(build_graph(storage.checkpointer), llm=llm)
|
||||
service = AccidentAppService(backend, storage.repository, locks=storage.locks)
|
||||
await service.chat(ChatInput('session', 'start'))
|
||||
async with aclosing(service.stream(ChatInput('session', '继续', request_key='interrupted'))) as source:
|
||||
while not isinstance(await asyncio.wait_for(anext(source), 3), ReplyFinished):
|
||||
pass
|
||||
assert llm.extraction_cancelled.is_set()
|
||||
assert not storage.repository.database.connection.in_transaction
|
||||
assert not storage.checkpointer.conn.in_transaction
|
||||
async with storage.repository.database.transaction() as conn:
|
||||
row = await conn.one('SELECT version,graph_state FROM agent_sessions WHERE session_id=?', ('session',))
|
||||
assert row['version'] == 1
|
||||
assert json.loads(row['graph_state'])['turn_count'] == 1
|
||||
llm.release_extraction.set()
|
||||
await service.chat(ChatInput('session', '继续', request_key='interrupted'))
|
||||
assert (await storage.repository.load('session'))['ywrysw'] is False
|
||||
assert not service._locks._entries
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_commit_failure_after_early_reply_still_emits_error_without_form_update_or_done(monkeypatch):
|
||||
llm = ParallelLLM()
|
||||
service, repo = await service_for(llm)
|
||||
monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph')
|
||||
monkeypatch.setattr(endpoints, 'get_app_service', lambda: service)
|
||||
async def fail(*args):
|
||||
raise OSError('commit unavailable')
|
||||
monkeypatch.setattr(repo, 'apply_update', fail)
|
||||
response = await endpoints.chat(
|
||||
ProcessRequest_chat(sessionId='session', timeStamp='probe', text='继续',
|
||||
needFormUpdate=True, useTextChunk=True),
|
||||
stream=True, backend=service._backend, idempotency_key=None,
|
||||
)
|
||||
received = []
|
||||
async with aclosing(response.body_iterator) as source:
|
||||
while True:
|
||||
event = await asyncio.wait_for(anext(source), 3)
|
||||
received.append(event)
|
||||
if 'event: text_delta' in event:
|
||||
break
|
||||
llm.release_extraction.set()
|
||||
received.extend([event async for event in source])
|
||||
body = ''.join(received)
|
||||
assert 'event: text_delta' in body and 'event: error' in body
|
||||
assert 'event: done' not in body and 'event: formUpdate' not in body
|
||||
assert await repo.load('session') == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_chatopenai_native_events_do_not_leak_extraction_or_force_streaming(monkeypatch):
|
||||
requests = []
|
||||
def handler(request):
|
||||
payload = json.loads(request.content)
|
||||
requests.append(payload)
|
||||
extraction = '【字段提取】' in payload['messages'][0]['content']
|
||||
if extraction:
|
||||
assert payload['stream'] is False
|
||||
return httpx.Response(200, json={
|
||||
'id': 'extract', 'object': 'chat.completion', 'created': 0, 'model': 'probe',
|
||||
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': '{"has_injury":false}'}, 'finish_reason': 'stop'}],
|
||||
})
|
||||
assert payload['stream'] is True
|
||||
data = {'id': 'reply', 'object': 'chat.completion.chunk', 'created': 0, 'model': 'probe',
|
||||
'choices': [{'index': 0, 'delta': {'content': '<state>1002</state>请问您还在事故现场吗？'}, 'finish_reason': None}]}
|
||||
return httpx.Response(200, headers={'content-type': 'text/event-stream'},
|
||||
content='data: '+json.dumps(data)+'\n\ndata: [DONE]\n\n')
|
||||
monkeypatch.setattr(Config, 'LLM_HTTP_TRUST_ENV', False)
|
||||
llm = OpenAILLMService('probe', 'offline', base_url='http://probe.local/v1')
|
||||
llm._get_client().http_async_client._transport = httpx.MockTransport(handler)
|
||||
try:
|
||||
service, repo = await service_for(llm)
|
||||
events = [event async for event in service.stream(ChatInput('session', '继续', need_form_update=True))]
|
||||
text = ''.join(event.text for event in events if isinstance(event, TextDelta))
|
||||
assert text == '<state>1002</state>请问您还在事故现场吗？'
|
||||
assert len(requests) == 2
|
||||
assert sum(isinstance(event, ReplyFinished) for event in events) == 1
|
||||
assert isinstance(events[-1], FormUpdate)
|
||||
finally:
|
||||
await llm.aclose()
|
||||
Reference in New Issue
Block a user