feat(llm): migrate model calls to LangChain ChatOpenAI

This commit is contained in:
Eric Wang
2026-10-24 17:08:46 +08:00
parent 97de142e6d
commit d8d48c26a3
13 changed files with 489 additions and 96 deletions
+291
View File
@@ -0,0 +1,291 @@
"""通过真实 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()
+1 -3
View File
@@ -37,9 +37,7 @@ def _service_with_failing_llm(exc: Exception, max_retries: int = 0):
"""complete 固定抛 exc 的 OpenAILLMService(client 打桩)。"""
llm = OpenAILLMService(model="test", api_key="test", max_retries=max_retries)
create = AsyncMock(side_effect=exc)
llm._client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
)
llm._client = SimpleNamespace(ainvoke=create)
return llm, create
+13 -19
View File
@@ -6,7 +6,9 @@ mock 客户端返回 token 事件流;验证消息构造与 complete 一致、
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import Mock
from langchain_core.messages import AIMessageChunk
import httpx
import pytest
@@ -22,9 +24,7 @@ from src.agent.llm import (
def _event(content):
return SimpleNamespace(
choices=[SimpleNamespace(delta=SimpleNamespace(content=content))]
)
return AIMessageChunk(content=content)
def _service_with_stream(events):
@@ -44,10 +44,8 @@ def _service_with_stream(events):
for item in self._items:
yield item
create = AsyncMock(return_value=_Stream(events))
llm._client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
)
create = Mock(return_value=_Stream(events))
llm._client = SimpleNamespace(astream=create)
return llm, create
@@ -60,13 +58,11 @@ async def test_stream_tokens_yields_incrementally_and_forwards_kwargs():
]
assert tokens == ["你", "好"] # 空 delta 被跳过
kwargs = create.call_args.kwargs
assert kwargs["model"] == "test-model"
assert kwargs["stream"] is True
assert kwargs["temperature"] == 0
assert kwargs["extra_body"]["reasoning_effort"] == "none" # extra_body 透传
assert kwargs["messages"][0] == {"role": "system", "content": "系统"}
assert kwargs["messages"][-1] == {"role": "user", "content": "输入"}
messages = create.call_args.args[0]
assert messages[0] == {"role": "system", "content": "系统"}
assert messages[1] == {"role": "assistant", "content": "h"}
assert messages[-1] == {"role": "user", "content": "输入"}
assert create.call_args.kwargs["config"]["run_name"] == "llm.stream"
@pytest.mark.asyncio
@@ -107,10 +103,8 @@ async def test_stream_tokens_times_out_on_stale_stream():
yield _event("好")
llm = OpenAILLMService(model="m", api_key="k", timeout_seconds=0.2)
create = AsyncMock(return_value=_HangingStream())
llm._client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
)
create = Mock(return_value=_HangingStream())
llm._client = SimpleNamespace(astream=create)
tokens = []
with pytest.raises(LLMTimeoutError):
async for token in llm.stream_tokens("系统", "输入"):
+5 -5
View File
@@ -104,14 +104,14 @@ async def test_openai_adapter_preserves_message_roles():
from unittest.mock import AsyncMock
from src.agent.llm import OpenAILLMService
create = AsyncMock(return_value=SimpleNamespace(choices=[
SimpleNamespace(message=SimpleNamespace(content="<state>1002</state>收到。"))
]))
from langchain_core.messages import AIMessage
create = AsyncMock(return_value=AIMessage(content="<state>1002</state>收到。"))
llm = OpenAILLMService(model="test", api_key="test", max_retries=0)
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
llm._client = SimpleNamespace(ainvoke=create)
history = [{"role": "assistant", "content": "之前的问题"}]
await llm.complete("系统提示词", "本轮输入", history=history)
assert create.call_args.kwargs["messages"] == [
assert create.call_args.args[0] == [
{"role": "system", "content": "系统提示词"},
{"role": "assistant", "content": "之前的问题"},
{"role": "user", "content": "本轮输入"},
+26 -12
View File
@@ -2,11 +2,13 @@ import asyncio
import json
from contextlib import aclosing
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, Mock
import httpx
import pytest
from openai import AsyncOpenAI
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
@@ -72,7 +74,10 @@ async def test_commit_is_not_lost_when_consumer_closes_after_graph_completes():
async def test_factory_disables_sdk_retries():
llm = OpenAILLMService(model='m', api_key='audit')
try:
assert llm._get_client().max_retries == 0
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()
@@ -101,8 +106,13 @@ async def test_actual_http_attempt_count_is_not_multiplied_by_sdk(stream, status
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._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')
@@ -120,19 +130,19 @@ async def test_stream_failure_after_first_token_closes_response_without_retry():
def __init__(self): self.closed = False
def __aiter__(self): return self.generate()
async def generate(self):
yield SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content='partial'))])
yield AIMessageChunk(content='partial')
raise OSError('broken socket')
async def close(self): self.closed = True
async def aclose(self): self.closed = True
response = Response()
create = AsyncMock(return_value=response)
create = Mock(return_value=response)
llm = OpenAILLMService(model='m', api_key='audit')
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
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.await_count == 1
assert create.call_count == 1
assert response.closed
@@ -145,8 +155,12 @@ async def test_retry_budget_exhaustion_is_exact(stream):
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)))
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:
+10 -10
View File
@@ -1,7 +1,7 @@
import asyncio
from contextlib import aclosing
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, Mock
import pytest
@@ -69,16 +69,16 @@ async def test_timeout_covers_whole_attempt_and_retries_are_bounded(stream):
async def generate(self):
await asyncio.Future()
yield None
async def close(self): self.closed = True
async def create(**kwargs):
if stream:
response = HangingResponse()
responses.append(response)
return response
async def aclose(self): self.closed = True
async def invoke(*args, **kwargs):
await asyncio.Future()
def create_stream(*args, **kwargs):
response = HangingResponse()
responses.append(response)
return response
llm = OpenAILLMService(model='m', api_key='audit', timeout_seconds=0.01, max_retries=2)
call = AsyncMock(side_effect=create)
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=call)))
call = Mock(side_effect=create_stream) if stream else AsyncMock(side_effect=invoke)
llm._client = SimpleNamespace(astream=call) if stream else SimpleNamespace(ainvoke=call)
llm._retry_delay = AsyncMock()
# 先加载 SDK 错误类;外层看门狗不测模块惰性导入耗时。
assert llm._retryable(TimeoutError())
@@ -89,6 +89,6 @@ async def test_timeout_covers_whole_attempt_and_retries_are_bounded(stream):
async for _ in llm.stream_tokens('system', 'user'): pass
else:
await llm.complete('system', 'user')
assert call.await_count == 3
assert call.call_count == 3
assert llm._retry_delay.await_count == 2
assert all(response.closed for response in responses)