70 lines
3.2 KiB
Python
70 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Run inside a separate LiteLLM container exec: fake OAuth failures, real API backup."""
|
|
import asyncio
|
|
from copy import deepcopy
|
|
from functools import wraps
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import litellm
|
|
import yaml
|
|
|
|
litellm.suppress_debug_info = True
|
|
|
|
|
|
async def main():
|
|
config = yaml.safe_load(Path('/tmp/litellm.yaml').read_text())
|
|
primary, backup = 'codex/gpt-luna-latest', 'openai/gpt-luna-latest'
|
|
routes = [deepcopy(r) for r in config['model_list'] if r['model_name'] in (primary, backup)]
|
|
assert len(routes) == 2
|
|
assert {primary: [backup]} in config['router_settings']['fallbacks']
|
|
for route in routes:
|
|
if route['model_name'] == backup:
|
|
route['litellm_params']['api_key'] = os.environ['OPENAI_API_KEY']
|
|
actual_chat, actual_responses = litellm.acompletion, litellm.aresponses
|
|
for protocol in ('chat', 'responses'):
|
|
for error_type, expected in ((litellm.RateLimitError, 429),
|
|
(litellm.AuthenticationError, 401),
|
|
(litellm.ServiceUnavailableError, 503)):
|
|
calls = []
|
|
|
|
async def invoke(real, *args, **kwargs):
|
|
model = kwargs.get('model', '')
|
|
if model.startswith('chatgpt/'):
|
|
calls.append('subscription-rejected')
|
|
raise error_type(message='Synthetic OAuth failure; isolated test only',
|
|
llm_provider='chatgpt', model=model)
|
|
assert model.startswith('openai/')
|
|
calls.append('official-api')
|
|
return await real(*args, **kwargs)
|
|
|
|
@wraps(actual_chat)
|
|
async def chat(*args, **kwargs):
|
|
return await invoke(actual_chat, *args, **kwargs)
|
|
|
|
@wraps(actual_responses)
|
|
async def responses(*args, **kwargs):
|
|
return await invoke(actual_responses, *args, **kwargs)
|
|
|
|
with patch.object(litellm, 'acompletion', chat), patch.object(litellm, 'aresponses', responses):
|
|
router = litellm.Router(model_list=deepcopy(routes), fallbacks=[{primary: [backup]}],
|
|
num_retries=0, timeout=60)
|
|
if protocol == 'chat':
|
|
result = await router.acompletion(model=primary, messages=[
|
|
{'role': 'user', 'content': 'Reply only OK. Synthetic test.'}], max_tokens=32)
|
|
assert result.choices and result.choices[0].message.content
|
|
else:
|
|
result = await router.aresponses(model=primary, input=[
|
|
{'role': 'user', 'content': 'Reply only OK. Synthetic test.'}],
|
|
max_output_tokens=32, store=False)
|
|
assert result.status == 'completed'
|
|
assert calls[0] == 'subscription-rejected' and calls[-1] == 'official-api', calls
|
|
print(json.dumps({'protocol': protocol, 'injected_error': expected,
|
|
'calls': calls, 'backup': backup, 'passed': True}), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|