mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-22 16:35:20 +08:00
perf(llm): harden request policy and optimize connections, reasoning and prompt caching (#2661)
* fix(llm): centralize request policy and preserve stream diagnostics * test(llm): dispose SQLite engine after migration regression * perf(llm): reuse connections and unify reasoning and prompt cache policy * fix(llm): preserve explicit native prompt cache options * perf(llm): fix measured TLS setup and streaming reuse overhead * fix(llm): resolve plugin request-chain regressions found by E2E --------- Co-authored-by: Hyu <[email protected]>
This commit is contained in:
RockChinQ
@@ -10,6 +10,27 @@ from __future__ import annotations
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from importlib import import_module
|
||||
from types import SimpleNamespace
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'call_args',
|
||||
[
|
||||
{'num_retries': 0},
|
||||
{'max_retries': 0},
|
||||
{'timeout': 3},
|
||||
{'request_policy': {'stream_idle_timeout': 2}},
|
||||
],
|
||||
)
|
||||
def test_llm_bridge_keeps_model_and_call_policy_layers_separate(call_args):
|
||||
from langbot.pkg.provider.modelmgr.request_policy import RequestPolicy
|
||||
|
||||
model_args = {'temperature': 0.2, 'request_policy': {'retries': 2, 'read_timeout': 90}}
|
||||
model = SimpleNamespace(model_entity=SimpleNamespace(extra_args=model_args))
|
||||
merged = get_handler_module()._merge_model_extra_args(model, call_args, llm_policy=True)
|
||||
assert RequestPolicy.resolve(model_args, merged) == RequestPolicy.resolve(model_args, call_args)
|
||||
assert merged['temperature'] == 0.2
|
||||
assert model_args['request_policy'] == {'retries': 2, 'read_timeout': 90}
|
||||
|
||||
|
||||
def get_handler_module():
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Replay synthetic HTTP/SSE traffic through the real Codex requester."""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
@@ -16,6 +17,32 @@ TOKENS = {'access_token': 'access-secret', 'account_id': 'account', 'connection_
|
||||
MODEL = SimpleNamespace(model_entity=SimpleNamespace(name='codex-test', extra_args={}, reasoning_config=None))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_headers_obey_host_first_event_deadline(monkeypatch):
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
|
||||
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def handler(request):
|
||||
try:
|
||||
await asyncio.sleep(10)
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
obj = requester(monkeypatch, handler)
|
||||
state = CallState(RequestPolicy(first_event_timeout=0.02))
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
with pytest.raises(RequesterError) as caught:
|
||||
await obj.invoke_llm(None, MODEL, [])
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
assert caught.value.code == 'timeout'
|
||||
assert caught.value.phase == 'first_event'
|
||||
assert cancelled.is_set()
|
||||
|
||||
|
||||
def requester(monkeypatch, handler):
|
||||
real_client = httpx.AsyncClient
|
||||
monkeypatch.setattr(
|
||||
@@ -198,3 +225,51 @@ def test_advanced_parameters_cannot_override_transport(monkeypatch, key):
|
||||
obj = requester(monkeypatch, lambda request: httpx.Response(200))
|
||||
with pytest.raises(ValueError, match='Unsupported Codex advanced'):
|
||||
obj._body(None, MODEL, [], None, {key: 'secret'}, TOKENS)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_early_close_releases_nested_http_stream(monkeypatch):
|
||||
closed = asyncio.Event()
|
||||
|
||||
class Body(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'data: {"type":"response.output_text.delta","delta":"first"}\n\n'
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
closed.set()
|
||||
|
||||
obj = requester(monkeypatch, lambda request: httpx.Response(200, stream=Body()))
|
||||
stream = obj.invoke_llm_stream(None, MODEL, [])
|
||||
assert (await anext(stream)).content == 'first'
|
||||
await stream.aclose()
|
||||
assert closed.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_deadline_does_not_cancel_suspended_consumer(monkeypatch, runtime_provider, runtime_llm_model):
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
from tests.unit_tests.provider.conftest import TEST_EXECUTION_CONTEXT
|
||||
|
||||
runtime_provider.requester = requester(
|
||||
monkeypatch,
|
||||
lambda request: stream(
|
||||
[
|
||||
{'type': 'response.output_text.delta', 'delta': 'first'},
|
||||
{'type': 'response.completed', 'response': {'status': 'completed', 'output': []}},
|
||||
]
|
||||
),
|
||||
)
|
||||
runtime_llm_model.model_entity.extra_args = {}
|
||||
result = runtime_provider.invoke_llm_stream(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'total_timeout': 0.05}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
assert (await anext(result)).content == 'first'
|
||||
await asyncio.sleep(0.07)
|
||||
with pytest.raises(RequesterError) as caught:
|
||||
await anext(result)
|
||||
assert caught.value.phase == 'total'
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from langbot.pkg.provider.modelmgr.http_pool import _CompletedSSEStream, RequesterHTTPPool
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
|
||||
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
|
||||
from tests.unit_tests.provider.test_reasoning_control import _runtime_model
|
||||
from scripts.llm_latency_probe import SSEMeter, cache_tokens
|
||||
|
||||
|
||||
def test_request_timing_uses_high_resolution_monotonic_clock(monkeypatch):
|
||||
from langbot.pkg.provider.modelmgr import request_policy
|
||||
|
||||
ticks = iter([100.0, 100.0004])
|
||||
monkeypatch.setattr(request_policy.time, 'perf_counter', lambda: next(ticks))
|
||||
state = CallState(RequestPolicy())
|
||||
state.mark('request_prepared_ms')
|
||||
assert state.metrics['request_prepared_ms'] == 0.4
|
||||
|
||||
|
||||
def test_meter_accepts_split_utf8_and_usage_without_content():
|
||||
meter = SSEMeter()
|
||||
event = {'choices': [{'delta': {'content': '你好'}}]}
|
||||
data = ('data: ' + json.dumps(event, ensure_ascii=False) + '\r\n\r\n').encode()
|
||||
for byte in data:
|
||||
meter.feed(bytes([byte]), 1)
|
||||
meter.feed(b'data: {"choices":[],"usage":{"prompt_tokens_details":{"cached_tokens":12}}}\n\n', 2)
|
||||
meter.feed(b'data: [DONE]\n\n', 3)
|
||||
assert meter.first_content == 1 and meter.content_chars == 2
|
||||
assert cache_tokens(meter.usage) == 12 and meter.done
|
||||
assert cache_tokens(None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('provider,name', [('openai', 'gpt-4o'), ('openrouter', 'openai/gpt-4o')])
|
||||
@pytest.mark.parametrize('drain', [False, True])
|
||||
async def test_real_stream_reuses_socket_and_preserves_late_usage(monkeypatch, provider, name, drain):
|
||||
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
|
||||
sockets = []
|
||||
|
||||
async def handle(request):
|
||||
await request.read()
|
||||
sockets.append(request.transport)
|
||||
response = web.StreamResponse(headers={'Content-Type': 'text/event-stream'})
|
||||
await response.prepare(request)
|
||||
events = [
|
||||
{'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': 'hello'}, 'finish_reason': None}]},
|
||||
{'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]},
|
||||
{
|
||||
'choices': [],
|
||||
'usage': {
|
||||
'prompt_tokens': 14,
|
||||
'completion_tokens': 1,
|
||||
'total_tokens': 15,
|
||||
'prompt_tokens_details': {'cached_tokens': 12},
|
||||
},
|
||||
},
|
||||
]
|
||||
for event in events:
|
||||
await response.write(
|
||||
(
|
||||
'data: '
|
||||
+ json.dumps(
|
||||
{
|
||||
'id': 'r',
|
||||
'object': 'chat.completion.chunk',
|
||||
'created': 1,
|
||||
'model': name,
|
||||
**event,
|
||||
}
|
||||
)
|
||||
+ '\n\n'
|
||||
).encode()
|
||||
)
|
||||
await asyncio.sleep(0.02)
|
||||
await response.write(b'data: [DO')
|
||||
await response.write(b'NE]\n\n')
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post('/chat/completions', handle)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
adapter = LiteLLMRequester(
|
||||
SimpleNamespace(logger=Mock()),
|
||||
{
|
||||
'custom_llm_provider': provider,
|
||||
'base_url': f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}',
|
||||
},
|
||||
)
|
||||
await adapter.initialize()
|
||||
adapter._http_pool.drain_sse = drain
|
||||
model = _runtime_model(adapter, name=name)
|
||||
model.provider.token_mgr.get_token = lambda: 'test'
|
||||
try:
|
||||
for _ in range(2):
|
||||
state = CallState(RequestPolicy())
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
chunks = [chunk async for chunk in adapter.invoke_llm_stream(None, model, [])]
|
||||
assert ''.join(chunk.content or '' for chunk in chunks) == 'hello'
|
||||
assert state.usage['prompt_tokens_details']['cached_tokens'] == 12
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
# OpenRouter's generic HTTP adapter already consumes EOF; the OpenAI
|
||||
# SDK path is the regression that needs terminal draining.
|
||||
assert (sockets[0] is sockets[1]) == (drain or provider == 'openrouter')
|
||||
finally:
|
||||
await adapter.aclose()
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_drain_is_bounded_but_early_close_does_not_drain():
|
||||
class Stream(httpx.AsyncByteStream):
|
||||
def __init__(self, terminal):
|
||||
self.terminal = terminal
|
||||
self.tail_read = False
|
||||
self.closed = False
|
||||
|
||||
async def __aiter__(self):
|
||||
yield b'data: [DONE]\n\n' if self.terminal else b'data: {"choices":[]}\n\n'
|
||||
self.tail_read = True
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
self.closed = True
|
||||
|
||||
for terminal in (False, True):
|
||||
raw = Stream(terminal)
|
||||
wrapped = _CompletedSSEStream(raw)
|
||||
iterator = wrapped.__aiter__()
|
||||
async with asyncio.timeout(1):
|
||||
await anext(iterator)
|
||||
await wrapped.aclose()
|
||||
await iterator.aclose()
|
||||
assert raw.closed and raw.tail_read == terminal
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancellation_during_terminal_drain_is_not_swallowed():
|
||||
reading = asyncio.Event()
|
||||
|
||||
class Stream(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'data: [DONE]\n\n'
|
||||
reading.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
wrapped = _CompletedSSEStream(Stream())
|
||||
iterator = wrapped.__aiter__()
|
||||
pending = asyncio.create_task(anext(iterator))
|
||||
await reading.wait()
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await pending
|
||||
await iterator.aclose()
|
||||
await wrapped.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('retire', [False, True])
|
||||
async def test_concurrent_pool_initialization_loads_tls_once_off_loop(monkeypatch, retire):
|
||||
context = httpx.create_ssl_context()
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
main_thread = threading.get_ident()
|
||||
calls = []
|
||||
|
||||
def load():
|
||||
assert threading.get_ident() != main_thread
|
||||
calls.append(True)
|
||||
entered.set()
|
||||
assert release.wait(2)
|
||||
return context
|
||||
|
||||
monkeypatch.setattr(httpx, 'create_ssl_context', load)
|
||||
pool = RequesterHTTPPool()
|
||||
|
||||
async def lease():
|
||||
async with pool.lease() as client:
|
||||
return client
|
||||
|
||||
first, second = asyncio.create_task(lease()), asyncio.create_task(lease())
|
||||
await asyncio.to_thread(entered.wait, 2)
|
||||
if retire:
|
||||
await pool.aclose()
|
||||
release.set()
|
||||
results = await asyncio.gather(first, second, return_exceptions=True)
|
||||
assert len(calls) == 1
|
||||
if retire:
|
||||
assert all(isinstance(result, RequesterError) for result in results)
|
||||
assert pool.client is None
|
||||
else:
|
||||
assert results[0] is results[1] is pool.client
|
||||
assert not pool.client.is_closed
|
||||
await pool.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_tls_initialization_leaks_no_http_client(monkeypatch):
|
||||
context = httpx.create_ssl_context()
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
|
||||
def load():
|
||||
entered.set()
|
||||
assert release.wait(2)
|
||||
return context
|
||||
|
||||
monkeypatch.setattr(httpx, 'create_ssl_context', load)
|
||||
pool = RequesterHTTPPool()
|
||||
|
||||
async def lease():
|
||||
async with pool.lease():
|
||||
pytest.fail('Cancelled initializer must not publish a client')
|
||||
|
||||
pending = asyncio.create_task(lease())
|
||||
await asyncio.to_thread(entered.wait, 2)
|
||||
pending.cancel()
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await pending
|
||||
await pool.aclose()
|
||||
assert pool.client is None and pool.active == 0
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_openai_stream_close_disconnects_upstream(monkeypatch):
|
||||
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
|
||||
disconnected = asyncio.Event()
|
||||
|
||||
async def handle(request):
|
||||
await request.read()
|
||||
response = web.StreamResponse(headers={'Content-Type': 'text/event-stream'})
|
||||
await response.prepare(request)
|
||||
await response.write(
|
||||
b'data: {"id":"r","object":"chat.completion.chunk","created":1,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}\n\n'
|
||||
)
|
||||
for _ in range(100):
|
||||
if request.transport is None or request.transport.is_closing():
|
||||
disconnected.set()
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
return response
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post('/chat/completions', handle)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
adapter = LiteLLMRequester(
|
||||
SimpleNamespace(logger=Mock()),
|
||||
{
|
||||
'custom_llm_provider': 'openai',
|
||||
'base_url': f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}',
|
||||
},
|
||||
)
|
||||
await adapter.initialize()
|
||||
model = _runtime_model(adapter, name='gpt-4o')
|
||||
model.provider.token_mgr.get_token = lambda: 'test'
|
||||
stream = adapter.invoke_llm_stream(None, model, [])
|
||||
try:
|
||||
assert (await anext(stream)).content == 'hello'
|
||||
await stream.aclose()
|
||||
await asyncio.wait_for(disconnected.wait(), 3)
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await adapter.aclose()
|
||||
await runner.cleanup()
|
||||
@@ -0,0 +1,516 @@
|
||||
"""Connection ownership, preparation reuse and provider-specific request policy."""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import threading
|
||||
import json
|
||||
from pathlib import Path
|
||||
from collections import OrderedDict
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from aiohttp import web
|
||||
|
||||
from langbot.pkg.provider.modelmgr import reasoning
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
from langbot.pkg.provider.modelmgr.http_pool import RequesterHTTPPool
|
||||
from langbot.pkg.provider.modelmgr.prompt_caching import apply_prompt_caching
|
||||
from langbot.pkg.provider.modelmgr.request_policy import CallState, RequestPolicy, current_call
|
||||
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
|
||||
from langbot.pkg.provider.modelmgr.requesters.litellmchat import _LeasedHTTPHandler
|
||||
from langbot.pkg.provider.modelmgr.requesters import codex
|
||||
from langbot_plugin.api.entities.builtin.provider.message import Message
|
||||
from tests.unit_tests.provider.test_reasoning_control import _runtime_model, _requester
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_retirement_waits_for_stream_lease():
|
||||
pool = RequesterHTTPPool()
|
||||
async with pool.lease() as first:
|
||||
async with pool.lease() as second:
|
||||
assert first is second
|
||||
await pool.aclose()
|
||||
assert not first.is_closed
|
||||
with pytest.raises(RequesterError, match='retired'):
|
||||
async with pool.lease():
|
||||
pass
|
||||
assert not first.is_closed
|
||||
assert first.is_closed
|
||||
await pool.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('provider,name', [('openai', 'gpt-4o'), ('openrouter', 'openai/gpt-4o')])
|
||||
async def test_real_http_reuses_socket_with_rotating_key_and_timeout(monkeypatch, provider, name):
|
||||
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
|
||||
sockets, headers, bodies = [], [], []
|
||||
|
||||
async def handle(request):
|
||||
sockets.append(request.transport)
|
||||
headers.append(request.headers['Authorization'])
|
||||
bodies.append(await request.json())
|
||||
return web.json_response(
|
||||
{
|
||||
'id': 'r',
|
||||
'object': 'chat.completion',
|
||||
'created': 1,
|
||||
'model': 'gpt-4o',
|
||||
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
|
||||
'usage': {'prompt_tokens': 1, 'completion_tokens': 1, 'total_tokens': 2},
|
||||
}
|
||||
)
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post('/v1/chat/completions', handle)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
adapter = LiteLLMRequester(Mock(), {'custom_llm_provider': provider, 'base_url': f'http://127.0.0.1:{port}/v1'})
|
||||
await adapter.initialize()
|
||||
model = _runtime_model(adapter, name=name)
|
||||
model.provider.token_mgr.get_token = Mock(side_effect=['first-key', 'second-key'])
|
||||
try:
|
||||
for seconds in (10, 20):
|
||||
state = CallState(RequestPolicy(read_timeout=seconds), 'lb-scoped')
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
message, _ = await adapter.invoke_llm(
|
||||
None,
|
||||
model,
|
||||
[Message(role='user', content='Hi')],
|
||||
extra_args={'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'}},
|
||||
)
|
||||
assert message.content == 'ok'
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
assert sockets[0] is sockets[1]
|
||||
assert headers == ['Bearer first-key', 'Bearer second-key']
|
||||
assert bodies[0]['prompt_cache_key'] == bodies[1]['prompt_cache_key']
|
||||
assert bodies[0]['prompt_cache_retention'] == 'in_memory'
|
||||
assert all('prompt_caching' not in body and 'reasoning_config' not in body for body in bodies)
|
||||
finally:
|
||||
await adapter.aclose()
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('cancel', [False, True])
|
||||
async def test_stream_holds_pool_lease_until_consumer_close(cancel):
|
||||
adapter = _requester('openai')
|
||||
await adapter.initialize()
|
||||
model = _runtime_model(adapter, name='gpt-4o')
|
||||
model.provider.token_mgr.get_token = Mock(return_value='test')
|
||||
closed = asyncio.Event()
|
||||
|
||||
async def raw():
|
||||
try:
|
||||
yield SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'content': 'hi'}), finish_reason=None)
|
||||
]
|
||||
)
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
closed.set()
|
||||
|
||||
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=AsyncMock(return_value=raw())):
|
||||
stream = adapter.invoke_llm_stream(None, model, [])
|
||||
await anext(stream)
|
||||
client = adapter._http_pool.client
|
||||
await adapter.aclose()
|
||||
assert not client.is_closed
|
||||
if cancel:
|
||||
pending = asyncio.create_task(anext(stream))
|
||||
await asyncio.sleep(0)
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await pending
|
||||
else:
|
||||
await stream.aclose()
|
||||
assert closed.is_set() and client.is_closed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preparation_reused_only_within_call_and_is_not_mutated():
|
||||
adapter = _requester('openai')
|
||||
adapter.ap.tool_mgr = SimpleNamespace(
|
||||
generate_tools_for_openai=AsyncMock(
|
||||
return_value=[{'type': 'function', 'function': {'name': 'tool', 'parameters': {'type': 'object'}}}]
|
||||
)
|
||||
)
|
||||
model = _runtime_model(adapter, name='gpt-4o')
|
||||
model.provider.token_mgr.get_token = Mock(side_effect=['key-1', 'key-2', 'key-3'])
|
||||
messages, funcs = [Message(role='user', content='original')], [object()]
|
||||
state = CallState(RequestPolicy(retries=1))
|
||||
with patch.object(adapter, '_convert_messages', wraps=adapter._convert_messages) as convert:
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
first = await adapter._build_completion_args(model, messages, funcs)
|
||||
first['messages'][0]['content'] = 'mutated'
|
||||
first['tools'][0]['function']['name'] = 'mutated'
|
||||
second = await adapter._build_completion_args(model, messages, funcs)
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
assert convert.call_count == 1
|
||||
assert adapter.ap.tool_mgr.generate_tools_for_openai.await_count == 1
|
||||
assert second['messages'][0]['content'] == 'original'
|
||||
assert second['tools'][0]['function']['name'] == 'tool'
|
||||
assert second['api_key'] == 'key-2'
|
||||
await adapter._build_completion_args(model, messages, funcs)
|
||||
assert convert.call_count == 2
|
||||
assert messages[0].content == 'original'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_counter_runs_off_event_loop():
|
||||
adapter = _requester('openai')
|
||||
model = _runtime_model(adapter, name='gpt-4o')
|
||||
model.provider.token_mgr.get_token = Mock(return_value='test')
|
||||
main_thread = threading.get_ident()
|
||||
|
||||
def count(**kwargs):
|
||||
assert threading.get_ident() != main_thread
|
||||
return 12
|
||||
|
||||
with patch('litellm.token_counter', new=count):
|
||||
assert await adapter.count_tokens(model, []) == 12
|
||||
|
||||
|
||||
def test_metadata_cache_is_bounded_and_returns_copies():
|
||||
adapter = _requester('openai')
|
||||
compute = Mock(return_value={'levels': ['low']})
|
||||
first = adapter._cached_metadata('model', compute)
|
||||
first['levels'].append('bad')
|
||||
assert adapter._cached_metadata('model', compute) == {'levels': ['low']}
|
||||
assert compute.call_count == 1
|
||||
for index in range(150):
|
||||
adapter._cached_metadata(index, lambda: index)
|
||||
assert len(adapter._metadata_cache) == 128
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'config',
|
||||
[
|
||||
{'budget_tokens': True},
|
||||
{'budget_tokens': 0},
|
||||
{'budget_tokens': -1},
|
||||
{'budget_tokens': 1.5},
|
||||
{'level': 'high', 'budget_tokens': 1024},
|
||||
{'level': 'disabled', 'budget_tokens': 1024},
|
||||
],
|
||||
)
|
||||
def test_invalid_reasoning_budget(config):
|
||||
with pytest.raises(ValueError):
|
||||
reasoning.normalize_reasoning_config(config)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'provider,name,budget,expected',
|
||||
[
|
||||
('anthropic', 'claude-sonnet-4-5', 2048, {'thinking': {'type': 'enabled', 'budget_tokens': 2048}}),
|
||||
('openai', 'qwen3-32b', 2048, {'extra_body': {'enable_thinking': True, 'thinking_budget': 2048}}),
|
||||
('gemini', 'gemini-2.5-flash', 2048, {'thinking': {'type': 'enabled', 'budget_tokens': 2048}}),
|
||||
('openrouter', 'anthropic/claude-sonnet-4.5', 2048, {'extra_body': {'reasoning': {'max_tokens': 2048}}}),
|
||||
],
|
||||
)
|
||||
async def test_budget_mapping_through_existing_extra_args(provider, name, budget, expected):
|
||||
adapter = _requester(provider)
|
||||
model = _runtime_model(adapter, name=name)
|
||||
model.provider.token_mgr.get_token = Mock(return_value='test')
|
||||
args = await adapter._build_completion_args(
|
||||
model, [], extra_args={'reasoning_config': {'budget_tokens': budget}, 'max_tokens': 4096}
|
||||
)
|
||||
for key, value in expected.items():
|
||||
assert args[key] == value
|
||||
assert 'reasoning_config' not in args
|
||||
assert model.reasoning_config_override is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'provider,name,budget',
|
||||
[
|
||||
('openai', 'gpt-5', 2048),
|
||||
('anthropic', 'claude-opus-4-7', 2048),
|
||||
('anthropic', 'claude-sonnet-4-5', 512),
|
||||
('gemini', 'gemini-2.5-flash', 30000),
|
||||
('gemini', 'gemini-2.5-flash-lite', 128),
|
||||
],
|
||||
)
|
||||
async def test_budget_not_silently_dropped_on_unsupported_model(provider, name, budget):
|
||||
adapter = _requester(provider)
|
||||
model = _runtime_model(adapter, name=name)
|
||||
model.provider.token_mgr.get_token = Mock(return_value='test')
|
||||
with pytest.raises(RequesterError, match='budget'):
|
||||
await adapter._build_completion_args(model, [], extra_args={'reasoning_config': {'budget_tokens': budget}})
|
||||
|
||||
|
||||
def test_runner_level_replaces_call_and_model_budget():
|
||||
adapter = _requester('anthropic')
|
||||
model = _runtime_model(adapter, name='claude-sonnet-4-5')
|
||||
model.model_entity.reasoning_config = {'budget_tokens': 2048}
|
||||
call = {'reasoning_config': {'budget_tokens': 4096}}
|
||||
assert reasoning.resolve_reasoning_config(model, call)['budget_tokens'] == 4096
|
||||
model.reasoning_config_override = {'level': 'low'}
|
||||
assert reasoning.resolve_reasoning_config(model, call) == {'level': 'low'}
|
||||
|
||||
|
||||
def test_anthropic_cache_boundary_stable_and_input_unchanged():
|
||||
messages = [{'role': 'system', 'content': 'stable instructions'}, {'role': 'user', 'content': 'first'}]
|
||||
original = copy.deepcopy(messages)
|
||||
args = {'messages': messages, 'prompt_caching': {'mode': 'anthropic', 'ttl': '1h'}}
|
||||
apply_prompt_caching(args, None, provider='anthropic', family='anthropic')
|
||||
prefix = args['messages'][0]
|
||||
assert prefix['content'][-1]['cache_control'] == {'type': 'ephemeral', 'ttl': '1h'}
|
||||
assert messages == original
|
||||
assert 'cache_control' not in args['messages'][1]
|
||||
args['messages'].append({'role': 'user', 'content': 'second'})
|
||||
args['prompt_caching'] = {'mode': 'anthropic'}
|
||||
apply_prompt_caching(args, None, provider='anthropic', family='anthropic')
|
||||
assert args['messages'][0] == prefix
|
||||
|
||||
|
||||
def test_explicit_cache_settings_win_and_default_does_nothing():
|
||||
args = {
|
||||
'model': 'openai/gpt-4o',
|
||||
'extra_body': {'prompt_cache_key': 'explicit', 'prompt_cache_retention': '24h'},
|
||||
'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'},
|
||||
}
|
||||
apply_prompt_caching(args, CallState(RequestPolicy(), 'lb-scoped'), provider='openai', family='openai')
|
||||
assert 'prompt_cache_key' not in args and 'prompt_cache_retention' not in args
|
||||
args = {'messages': [{'role': 'system', 'content': 'same'}], 'prompt_caching': {'mode': 'provider_default'}}
|
||||
apply_prompt_caching(args, None, provider='anthropic', family='anthropic')
|
||||
assert args == {'messages': [{'role': 'system', 'content': 'same'}]}
|
||||
|
||||
|
||||
def test_native_cache_options_bypass_model_filter_without_legacy_hints():
|
||||
args = {
|
||||
'model': 'openai/gpt-test',
|
||||
'prompt_cache_options': {'mode': 'implicit', 'ttl': '30m'},
|
||||
'prompt_caching': {'mode': 'openai', 'retention': '24h'},
|
||||
}
|
||||
apply_prompt_caching(args, CallState(RequestPolicy(), 'lb-scoped'), provider='openai', family='openai')
|
||||
assert args == {
|
||||
'model': 'openai/gpt-test',
|
||||
'extra_body': {'prompt_cache_options': {'mode': 'implicit', 'ttl': '30m'}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'config',
|
||||
[
|
||||
{'mode': 'typo'},
|
||||
{'mode': 'anthropic', 'ttl': '24h'},
|
||||
{'mode': 'openai', 'retention': 'forever'},
|
||||
{'mode': 'provider_default', 'ttl': '1h'},
|
||||
],
|
||||
)
|
||||
def test_invalid_cache_config(config):
|
||||
with pytest.raises(RequesterError):
|
||||
apply_prompt_caching(
|
||||
{'model': 'test', 'messages': [], 'prompt_caching': config}, None, provider='openai', family='openai'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_transport_never_retries_or_closes_borrowed_client():
|
||||
calls = 0
|
||||
|
||||
async def fail(request):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
raise httpx.ConnectError('unavailable', request=request)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(fail)) as client:
|
||||
handler = _LeasedHTTPHandler(client, 10)
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await handler.post('https://example.com', json={})
|
||||
await handler.close()
|
||||
assert calls == 1 and not client.is_closed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('provider', ['anthropic', 'openrouter', 'openrouter-metadata'])
|
||||
async def test_anthropic_cache_and_budget_reach_upstream(provider):
|
||||
bodies = []
|
||||
|
||||
async def handle(request):
|
||||
bodies.append(json.loads(request.content))
|
||||
if provider == 'anthropic':
|
||||
payload = {
|
||||
'id': 'msg_test',
|
||||
'type': 'message',
|
||||
'role': 'assistant',
|
||||
'model': 'claude-sonnet-4-5',
|
||||
'content': [{'type': 'text', 'text': 'ok'}],
|
||||
'stop_reason': 'end_turn',
|
||||
'usage': {'input_tokens': 1, 'output_tokens': 1, 'cache_read_input_tokens': 50},
|
||||
}
|
||||
else:
|
||||
payload = {
|
||||
'id': 'r',
|
||||
'object': 'chat.completion',
|
||||
'created': 1,
|
||||
'model': 'anthropic/claude-sonnet-4.5',
|
||||
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
|
||||
'usage': {
|
||||
'prompt_tokens': 51,
|
||||
'completion_tokens': 1,
|
||||
'total_tokens': 52,
|
||||
'prompt_tokens_details': {'cached_tokens': 50},
|
||||
},
|
||||
}
|
||||
return httpx.Response(200, json=payload)
|
||||
|
||||
adapter = _requester('openai' if provider == 'openrouter-metadata' else provider)
|
||||
if provider == 'openrouter-metadata':
|
||||
metadata = yaml.safe_load(
|
||||
(
|
||||
Path(__file__).parents[3] / 'src/langbot/pkg/provider/modelmgr/requesters/openrouterchatcmpl.yaml'
|
||||
).read_text(encoding='utf-8')
|
||||
)
|
||||
adapter.requester_cfg.update(
|
||||
requester_name=metadata['metadata']['name'], custom_llm_provider=metadata['spec']['litellm_provider']
|
||||
)
|
||||
name = 'claude-sonnet-4-5' if provider == 'anthropic' else 'anthropic/claude-sonnet-4.5'
|
||||
model = _runtime_model(adapter, name=name)
|
||||
model.provider.token_mgr.get_token = Mock(return_value='test')
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
|
||||
result, usage = await adapter.invoke_llm(
|
||||
None,
|
||||
model,
|
||||
[
|
||||
Message(role='system', content='stable'),
|
||||
Message(role='user', content='hello'),
|
||||
],
|
||||
extra_args={
|
||||
'client': _LeasedHTTPHandler(client, 10),
|
||||
'reasoning_config': {'budget_tokens': 2048},
|
||||
'max_tokens': 4096,
|
||||
'prompt_caching': {'mode': 'anthropic', 'ttl': '1h'},
|
||||
},
|
||||
)
|
||||
assert result.content == 'ok'
|
||||
assert usage
|
||||
body = bodies[0]
|
||||
if provider == 'anthropic':
|
||||
assert body['system'][-1]['cache_control'] == {'type': 'ephemeral', 'ttl': '1h'}
|
||||
assert body['thinking']['budget_tokens'] == 2048
|
||||
else:
|
||||
assert body['messages'][0]['content'][-1]['cache_control'] == {'type': 'ephemeral', 'ttl': '1h'}
|
||||
assert body['reasoning']['max_tokens'] == 2048
|
||||
assert 'reasoning_config' not in body and 'prompt_caching' not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_codex_reuses_real_connection(monkeypatch):
|
||||
monkeypatch.setenv('NO_PROXY', '127.0.0.1,localhost')
|
||||
sockets = []
|
||||
|
||||
async def handle(request):
|
||||
sockets.append(request.transport)
|
||||
await request.read()
|
||||
events = [
|
||||
{'type': 'response.output_text.delta', 'delta': 'ok'},
|
||||
{'type': 'response.completed', 'response': {'id': 'r', 'status': 'completed', 'output': []}},
|
||||
]
|
||||
return web.Response(
|
||||
text=''.join('data: ' + json.dumps(event) + '\n\n' for event in events), content_type='text/event-stream'
|
||||
)
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post('/responses', handle)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
monkeypatch.setattr(codex, 'BASE_URL', f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}')
|
||||
adapter = object.__new__(codex.CodexRequester)
|
||||
adapter.workspace, adapter.provider = 'w', 'p'
|
||||
adapter._replay = OrderedDict()
|
||||
adapter._http_pool = RequesterHTTPPool()
|
||||
adapter.auth = SimpleNamespace(
|
||||
access=AsyncMock(
|
||||
return_value={
|
||||
'access_token': 'test',
|
||||
'account_id': 'a',
|
||||
'connection_id': 'c',
|
||||
}
|
||||
)
|
||||
)
|
||||
model = SimpleNamespace(model_entity=SimpleNamespace(name='test', extra_args={}, reasoning_config=None))
|
||||
try:
|
||||
for _ in range(2):
|
||||
message, _ = await adapter.invoke_llm(None, model, [])
|
||||
assert message.content == 'ok'
|
||||
assert sockets[0] is sockets[1]
|
||||
finally:
|
||||
await adapter.aclose()
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_client_and_redirected_provider_are_preserved():
|
||||
adapter = _requester('openai')
|
||||
await adapter.initialize()
|
||||
explicit = object()
|
||||
for args in ({'model': 'openai/gpt-4o', 'client': explicit}, {'model': 'anthropic/claude-sonnet-4-5'}):
|
||||
async with adapter._client_args(args) as result:
|
||||
assert result is args
|
||||
assert adapter._http_pool.client is None
|
||||
await adapter.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_codex_terminal_tail_is_bounded_and_releases_stream():
|
||||
closed = asyncio.Event()
|
||||
|
||||
class HangingTail(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'data: {"type":"response.completed"}\n\n'
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
closed.set()
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=HangingTail()))
|
||||
) as client:
|
||||
async with client.stream('GET', 'https://example.com') as response:
|
||||
events = [event async for event in codex.sse_events(response)]
|
||||
assert events == [{'type': 'response.completed'}]
|
||||
assert closed.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_tls_keeps_litellm_transport(monkeypatch):
|
||||
import litellm
|
||||
|
||||
adapter = _requester('openai')
|
||||
await adapter.initialize()
|
||||
monkeypatch.setattr(litellm, 'ssl_verify', '/custom/ca.pem')
|
||||
args = {'model': 'openai/gpt-4o'}
|
||||
async with adapter._client_args(args) as result:
|
||||
assert result is args
|
||||
assert adapter._http_pool.client is None
|
||||
await adapter.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pooled_sdk_preserves_environment_fallbacks(monkeypatch):
|
||||
monkeypatch.setenv('OPENAI_BASE_URL', 'https://gateway.example/v1')
|
||||
monkeypatch.setenv('OPENAI_API_KEY', 'env-test-key')
|
||||
adapter = _requester('openai')
|
||||
await adapter.initialize()
|
||||
try:
|
||||
async with adapter._client_args({'model': 'openai/gpt-4o', 'api_key': ''}) as args:
|
||||
assert str(args['client'].base_url) == 'https://gateway.example/v1/'
|
||||
assert args['client'].api_key == 'env-test-key'
|
||||
finally:
|
||||
await adapter.aclose()
|
||||
@@ -0,0 +1,549 @@
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.provider.modelmgr.errors import RequesterError
|
||||
from langbot.pkg.provider.modelmgr.request_policy import (
|
||||
CallState,
|
||||
RequestPolicy,
|
||||
current_call,
|
||||
error_metadata,
|
||||
opaque_conversation_id,
|
||||
)
|
||||
from langbot.pkg.provider.modelmgr.requesters.litellmchat import LiteLLMRequester
|
||||
from langbot_plugin.api.entities.builtin.provider.message import Message, MessageChunk
|
||||
from tests.unit_tests.provider.conftest import TEST_EXECUTION_CONTEXT
|
||||
from tests.unit_tests.provider.test_litellmchat import MockRuntimeModel
|
||||
|
||||
|
||||
def test_policy_precedence_zero_and_no_mutation():
|
||||
model = {'num_retries': 3, 'request_policy': {'read_timeout': 42}}
|
||||
policy = RequestPolicy.resolve({'request_policy': {'retries': 2}}, model, {'num_retries': 0})
|
||||
assert policy.retries == 0
|
||||
assert policy.read_timeout == 42
|
||||
assert model['num_retries'] == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'policy',
|
||||
[
|
||||
{'retries': True},
|
||||
{'retries': -1},
|
||||
{'retries': 6},
|
||||
{'read_timeout': 0},
|
||||
{'total_timeout': float('nan')},
|
||||
{'connect_timeout': float('inf')},
|
||||
{'attribution': 'false'},
|
||||
{'conversation_header': 'authorization'},
|
||||
{'typo': 1},
|
||||
],
|
||||
)
|
||||
def test_invalid_policy(policy):
|
||||
with pytest.raises(RequesterError, match='模型请求失败'):
|
||||
RequestPolicy.resolve({'request_policy': policy})
|
||||
|
||||
|
||||
def test_retry_metadata_does_not_include_upstream_secrets():
|
||||
response = httpx.Response(429, headers={'retry-after': '2', 'x-request-id': 'req-123'})
|
||||
meta = error_metadata(SimpleNamespace(status_code=429, response=response))
|
||||
assert meta == {'status_code': 429, 'retry_after': 2, 'request_id': 'req-123'}
|
||||
adapter = LiteLLMRequester(Mock(), {})
|
||||
exc = litellm.RateLimitError('secret-key-and-prompt', llm_provider='openai', model='test', response=response)
|
||||
with pytest.raises(RequesterError) as caught:
|
||||
adapter._handle_litellm_error(exc)
|
||||
assert caught.value.retryable
|
||||
assert caught.value.request_id == 'req-123'
|
||||
assert 'secret-key' not in str(caught.value)
|
||||
|
||||
|
||||
def test_quota_is_not_retryable():
|
||||
adapter = LiteLLMRequester(Mock(), {})
|
||||
exc = litellm.RateLimitError('quota', llm_provider='openai', model='test')
|
||||
exc.body = {'error': {'code': 'insufficient_quota'}}
|
||||
with pytest.raises(RequesterError) as caught:
|
||||
adapter._handle_litellm_error(exc)
|
||||
assert caught.value.code == 'quota_exceeded'
|
||||
assert not caught.value.retryable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_headers_and_sdk_retry_ownership():
|
||||
adapter = LiteLLMRequester(Mock(), {'base_url': 'https://openrouter.ai/api/v1'})
|
||||
state = CallState(RequestPolicy(retries=2), 'lb-opaque')
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
args = await adapter._build_completion_args(
|
||||
MockRuntimeModel(),
|
||||
[],
|
||||
extra_args={
|
||||
'request_policy': {'retries': 2},
|
||||
'num_retries': 2,
|
||||
'extra_headers': {'x-openrouter-title': 'Custom app', 'x-other': 'keep'},
|
||||
},
|
||||
)
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
assert 'request_policy' not in args
|
||||
assert args['num_retries'] == args['max_retries'] == 0
|
||||
assert isinstance(args['timeout'], httpx.Timeout)
|
||||
assert args['extra_headers'] == {
|
||||
'x-openrouter-title': 'Custom app',
|
||||
'x-other': 'keep',
|
||||
'HTTP-Referer': 'https://langbot.app',
|
||||
'x-session-id': 'lb-opaque',
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('url', ['https://gateway.example/v1', 'https://openrouter.ai.evil.example/v1'])
|
||||
async def test_no_vendor_headers_on_unrelated_hosts(url):
|
||||
adapter = LiteLLMRequester(Mock(), {'base_url': url})
|
||||
args = await adapter._build_completion_args(MockRuntimeModel(), [])
|
||||
assert not args.get('extra_headers')
|
||||
assert args['num_retries'] == 0
|
||||
|
||||
|
||||
def test_conversation_is_stable_and_scoped():
|
||||
query = SimpleNamespace(session=None, launcher_type='person', launcher_id='123', bot_uuid='bot', pipeline_uuid='p')
|
||||
first = opaque_conversation_id(query, TEST_EXECUTION_CONTEXT)
|
||||
assert first == opaque_conversation_id(query, TEST_EXECUTION_CONTEXT)
|
||||
query.bot_uuid = 'another-bot'
|
||||
assert first != opaque_conversation_id(query, TEST_EXECUTION_CONTEXT)
|
||||
assert opaque_conversation_id(None, TEST_EXECUTION_CONTEXT) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_retries_only_before_output(runtime_provider, runtime_llm_model):
|
||||
provider = runtime_provider
|
||||
attempts = 0
|
||||
closed = []
|
||||
|
||||
async def stream(**kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
try:
|
||||
if attempts == 1:
|
||||
raise RequesterError('transient', retryable=True)
|
||||
yield MessageChunk(role='assistant', content='first')
|
||||
raise RequesterError('after output', retryable=True)
|
||||
finally:
|
||||
closed.append(attempts)
|
||||
|
||||
provider.requester.invoke_llm_stream = stream
|
||||
result = []
|
||||
with pytest.raises(RequesterError, match='after output'):
|
||||
async for chunk in provider.invoke_llm_stream(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'retries': 3, 'backoff_base': 0.001}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
):
|
||||
result.append(chunk.content)
|
||||
assert result == ['first']
|
||||
assert attempts == 2
|
||||
assert closed == [1, 2]
|
||||
assert current_call.get() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_deadline_and_close(runtime_provider, runtime_llm_model):
|
||||
closed = asyncio.Event()
|
||||
|
||||
async def stream(**kwargs):
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
yield MessageChunk(role='assistant', content='never')
|
||||
finally:
|
||||
closed.set()
|
||||
|
||||
runtime_provider.requester.invoke_llm_stream = stream
|
||||
with pytest.raises(RequesterError) as caught:
|
||||
async for _ in runtime_provider.invoke_llm_stream(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'first_event_timeout': 0.02}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
):
|
||||
pass
|
||||
assert caught.value.phase == 'first_event'
|
||||
assert closed.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_after_cannot_extend_total_budget():
|
||||
state = CallState(RequestPolicy(retries=5, total_timeout=0.1))
|
||||
state.metrics['attempts'] = 1
|
||||
assert not await state.retry(RequesterError('busy', retryable=True, retry_after=30))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancellation_is_not_retried(runtime_provider, runtime_llm_model):
|
||||
invoke = AsyncMock(side_effect=asyncio.CancelledError)
|
||||
runtime_provider.requester.invoke_llm = invoke
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await runtime_provider.invoke_llm(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'retries': 3}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
assert invoke.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_usage_and_timings_persist(runtime_provider, runtime_llm_model):
|
||||
from langbot_plugin.api.entities.builtin.pipeline.query import Query
|
||||
|
||||
query = Query.model_construct(
|
||||
query_id=1, launcher_type='person', launcher_id=1, variables={}, bot_uuid='bot', pipeline_uuid='p', session=None
|
||||
)
|
||||
object.__setattr__(query, '_execution_context', TEST_EXECUTION_CONTEXT)
|
||||
usage = {
|
||||
'prompt_tokens': 20,
|
||||
'completion_tokens': 10,
|
||||
'total_tokens': 30,
|
||||
'prompt_tokens_details': {'cached_tokens': 15},
|
||||
'completion_tokens_details': {'reasoning_tokens': 8},
|
||||
'cost': 0.01,
|
||||
}
|
||||
runtime_provider.requester.invoke_llm = AsyncMock(return_value=(Message(role='assistant', content='ok'), usage))
|
||||
with patch(
|
||||
'langbot.pkg.pipeline.monitoring_helper.MonitoringHelper.record_llm_call', new_callable=AsyncMock
|
||||
) as record:
|
||||
result = await runtime_provider.invoke_llm(query, runtime_llm_model, [])
|
||||
assert result.content == 'ok'
|
||||
details = record.call_args.kwargs['details']
|
||||
assert details['usage'] == usage
|
||||
assert details['timing']['attempts'] == 1
|
||||
assert details['timing']['total_ms'] >= details['timing']['response_completed_ms']
|
||||
assert query.variables['_llm_usage'] == usage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hidden_reasoning_marks_raw_event_and_usage():
|
||||
state = CallState(RequestPolicy())
|
||||
adapter = LiteLLMRequester(Mock(), {})
|
||||
|
||||
async def raw():
|
||||
yield SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(
|
||||
delta=SimpleNamespace(model_dump=lambda: {'reasoning_content': 'hidden'}), finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield SimpleNamespace(choices=[], usage={'prompt_tokens': 3, 'completion_tokens': 2, 'total_tokens': 5})
|
||||
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
with patch(
|
||||
'langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=AsyncMock(return_value=raw())
|
||||
):
|
||||
chunks = [c async for c in adapter.invoke_llm_stream(None, MockRuntimeModel(), [], remove_think=True)]
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
assert state.raw_started
|
||||
assert 'first_event_ms' in state.metrics and 'first_reasoning_ms' in state.metrics
|
||||
assert 'first_content_ms' not in state.metrics
|
||||
assert state.usage['total_tokens'] == 5
|
||||
assert not any(c.content for c in chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_http_retry_count_and_request_body(runtime_provider, runtime_llm_model):
|
||||
"""Exercise the pinned LiteLLM/OpenAI SDK against a local HTTP server."""
|
||||
from aiohttp import web
|
||||
|
||||
received = []
|
||||
|
||||
async def handler(request):
|
||||
received.append(await request.json())
|
||||
if len(received) == 1:
|
||||
return web.json_response({'error': {'message': 'busy', 'type': 'server_error'}}, status=503)
|
||||
return web.json_response(
|
||||
{
|
||||
'id': 'chatcmpl-local',
|
||||
'object': 'chat.completion',
|
||||
'created': 1,
|
||||
'model': 'gpt-4o',
|
||||
'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'ok'}, 'finish_reason': 'stop'}],
|
||||
'usage': {'prompt_tokens': 4, 'completion_tokens': 1, 'total_tokens': 5},
|
||||
}
|
||||
)
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post('/v1/chat/completions', handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
runtime_provider.requester = LiteLLMRequester(
|
||||
Mock(),
|
||||
{
|
||||
'base_url': f'http://127.0.0.1:{port}/v1',
|
||||
'custom_llm_provider': 'openai',
|
||||
},
|
||||
)
|
||||
runtime_llm_model.model_entity.name = 'gpt-4o'
|
||||
runtime_llm_model.model_entity.extra_args = {}
|
||||
try:
|
||||
with pytest.raises(RequesterError):
|
||||
await runtime_provider.invoke_llm(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[Message(role='user', content='hello')],
|
||||
extra_args={'num_retries': 0},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
assert len(received) == 1
|
||||
received.clear()
|
||||
result = await runtime_provider.invoke_llm(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[Message(role='user', content='hello')],
|
||||
extra_args={'request_policy': {'retries': 1, 'backoff_base': 0.001}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
assert result.content == 'ok'
|
||||
assert len(received) == 2
|
||||
assert all('request_policy' not in body and 'num_retries' not in body for body in received)
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_idle_timeout_after_output(runtime_provider, runtime_llm_model):
|
||||
async def stream(**kwargs):
|
||||
yield MessageChunk(role='assistant', content='first')
|
||||
await asyncio.Event().wait()
|
||||
|
||||
runtime_provider.requester.invoke_llm_stream = stream
|
||||
with pytest.raises(RequesterError) as caught:
|
||||
async for _ in runtime_provider.invoke_llm_stream(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'stream_idle_timeout': 0.02}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
):
|
||||
pass
|
||||
assert caught.value.phase == 'stream_idle'
|
||||
|
||||
|
||||
def test_details_migration_preserves_existing_records_and_is_idempotent():
|
||||
import importlib
|
||||
import sqlalchemy as sa
|
||||
from alembic.migration import MigrationContext
|
||||
from alembic.operations import Operations
|
||||
|
||||
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0038_llm_call_details')
|
||||
engine = sa.create_engine('sqlite://')
|
||||
try:
|
||||
with engine.begin() as connection:
|
||||
connection.exec_driver_sql('CREATE TABLE monitoring_llm_calls (id TEXT PRIMARY KEY)')
|
||||
connection.exec_driver_sql("INSERT INTO monitoring_llm_calls (id) VALUES ('old')")
|
||||
with Operations.context(MigrationContext.configure(connection)):
|
||||
migration.upgrade()
|
||||
migration.upgrade()
|
||||
assert connection.exec_driver_sql('SELECT id, details FROM monitoring_llm_calls').one() == ('old', None)
|
||||
table = sa.Table('monitoring_llm_calls', sa.MetaData(), autoload_with=connection)
|
||||
payload = {'usage': {'completion_tokens_details': {'reasoning_tokens': 8}}, 'timing': {'first_event_ms': 1}}
|
||||
connection.execute(table.insert().values(id='new', details=payload))
|
||||
assert connection.execute(sa.select(table.c.details).where(table.c.id == 'new')).scalar() == payload
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hidden_raw_events_reset_idle_not_first_visible_deadline(runtime_provider, runtime_llm_model):
|
||||
adapter = LiteLLMRequester(Mock(), {})
|
||||
runtime_provider.requester = adapter
|
||||
runtime_llm_model.model_entity.extra_args = {}
|
||||
|
||||
async def raw():
|
||||
for _ in range(5):
|
||||
await asyncio.sleep(0.3)
|
||||
yield SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'role': 'assistant'}), finish_reason=None)
|
||||
]
|
||||
)
|
||||
yield SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'content': 'answer'}), finish_reason='stop')
|
||||
]
|
||||
)
|
||||
|
||||
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=AsyncMock(return_value=raw())):
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in runtime_provider.invoke_llm_stream(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'first_event_timeout': 1, 'stream_idle_timeout': 2}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
]
|
||||
assert ''.join(chunk.content or '' for chunk in chunks) == 'answer'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitoring_service_persists_full_details():
|
||||
from langbot.pkg.api.http.service.monitoring import MonitoringService
|
||||
|
||||
persistence = SimpleNamespace(execute_async=AsyncMock())
|
||||
service = MonitoringService(SimpleNamespace(persistence_mgr=persistence))
|
||||
details = {
|
||||
'usage': {
|
||||
'prompt_tokens': 4,
|
||||
'completion_tokens': 3,
|
||||
'total_tokens': 9,
|
||||
'cache_read_input_tokens': 2,
|
||||
'provider_extra': {'cost': 0.01},
|
||||
},
|
||||
'timing': {'attempts': 2},
|
||||
'error': None,
|
||||
}
|
||||
await service.record_llm_call(
|
||||
TEST_EXECUTION_CONTEXT,
|
||||
bot_id='b',
|
||||
bot_name='b',
|
||||
pipeline_id='p',
|
||||
pipeline_name='p',
|
||||
session_id='s',
|
||||
model_name='m',
|
||||
input_tokens=4,
|
||||
output_tokens=3,
|
||||
duration=10,
|
||||
details=details,
|
||||
)
|
||||
values = persistence.execute_async.call_args.args[0].compile().params
|
||||
assert values['details'] == details
|
||||
assert values['total_tokens'] == 9
|
||||
|
||||
|
||||
def test_codex_usage_preserves_provider_fields():
|
||||
from langbot.pkg.provider.modelmgr.requesters.codex import _usage
|
||||
|
||||
raw = {
|
||||
'input_tokens': 4,
|
||||
'output_tokens': 2,
|
||||
'cost': 0.02,
|
||||
'output_tokens_details': {'reasoning_tokens': 1},
|
||||
'future_counter': 7,
|
||||
}
|
||||
normalized = _usage({'usage': raw})
|
||||
assert all(normalized[key] == value for key, value in raw.items())
|
||||
assert normalized['total_tokens'] == 6
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nonstream_uses_response_budget_not_first_event():
|
||||
adapter = LiteLLMRequester(Mock(), {})
|
||||
state = CallState(RequestPolicy(first_event_timeout=0.001, total_timeout=1))
|
||||
|
||||
async def completion(**kwargs):
|
||||
await asyncio.sleep(0.02)
|
||||
return 'complete'
|
||||
|
||||
token = current_call.set(state)
|
||||
try:
|
||||
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=completion):
|
||||
assert await adapter._completion({'stream': False}) == 'complete'
|
||||
finally:
|
||||
current_call.reset(token)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('cancel_read', [False, True])
|
||||
async def test_litellm_stream_close_and_cancel_release_upstream(runtime_provider, runtime_llm_model, cancel_read):
|
||||
adapter = LiteLLMRequester(Mock(), {})
|
||||
runtime_provider.requester = adapter
|
||||
runtime_llm_model.model_entity.extra_args = {}
|
||||
closed = asyncio.Event()
|
||||
reading = asyncio.Event()
|
||||
|
||||
async def raw():
|
||||
try:
|
||||
yield SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(delta=SimpleNamespace(model_dump=lambda: {'content': 'first'}), finish_reason=None)
|
||||
]
|
||||
)
|
||||
reading.set()
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
closed.set()
|
||||
|
||||
completion = AsyncMock(return_value=raw())
|
||||
with patch('langbot.pkg.provider.modelmgr.requesters.litellmchat.acompletion', new=completion):
|
||||
stream = runtime_provider.invoke_llm_stream(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'retries': 3}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
assert (await anext(stream)).content == 'first'
|
||||
if cancel_read:
|
||||
task = asyncio.create_task(anext(stream))
|
||||
await reading.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
else:
|
||||
await stream.aclose()
|
||||
assert closed.is_set()
|
||||
assert completion.await_count == 1
|
||||
assert current_call.get() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_backoff_does_not_start_another_attempt(runtime_provider, runtime_llm_model):
|
||||
runtime_provider.requester.invoke_llm = AsyncMock(side_effect=RequesterError('busy', retryable=True))
|
||||
sleeping = asyncio.Event()
|
||||
|
||||
async def backoff(delay):
|
||||
sleeping.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
with patch('langbot.pkg.provider.modelmgr.request_policy.asyncio.sleep', new=backoff):
|
||||
task = asyncio.create_task(
|
||||
runtime_provider.invoke_llm(
|
||||
None,
|
||||
runtime_llm_model,
|
||||
[],
|
||||
extra_args={'request_policy': {'retries': 3}},
|
||||
execution_context=TEST_EXECUTION_CONTEXT,
|
||||
)
|
||||
)
|
||||
await sleeping.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert runtime_provider.requester.invoke_llm.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_none_chunk_is_ignored(runtime_provider, runtime_llm_model):
|
||||
async def stream(**kwargs):
|
||||
yield None
|
||||
yield MessageChunk(role='assistant', content='ok')
|
||||
|
||||
runtime_provider.requester.invoke_llm_stream = stream
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in runtime_provider.invoke_llm_stream(
|
||||
None, runtime_llm_model, [], execution_context=TEST_EXECUTION_CONTEXT
|
||||
)
|
||||
]
|
||||
assert [chunk.content for chunk in chunks] == ['ok']
|
||||
Reference in New Issue
Block a user