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
@@ -0,0 +1,442 @@
|
||||
"""Real HTTP -> Core -> Plugin Runtime -> Runner -> Core -> HTTP provider tests.
|
||||
|
||||
Only the remote provider is deterministic. No requester, SDK action transport,
|
||||
model manager, persistence or monitoring implementation is mocked.
|
||||
"""
|
||||
|
||||
import json
|
||||
import importlib.metadata
|
||||
import os
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import select
|
||||
import socket
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
import httpx
|
||||
from packaging.version import Version
|
||||
|
||||
from tests.e2e import test_runner_plugin_runtime as runtime
|
||||
|
||||
# Reuse the real process lifecycle; override the package/config and port below.
|
||||
runner_runtime_ports = runtime.runner_runtime_ports
|
||||
runner_runtime_process = runtime.runner_runtime_process
|
||||
runner_langbot_process = runtime.runner_langbot_process
|
||||
runner_client = runtime.runner_client
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
RUNNER = runtime.QA_RUNNER_ID
|
||||
USAGE = {
|
||||
'prompt_tokens': 24,
|
||||
'completion_tokens': 4,
|
||||
'total_tokens': 28,
|
||||
'prompt_tokens_details': {'cached_tokens': 16},
|
||||
'completion_tokens_details': {'reasoning_tokens': 2},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def runner_e2e_tmpdir(tmp_path_factory):
|
||||
# Retain logs/SQLite in pytest's bounded temp history for failure diagnosis.
|
||||
return tmp_path_factory.mktemp('llm-chain')
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def runner_e2e_port():
|
||||
return runtime._free_port()
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def runner_e2e_config_path(runner_e2e_tmpdir, runner_e2e_port, runner_runtime_ports):
|
||||
path = runtime.runner_e2e_config_path.__wrapped__(runner_e2e_tmpdir, runner_e2e_port, runner_runtime_ports)
|
||||
root = runner_e2e_tmpdir / 'agent-runner-qa-package'
|
||||
manifest = root / 'components/runner/default.yaml'
|
||||
data = yaml.safe_load(manifest.read_text(encoding='utf-8'))
|
||||
data['spec']['permissions'] = {'models': ['invoke', 'stream', 'count_tokens']}
|
||||
data['spec']['config'] = [
|
||||
{'name': 'model', 'label': {'en_US': 'Model'}, 'type': 'llm-model-selector', 'required': True},
|
||||
]
|
||||
manifest.write_text(yaml.safe_dump(data), encoding='utf-8')
|
||||
shutil.copyfile(Path(__file__).parent / 'utils/llm_probe_runner.py', root / 'components/runner/default.py')
|
||||
if legacy_root := os.environ.get('LLM_E2E_LEGACY_API_DIR'):
|
||||
for filename in ('legacy_api.py', 'legacy_message.py'):
|
||||
shutil.copyfile(Path(legacy_root) / filename, root / filename)
|
||||
shutil.make_archive(str(runner_e2e_tmpdir / 'agent-runner-qa'), 'zip', root_dir=root)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope='module')
|
||||
def provider_wire():
|
||||
records = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
protocol_version = 'HTTP/1.1'
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
body = json.loads(self.rfile.read(int(self.headers['Content-Length'])))
|
||||
scenario = self.headers.get('x-e2e-case', 'normal')
|
||||
row = {'body': body, 'headers': dict(self.headers), 'case': scenario, 'closed': False}
|
||||
records.append(row)
|
||||
attempt = sum(x['case'] == scenario for x in records)
|
||||
status = 503 if scenario.startswith('retry') and attempt == 1 else 200
|
||||
if scenario == 'auth':
|
||||
status = 401
|
||||
if status != 200:
|
||||
data = json.dumps({'error': {'message': 'injected failure', 'type': 'api_error'}}).encode()
|
||||
self.send_response(status)
|
||||
self.send_header('Content-Type', 'application/json')
|
||||
self.send_header('Content-Length', str(len(data)))
|
||||
self.send_header('Retry-After', '0')
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
if scenario in ('first-timeout', 'total-timeout'):
|
||||
time.sleep(1)
|
||||
try:
|
||||
if not body.get('stream'):
|
||||
data = json.dumps(
|
||||
{
|
||||
'id': 'e2e',
|
||||
'object': 'chat.completion',
|
||||
'created': 1,
|
||||
'model': body['model'],
|
||||
'choices': [
|
||||
{
|
||||
'index': 0,
|
||||
'message': {'role': 'assistant', 'content': 'hello world'},
|
||||
'finish_reason': 'stop',
|
||||
}
|
||||
],
|
||||
'usage': USAGE,
|
||||
}
|
||||
).encode()
|
||||
self.send_response(200)
|
||||
self.send_header('Content-Type', 'application/json')
|
||||
self.send_header('Content-Length', str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
self.send_response(200)
|
||||
self.send_header('Content-Type', 'text/event-stream')
|
||||
self.send_header('Transfer-Encoding', 'chunked')
|
||||
self.end_headers()
|
||||
|
||||
def send(event):
|
||||
data = ('data: ' + (event if isinstance(event, str) else json.dumps(event)) + '\n\n').encode()
|
||||
self.wfile.write(f'{len(data):X}\r\n'.encode() + data + b'\r\n')
|
||||
self.wfile.flush()
|
||||
|
||||
base = {'id': 'e2e', 'object': 'chat.completion.chunk', 'created': 1, 'model': body['model']}
|
||||
if scenario.startswith('reasoning-'):
|
||||
send(
|
||||
{
|
||||
**base,
|
||||
'choices': [
|
||||
{'index': 0, 'delta': {'reasoning_content': 'private thought'}, 'finish_reason': None}
|
||||
],
|
||||
}
|
||||
)
|
||||
for index, content in enumerate(['hello ', 'world']):
|
||||
send({**base, 'choices': [{'index': 0, 'delta': {'content': content}, 'finish_reason': None}]})
|
||||
if index == 0 and (scenario in ('idle-timeout', 'disconnect') or scenario.startswith('cancel-')):
|
||||
if scenario == 'disconnect':
|
||||
self.close_connection = True
|
||||
return
|
||||
deadline = time.monotonic() + 1
|
||||
while time.monotonic() < deadline:
|
||||
if select.select([self.connection], [], [], 0.02)[0]:
|
||||
if not self.connection.recv(1, socket.MSG_PEEK):
|
||||
row['closed'] = True
|
||||
return
|
||||
send({**base, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})
|
||||
send({**base, 'choices': [], 'usage': USAGE})
|
||||
send('[DONE]')
|
||||
self.wfile.write(b'0\r\n\r\n')
|
||||
self.wfile.flush()
|
||||
except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
|
||||
row['closed'] = True
|
||||
|
||||
server = ThreadingHTTPServer(('127.0.0.1', 0), Handler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield f'http://127.0.0.1:{server.server_port}/v1', records
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(scope='module')
|
||||
def chain(request, provider_wire):
|
||||
request.getfixturevalue('runner_langbot_process')
|
||||
port = request.getfixturevalue('runner_e2e_port')
|
||||
client = httpx.Client(base_url=f'http://127.0.0.1:{port}', timeout=90, trust_env=False)
|
||||
request.addfinalizer(client.close)
|
||||
tmpdir = request.getfixturevalue('runner_e2e_tmpdir')
|
||||
token = runtime._init_and_auth(client)
|
||||
runtime._ensure_qa_plugin(client, token, tmpdir / 'agent-runner-qa.zip')
|
||||
client.headers['Authorization'] = f'Bearer {token}'
|
||||
|
||||
def post(path, body):
|
||||
response = client.post('/api/v1/' + path, json=body)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload['code'] == 0, payload
|
||||
return payload['data']
|
||||
|
||||
provider = post(
|
||||
'provider/providers',
|
||||
{
|
||||
'name': 'LLM chain E2E',
|
||||
'requester': 'openai-chat-completions',
|
||||
'base_url': provider_wire[0],
|
||||
'api_keys': ['e2e-not-a-secret'],
|
||||
},
|
||||
)['uuid']
|
||||
model = post(
|
||||
'provider/models/llm',
|
||||
{
|
||||
'name': 'gpt-4o',
|
||||
'provider_uuid': provider,
|
||||
'abilities': ['func_call', 'vision'],
|
||||
'extra_args': {
|
||||
'temperature': 0.1,
|
||||
'request_policy': {'retries': 1, 'backoff_base': 0.01, 'backoff_max': 0.01},
|
||||
},
|
||||
},
|
||||
)['uuid']
|
||||
agent = post(
|
||||
'agents',
|
||||
{
|
||||
'kind': 'agent',
|
||||
'name': 'LLM chain E2E',
|
||||
'config': {
|
||||
'runner': {'id': RUNNER},
|
||||
'runner_config': {RUNNER: {'model': model}},
|
||||
},
|
||||
},
|
||||
)['uuid']
|
||||
|
||||
def invoke(mode, case='normal', extra=None, *, configuration=None, options=None):
|
||||
selected_model, selected_agent = model, agent
|
||||
if configuration:
|
||||
configured_provider = post(
|
||||
'provider/providers',
|
||||
{
|
||||
'name': case,
|
||||
'requester': configuration.get('requester', 'openai-chat-completions'),
|
||||
'base_url': provider_wire[0],
|
||||
'api_keys': ['e2e-not-a-secret'],
|
||||
},
|
||||
)['uuid']
|
||||
selected_model = post(
|
||||
'provider/models/llm',
|
||||
{
|
||||
'name': configuration.get('name', 'gpt-4o'),
|
||||
'provider_uuid': configured_provider,
|
||||
'abilities': ['func_call', 'vision'],
|
||||
'extra_args': configuration.get('extra_args', {}),
|
||||
},
|
||||
)['uuid']
|
||||
selected_agent = post(
|
||||
'agents',
|
||||
{
|
||||
'kind': 'agent',
|
||||
'name': case,
|
||||
'config': {
|
||||
'runner': {'id': RUNNER},
|
||||
'runner_config': {RUNNER: {'model': selected_model}},
|
||||
},
|
||||
},
|
||||
)['uuid']
|
||||
result = post(
|
||||
f'agents/{selected_agent}/debug',
|
||||
{
|
||||
'text': json.dumps(
|
||||
{'mode': mode, 'model': selected_model, 'case': case, 'extra': extra or {}, **(options or {})}
|
||||
),
|
||||
'conversation_id': 'e2e-' + case,
|
||||
'event_type': 'message.received',
|
||||
},
|
||||
)
|
||||
assert result.get('final_text'), result
|
||||
return json.loads(result['final_text'])
|
||||
|
||||
return invoke, provider_wire[1], tmpdir
|
||||
|
||||
|
||||
def test_multimodal_and_tool_definitions_survive_plugin_transport(chain):
|
||||
invoke, records, _ = chain
|
||||
image = 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII='
|
||||
messages = [
|
||||
{
|
||||
'role': 'user',
|
||||
'content': [
|
||||
{'type': 'text', 'text': 'Describe this image'},
|
||||
{'type': 'image_url', 'image_url': {'url': image}},
|
||||
],
|
||||
}
|
||||
]
|
||||
result = invoke('runner-events', 'multimodal', options={'messages': messages, 'tool': True})
|
||||
assert result['text'] == 'hello world', result
|
||||
wire = records[-1]['body']
|
||||
assert wire['messages'][0]['content'][1]['image_url']['url'] == image
|
||||
assert wire['tools'][0]['function']['name'] == 'e2e_lookup'
|
||||
|
||||
|
||||
def test_openrouter_budget_cache_and_explicit_attribution(chain):
|
||||
invoke, records, _ = chain
|
||||
result = invoke(
|
||||
'runner-events',
|
||||
'openrouter-budget-cache',
|
||||
configuration={
|
||||
'requester': 'openrouter-chat-completions',
|
||||
'name': 'anthropic/claude-sonnet-4',
|
||||
'extra_args': {
|
||||
'reasoning_config': {'budget_tokens': 2048},
|
||||
'prompt_caching': {'mode': 'anthropic'},
|
||||
'max_tokens': 4096,
|
||||
},
|
||||
},
|
||||
extra={
|
||||
'extra_headers': {'HTTP-Referer': 'https://langbot.app', 'X-OpenRouter-Title': 'LangBot'},
|
||||
'request_policy': {'conversation_header': 'x-langbot-conversation-id'},
|
||||
},
|
||||
options={'messages': [{'role': 'system', 'content': 'Stable prefix'}, {'role': 'user', 'content': 'Hello'}]},
|
||||
)
|
||||
assert result['text'] == 'hello world', result
|
||||
wire = records[-1]
|
||||
assert wire['body']['reasoning']['max_tokens'] == 2048
|
||||
assert wire['body']['messages'][0]['content'][0]['cache_control']['type'] == 'ephemeral'
|
||||
headers = {k.lower(): v for k, v in wire['headers'].items()}
|
||||
assert headers['http-referer'] == 'https://langbot.app'
|
||||
assert headers['x-openrouter-title'] == 'LangBot'
|
||||
assert headers['x-langbot-conversation-id'].startswith('lb-')
|
||||
assert 'prompt_caching' not in wire['body'] and 'reasoning_config' not in wire['body']
|
||||
|
||||
|
||||
def test_openai_cache_options_and_call_override(chain):
|
||||
invoke, records, _ = chain
|
||||
result = invoke(
|
||||
'runner-usage',
|
||||
'openai-cache',
|
||||
configuration={
|
||||
'extra_args': {'temperature': 0.1, 'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'}},
|
||||
},
|
||||
extra={'temperature': 0.7, 'prompt_cache_key': 'explicit-test-cache'},
|
||||
)
|
||||
assert result['text'] == 'hello world', result
|
||||
assert records[-1]['body']['temperature'] == 0.7
|
||||
assert records[-1]['body']['prompt_cache_key'] == 'explicit-test-cache'
|
||||
assert records[-1]['body']['prompt_cache_retention'] == 'in_memory'
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'mode',
|
||||
[
|
||||
'runner-message',
|
||||
'runner-usage',
|
||||
'runner-stream',
|
||||
'runner-events',
|
||||
'plugin-message',
|
||||
'plugin-usage',
|
||||
'plugin-stream',
|
||||
'plugin-events',
|
||||
],
|
||||
)
|
||||
def test_public_api_shapes_and_usage(chain, mode):
|
||||
invoke, records, _ = chain
|
||||
result = invoke(mode, mode)
|
||||
assert result['text'] == 'hello world', result
|
||||
assert not result.get('error'), result
|
||||
if mode.endswith(('usage', 'events')):
|
||||
assert result['usage']['prompt_tokens_details']['cached_tokens'] == 16
|
||||
assert result['usage']['completion_tokens_details']['reasoning_tokens'] == 2
|
||||
wire = records[-1]['body']
|
||||
assert wire['temperature'] == 0.1
|
||||
assert 'request_policy' not in wire
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'case,extra,attempts,error',
|
||||
[
|
||||
('retry-once', {}, 2, False),
|
||||
('retry-disabled', {'num_retries': 0}, 1, True),
|
||||
('retry-max-disabled', {'max_retries': 0}, 1, True),
|
||||
('auth', {}, 1, True),
|
||||
('first-timeout', {'request_policy': {'first_event_timeout': 0.2, 'retries': 0}}, 1, True),
|
||||
('total-timeout', {'request_policy': {'total_timeout': 0.2, 'retries': 0}}, 1, True),
|
||||
('idle-timeout', {'request_policy': {'stream_idle_timeout': 0.2}}, 1, True),
|
||||
('disconnect', {}, 1, True),
|
||||
],
|
||||
)
|
||||
def test_failures_retry_boundaries_and_recovery(chain, case, extra, attempts, error):
|
||||
invoke, records, _ = chain
|
||||
result = invoke('runner-events', case, extra)
|
||||
assert bool(result.get('error')) == error, result
|
||||
assert sum(row['case'] == case for row in records) == attempts
|
||||
assert invoke('runner-events', 'recover-' + case)['text'] == 'hello world'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('surface', ['runner-events', 'runner-stream', 'plugin-events', 'plugin-stream'])
|
||||
def test_cancellation_releases_upstream_and_next_request_works(chain, surface):
|
||||
invoke, records, tmpdir = chain
|
||||
case = 'cancel-' + surface
|
||||
result = invoke(case, case)
|
||||
assert result['text'] == 'hello ', result
|
||||
time.sleep(1.5)
|
||||
cancelled = [row for row in records if row['case'] == case]
|
||||
with sqlite3.connect(tmpdir / 'data/langbot.db') as db:
|
||||
diagnostics = db.execute('SELECT session_id, details FROM monitoring_llm_calls').fetchall()
|
||||
assert len(cancelled) == 1
|
||||
assert invoke('runner-events', 'after-' + case)['text'] == 'hello world'
|
||||
if (
|
||||
not cancelled[0]['closed']
|
||||
and Version(importlib.metadata.version('langbot-plugin')) <= Version('0.7.11')
|
||||
and not os.environ.get('LLM_E2E_REQUIRE_STREAM_CLOSE')
|
||||
):
|
||||
pytest.xfail('Published SDK <=0.7.11 omits aclosing in public proxy wrappers; requires companion SDK fix')
|
||||
assert cancelled[0]['closed'], diagnostics[-2:]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('remove_think', [True, False])
|
||||
def test_reasoning_visibility_through_runner(chain, remove_think):
|
||||
result = chain[0]('runner-events', f'reasoning-{remove_think}', options={'remove_think': remove_think})
|
||||
assert result['text'].endswith('hello world'), result
|
||||
assert ('private thought' in result['text']) == (not remove_think)
|
||||
|
||||
|
||||
def test_invalid_budget_fails_before_network_and_recovers(chain):
|
||||
invoke, records, _ = chain
|
||||
result = invoke('runner-events', 'invalid-budget', {'reasoning_config': {'budget_tokens': 1}})
|
||||
assert result.get('error'), result
|
||||
assert not any(row['case'] == 'invalid-budget' for row in records)
|
||||
assert invoke('runner-events', 'after-invalid-budget')['text'] == 'hello world'
|
||||
|
||||
|
||||
def test_monitoring_persists_complete_usage_and_structured_errors(chain):
|
||||
invoke, _, tmpdir = chain
|
||||
invoke('runner-events', 'monitoring')
|
||||
with sqlite3.connect(tmpdir / 'data/langbot.db') as db:
|
||||
rows = db.execute('SELECT details FROM monitoring_llm_calls').fetchall()
|
||||
details = [json.loads(row[0]) for row in rows if row[0]]
|
||||
assert any(x.get('usage', {}).get('prompt_tokens_details', {}).get('cached_tokens') == 16 for x in details)
|
||||
assert any(x.get('timing', {}).get('transport_setup_ms') is not None for x in details)
|
||||
assert any(x.get('error', {}).get('code') for x in details if x.get('error'))
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get('LLM_E2E_LEGACY_API_DIR'), reason='Historical SDK sources not supplied')
|
||||
@pytest.mark.parametrize('mode', ['legacy-message', 'legacy-usage', 'legacy-stream', 'legacy-events'])
|
||||
def test_historical_sdk_proxy_and_message_entities(chain, mode):
|
||||
result = chain[0](mode, mode)
|
||||
assert result['text'] == 'hello world', result
|
||||
assert not result.get('error'), result
|
||||
if mode.endswith(('usage', 'events')):
|
||||
assert result['usage']['prompt_tokens_details']['cached_tokens'] == 16
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Packaged into the isolated E2E plugin; calls only public SDK LLM APIs."""
|
||||
|
||||
from contextlib import aclosing
|
||||
import json
|
||||
|
||||
from langbot_plugin.api.definition.components.runner.runner import Runner
|
||||
from langbot_plugin.api.entities.builtin.provider.message import Message
|
||||
from langbot_plugin.api.entities.builtin.runner.result import RunnerResult
|
||||
from langbot_plugin.api.proxies.langbot_api import LangBotAPIProxy
|
||||
|
||||
|
||||
class DefaultRunner(Runner):
|
||||
async def run(self, ctx):
|
||||
options = json.loads(ctx.input.to_text())
|
||||
mode = options['mode']
|
||||
cancel = mode.startswith('cancel-')
|
||||
if cancel:
|
||||
mode = mode.removeprefix('cancel-')
|
||||
api = LangBotAPIProxy(self._plugin_runtime_handler) if mode.startswith('plugin-') else ctx.api
|
||||
message_type = Message
|
||||
if mode.startswith('legacy-'):
|
||||
# Historical files are supplied verbatim from a released SDK tag.
|
||||
# Its proxy and serializers use the real current runtime transport.
|
||||
import legacy_api
|
||||
import legacy_message
|
||||
|
||||
legacy_api.provider_message = legacy_message
|
||||
api = legacy_api.LangBotAPIProxy(self._plugin_runtime_handler)
|
||||
message_type = legacy_message.Message
|
||||
extra = {
|
||||
**options['extra'],
|
||||
'extra_headers': {
|
||||
**options['extra'].get('extra_headers', {}),
|
||||
'x-e2e-case': options['case'],
|
||||
},
|
||||
}
|
||||
kwargs = {
|
||||
'llm_model_uuid': options['model'],
|
||||
'messages': [
|
||||
message_type.model_validate(m)
|
||||
for m in options.get(
|
||||
'messages',
|
||||
[
|
||||
{'role': 'user', 'content': 'Hello'},
|
||||
],
|
||||
)
|
||||
],
|
||||
'extra_args': extra,
|
||||
}
|
||||
if options.get('tool'):
|
||||
from langbot_plugin.api.entities.builtin.resource.tool import LLMTool
|
||||
|
||||
async def lookup(**kwargs):
|
||||
return 'result'
|
||||
|
||||
kwargs['funcs'] = [
|
||||
LLMTool(
|
||||
name='e2e_lookup',
|
||||
human_desc='Lookup',
|
||||
description='Lookup',
|
||||
parameters={'type': 'object', 'properties': {}},
|
||||
func=lookup,
|
||||
)
|
||||
]
|
||||
result = {'text': '', 'usage': None}
|
||||
if 'remove_think' in options:
|
||||
kwargs['remove_think'] = options['remove_think']
|
||||
try:
|
||||
if mode.endswith('message'):
|
||||
message = await api.invoke_llm(**kwargs)
|
||||
assert isinstance(message, message_type)
|
||||
result['text'] = message.content
|
||||
elif mode.endswith('usage'):
|
||||
response = await api.invoke_llm_with_usage(**kwargs)
|
||||
result.update(
|
||||
text=response.message.content, usage=response.usage.model_dump() if response.usage else None
|
||||
)
|
||||
else:
|
||||
events = mode.endswith('events')
|
||||
method = api.invoke_llm_stream_events if events else api.invoke_llm_stream
|
||||
async with aclosing(method(**kwargs)) as stream:
|
||||
async for item in stream:
|
||||
chunk = item.chunk if events else item
|
||||
if chunk:
|
||||
result['text'] += chunk.content or ''
|
||||
if events and item.usage:
|
||||
result['usage'] = item.usage.model_dump()
|
||||
if cancel and result['text']:
|
||||
break
|
||||
except Exception as exc:
|
||||
result['error'] = type(exc).__name__ + ': ' + str(exc)
|
||||
yield RunnerResult.message_completed(ctx.run_id, Message(role='assistant', content=json.dumps(result)))
|
||||
yield RunnerResult.run_completed(ctx.run_id, finish_reason='stop')
|
||||
@@ -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