292 lines
11 KiB
Python
292 lines
11 KiB
Python
"""通过真实 ChatOpenAI/OpenAI SDK + HTTP mock 验证迁移边界。"""
|
|
|
|
import asyncio
|
|
import json
|
|
from contextlib import aclosing
|
|
from typing import TypedDict
|
|
from unittest.mock import AsyncMock
|
|
|
|
import httpx
|
|
import pytest
|
|
from langchain_core.callbacks import AsyncCallbackHandler
|
|
from langchain_core.messages import AIMessage
|
|
from langgraph.graph import END, START, StateGraph
|
|
|
|
from src.agent.llm import LLMAPIError, LLMTimeoutError, OpenAILLMService
|
|
from src.core.config import Config
|
|
|
|
|
|
def _chunk(text):
|
|
data = {
|
|
'id': 'response', 'object': 'chat.completion.chunk', 'created': 0, 'model': 'probe',
|
|
'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}],
|
|
}
|
|
return ('data: ' + json.dumps(data) + '\n\n').encode()
|
|
|
|
|
|
def _response(text, *, stream=False):
|
|
if stream:
|
|
return httpx.Response(200, headers={'content-type': 'text/event-stream'},
|
|
content=_chunk(text) + b'data: [DONE]\n\n')
|
|
return httpx.Response(200, json={
|
|
'id': 'response', 'object': 'chat.completion', 'created': 0, 'model': 'probe',
|
|
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': text},
|
|
'finish_reason': 'stop'}],
|
|
'usage': {'prompt_tokens': 7, 'completion_tokens': 2, 'total_tokens': 9},
|
|
})
|
|
|
|
|
|
def _service(monkeypatch, handler, **kwargs):
|
|
monkeypatch.setattr(Config, 'LLM_HTTP_TRUST_ENV', False)
|
|
service = OpenAILLMService(
|
|
model='probe', api_key='offline-test', base_url='http://probe.local/v1',
|
|
extra_body={'reasoning_effort': 'none'}, **kwargs,
|
|
)
|
|
# 只替换 HTTP 传输，模型工厂、消息转换、callbacks 和 SDK 全部真实执行。
|
|
model = service._get_client()
|
|
model.http_async_client._transport = httpx.MockTransport(handler)
|
|
return service
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('stream', [False, True])
|
|
@pytest.mark.parametrize('temperature', [None, 0.0])
|
|
async def test_wire_contract_preserves_profile_messages_and_chat_completions(monkeypatch, stream, temperature):
|
|
monkeypatch.setenv('OPENAI_PROXY', 'http://invalid-proxy.local:1')
|
|
monkeypatch.setenv('HTTP_PROXY', 'http://invalid-proxy.local:1')
|
|
monkeypatch.setenv('HTTPS_PROXY', 'http://invalid-proxy.local:1')
|
|
requests = []
|
|
def handler(request):
|
|
requests.append(request)
|
|
return _response('收到', stream=stream)
|
|
llm = _service(monkeypatch, handler, temperature=temperature)
|
|
http_async = llm._client.http_async_client
|
|
http_sync = llm._client.http_client
|
|
history = [{'role': 'user', 'content': '之前的输入'}, {'role': 'assistant', 'content': '之前的回复'}]
|
|
try:
|
|
if stream:
|
|
result = ''.join([t async for t in llm.stream_tokens('系统', '本轮', history=history)])
|
|
else:
|
|
result = await llm.complete('系统', '本轮', history=history)
|
|
assert result == '收到'
|
|
assert len(requests) == 1
|
|
request = requests[0]
|
|
payload = json.loads(request.content)
|
|
assert request.url.path == '/v1/chat/completions'
|
|
assert payload['model'] == 'probe'
|
|
assert payload['reasoning_effort'] == 'none'
|
|
assert payload['stream'] is stream
|
|
assert 'stream_options' not in payload
|
|
if temperature is None:
|
|
assert 'temperature' not in payload
|
|
else:
|
|
assert payload['temperature'] == temperature
|
|
assert payload['messages'] == [
|
|
{'role': 'system', 'content': '系统'}, *history, {'role': 'user', 'content': '本轮'},
|
|
]
|
|
assert history == [{'role': 'user', 'content': '之前的输入'}, {'role': 'assistant', 'content': '之前的回复'}]
|
|
finally:
|
|
await llm.aclose()
|
|
assert http_async.is_closed and http_sync.is_closed
|
|
assert llm._client is None
|
|
await llm.aclose()
|
|
|
|
|
|
class TrackingBody(httpx.AsyncByteStream):
|
|
def __init__(self, text='partial', *, fail=False, emit=True):
|
|
self.text = text
|
|
self.fail = fail
|
|
self.emit = emit
|
|
self.closed = False
|
|
self.started = asyncio.Event()
|
|
self.release = asyncio.Event()
|
|
|
|
async def __aiter__(self):
|
|
self.started.set()
|
|
if self.emit:
|
|
yield _chunk(self.text)
|
|
if self.fail:
|
|
raise httpx.ReadError('broken transport')
|
|
await self.release.wait()
|
|
yield _chunk('tail')
|
|
yield b'data: [DONE]\n\n'
|
|
|
|
async def aclose(self):
|
|
self.closed = True
|
|
|
|
|
|
def _stream_response(body):
|
|
return httpx.Response(200, headers={'content-type': 'text/event-stream'}, stream=body)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_early_close_is_immediate_and_does_not_close_another_session(monkeypatch):
|
|
bodies = {}
|
|
def handler(request):
|
|
text = json.loads(request.content)['messages'][-1]['content']
|
|
body = bodies[text] = TrackingBody(text)
|
|
return _stream_response(body)
|
|
llm = _service(monkeypatch, handler)
|
|
first = llm.stream_tokens('system', 'first')
|
|
second = llm.stream_tokens('system', 'second')
|
|
try:
|
|
# anext 运行在其他任务中，aclose 在当前任务中；上下文也必须安全复原。
|
|
assert await asyncio.wait_for(anext(first), 5) == 'first'
|
|
assert await asyncio.wait_for(anext(second), 5) == 'second'
|
|
await first.aclose()
|
|
assert bodies['first'].closed
|
|
assert not bodies['second'].closed
|
|
bodies['second'].release.set()
|
|
assert [t async for t in second] == ['tail']
|
|
assert bodies['second'].closed
|
|
finally:
|
|
await first.aclose()
|
|
await second.aclose()
|
|
await llm.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transport_failure_after_text_closes_immediately_without_retry(monkeypatch):
|
|
body = TrackingBody(fail=True)
|
|
calls = []
|
|
def handler(request):
|
|
calls.append(request)
|
|
return _stream_response(body)
|
|
llm = _service(monkeypatch, handler)
|
|
llm._retry_delay = AsyncMock()
|
|
tokens = []
|
|
try:
|
|
with pytest.raises(LLMAPIError):
|
|
async for token in llm.stream_tokens('system', 'user'):
|
|
tokens.append(token)
|
|
assert tokens == ['partial']
|
|
assert body.closed
|
|
assert len(calls) == 1
|
|
llm._retry_delay.assert_not_awaited()
|
|
finally:
|
|
await llm.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_cancellation_closes_http_response_before_returning(monkeypatch):
|
|
body = TrackingBody(emit=False)
|
|
llm = _service(monkeypatch, lambda request: _stream_response(body))
|
|
llm._retry_delay = AsyncMock()
|
|
async def consume():
|
|
async with aclosing(llm.stream_tokens('system', 'user')) as stream:
|
|
async for _ in stream:
|
|
pass
|
|
task = asyncio.create_task(consume())
|
|
try:
|
|
await asyncio.wait_for(body.started.wait(), 5)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert body.closed
|
|
llm._retry_delay.assert_not_awaited()
|
|
finally:
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
await llm.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_timeout_closes_each_attempt_before_retry(monkeypatch):
|
|
bodies = []
|
|
def handler(request):
|
|
assert all(body.closed for body in bodies)
|
|
body = TrackingBody(emit=False)
|
|
bodies.append(body)
|
|
return _stream_response(body)
|
|
llm = _service(monkeypatch, handler, timeout_seconds=0.1, max_retries=2)
|
|
llm._retry_delay = AsyncMock()
|
|
try:
|
|
with pytest.raises(LLMTimeoutError):
|
|
async for _ in llm.stream_tokens('system', 'user'):
|
|
pass
|
|
assert len(bodies) == 3
|
|
assert all(body.closed for body in bodies)
|
|
assert llm._retry_delay.await_count == 2
|
|
finally:
|
|
await llm.aclose()
|
|
|
|
|
|
class TraceRecorder(AsyncCallbackHandler):
|
|
def __init__(self):
|
|
self.chains = {}
|
|
self.models = []
|
|
self.results = []
|
|
|
|
async def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
|
|
self.chains[run_id] = kwargs.get('name')
|
|
|
|
async def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
|
|
self.models.append({'run_id': run_id, **kwargs})
|
|
|
|
async def on_llm_end(self, response, **kwargs):
|
|
self.results.append(response)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('stream', [False, True])
|
|
async def test_model_callbacks_inherit_graph_parent_metadata_and_tags(monkeypatch, stream):
|
|
llm = _service(monkeypatch, lambda request: _response('ok', stream=stream))
|
|
recorder = TraceRecorder()
|
|
class State(TypedDict):
|
|
text: str
|
|
async def dialogue(state):
|
|
if stream:
|
|
text = ''.join([t async for t in llm.stream_tokens('system', state['text'])])
|
|
else:
|
|
text = await llm.complete('system', state['text'])
|
|
return {'text': text}
|
|
graph = StateGraph(State)
|
|
graph.add_node('dialogue', dialogue)
|
|
graph.add_edge(START, 'dialogue')
|
|
graph.add_edge('dialogue', END)
|
|
try:
|
|
result = await graph.compile().ainvoke({'text': 'hello'}, config={
|
|
'callbacks': [recorder], 'metadata': {'session_id': 'trace-session'}, 'tags': ['probe'],
|
|
})
|
|
assert result['text'] == 'ok'
|
|
assert len(recorder.models) == 1
|
|
model = recorder.models[0]
|
|
assert recorder.chains[model['parent_run_id']] == 'dialogue'
|
|
assert model['name'] == ('llm.stream' if stream else 'llm.complete')
|
|
assert model['metadata']['session_id'] == 'trace-session'
|
|
assert model['metadata']['llm_attempt'] == 1
|
|
assert 'probe' in model['tags']
|
|
assert len(recorder.results) == 1
|
|
if not stream:
|
|
usage = recorder.results[0].generations[0][0].message.usage_metadata
|
|
assert usage['input_tokens'] == 7 and usage['output_tokens'] == 2
|
|
finally:
|
|
await llm.aclose()
|
|
|
|
|
|
@pytest.mark.parametrize('content, expected', [
|
|
('正文', '正文'),
|
|
([{'type': 'reasoning', 'reasoning': '内部内容'}, {'type': 'text', 'text': '正文'}], '正文'),
|
|
([{'type': 'reasoning', 'reasoning': '内部内容'}], ''),
|
|
])
|
|
def test_only_text_content_is_forwarded(content, expected):
|
|
assert OpenAILLMService._text(AIMessage(content=content)) == expected
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shared_model_serves_concurrent_calls_without_mixing_messages(monkeypatch):
|
|
async def handler(request):
|
|
payload = json.loads(request.content)
|
|
await asyncio.sleep(0)
|
|
return _response(payload['messages'][-1]['content'], stream=payload['stream'])
|
|
llm = _service(monkeypatch, handler)
|
|
async def request(index):
|
|
text = f'input-{index}'
|
|
if index % 2:
|
|
return ''.join([t async for t in llm.stream_tokens('system', text)])
|
|
return await llm.complete('system', text)
|
|
try:
|
|
assert await asyncio.gather(*(request(i) for i in range(12))) == [f'input-{i}' for i in range(12)]
|
|
finally:
|
|
await llm.aclose()
|