mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
* fix(ci): stop five stale or flaky CI reds and retry CyberArk policy-load conflicts The Langfuse redaction unit test exports to a local OTLP capture instead of polling Langfuse Cloud through a recorded lookup. The passthrough worker-kill test only requires spend rows for requests the surviving worker served. The spend-routes sweep treats the intentional /spend/capture_rate 503 as expected. CyberArk retries a 409 policy load in Python, Rust and the e2e Conjur helper instead of reading it as "variable exists". The integration egress guard now matches the script's own cgroup, so it no longer blocks the CircleCI agent, which runs as the same user. * fix(ci): keep the policy-load backoff typed as float * fix(ci): retry CyberArk policy loads without blocking the event loop and tighten the worker-kill and Langfuse tests * fix(secrets): load CyberArk policy one request at a time per manager * test(secrets): pin that non-conflict CyberArk policy failures are not retried * test(unit): run tests/unit with only an allowlisted host environment CircleCI's unit job inherits every project env var, so real provider keys, REDIS_HOST, DATABASE_URL and AWS or Azure credentials reached tests that assume none are set. Locally, litellm's import-time load_dotenv did the same from any .env up the tree. The unit conftest now drops every variable outside a small allowlist and disables dotenv before litellm is imported. * test(e2e): name a failed search and the stuck batch status instead of misattributing them The websearch session test read an empty web_search_tool_result_error block as a successful search, so a failing search tool surfaced as a session billing bug. The batch cancellation timeout now reports the last status the proxy returned. * fix(ci): scrub the host environment per unit test instead of for the whole pytest process GHA shards run tests/unit next to other suites in one process, so the import-time scrub deleted MCP_TEST_PEER_PYTHON before tests/mcp_tests read it and the MCP upstream fell back to the SDK2 interpreter. The two websearch tests that called OpenAI and Perplexity live are removed: tests/unit no longer sees their keys. * fix(ci): scrub only the host variables present before litellm is imported The per-test scrub also deleted TIKTOKEN_CACHE_DIR, which litellm sets at import to its bundled encodings, so tokenizer paths tried to download them and hit the socket guard. The prisma setup test now passes its own database URL instead of reading one another test leaked into the process environment. * fix(ci): stop the order-dependent unit reds and settle logging tasks on their own queue LoggingWorker marked a task done on whichever queue was current when the callback finished, so a callback that outlived an event-loop change raised "task_done() called too many times" or undercounted the new loop's queue. It now settles the queue the task came from. The rest are test isolation fixes for failures that only appeared when another file ran first on the same xdist worker: a replaced user_api_key_cache, breaker metrics unregistered by prometheus tests, semantic_router's health-check filter on uvicorn.access, logging tasks carried over from bedrock tests, a Router-written model_cost entry, and a stray post captured by the langflow test. The token counter check now asserts bounded chunking instead of wall-clock time. * test(e2e/ui): wait for the logout redirect before visiting a protected page Logout revokes the session server-side before clearing cookies and navigating, so an immediate page.goto either ran with the cookie still set or was aborted by the logout redirect (net::ERR_ABORTED). * test(unit): restore the prometheus metrics config per test and settle logs carried from earlier tests in the a2a cost tests * test(router): pin the router clock in the usage counter tests so a minute rollover cannot empty the read * test(e2e/ui): wait for logout to clear the token cookie instead of for a login redirect * test(integration/mcp): answer the model-info probe another test's proxy sends to the model double
3174 lines
111 KiB
Python
3174 lines
111 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import traceback
|
|
from dotenv import load_dotenv
|
|
from fastapi import Request
|
|
from datetime import datetime, timezone
|
|
|
|
from litellm import Router
|
|
import pytest
|
|
import litellm
|
|
from unittest.mock import patch, MagicMock, AsyncMock
|
|
from create_mock_standard_logging_payload import create_standard_logging_payload
|
|
from litellm.types.utils import ModelResponse, StandardLoggingPayload
|
|
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
|
from litellm.caching.dual_cache import DualCache
|
|
from litellm.caching.in_memory_cache import InMemoryCache
|
|
from litellm.types.caching import RedisPipelineIncrementOperation
|
|
from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
|
|
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
|
|
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
|
|
|
|
|
|
@pytest.fixture
|
|
def model_list():
|
|
return [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"tpm": 1000, # Add TPM limit so async method doesn't return early
|
|
"rpm": 100, # Add RPM limit so async method doesn't return early
|
|
},
|
|
"model_info": {
|
|
"access_groups": ["group1", "group2"],
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-5.5",
|
|
"litellm_params": {
|
|
"model": "gpt-5.5",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-image-1",
|
|
"litellm_params": {
|
|
"model": "gpt-image-1",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "*",
|
|
"litellm_params": {
|
|
"model": "openai/*",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "claude-*",
|
|
"litellm_params": {
|
|
"model": "anthropic/*",
|
|
"api_key": os.getenv("ANTHROPIC_API_KEY"),
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
def test_validate_fallbacks(model_list):
|
|
router = Router(model_list=model_list, fallbacks=[{"gpt-5.5": "gpt-5-mini"}])
|
|
router.validate_fallbacks(fallback_param=[{"gpt-5.5": "gpt-5-mini"}])
|
|
|
|
|
|
def test_routing_strategy_init(model_list):
|
|
"""Test if all routing strategies are initialized correctly"""
|
|
from litellm.types.router import RoutingStrategy
|
|
|
|
router = Router(model_list=model_list)
|
|
for strategy in RoutingStrategy:
|
|
router.routing_strategy_init(
|
|
routing_strategy=strategy, routing_strategy_args={}
|
|
)
|
|
|
|
|
|
def test_routing_strategy_init_invalid_strategy(model_list):
|
|
"""Test that invalid routing_strategy raises ValueError with helpful message.
|
|
|
|
See: https://github.com/BerriAI/litellm/issues/11330
|
|
Invalid strategies like 'simple' (without '-shuffle') should fail fast
|
|
with a clear error, not silently cause 'No deployments available' errors.
|
|
"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Test common mistake: "simple" instead of "simple-shuffle"
|
|
with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info:
|
|
router.routing_strategy_init(
|
|
routing_strategy="simple", routing_strategy_args={}
|
|
)
|
|
|
|
# Verify error message is helpful
|
|
error_msg = str(exc_info.value)
|
|
assert "Invalid routing_strategy" in error_msg
|
|
assert "simple" in error_msg
|
|
assert "simple-shuffle" in error_msg # Suggests the correct option
|
|
# Verify error message tells user WHERE to fix it
|
|
assert "config.yaml" in error_msg
|
|
assert "router_settings.routing_strategy" in error_msg
|
|
assert "Router SDK" in error_msg
|
|
|
|
# Test completely invalid strategy
|
|
with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info:
|
|
router.routing_strategy_init(
|
|
routing_strategy="not-a-real-strategy", routing_strategy_args={}
|
|
)
|
|
assert "Invalid routing_strategy" in str(exc_info.value)
|
|
|
|
|
|
def test_routing_strategy_init_valid_string_strategies(model_list):
|
|
"""Test that all valid string routing strategies work without error.
|
|
|
|
Valid strategies are derived from RoutingStrategy enum values plus 'simple-shuffle'.
|
|
"""
|
|
from litellm.types.router import RoutingStrategy
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
# All strategies from enum + simple-shuffle (default, not in enum)
|
|
valid_strategies = ["simple-shuffle"] + [s.value for s in RoutingStrategy]
|
|
|
|
for strategy in valid_strategies:
|
|
# Should not raise
|
|
router.routing_strategy_init(
|
|
routing_strategy=strategy, routing_strategy_args={}
|
|
)
|
|
|
|
|
|
def test_print_deployment(model_list):
|
|
"""Test if the api key is masked correctly"""
|
|
|
|
router = Router(model_list=model_list)
|
|
deployment = {
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
}
|
|
printed_deployment = router.print_deployment(deployment)
|
|
assert 10 * "*" in printed_deployment["litellm_params"]["api_key"]
|
|
|
|
|
|
def test_print_deployment_with_redact_enabled(model_list):
|
|
"""Test if sensitive credentials are masked when redact_user_api_key_info is enabled"""
|
|
import litellm
|
|
|
|
router = Router(model_list=model_list)
|
|
deployment = {
|
|
"model_name": "bedrock-claude",
|
|
"litellm_params": {
|
|
"model": "bedrock/anthropic.claude-v2",
|
|
"aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
|
|
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
|
"aws_region_name": "us-west-2",
|
|
},
|
|
}
|
|
|
|
original_setting = litellm.redact_user_api_key_info
|
|
try:
|
|
litellm.redact_user_api_key_info = True
|
|
printed_deployment = router.print_deployment(deployment)
|
|
|
|
assert "*" in printed_deployment["litellm_params"]["aws_access_key_id"]
|
|
assert "*" in printed_deployment["litellm_params"]["aws_secret_access_key"]
|
|
assert "us-west-2" == printed_deployment["litellm_params"]["aws_region_name"]
|
|
finally:
|
|
litellm.redact_user_api_key_info = original_setting
|
|
|
|
|
|
def test_completion(model_list):
|
|
"""Test if the completion function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
response = router._completion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.flaky(retries=6, delay=1)
|
|
@pytest.mark.asyncio
|
|
async def test_image_generation(model_list, sync_mode):
|
|
"""Test if the underlying '_image_generation' function is working correctly"""
|
|
from litellm.types.utils import ImageResponse
|
|
|
|
router = Router(model_list=model_list)
|
|
if sync_mode:
|
|
response = router._image_generation(
|
|
model="gpt-image-1",
|
|
prompt="A cute baby sea otter",
|
|
)
|
|
else:
|
|
response = await router._aimage_generation(
|
|
model="gpt-image-1",
|
|
prompt="A cute baby sea otter",
|
|
)
|
|
|
|
ImageResponse.model_validate(response)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_acompletion_util(model_list):
|
|
"""Test if the underlying '_acompletion' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
response = await router._acompletion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_abatch_completion_one_model_multiple_requests_util(model_list):
|
|
"""Test if the 'abatch_completion_one_model_multiple_requests' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
response = await router.abatch_completion_one_model_multiple_requests(
|
|
model="gpt-5-mini",
|
|
messages=[
|
|
[{"role": "user", "content": "Hello, how are you?"}],
|
|
[{"role": "user", "content": "Hello, how are you?"}],
|
|
],
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
print(response)
|
|
assert response[0]["choices"][0]["message"]["content"] == "I'm fine, thank you!"
|
|
assert response[1]["choices"][0]["message"]["content"] == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_schedule_acompletion(model_list):
|
|
"""Test if the 'schedule_acompletion' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
response = await router.schedule_acompletion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response="I'm fine, thank you!",
|
|
priority=1,
|
|
)
|
|
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_schedule_atext_completion(model_list):
|
|
"""Test if the 'schedule_atext_completion' function is working correctly"""
|
|
from litellm.types.utils import TextCompletionResponse
|
|
|
|
router = Router(model_list=model_list)
|
|
with patch.object(
|
|
router, "_atext_completion", AsyncMock()
|
|
) as mock_atext_completion:
|
|
mock_atext_completion.return_value = TextCompletionResponse()
|
|
response = await router.atext_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
priority=1,
|
|
)
|
|
mock_atext_completion.assert_awaited_once()
|
|
assert "priority" not in mock_atext_completion.call_args.kwargs
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_schedule_factory(model_list):
|
|
"""Test if the 'schedule_atext_completion' function is working correctly"""
|
|
from litellm.types.utils import TextCompletionResponse
|
|
|
|
router = Router(model_list=model_list)
|
|
with patch.object(
|
|
router, "_atext_completion", AsyncMock()
|
|
) as mock_atext_completion:
|
|
mock_atext_completion.return_value = TextCompletionResponse()
|
|
response = await router._schedule_factory(
|
|
model="gpt-5-mini",
|
|
args=(
|
|
"gpt-5-mini",
|
|
"Hello, how are you?",
|
|
),
|
|
priority=1,
|
|
kwargs={},
|
|
original_function=router.atext_completion,
|
|
)
|
|
mock_atext_completion.assert_awaited_once()
|
|
assert "priority" not in mock_atext_completion.call_args.kwargs
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_router_function_with_fallbacks(model_list, sync_mode):
|
|
"""Test if the router 'async_function_with_fallbacks' + 'function_with_fallbacks' are working correctly"""
|
|
router = Router(model_list=model_list)
|
|
data = {
|
|
"model": "gpt-5-mini",
|
|
"messages": [{"role": "user", "content": "Hello, how are you?"}],
|
|
"mock_response": "I'm fine, thank you!",
|
|
"num_retries": 0,
|
|
}
|
|
if sync_mode:
|
|
response = router.function_with_fallbacks(
|
|
original_function=router._completion,
|
|
**data,
|
|
)
|
|
else:
|
|
response = await router.async_function_with_fallbacks(
|
|
original_function=router._acompletion,
|
|
**data,
|
|
)
|
|
assert response.choices[0].message.content == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_router_function_with_retries(model_list, sync_mode):
|
|
"""Test if the router 'async_function_with_retries' + 'function_with_retries' are working correctly"""
|
|
router = Router(model_list=model_list)
|
|
data = {
|
|
"model": "gpt-5-mini",
|
|
"messages": [{"role": "user", "content": "Hello, how are you?"}],
|
|
"mock_response": "I'm fine, thank you!",
|
|
"num_retries": 0,
|
|
}
|
|
response = await router.async_function_with_retries(
|
|
original_function=router._acompletion,
|
|
**data,
|
|
)
|
|
|
|
assert response.choices[0].message.content == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_make_call(model_list):
|
|
"""Test if the router 'make_call' function is working correctly"""
|
|
|
|
## ACOMPLETION
|
|
router = Router(model_list=model_list)
|
|
response = await router.make_call(
|
|
original_function=router._acompletion,
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
assert response.choices[0].message.content == "I'm fine, thank you!"
|
|
|
|
## ATEXT_COMPLETION
|
|
response = await router.make_call(
|
|
original_function=router._atext_completion,
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
assert response.choices[0].text == "I'm fine, thank you!"
|
|
|
|
## AEMBEDDING
|
|
response = await router.make_call(
|
|
original_function=router._aembedding,
|
|
model="gpt-5-mini",
|
|
input="Hello, how are you?",
|
|
mock_response=[0.1, 0.2, 0.3],
|
|
)
|
|
assert response.data[0].embedding == [0.1, 0.2, 0.3]
|
|
|
|
## AIMAGE_GENERATION
|
|
response = await router.make_call(
|
|
original_function=router._aimage_generation,
|
|
model="gpt-image-1",
|
|
prompt="A cute baby sea otter",
|
|
mock_response="https://example.com/image.png",
|
|
)
|
|
assert response.data[0].url == "https://example.com/image.png"
|
|
|
|
|
|
def test_update_kwargs_with_deployment(model_list):
|
|
"""Test if the '_update_kwargs_with_deployment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
kwargs: dict = {"metadata": {}}
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
router._update_kwargs_with_deployment(
|
|
deployment=deployment,
|
|
kwargs=kwargs,
|
|
)
|
|
set_fields = ["deployment", "api_base", "model_info"]
|
|
assert all(field in kwargs["metadata"] for field in set_fields)
|
|
|
|
|
|
def test_update_kwargs_with_default_litellm_params(model_list):
|
|
"""Test if the '_update_kwargs_with_default_litellm_params' function is working correctly"""
|
|
router = Router(
|
|
model_list=model_list,
|
|
default_litellm_params={"api_key": "test", "metadata": {"key": "value"}},
|
|
)
|
|
kwargs: dict = {"metadata": {"key2": "value2"}}
|
|
router._update_kwargs_with_default_litellm_params(kwargs=kwargs)
|
|
assert kwargs["api_key"] == "test"
|
|
assert kwargs["metadata"]["key"] == "value"
|
|
assert kwargs["metadata"]["key2"] == "value2"
|
|
|
|
|
|
def test_get_timeout(model_list):
|
|
"""Test if the '_get_timeout' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
timeout = router._get_timeout(kwargs={}, data={"timeout": 100})
|
|
assert timeout == 100
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"fallback_kwarg, expected_error",
|
|
[
|
|
("mock_testing_fallbacks", litellm.InternalServerError),
|
|
("mock_testing_context_fallbacks", litellm.ContextWindowExceededError),
|
|
("mock_testing_content_policy_fallbacks", litellm.ContentPolicyViolationError),
|
|
],
|
|
)
|
|
def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_error):
|
|
"""Test if the '_handle_mock_testing_fallbacks' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
data = {
|
|
fallback_kwarg: True,
|
|
}
|
|
|
|
with pytest.raises(expected_error):
|
|
router._handle_mock_testing_fallbacks(
|
|
kwargs=data,
|
|
)
|
|
|
|
|
|
def test_handle_mock_testing_rate_limit_error(model_list):
|
|
"""Test if the '_handle_mock_testing_rate_limit_error' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
data = {
|
|
"mock_testing_rate_limit_error": True,
|
|
}
|
|
|
|
with pytest.raises(litellm.RateLimitError):
|
|
router._handle_mock_testing_rate_limit_error(
|
|
kwargs=data,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_callback_on_success(sync_mode):
|
|
"""Test if the '_deployment_callback_on_success' function is working correctly"""
|
|
import time
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
"rpm": 100,
|
|
},
|
|
"model_info": {"id": "100"},
|
|
}
|
|
]
|
|
router = Router(model_list=model_list)
|
|
# Get the actual deployment ID that was generated
|
|
gpt_deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
deployment_id = gpt_deployment["model_info"]["id"]
|
|
|
|
standard_logging_payload = create_standard_logging_payload()
|
|
standard_logging_payload["total_tokens"] = 100
|
|
standard_logging_payload["model_id"] = "100"
|
|
kwargs = {
|
|
"litellm_params": {
|
|
"metadata": {
|
|
"model_group": "gpt-5-mini",
|
|
},
|
|
"model_info": {"id": deployment_id},
|
|
},
|
|
"standard_logging_object": standard_logging_payload,
|
|
}
|
|
response = litellm.ModelResponse(
|
|
model="gpt-5-mini",
|
|
usage={"total_tokens": 100},
|
|
)
|
|
if sync_mode:
|
|
tpm_key = router.sync_deployment_callback_on_success(
|
|
kwargs=kwargs,
|
|
completion_response=response,
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
else:
|
|
tpm_key = await router.deployment_callback_on_success(
|
|
kwargs=kwargs,
|
|
completion_response=response,
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
assert tpm_key is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
|
|
"""
|
|
An IO-limited deployment (itpm/otpm, no tpm/rpm) must still record TPM usage
|
|
in the router's routing counter so TPM-aware routing strategies see its real
|
|
load in mixed model groups; its itpm/otpm enforcement runs separately.
|
|
"""
|
|
import time
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "opus",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-4o-mini",
|
|
"api_key": "sk-fake",
|
|
"itpm": 1000,
|
|
},
|
|
"model_info": {"id": "io-100"},
|
|
}
|
|
]
|
|
router = Router(model_list=model_list)
|
|
|
|
standard_logging_payload = create_standard_logging_payload()
|
|
standard_logging_payload["total_tokens"] = 100
|
|
standard_logging_payload["model_id"] = "io-100"
|
|
kwargs = {
|
|
"litellm_params": {
|
|
"metadata": {
|
|
"deployment": "openai/gpt-4o-mini",
|
|
"model_group": "opus",
|
|
},
|
|
"model_info": {"id": "io-100"},
|
|
},
|
|
"standard_logging_object": standard_logging_payload,
|
|
}
|
|
response = litellm.ModelResponse(model="openai/gpt-4o-mini", usage={"total_tokens": 100})
|
|
|
|
tpm_key = await router.deployment_callback_on_success(
|
|
kwargs=kwargs,
|
|
completion_response=response,
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
|
|
# The IO deployment is no longer skipped: its TPM routing counter is tracked.
|
|
assert tpm_key is not None
|
|
assert await router.cache.async_get_cache(key=tpm_key) == 100
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_callback_on_failure(model_list):
|
|
"""Test if the '_deployment_callback_on_failure' function is working correctly"""
|
|
import time
|
|
|
|
router = Router(model_list=model_list)
|
|
kwargs = {
|
|
"litellm_params": {
|
|
"metadata": {
|
|
"model_group": "gpt-5-mini",
|
|
},
|
|
"model_info": {"id": 100},
|
|
},
|
|
}
|
|
result = router.deployment_callback_on_failure(
|
|
kwargs=kwargs,
|
|
completion_response=None,
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
assert isinstance(result, bool)
|
|
assert result is False
|
|
|
|
model_response = router.completion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
result = await router.async_deployment_callback_on_failure(
|
|
kwargs=kwargs,
|
|
completion_response=model_response,
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
|
|
|
|
def test_deployment_callback_respects_cooldown_time(model_list):
|
|
"""Ensure per-model cooldown_time is honored even when exception headers are present."""
|
|
import httpx
|
|
import time
|
|
from unittest.mock import patch
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
class FakeException(Exception):
|
|
def __init__(self):
|
|
self.status_code = 429
|
|
self.headers = httpx.Headers({"x-test": "1"})
|
|
|
|
kwargs = {
|
|
"exception": FakeException(),
|
|
"litellm_params": {
|
|
"metadata": {"model_group": "gpt-5-mini"},
|
|
"model_info": {"id": 100},
|
|
"cooldown_time": 0,
|
|
},
|
|
}
|
|
|
|
with patch("litellm.router._set_cooldown_deployments") as mock_set:
|
|
router.deployment_callback_on_failure(
|
|
kwargs=kwargs,
|
|
completion_response=None,
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
|
|
mock_set.assert_called_once()
|
|
assert mock_set.call_args.kwargs["time_to_cooldown"] == 0
|
|
|
|
|
|
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
|
def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> None:
|
|
"""log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the
|
|
request metadata into it, counts every failed attempt of the request independently of the
|
|
per-hop attempted_retries, and never trusts a negative count planted before the first failure"""
|
|
router = Router(model_list=model_list)
|
|
rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo")
|
|
new_kwargs = router.log_retry(
|
|
kwargs={
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "sk-must-not-be-recorded",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
|
|
},
|
|
e=rate_limit_error,
|
|
)
|
|
assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [
|
|
{
|
|
"model_group": "gpt-3.5-turbo",
|
|
"deployment_id": "deployment-1",
|
|
"exception_type": "RateLimitError",
|
|
"exception_string": "litellm.RateLimitError: slow down",
|
|
"attempted_retries": 2,
|
|
}
|
|
]
|
|
assert new_kwargs[metadata_key]["request_retry_count"] == 1
|
|
assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2
|
|
planted_kwargs = {"model": "gpt-3.5-turbo", metadata_key: {"request_retry_count": -100}}
|
|
assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1
|
|
|
|
|
|
def test_update_usage(model_list):
|
|
"""Test if the '_update_usage' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
deployment_id = deployment["model_info"]["id"]
|
|
request_count = router._update_usage(
|
|
deployment_id=deployment_id, parent_otel_span=None
|
|
)
|
|
assert request_count == 1
|
|
|
|
request_count = router._update_usage(
|
|
deployment_id=deployment_id, parent_otel_span=None
|
|
)
|
|
|
|
assert request_count == 2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"finish_reason, expected_fallback", [("content_filter", True), ("stop", False)]
|
|
)
|
|
@pytest.mark.parametrize("fallback_type", ["model-specific", "default"])
|
|
def test_should_raise_content_policy_error(
|
|
model_list, finish_reason, expected_fallback, fallback_type
|
|
):
|
|
"""Test if the '_should_raise_content_policy_error' function is working correctly"""
|
|
router = Router(
|
|
model_list=model_list,
|
|
default_fallbacks=["gpt-5.5"] if fallback_type == "default" else None,
|
|
)
|
|
|
|
assert (
|
|
router._should_raise_content_policy_error(
|
|
model="gpt-5-mini",
|
|
response=litellm.ModelResponse(
|
|
model="gpt-5-mini",
|
|
choices=[
|
|
{
|
|
"finish_reason": finish_reason,
|
|
"message": {"content": "I'm fine, thank you!"},
|
|
}
|
|
],
|
|
usage={"total_tokens": 100},
|
|
),
|
|
kwargs={
|
|
"content_policy_fallbacks": (
|
|
[{"gpt-5-mini": "gpt-5.5"}]
|
|
if fallback_type == "model-specific"
|
|
else None
|
|
)
|
|
},
|
|
)
|
|
is expected_fallback
|
|
)
|
|
|
|
|
|
def test_get_healthy_deployments(model_list):
|
|
"""Test if the '_get_healthy_deployments' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployments = router._get_healthy_deployments(
|
|
model="gpt-5-mini", parent_otel_span=None
|
|
)
|
|
assert len(deployments) > 0
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_routing_strategy_pre_call_checks(model_list, sync_mode):
|
|
"""Test if the '_routing_strategy_pre_call_checks' function is working correctly"""
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
|
|
callback = CustomLogger()
|
|
litellm.callbacks = [callback]
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
|
|
litellm_logging_obj = Logging(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=False,
|
|
call_type="acompletion",
|
|
litellm_call_id="1234",
|
|
start_time=datetime.now(),
|
|
function_id="1234",
|
|
)
|
|
if sync_mode:
|
|
router.routing_strategy_pre_call_checks(deployment)
|
|
else:
|
|
## NO EXCEPTION
|
|
await router.async_routing_strategy_pre_call_checks(
|
|
deployment, litellm_logging_obj
|
|
)
|
|
|
|
## WITH EXCEPTION - rate limit error
|
|
with patch.object(
|
|
callback,
|
|
"async_pre_call_check",
|
|
AsyncMock(
|
|
side_effect=litellm.RateLimitError(
|
|
message="Rate limit error",
|
|
llm_provider="openai",
|
|
model="gpt-5-mini",
|
|
)
|
|
),
|
|
):
|
|
with pytest.raises(litellm.RateLimitError):
|
|
await router.async_routing_strategy_pre_call_checks(
|
|
deployment, litellm_logging_obj
|
|
)
|
|
|
|
## WITH EXCEPTION - generic error
|
|
with patch.object(
|
|
callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))
|
|
):
|
|
with pytest.raises(Exception, match="Error"):
|
|
await router.async_routing_strategy_pre_call_checks(
|
|
deployment, litellm_logging_obj
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_supported_environments, supported_environments, is_supported",
|
|
[(True, ["staging"], True), (False, None, True), (True, ["development"], False)],
|
|
)
|
|
def test_create_deployment(
|
|
model_list, set_supported_environments, supported_environments, is_supported
|
|
):
|
|
"""Test if the '_create_deployment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
|
|
if set_supported_environments:
|
|
os.environ["LITELLM_ENVIRONMENT"] = "staging"
|
|
deployment = router._create_deployment(
|
|
deployment_info={},
|
|
_model_name="gpt-5-mini",
|
|
_litellm_params={
|
|
"model": "gpt-5-mini",
|
|
"api_key": "test",
|
|
"custom_llm_provider": "openai",
|
|
},
|
|
_model_info={
|
|
"id": 100,
|
|
"supported_environments": supported_environments,
|
|
},
|
|
)
|
|
if is_supported:
|
|
assert deployment is not None
|
|
else:
|
|
assert deployment is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_supported_environments, supported_environments, is_supported",
|
|
[(True, ["staging"], True), (False, None, True), (True, ["development"], False)],
|
|
)
|
|
def test_deployment_is_active_for_environment(
|
|
model_list, set_supported_environments, supported_environments, is_supported
|
|
):
|
|
"""Test if the '_deployment_is_active_for_environment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
if set_supported_environments:
|
|
os.environ["LITELLM_ENVIRONMENT"] = "staging"
|
|
deployment["model_info"]["supported_environments"] = supported_environments
|
|
if is_supported:
|
|
assert (
|
|
router.deployment_is_active_for_environment(deployment=deployment) is True
|
|
)
|
|
else:
|
|
assert (
|
|
router.deployment_is_active_for_environment(deployment=deployment) is False
|
|
)
|
|
|
|
|
|
def test_set_model_list(model_list):
|
|
"""Test if the '_set_model_list' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
router.set_model_list(model_list=model_list)
|
|
assert len(router.model_list) == len(model_list)
|
|
|
|
|
|
def test_add_deployment(model_list):
|
|
"""Test if the '_add_deployment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
deployment["model_info"]["id"] = "100"
|
|
## Test 1: call user facing function
|
|
router.add_deployment(deployment=deployment)
|
|
|
|
## Test 2: call internal function
|
|
router._add_deployment(deployment=deployment)
|
|
assert len(router.model_list) == len(model_list) + 1
|
|
|
|
|
|
def test_upsert_deployment(model_list):
|
|
"""Test if the 'upsert_deployment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
print("model list", len(router.model_list))
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
deployment.litellm_params.model = "gpt-5.5"
|
|
router.upsert_deployment(deployment=deployment)
|
|
assert len(router.model_list) == len(model_list)
|
|
|
|
|
|
def test_delete_deployment(model_list):
|
|
"""Test if the 'delete_deployment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
router.delete_deployment(id=deployment["model_info"]["id"])
|
|
assert len(router.model_list) == len(model_list) - 1
|
|
|
|
|
|
def test_get_model_info(model_list):
|
|
"""Test if the 'get_model_info' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
model_info = router.get_model_info(id=deployment["model_info"]["id"])
|
|
assert model_info is not None
|
|
|
|
|
|
def test_get_model_group(model_list):
|
|
"""Test if the 'get_model_group' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
)
|
|
model_group = router.get_model_group(id=deployment["model_info"]["id"])
|
|
assert model_group is not None
|
|
assert model_group[0]["model_name"] == "gpt-5-mini"
|
|
|
|
|
|
@pytest.mark.parametrize("user_facing_model_group_name", ["gpt-5-mini", "gpt-5.5"])
|
|
def test_set_model_group_info(model_list, user_facing_model_group_name):
|
|
"""Test if the 'set_model_group_info' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
resp = router._set_model_group_info(
|
|
model_group="gpt-5-mini",
|
|
user_facing_model_group_name=user_facing_model_group_name,
|
|
)
|
|
assert resp is not None
|
|
assert resp.model_group == user_facing_model_group_name
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers(model_list):
|
|
"""Test if the 'set_response_headers' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
resp = await router.set_response_headers(response=None, model_group=None)
|
|
assert resp is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers_passes_through_post_increment_counters(model_list):
|
|
from pydantic import BaseModel
|
|
|
|
class _Usage(BaseModel):
|
|
total_tokens: int = 42
|
|
|
|
class _Resp(BaseModel):
|
|
usage: _Usage = _Usage()
|
|
_hidden_params: dict = {}
|
|
|
|
router = Router(model_list=model_list)
|
|
router.get_remaining_model_group_usage = AsyncMock(
|
|
return_value={
|
|
"x-ratelimit-remaining-tokens": 958,
|
|
"x-ratelimit-limit-tokens": 1000,
|
|
"x-ratelimit-remaining-requests": 99,
|
|
"x-ratelimit-limit-requests": 100,
|
|
"x-ratelimit-remaining-input-tokens": 1000,
|
|
"x-ratelimit-remaining-output-tokens": 500,
|
|
}
|
|
)
|
|
|
|
resp = _Resp()
|
|
resp._hidden_params = {}
|
|
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
|
|
|
headers = resp._hidden_params["additional_headers"]
|
|
assert headers["x-ratelimit-remaining-tokens"] == 958
|
|
assert headers["x-ratelimit-remaining-requests"] == 99
|
|
assert headers["x-ratelimit-limit-tokens"] == 1000
|
|
assert headers["x-ratelimit-limit-requests"] == 100
|
|
assert headers["x-ratelimit-remaining-input-tokens"] == 1000
|
|
assert headers["x-ratelimit-remaining-output-tokens"] == 500
|
|
|
|
|
|
def _rpm_tpm_router(model_id: str) -> Router:
|
|
return Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "tpm": 1000, "rpm": 100},
|
|
"model_info": {"id": model_id},
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def router_minute_pinned(monkeypatch):
|
|
pinned = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
|
|
monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned)
|
|
|
|
|
|
def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]:
|
|
return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.usefixtures("router_minute_pinned")
|
|
async def test_acompletion_headers_read_post_increment_counter_and_count_once():
|
|
router = _rpm_tpm_router("lit-3058-async")
|
|
|
|
response = await router.acompletion(
|
|
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong"
|
|
)
|
|
total_tokens = response.usage.total_tokens
|
|
assert total_tokens > 0
|
|
|
|
headers = _ratelimit_headers(response)
|
|
assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens
|
|
assert headers["x-ratelimit-remaining-requests"] == 99
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
|
|
|
|
await asyncio.sleep(0.5)
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "openai/*",
|
|
"litellm_params": {"model": "openai/*", "api_key": "sk-fake", "tpm": 1000, "rpm": 100},
|
|
"model_info": {"id": "lit-3058-wildcard"},
|
|
}
|
|
]
|
|
)
|
|
|
|
response = await router.acompletion(
|
|
model="openai/gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong"
|
|
)
|
|
total_tokens = response.usage.total_tokens
|
|
|
|
headers = _ratelimit_headers(response)
|
|
assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens
|
|
assert headers["x-ratelimit-remaining-requests"] == 99
|
|
assert await router.get_model_group_usage("openai/gpt-5-mini") == (total_tokens, 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.usefixtures("router_minute_pinned")
|
|
async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion():
|
|
router = _rpm_tpm_router("lit-3058-stream")
|
|
|
|
stream = await router.acompletion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
mock_response="pong pong pong",
|
|
stream=True,
|
|
stream_options={"include_usage": True},
|
|
)
|
|
headers = _ratelimit_headers(stream)
|
|
assert headers["x-ratelimit-remaining-tokens"] == 1000
|
|
assert headers["x-ratelimit-remaining-requests"] == 99
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (0, 1)
|
|
|
|
chunks = [chunk async for chunk in stream]
|
|
total_tokens = chunks[-1].usage.total_tokens
|
|
assert total_tokens > 0
|
|
|
|
await asyncio.sleep(0.5)
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deployment_callback_on_success_adds_only_uncounted_tokens():
|
|
import time
|
|
|
|
router = _rpm_tpm_router("lit-3058-callback")
|
|
standard_logging_payload = create_standard_logging_payload()
|
|
standard_logging_payload["total_tokens"] = 100
|
|
kwargs = {
|
|
"litellm_params": {
|
|
"metadata": {
|
|
"deployment": "gpt-5-mini",
|
|
"model_group": "gpt-5-mini",
|
|
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 60,
|
|
},
|
|
"model_info": {"id": "lit-3058-callback"},
|
|
},
|
|
"standard_logging_object": standard_logging_payload,
|
|
}
|
|
|
|
tpm_key = await router.deployment_callback_on_success(
|
|
kwargs=kwargs,
|
|
completion_response=litellm.ModelResponse(model="gpt-5-mini", usage={"total_tokens": 100}),
|
|
start_time=time.time(),
|
|
end_time=time.time(),
|
|
)
|
|
|
|
assert tpm_key is not None
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (40, 0)
|
|
|
|
|
|
class _GatedIncrementCache(DualCache):
|
|
def __init__(self) -> None:
|
|
super().__init__(in_memory_cache=InMemoryCache())
|
|
self.first_increment_started = asyncio.Event()
|
|
self.release_first_increment = asyncio.Event()
|
|
self.increment_calls = 0
|
|
|
|
async def async_increment_cache_pipeline(
|
|
self,
|
|
increment_list: list[RedisPipelineIncrementOperation],
|
|
local_only: bool = False,
|
|
parent_otel_span: object = None,
|
|
**kwargs: object,
|
|
) -> list[float] | None:
|
|
self.increment_calls += 1
|
|
if self.increment_calls == 1:
|
|
self.first_increment_started.set()
|
|
await self.release_first_increment.wait()
|
|
return await super().async_increment_cache_pipeline(
|
|
increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_success_callback_running_during_pre_header_increment_does_not_double_count():
|
|
router = _rpm_tpm_router("lit-3058-race")
|
|
cache = _GatedIncrementCache()
|
|
router.cache = cache
|
|
|
|
request = asyncio.ensure_future(
|
|
router.acompletion(model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong")
|
|
)
|
|
await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5)
|
|
for _ in range(50):
|
|
if get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1:
|
|
break
|
|
await asyncio.sleep(0.1)
|
|
assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1
|
|
assert cache.increment_calls == 1
|
|
|
|
cache.release_first_increment.set()
|
|
response = await request
|
|
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
|
|
|
|
|
|
class _UnavailableIncrementCache(DualCache):
|
|
def __init__(self) -> None:
|
|
super().__init__(in_memory_cache=InMemoryCache())
|
|
self.first_increment_started = asyncio.Event()
|
|
self.release_first_increment = asyncio.Event()
|
|
self.increment_calls = 0
|
|
|
|
async def async_increment_cache_pipeline(
|
|
self,
|
|
increment_list: list[RedisPipelineIncrementOperation],
|
|
local_only: bool = False,
|
|
parent_otel_span: object = None,
|
|
**kwargs: object,
|
|
) -> list[float] | None:
|
|
self.increment_calls += 1
|
|
if self.increment_calls == 1:
|
|
self.first_increment_started.set()
|
|
await self.release_first_increment.wait()
|
|
raise RuntimeError("cache unavailable")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callback_observing_stamp_before_pre_header_increment_fails_leaves_no_stamp_behind():
|
|
router = _rpm_tpm_router("lit-3058-fail")
|
|
cache = _UnavailableIncrementCache()
|
|
router.cache = cache
|
|
metadata: dict[str, object] = {}
|
|
|
|
request = asyncio.ensure_future(
|
|
router.acompletion(
|
|
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong", metadata=metadata
|
|
)
|
|
)
|
|
await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5)
|
|
assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30
|
|
for _ in range(50):
|
|
if get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1:
|
|
break
|
|
await asyncio.sleep(0.1)
|
|
assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1
|
|
assert cache.increment_calls == 1
|
|
|
|
cache.release_first_increment.set()
|
|
response = await request
|
|
|
|
assert response.usage.total_tokens == 30
|
|
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata
|
|
assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 100
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_increment_deployment_usage_for_response_skips_session_wrappers():
|
|
router = _rpm_tpm_router("lit-3058-ws")
|
|
request_kwargs = {
|
|
"model": "gpt-5-mini",
|
|
"litellm_metadata": {"model_group": "gpt-5-mini", "model_info": {"id": "lit-3058-ws"}},
|
|
}
|
|
|
|
await router.increment_deployment_usage_for_response(response=None, request_kwargs=request_kwargs)
|
|
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
|
|
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_increment_deployment_usage_writes_only_positive_deltas_for_limited_deployments():
|
|
router = _rpm_tpm_router("lit-3058-delta")
|
|
unlimited = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"},
|
|
"model_info": {"id": "lit-3058-unlimited"},
|
|
}
|
|
]
|
|
)
|
|
|
|
tpm_key = await router._increment_deployment_usage(
|
|
deployment_id="lit-3058-delta",
|
|
deployment_name="gpt-5-mini",
|
|
model_group="gpt-5-mini",
|
|
total_tokens=25,
|
|
rpm_increment=1,
|
|
parent_otel_span=None,
|
|
)
|
|
assert tpm_key is not None
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (25, 1)
|
|
|
|
assert (
|
|
await router._increment_deployment_usage(
|
|
deployment_id="lit-3058-delta",
|
|
deployment_name="gpt-5-mini",
|
|
model_group="gpt-5-mini",
|
|
total_tokens=0,
|
|
rpm_increment=0,
|
|
parent_otel_span=None,
|
|
)
|
|
is None
|
|
)
|
|
assert await router.get_model_group_usage("gpt-5-mini") == (25, 1)
|
|
|
|
assert (
|
|
await unlimited._increment_deployment_usage(
|
|
deployment_id="lit-3058-unlimited",
|
|
deployment_name="gpt-5-mini",
|
|
model_group="gpt-5-mini",
|
|
total_tokens=25,
|
|
rpm_increment=1,
|
|
parent_otel_span=None,
|
|
)
|
|
is None
|
|
)
|
|
assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None)
|
|
|
|
|
|
def _shared_redis_stub(store: dict) -> MagicMock:
|
|
from litellm.caching.redis_cache import RedisCache
|
|
|
|
async def increment_pipeline(increment_list, **kwargs):
|
|
for op in increment_list:
|
|
store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"]
|
|
return [store[op["key"]] for op in increment_list]
|
|
|
|
async def batch_get(keys, **kwargs):
|
|
return {key: store.get(key) for key in keys}
|
|
|
|
redis_stub = MagicMock(spec=RedisCache)
|
|
redis_stub.async_increment_pipeline = increment_pipeline
|
|
redis_stub.async_batch_get_cache = batch_get
|
|
return redis_stub
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
|
|
store: dict = {}
|
|
worker_a = _rpm_tpm_router("lit-3058-workers")
|
|
worker_b = _rpm_tpm_router("lit-3058-workers")
|
|
worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
|
|
worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
|
|
|
|
messages = [{"role": "user", "content": "hi"}]
|
|
tokens_on_a = 0
|
|
for _ in range(3):
|
|
response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
|
tokens_on_a += response.usage.total_tokens
|
|
|
|
response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
|
headers = _ratelimit_headers(response)
|
|
assert headers["x-ratelimit-remaining-requests"] == 96
|
|
assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens
|
|
|
|
counted_tokens = tokens_on_a + response.usage.total_tokens
|
|
for _ in range(2):
|
|
response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
|
counted_tokens += response.usage.total_tokens
|
|
|
|
stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True)
|
|
stream_headers = _ratelimit_headers(stream)
|
|
assert stream_headers["x-ratelimit-remaining-requests"] == 93
|
|
assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens
|
|
assert [chunk async for chunk in stream]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_model_group_io_token_usage_sums_across_deployments():
|
|
"""
|
|
get_model_group_io_token_usage must sum ITPM/OTPM across every deployment
|
|
in the model group (not just the first), reading the same per-deployment
|
|
cache keys the pre-call reservation writes to.
|
|
"""
|
|
from litellm.types.router import RouterCacheEnum
|
|
from litellm.utils import get_utc_datetime
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "opus",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-4o-mini",
|
|
"itpm": 1000,
|
|
"otpm": 500,
|
|
},
|
|
"model_info": {"id": "io-usage-dep-1"},
|
|
},
|
|
{
|
|
"model_name": "opus",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-4o",
|
|
"itpm": 1000,
|
|
"otpm": 500,
|
|
},
|
|
"model_info": {"id": "io-usage-dep-2"},
|
|
},
|
|
]
|
|
)
|
|
|
|
minute = get_utc_datetime().strftime("%H-%M")
|
|
keys_and_values = [
|
|
(
|
|
RouterCacheEnum.ITPM.value.format(
|
|
id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute
|
|
),
|
|
30,
|
|
),
|
|
(
|
|
RouterCacheEnum.OTPM.value.format(
|
|
id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute
|
|
),
|
|
10,
|
|
),
|
|
(
|
|
RouterCacheEnum.ITPM.value.format(
|
|
id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute
|
|
),
|
|
70,
|
|
),
|
|
(
|
|
RouterCacheEnum.OTPM.value.format(
|
|
id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute
|
|
),
|
|
20,
|
|
),
|
|
]
|
|
for key, value in keys_and_values:
|
|
await router.cache.async_increment_cache(key=key, value=value, ttl=60)
|
|
|
|
current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus")
|
|
|
|
assert current_itpm == 100
|
|
assert current_otpm == 30
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_model_group_io_token_usage_no_deployments_returns_none():
|
|
router = Router(model_list=[])
|
|
current_itpm, current_otpm = await router.get_model_group_io_token_usage(
|
|
"nonexistent-group"
|
|
)
|
|
assert current_itpm is None
|
|
assert current_otpm is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(model_list):
|
|
"""
|
|
A model group with both itpm/otpm and tpm/rpm limits must expose the
|
|
standard remaining-tokens/requests headers alongside the input/output token
|
|
headers, so clients and prometheus gauges relying on either still get data.
|
|
"""
|
|
from unittest.mock import Mock
|
|
|
|
from litellm.types.router import ModelGroupInfo
|
|
|
|
router = Router(model_list=model_list)
|
|
router._cached_get_model_group_info = Mock(
|
|
return_value=ModelGroupInfo(
|
|
model_group="gpt-3.5-turbo",
|
|
providers=["openai"],
|
|
itpm=2000,
|
|
otpm=1000,
|
|
tpm=5000,
|
|
rpm=50,
|
|
)
|
|
)
|
|
router.get_model_group_io_token_usage = AsyncMock(return_value=(100, 40))
|
|
router.get_model_group_usage = AsyncMock(return_value=(500, 5))
|
|
|
|
headers = await router.get_remaining_model_group_usage("gpt-3.5-turbo")
|
|
|
|
assert headers["x-ratelimit-remaining-input-tokens"] == 1900
|
|
assert headers["x-ratelimit-remaining-output-tokens"] == 960
|
|
assert headers["x-ratelimit-remaining-tokens"] == 4500
|
|
assert headers["x-ratelimit-remaining-requests"] == 45
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(model_list):
|
|
"""
|
|
A provider that natively returns `x-ratelimit-remaining-input-tokens` must
|
|
not suppress the router's own remaining-tokens/requests headers for a
|
|
non-IO model group.
|
|
"""
|
|
from pydantic import BaseModel
|
|
|
|
class _Usage(BaseModel):
|
|
total_tokens: int = 42
|
|
|
|
class _Resp(BaseModel):
|
|
usage: _Usage = _Usage()
|
|
_hidden_params: dict = {}
|
|
|
|
router = Router(model_list=model_list)
|
|
router.get_remaining_model_group_usage = AsyncMock(
|
|
return_value={
|
|
"x-ratelimit-remaining-tokens": 1000,
|
|
"x-ratelimit-remaining-requests": 100,
|
|
}
|
|
)
|
|
|
|
resp = _Resp()
|
|
resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-input-tokens": 5}}
|
|
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
|
|
|
headers = resp._hidden_params["additional_headers"]
|
|
assert headers["x-ratelimit-remaining-tokens"] == 1000
|
|
assert headers["x-ratelimit-remaining-requests"] == 100
|
|
# the provider's native header is left untouched
|
|
assert headers["x-ratelimit-remaining-input-tokens"] == 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(model_list):
|
|
from pydantic import BaseModel
|
|
|
|
class _Usage(BaseModel):
|
|
total_tokens: int = 42
|
|
|
|
class _Resp(BaseModel):
|
|
usage: _Usage = _Usage()
|
|
_hidden_params: dict = {}
|
|
|
|
router = Router(model_list=model_list)
|
|
router.get_remaining_model_group_usage = AsyncMock(
|
|
return_value={
|
|
"x-ratelimit-remaining-tokens": 1000,
|
|
"x-ratelimit-remaining-requests": 100,
|
|
"x-ratelimit-remaining-input-tokens": 900,
|
|
"x-ratelimit-remaining-output-tokens": 450,
|
|
}
|
|
)
|
|
|
|
resp = _Resp()
|
|
resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-tokens": 5}}
|
|
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
|
|
|
headers = resp._hidden_params["additional_headers"]
|
|
assert headers["x-ratelimit-remaining-tokens"] == 5
|
|
assert headers["x-ratelimit-remaining-requests"] == 100
|
|
assert headers["x-ratelimit-remaining-input-tokens"] == 900
|
|
assert headers["x-ratelimit-remaining-output-tokens"] == 450
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers_handles_missing_usage(model_list):
|
|
"""
|
|
Streaming chunks and some response shapes may lack a `usage` attribute or
|
|
populated `total_tokens`. Header composition must not depend on usage and never raise.
|
|
"""
|
|
from pydantic import BaseModel
|
|
|
|
class _Resp(BaseModel):
|
|
_hidden_params: dict = {}
|
|
|
|
router = Router(model_list=model_list)
|
|
router.get_remaining_model_group_usage = AsyncMock(
|
|
return_value={
|
|
"x-ratelimit-remaining-tokens": 1000,
|
|
"x-ratelimit-remaining-requests": 100,
|
|
}
|
|
)
|
|
|
|
resp = _Resp()
|
|
resp._hidden_params = {}
|
|
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
|
|
|
|
headers = resp._hidden_params["additional_headers"]
|
|
assert headers["x-ratelimit-remaining-tokens"] == 1000
|
|
assert headers["x-ratelimit-remaining-requests"] == 100
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers_dict_anthropic_messages_response(model_list):
|
|
"""Anthropic /v1/messages returns a dict; IO rate-limit headers must attach."""
|
|
router = Router(model_list=model_list)
|
|
router.get_remaining_model_group_usage = AsyncMock(
|
|
return_value={
|
|
"x-ratelimit-limit-input-tokens": 25,
|
|
"x-ratelimit-remaining-input-tokens": 20,
|
|
"x-ratelimit-limit-output-tokens": 100,
|
|
"x-ratelimit-remaining-output-tokens": 95,
|
|
}
|
|
)
|
|
|
|
resp = {
|
|
"id": "msg_123",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "hi"}],
|
|
"usage": {"input_tokens": 5, "output_tokens": 1},
|
|
}
|
|
await router.set_response_headers(response=resp, model_group="io-itpm-strict")
|
|
|
|
assert "_hidden_params" in resp
|
|
headers = resp["_hidden_params"]["additional_headers"]
|
|
assert headers["x-litellm-model-group"] == "io-itpm-strict"
|
|
assert headers["x-ratelimit-limit-input-tokens"] == 25
|
|
assert headers["x-ratelimit-remaining-input-tokens"] == 20
|
|
assert headers["x-ratelimit-remaining-output-tokens"] == 95
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_response_headers_wraps_bare_async_generator(model_list):
|
|
"""
|
|
Streaming responses that never go through Router.make_call's usual
|
|
object-based wrappers (e.g. the Anthropic /v1/messages -> Responses API
|
|
bridge, which yields a raw async generator with no `_hidden_params` slot)
|
|
must still get IO rate-limit headers attached via a thin wrapper.
|
|
"""
|
|
|
|
async def _raw_generator():
|
|
yield {"type": "message_start"}
|
|
yield {"type": "message_stop"}
|
|
|
|
router = Router(model_list=model_list)
|
|
router.get_remaining_model_group_usage = AsyncMock(
|
|
return_value={
|
|
"x-ratelimit-limit-input-tokens": 25,
|
|
"x-ratelimit-remaining-input-tokens": 20,
|
|
}
|
|
)
|
|
|
|
wrapped = await router.set_response_headers(response=_raw_generator(), model_group="io-itpm-strict")
|
|
|
|
assert hasattr(wrapped, "_hidden_params")
|
|
headers = wrapped._hidden_params["additional_headers"]
|
|
assert headers["x-litellm-model-group"] == "io-itpm-strict"
|
|
assert headers["x-ratelimit-limit-input-tokens"] == 25
|
|
assert headers["x-ratelimit-remaining-input-tokens"] == 20
|
|
|
|
from collections.abc import AsyncIterator
|
|
|
|
assert isinstance(wrapped, AsyncIterator)
|
|
chunks = [chunk async for chunk in wrapped]
|
|
assert chunks == [{"type": "message_start"}, {"type": "message_stop"}]
|
|
|
|
|
|
def test_get_all_deployments(model_list):
|
|
"""Test if the 'get_all_deployments' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployments = router._get_all_deployments(
|
|
model_name="gpt-5-mini", model_alias="gpt-5-mini"
|
|
)
|
|
assert len(deployments) > 0
|
|
|
|
|
|
def test_get_model_access_groups(model_list):
|
|
"""Test if the 'get_model_access_groups' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
access_groups = router.get_model_access_groups()
|
|
assert len(access_groups) == 2
|
|
|
|
|
|
def test_update_settings(model_list):
|
|
"""Test if the 'update_settings' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
pre_update_allowed_fails = router.allowed_fails
|
|
router.update_settings(**{"allowed_fails": 20})
|
|
assert router.allowed_fails != pre_update_allowed_fails
|
|
assert router.allowed_fails == 20
|
|
|
|
|
|
def test_common_checks_available_deployment(model_list):
|
|
"""Test if the 'common_checks_available_deployment' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
_, available_deployments = router._common_checks_available_deployment(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
input="hi",
|
|
specific_deployment=False,
|
|
)
|
|
|
|
assert len(available_deployments) > 0
|
|
|
|
|
|
def test_filter_cooldown_deployments(model_list):
|
|
"""Test if the 'filter_cooldown_deployments' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployments = router._filter_cooldown_deployments(
|
|
healthy_deployments=router._get_all_deployments(model_name="gpt-5-mini"), # type: ignore
|
|
cooldown_deployments=[],
|
|
)
|
|
assert len(deployments) == len(router._get_all_deployments(model_name="gpt-5-mini"))
|
|
|
|
|
|
def test_track_deployment_metrics(model_list):
|
|
"""Test if the 'track_deployment_metrics' function is working correctly"""
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
router = Router(model_list=model_list)
|
|
router._track_deployment_metrics(
|
|
deployment=router.get_deployment_by_model_group_name(
|
|
model_group_name="gpt-5-mini"
|
|
),
|
|
response=ModelResponse(
|
|
model="gpt-5-mini",
|
|
usage={"total_tokens": 100},
|
|
),
|
|
parent_otel_span=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"exception_type, exception_name, num_retries",
|
|
[
|
|
(litellm.exceptions.BadRequestError, "BadRequestError", 3),
|
|
(litellm.exceptions.AuthenticationError, "AuthenticationError", 4),
|
|
(litellm.exceptions.RateLimitError, "RateLimitError", 6),
|
|
(
|
|
litellm.exceptions.ContentPolicyViolationError,
|
|
"ContentPolicyViolationError",
|
|
7,
|
|
),
|
|
],
|
|
)
|
|
def test_get_num_retries_from_retry_policy(
|
|
model_list, exception_type, exception_name, num_retries
|
|
):
|
|
"""Test if the 'get_num_retries_from_retry_policy' function is working correctly"""
|
|
from litellm.router import RetryPolicy
|
|
|
|
data = {exception_name + "Retries": num_retries}
|
|
print("data", data)
|
|
router = Router(
|
|
model_list=model_list,
|
|
retry_policy=RetryPolicy(**data),
|
|
)
|
|
print("exception_type", exception_type)
|
|
calc_num_retries = router.get_num_retries_from_retry_policy(
|
|
exception=exception_type(
|
|
message="test", llm_provider="openai", model="gpt-5-mini"
|
|
)
|
|
)
|
|
assert calc_num_retries == num_retries
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"exception_type, exception_name, allowed_fails",
|
|
[
|
|
(litellm.exceptions.BadRequestError, "BadRequestError", 3),
|
|
(litellm.exceptions.AuthenticationError, "AuthenticationError", 4),
|
|
(litellm.exceptions.RateLimitError, "RateLimitError", 6),
|
|
(
|
|
litellm.exceptions.ContentPolicyViolationError,
|
|
"ContentPolicyViolationError",
|
|
7,
|
|
),
|
|
],
|
|
)
|
|
def test_get_allowed_fails_from_policy(
|
|
model_list, exception_type, exception_name, allowed_fails
|
|
):
|
|
"""Test if the 'get_allowed_fails_from_policy' function is working correctly"""
|
|
from litellm.types.router import AllowedFailsPolicy
|
|
|
|
data = {exception_name + "AllowedFails": allowed_fails}
|
|
router = Router(
|
|
model_list=model_list, allowed_fails_policy=AllowedFailsPolicy(**data)
|
|
)
|
|
calc_allowed_fails = router.get_allowed_fails_from_policy(
|
|
exception=exception_type(
|
|
message="test", llm_provider="openai", model="gpt-5-mini"
|
|
)
|
|
)
|
|
assert calc_allowed_fails == allowed_fails
|
|
|
|
|
|
def test_initialize_alerting(model_list):
|
|
"""Test if the 'initialize_alerting' function is working correctly"""
|
|
from litellm.types.router import AlertingConfig
|
|
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
|
|
|
router = Router(
|
|
model_list=model_list, alerting_config=AlertingConfig(webhook_url="test")
|
|
)
|
|
router._initialize_alerting()
|
|
|
|
callback_added = False
|
|
for callback in litellm.callbacks:
|
|
if isinstance(callback, SlackAlerting):
|
|
callback_added = True
|
|
assert callback_added is True
|
|
|
|
|
|
def test_flush_cache(model_list):
|
|
"""Test if the 'flush_cache' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
router.cache.set_cache("test", "test")
|
|
assert router.cache.get_cache("test") == "test"
|
|
router.flush_cache()
|
|
assert router.cache.get_cache("test") is None
|
|
|
|
|
|
def test_discard(model_list):
|
|
"""
|
|
Test that discard properly removes a Router from the callback lists
|
|
"""
|
|
litellm.callbacks = []
|
|
litellm.success_callback = []
|
|
litellm._async_success_callback = []
|
|
litellm.failure_callback = []
|
|
litellm._async_failure_callback = []
|
|
litellm.input_callback = []
|
|
litellm.service_callback = []
|
|
|
|
router = Router(model_list=model_list)
|
|
router.discard()
|
|
|
|
# Verify all callback lists are empty
|
|
assert len(litellm.callbacks) == 0
|
|
assert len(litellm.success_callback) == 0
|
|
assert len(litellm.failure_callback) == 0
|
|
assert len(litellm._async_success_callback) == 0
|
|
assert len(litellm._async_failure_callback) == 0
|
|
assert len(litellm.input_callback) == 0
|
|
assert len(litellm.service_callback) == 0
|
|
|
|
|
|
def test_initialize_assistants_endpoint(model_list):
|
|
"""Test if the 'initialize_assistants_endpoint' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
router.initialize_assistants_endpoint()
|
|
assert router.acreate_assistants is not None
|
|
assert router.adelete_assistant is not None
|
|
assert router.aget_assistants is not None
|
|
assert router.acreate_thread is not None
|
|
assert router.aget_thread is not None
|
|
assert router.arun_thread is not None
|
|
assert router.aget_messages is not None
|
|
assert router.a_add_message is not None
|
|
|
|
|
|
def test_pass_through_assistants_endpoint_factory(model_list):
|
|
"""Test if the 'pass_through_assistants_endpoint_factory' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
router._pass_through_assistants_endpoint_factory(
|
|
original_function=litellm.acreate_assistants,
|
|
custom_llm_provider="openai",
|
|
client=None,
|
|
**{},
|
|
)
|
|
|
|
|
|
def test_factory_function(model_list):
|
|
"""Test if the 'factory_function' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
router.factory_function(litellm.acreate_assistants)
|
|
|
|
|
|
def test_get_model_from_alias(model_list):
|
|
"""Test if the 'get_model_from_alias' function is working correctly"""
|
|
router = Router(
|
|
model_list=model_list,
|
|
model_group_alias={"gpt-5.5": "gpt-5-mini"},
|
|
)
|
|
model = router._get_model_from_alias(model="gpt-5.5")
|
|
assert model == "gpt-5-mini"
|
|
|
|
|
|
def test_get_deployment_by_litellm_model(model_list):
|
|
"""Test if the 'get_deployment_by_litellm_model' function is working correctly"""
|
|
router = Router(model_list=model_list)
|
|
deployment = router._get_deployment_by_litellm_model(model="gpt-5-mini")
|
|
assert deployment is not None
|
|
|
|
|
|
def test_get_pattern(model_list):
|
|
router = Router(model_list=model_list)
|
|
pattern = router.pattern_router.get_pattern(model="claude-3")
|
|
assert pattern is not None
|
|
|
|
|
|
def test_deployments_by_pattern(model_list):
|
|
router = Router(model_list=model_list)
|
|
deployments = router.pattern_router.get_deployments_by_pattern(model="claude-3")
|
|
assert deployments is not None
|
|
|
|
|
|
# def test_pattern_match_deployments(model_list):
|
|
# from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
|
# import re
|
|
|
|
# patter_router = PatternMatchRouter()
|
|
|
|
# request = "fo::hi::static::hello"
|
|
# model_name = "fo::*:static::*"
|
|
|
|
# model_name_regex = patter_router._pattern_to_regex(model_name)
|
|
|
|
# # Match against the request
|
|
# match = re.match(model_name_regex, request)
|
|
|
|
# print(f"match: {match}")
|
|
# print(f"match.end: {match.end()}")
|
|
# if match is None:
|
|
# raise ValueError("Match not found")
|
|
# updated_model = patter_router.set_deployment_model_name(
|
|
# matched_pattern=match, litellm_deployment_litellm_model="openai/*"
|
|
# )
|
|
# assert updated_model == "openai/fo::hi:static::hello"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"user_request_model, model_name, litellm_model, expected_model",
|
|
[
|
|
("llmengine/foo", "llmengine/*", "openai/foo", "openai/foo"),
|
|
("llmengine/foo", "llmengine/*", "openai/*", "openai/foo"),
|
|
(
|
|
"fo::hi::static::hello",
|
|
"fo::*::static::*",
|
|
"openai/fo::*:static::*",
|
|
"openai/fo::hi:static::hello",
|
|
),
|
|
(
|
|
"fo::hi::static::hello",
|
|
"fo::*::static::*",
|
|
"openai/gpt-5-mini",
|
|
"openai/gpt-5-mini",
|
|
),
|
|
(
|
|
"bedrock/meta.llama3-70b",
|
|
"*meta.llama3*",
|
|
"bedrock/meta.llama3-*",
|
|
"bedrock/meta.llama3-70b",
|
|
),
|
|
(
|
|
"meta.llama3-70b",
|
|
"*meta.llama3*",
|
|
"bedrock/meta.llama3-*",
|
|
"meta.llama3-70b",
|
|
),
|
|
],
|
|
)
|
|
def test_pattern_match_deployment_set_model_name(
|
|
user_request_model, model_name, litellm_model, expected_model
|
|
):
|
|
from re import Match
|
|
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
|
|
|
pattern_router = PatternMatchRouter()
|
|
|
|
import re
|
|
|
|
# Convert model_name into a proper regex
|
|
model_name_regex = pattern_router._pattern_to_regex(model_name)
|
|
|
|
# Match against the request
|
|
match = re.match(model_name_regex, user_request_model)
|
|
|
|
if match is None:
|
|
raise ValueError("Match not found")
|
|
|
|
# Call the set_deployment_model_name function
|
|
updated_model = pattern_router.set_deployment_model_name(match, litellm_model)
|
|
|
|
print(updated_model) # Expected output: "openai/fo::hi:static::hello"
|
|
assert updated_model == expected_model
|
|
|
|
updated_models = pattern_router._return_pattern_matched_deployments(
|
|
match,
|
|
deployments=[
|
|
{
|
|
"model_name": model_name,
|
|
"litellm_params": {"model": litellm_model},
|
|
}
|
|
],
|
|
)
|
|
|
|
for model in updated_models:
|
|
assert model["litellm_params"]["model"] == expected_model
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pass_through_moderation_endpoint_factory(model_list):
|
|
router = Router(model_list=model_list)
|
|
response = await router._pass_through_moderation_endpoint_factory(
|
|
original_function=litellm.amoderation,
|
|
input="this is valid good text",
|
|
model=None,
|
|
)
|
|
assert response is not None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"has_default_fallbacks, expected_result",
|
|
[(True, True), (False, False)],
|
|
)
|
|
def test_has_default_fallbacks(model_list, has_default_fallbacks, expected_result):
|
|
router = Router(
|
|
model_list=model_list,
|
|
default_fallbacks=(
|
|
["my-default-fallback-model"] if has_default_fallbacks else None
|
|
),
|
|
)
|
|
assert router._has_default_fallbacks() is expected_result
|
|
|
|
|
|
def test_add_optional_pre_call_checks(model_list):
|
|
router = Router(model_list=model_list)
|
|
|
|
router.add_optional_pre_call_checks(["prompt_caching"])
|
|
assert len(litellm.callbacks) > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_callback_filter_deployments(model_list):
|
|
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
healthy_deployments = router.get_model_list(model_name="gpt-5-mini")
|
|
|
|
new_healthy_deployments = await router.async_callback_filter_deployments(
|
|
model="gpt-5-mini",
|
|
healthy_deployments=healthy_deployments,
|
|
messages=[],
|
|
parent_otel_span=None,
|
|
)
|
|
|
|
assert len(new_healthy_deployments) == len(healthy_deployments)
|
|
|
|
|
|
def test_cached_get_model_group_info(model_list):
|
|
"""Test if the '_cached_get_model_group_info' function is working correctly with LRU cache"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# First call - should hit the actual function
|
|
result1 = router._cached_get_model_group_info("gpt-5-mini")
|
|
|
|
# Second call with same argument - should hit the cache
|
|
result2 = router._cached_get_model_group_info("gpt-5-mini")
|
|
|
|
# Verify results are the same
|
|
assert result1 == result2
|
|
|
|
# Verify the cache info shows hits
|
|
cache_info = router._cached_get_model_group_info.cache_info()
|
|
assert cache_info.hits > 0 # Should have at least one cache hit
|
|
|
|
|
|
def test_init_responses_api_endpoints(model_list):
|
|
"""Test if the '_init_responses_api_endpoints' function is working correctly"""
|
|
from typing import Callable
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
assert router.aget_responses is not None
|
|
assert isinstance(router.aget_responses, Callable)
|
|
assert router.adelete_responses is not None
|
|
assert isinstance(router.adelete_responses, Callable)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"mock_testing_fallbacks, mock_testing_context_fallbacks, mock_testing_content_policy_fallbacks, expected_fallbacks, expected_context, expected_content_policy",
|
|
[
|
|
# Test string to bool conversion
|
|
("true", "false", "True", True, False, True),
|
|
("TRUE", "FALSE", "False", True, False, False),
|
|
("false", "true", "false", False, True, False),
|
|
# Test actual boolean values (should pass through unchanged)
|
|
(True, False, True, True, False, True),
|
|
(False, True, False, False, True, False),
|
|
# Test None values
|
|
(None, None, None, None, None, None),
|
|
# Test mixed types
|
|
("true", False, None, True, False, None),
|
|
],
|
|
)
|
|
def test_mock_router_testing_params_str_to_bool_conversion(
|
|
mock_testing_fallbacks,
|
|
mock_testing_context_fallbacks,
|
|
mock_testing_content_policy_fallbacks,
|
|
expected_fallbacks,
|
|
expected_context,
|
|
expected_content_policy,
|
|
):
|
|
"""Test if MockRouterTestingParams.from_kwargs correctly converts string values to booleans using str_to_bool"""
|
|
from litellm.types.router import MockRouterTestingParams
|
|
|
|
kwargs = {
|
|
"mock_testing_fallbacks": mock_testing_fallbacks,
|
|
"mock_testing_context_fallbacks": mock_testing_context_fallbacks,
|
|
"mock_testing_content_policy_fallbacks": mock_testing_content_policy_fallbacks,
|
|
"other_param": "should_remain", # This should not be affected
|
|
}
|
|
|
|
# Make a copy to verify kwargs are properly popped
|
|
original_kwargs = kwargs.copy()
|
|
|
|
mock_params = MockRouterTestingParams.from_kwargs(kwargs)
|
|
|
|
# Verify the converted values
|
|
assert mock_params.mock_testing_fallbacks == expected_fallbacks
|
|
assert mock_params.mock_testing_context_fallbacks == expected_context
|
|
assert mock_params.mock_testing_content_policy_fallbacks == expected_content_policy
|
|
|
|
# Verify that the mock testing params were popped from kwargs
|
|
assert "mock_testing_fallbacks" not in kwargs
|
|
assert "mock_testing_context_fallbacks" not in kwargs
|
|
assert "mock_testing_content_policy_fallbacks" not in kwargs
|
|
|
|
# Verify other params remain unchanged
|
|
assert kwargs["other_param"] == "should_remain"
|
|
|
|
|
|
def test_is_auto_router_deployment(model_list):
|
|
"""Test if the '_is_auto_router_deployment' function correctly identifies auto-router deployments"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Test case 1: Model starts with "auto_router/" - should return True
|
|
litellm_params_auto = LiteLLM_Params(model="auto_router/my-auto-router")
|
|
assert router._is_auto_router_deployment(litellm_params_auto) is True
|
|
|
|
# Test case 2: Model doesn't start with "auto_router/" - should return False
|
|
litellm_params_regular = LiteLLM_Params(model="gpt-5-mini")
|
|
assert router._is_auto_router_deployment(litellm_params_regular) is False
|
|
|
|
# Test case 3: Model is empty string - should return False
|
|
litellm_params_empty = LiteLLM_Params(model="")
|
|
assert router._is_auto_router_deployment(litellm_params_empty) is False
|
|
|
|
# Test case 4: Model contains "auto_router/" but doesn't start with it - should return False
|
|
litellm_params_contains = LiteLLM_Params(model="prefix_auto_router/something")
|
|
assert router._is_auto_router_deployment(litellm_params_contains) is False
|
|
|
|
|
|
@patch("litellm.router_strategy.auto_router.auto_router.AutoRouter")
|
|
def test_init_auto_router_deployment_success(mock_auto_router, model_list):
|
|
"""Test if the 'init_auto_router_deployment' function successfully initializes auto-router when all params provided"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Create a mock AutoRouter instance
|
|
mock_auto_router_instance = MagicMock()
|
|
mock_auto_router.return_value = mock_auto_router_instance
|
|
|
|
# Test case: All required parameters provided
|
|
litellm_params = LiteLLM_Params(
|
|
model="auto_router/test",
|
|
auto_router_config_path="/path/to/config",
|
|
auto_router_default_model="gpt-5-mini",
|
|
auto_router_embedding_model="text-embedding-3-small",
|
|
)
|
|
deployment = Deployment(
|
|
model_name="test-auto-router",
|
|
litellm_params=litellm_params,
|
|
model_info={"id": "test-id"},
|
|
)
|
|
|
|
# Should not raise any exception
|
|
router.init_auto_router_deployment(deployment)
|
|
|
|
# Verify AutoRouter was called with correct parameters
|
|
mock_auto_router.assert_called_once_with(
|
|
model_name="test-auto-router",
|
|
auto_router_config_path="/path/to/config",
|
|
auto_router_config=None,
|
|
default_model="gpt-5-mini",
|
|
embedding_model="text-embedding-3-small",
|
|
litellm_router_instance=router,
|
|
max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS,
|
|
)
|
|
|
|
# Verify the auto-router was added to the router's auto_routers dict
|
|
assert "test-auto-router" in router.auto_routers
|
|
assert router.auto_routers["test-auto-router"][0].strategy == mock_auto_router_instance
|
|
|
|
|
|
@patch("litellm.router_strategy.auto_router.auto_router.AutoRouter")
|
|
def test_init_auto_router_deployment_duplicate_model_name(mock_auto_router, model_list):
|
|
"""Test if the 'init_auto_router_deployment' function raises ValueError when model_name already exists"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Create a mock AutoRouter instance
|
|
mock_auto_router_instance = MagicMock()
|
|
mock_auto_router.return_value = mock_auto_router_instance
|
|
|
|
# Add an existing auto-router
|
|
from litellm.types.router import TaggedPreRoutingStrategy
|
|
|
|
router.auto_routers["test-auto-router"] = [
|
|
TaggedPreRoutingStrategy(tags=(), strategy=mock_auto_router_instance)
|
|
]
|
|
|
|
# Try to add another auto-router with the same name
|
|
litellm_params = LiteLLM_Params(
|
|
model="auto_router/test",
|
|
auto_router_config_path="/path/to/config",
|
|
auto_router_default_model="gpt-5-mini",
|
|
auto_router_embedding_model="text-embedding-3-small",
|
|
)
|
|
deployment = Deployment(
|
|
model_name="test-auto-router",
|
|
litellm_params=litellm_params,
|
|
model_info={"id": "test-id"},
|
|
)
|
|
|
|
with pytest.raises(
|
|
ValueError, match=r"Auto-router deployment test-auto-router with tags .* already exists"
|
|
):
|
|
router.init_auto_router_deployment(deployment)
|
|
|
|
|
|
def testgenerate_model_id_with_deployment_model_name(model_list):
|
|
"""Test that generate_model_id works correctly with deployment model_name and handles None values properly"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Test case 1: Normal case with valid model_group and litellm_params
|
|
model_group = "gpt-4.1"
|
|
litellm_params = {
|
|
"model": "gpt-4.1",
|
|
"api_key": "test_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
}
|
|
|
|
try:
|
|
result = router.generate_model_id(
|
|
model_group=model_group, litellm_params=litellm_params
|
|
)
|
|
assert isinstance(result, str)
|
|
assert len(result) > 0
|
|
print(f"✓ Success with valid model_group: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"Failed with valid model_group: {e}")
|
|
|
|
# Test case 2: Edge case with None model_group (this should fail as expected - our fix prevents this from happening)
|
|
with pytest.raises(TypeError) as exc_info:
|
|
router.generate_model_id(model_group=None, litellm_params=litellm_params)
|
|
# After optimization, error message changed but still fails appropriately on None
|
|
error_str = str(exc_info.value)
|
|
assert (
|
|
"unsupported operand type(s) for +=" in error_str
|
|
or "expected str instance, NoneType found" in error_str
|
|
)
|
|
|
|
# Test case 3: Edge case with None key in litellm_params
|
|
litellm_params_with_none_key = {
|
|
"model": "gpt-4.1",
|
|
"api_key": "test_key",
|
|
None: "should_be_skipped", # This should be handled gracefully
|
|
}
|
|
|
|
try:
|
|
result = router.generate_model_id(
|
|
model_group=model_group, litellm_params=litellm_params_with_none_key
|
|
)
|
|
assert isinstance(result, str)
|
|
assert len(result) > 0
|
|
print(f"✓ Success with None key in litellm_params: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"Failed with None key in litellm_params: {e}")
|
|
|
|
# Test case 4: Edge case with empty litellm_params
|
|
try:
|
|
result = router.generate_model_id(model_group=model_group, litellm_params={})
|
|
assert isinstance(result, str)
|
|
assert len(result) > 0
|
|
print(f"✓ Success with empty litellm_params: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"Failed with empty litellm_params: {e}")
|
|
|
|
# Test case 5: Verify that the same inputs produce the same result (deterministic)
|
|
result1 = router.generate_model_id(
|
|
model_group=model_group, litellm_params=litellm_params
|
|
)
|
|
result2 = router.generate_model_id(
|
|
model_group=model_group, litellm_params=litellm_params
|
|
)
|
|
assert result1 == result2, "Model ID generation should be deterministic"
|
|
|
|
print("✓ All generate_model_id tests passed!")
|
|
|
|
|
|
def test_handle_clientside_credential_with_deployment_model_name(model_list):
|
|
"""Test that _handle_clientside_credential uses deployment model_name correctly"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Mock deployment with model_name
|
|
deployment = {
|
|
"model_name": "gpt-4.1",
|
|
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
|
|
}
|
|
|
|
# Mock kwargs with empty metadata (simulating the original issue)
|
|
kwargs = {
|
|
"metadata": {}, # Empty metadata, no model_group
|
|
"litellm_params": {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
},
|
|
}
|
|
|
|
# Mock dynamic_litellm_params that would be returned by get_dynamic_litellm_params
|
|
dynamic_litellm_params = {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
}
|
|
|
|
# Test that the method doesn't fail when metadata is empty
|
|
try:
|
|
# This would normally call generate_model_id internally
|
|
# We're testing that the fix prevents the TypeError
|
|
model_group = deployment["model_name"] # This is what our fix does
|
|
assert model_group == "gpt-4.1"
|
|
|
|
# Verify that generate_model_id works with this model_group
|
|
result = router.generate_model_id(
|
|
model_group=model_group, litellm_params=dynamic_litellm_params
|
|
)
|
|
assert isinstance(result, str)
|
|
assert len(result) > 0
|
|
|
|
print(f"✓ Success with deployment model_name: {result}")
|
|
except Exception as e:
|
|
pytest.fail(f"Failed with deployment model_name: {e}")
|
|
|
|
print("✓ _handle_clientside_credential test passed!")
|
|
|
|
|
|
def test_sync_generic_api_call_preserves_requested_model_group_in_logs():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "claude-sonnet-4-6",
|
|
"litellm_params": {
|
|
"model": "bedrock/global.anthropic.claude-sonnet-4-6",
|
|
"aws_access_key_id": "test-access-key",
|
|
"aws_secret_access_key": "test-secret-key",
|
|
"aws_region_name": "us-west-2",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
try:
|
|
captured_kwargs = {}
|
|
|
|
def mock_original_function(**kwargs):
|
|
captured_kwargs.update(kwargs)
|
|
return {"status": "ok"}
|
|
|
|
response = router._generic_api_call_with_fallbacks(
|
|
model="claude-sonnet-4-6",
|
|
original_function=mock_original_function,
|
|
)
|
|
|
|
assert response == {"status": "ok"}
|
|
assert captured_kwargs["model"] == "bedrock/global.anthropic.claude-sonnet-4-6"
|
|
assert captured_kwargs["litellm_metadata"]["model_group"] == "claude-sonnet-4-6"
|
|
assert (
|
|
captured_kwargs["litellm_metadata"]["deployment"]
|
|
== "bedrock/global.anthropic.claude-sonnet-4-6"
|
|
)
|
|
finally:
|
|
router.discard()
|
|
|
|
|
|
def test_sync_generic_api_call_uses_request_kwargs_for_deployment_selection():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "regional-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/us-model",
|
|
"api_key": "test-api-key",
|
|
"region_name": "us",
|
|
},
|
|
},
|
|
{
|
|
"model_name": "regional-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/eu-model",
|
|
"api_key": "test-api-key",
|
|
"region_name": "eu",
|
|
},
|
|
},
|
|
],
|
|
enable_pre_call_checks=True,
|
|
)
|
|
|
|
try:
|
|
captured_kwargs = {}
|
|
|
|
def mock_original_function(**kwargs):
|
|
captured_kwargs.update(kwargs)
|
|
return {"status": "ok"}
|
|
|
|
response = router._generic_api_call_with_fallbacks(
|
|
model="regional-model",
|
|
original_function=mock_original_function,
|
|
messages=[{"role": "user", "content": "Hello from Europe"}],
|
|
allowed_model_region="eu",
|
|
)
|
|
|
|
assert response == {"status": "ok"}
|
|
assert captured_kwargs["model"] == "anthropic/eu-model"
|
|
finally:
|
|
router.discard()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"function_name, expected_metadata_key",
|
|
[
|
|
("acompletion", "metadata"),
|
|
("_ageneric_api_call_with_fallbacks", "litellm_metadata"),
|
|
("batch", "litellm_metadata"),
|
|
("completion", "metadata"),
|
|
("acreate_file", "litellm_metadata"),
|
|
("aget_file", "litellm_metadata"),
|
|
],
|
|
)
|
|
def test_handle_clientside_credential_metadata_loading(
|
|
model_list, function_name, expected_metadata_key
|
|
):
|
|
"""Test that _handle_clientside_credential correctly loads metadata based on function name"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Mock deployment
|
|
deployment = {
|
|
"model_name": "gpt-4.1",
|
|
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
|
|
"model_info": {"id": "original-id-123"},
|
|
}
|
|
|
|
# Mock kwargs with clientside credentials and metadata
|
|
kwargs = {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
expected_metadata_key: {"model_group": "gpt-4.1", "custom_field": "test_value"},
|
|
}
|
|
|
|
# Call the function
|
|
result_deployment = router._handle_clientside_credential(
|
|
deployment=deployment, kwargs=kwargs, function_name=function_name
|
|
)
|
|
|
|
# Verify the result is a Deployment object
|
|
assert isinstance(result_deployment, Deployment)
|
|
|
|
# Verify the deployment has the correct model_name (should be the model_group from metadata)
|
|
assert result_deployment.model_name == "gpt-4.1"
|
|
|
|
# Verify the litellm_params contain the clientside credentials
|
|
assert result_deployment.litellm_params.api_key == "client_side_key"
|
|
assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1"
|
|
|
|
# Verify the model_info has been updated with a new ID
|
|
assert result_deployment.model_info.id != "original-id-123"
|
|
assert result_deployment.model_info.original_model_id == "original-id-123"
|
|
|
|
# The caller-supplied credential must stay scoped to this call: it must never be
|
|
# registered as a router deployment, or a later caller with no override of their
|
|
# own could be load-balanced onto it and reach the provider with this credential
|
|
# (see LIT-7811).
|
|
assert len(router.model_list) == len(model_list)
|
|
assert router.get_deployment(model_id=result_deployment.model_info.id) is None
|
|
|
|
# Test that the function correctly uses the right metadata key
|
|
# For acompletion, it should use "metadata"
|
|
# For _ageneric_api_call_with_fallbacks/batch, it should use "litellm_metadata"
|
|
if function_name == "acompletion":
|
|
assert "metadata" in kwargs
|
|
assert "litellm_metadata" not in kwargs
|
|
elif function_name in [
|
|
"_ageneric_api_call_with_fallbacks",
|
|
"batch",
|
|
"acreate_file",
|
|
"aget_file",
|
|
]:
|
|
assert "litellm_metadata" in kwargs
|
|
# Note: acompletion would not have litellm_metadata, but other functions might have both
|
|
|
|
print(
|
|
f"✓ Success with function_name '{function_name}' using '{expected_metadata_key}' metadata key"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"function_name, metadata_key",
|
|
[
|
|
("acompletion", "metadata"),
|
|
("_ageneric_api_call_with_fallbacks", "litellm_metadata"),
|
|
],
|
|
)
|
|
def test_handle_clientside_credential_metadata_variable_name(
|
|
model_list, function_name, metadata_key
|
|
):
|
|
"""Test that _handle_clientside_credential uses the correct metadata variable name based on function name"""
|
|
from litellm.router_utils.batch_utils import _get_router_metadata_variable_name
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
# Verify the metadata variable name is correct for each function
|
|
expected_metadata_key = _get_router_metadata_variable_name(
|
|
function_name=function_name
|
|
)
|
|
assert expected_metadata_key == metadata_key
|
|
|
|
# Mock deployment
|
|
deployment = {
|
|
"model_name": "gpt-4.1",
|
|
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
|
|
"model_info": {"id": "original-id-456"},
|
|
}
|
|
|
|
# Mock kwargs with clientside credentials and the correct metadata key
|
|
kwargs = {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
metadata_key: {"model_group": "gpt-4.1", "test_field": "test_value"},
|
|
}
|
|
|
|
# Call the function
|
|
result_deployment = router._handle_clientside_credential(
|
|
deployment=deployment, kwargs=kwargs, function_name=function_name
|
|
)
|
|
|
|
# Verify the function correctly extracted model_group from the right metadata key
|
|
assert result_deployment.model_name == "gpt-4.1"
|
|
|
|
# Verify the deployment was created with the correct metadata
|
|
assert result_deployment.litellm_params.api_key == "client_side_key"
|
|
assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1"
|
|
|
|
print(
|
|
f"✓ Success with function_name '{function_name}' correctly using '{metadata_key}' for metadata"
|
|
)
|
|
|
|
|
|
def test_handle_clientside_credential_no_metadata(model_list):
|
|
"""Test that _handle_clientside_credential handles cases where no metadata is provided"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Mock deployment
|
|
deployment = {
|
|
"model_name": "gpt-4.1",
|
|
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
|
|
"model_info": {"id": "original-id-789"},
|
|
}
|
|
|
|
# Mock kwargs with clientside credentials but NO metadata
|
|
kwargs = {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
# No metadata key at all
|
|
}
|
|
|
|
# This should fail because there's no model_group in metadata
|
|
# The function expects to find model_group in the metadata
|
|
try:
|
|
result_deployment = router._handle_clientside_credential(
|
|
deployment=deployment, kwargs=kwargs, function_name="acompletion"
|
|
)
|
|
# If we get here, the function should have used deployment.model_name as fallback
|
|
assert result_deployment.model_name == "gpt-4.1"
|
|
print("✓ Success with no metadata - used deployment.model_name as fallback")
|
|
except Exception as e:
|
|
# This is expected behavior - the function needs model_group to generate model_id
|
|
print(f"✓ Correctly handled no metadata case: {e}")
|
|
|
|
# Test with empty metadata
|
|
kwargs_with_empty_metadata = {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
"metadata": {}, # Empty metadata
|
|
}
|
|
|
|
try:
|
|
result_deployment = router._handle_clientside_credential(
|
|
deployment=deployment,
|
|
kwargs=kwargs_with_empty_metadata,
|
|
function_name="acompletion",
|
|
)
|
|
# Should fail because empty metadata has no model_group
|
|
pytest.fail("Expected failure with empty metadata")
|
|
except Exception as e:
|
|
print(f"✓ Correctly handled empty metadata case: {e}")
|
|
|
|
|
|
def test_handle_clientside_credential_with_responses_function(model_list):
|
|
"""Test that _handle_clientside_credential works correctly with responses function name"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Mock deployment
|
|
deployment = {
|
|
"model_name": "gpt-4.1",
|
|
"litellm_params": {"model": "gpt-4.1", "api_key": "test_key"},
|
|
"model_info": {"id": "original-id-responses"},
|
|
}
|
|
|
|
# Mock kwargs with clientside credentials and litellm_metadata (for responses function)
|
|
kwargs = {
|
|
"api_key": "client_side_key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
"litellm_metadata": {
|
|
"model_group": "gpt-4.1",
|
|
"responses_field": "responses_value",
|
|
},
|
|
}
|
|
|
|
# Call the function with _ageneric_api_call_with_fallbacks function name (which handles responses)
|
|
result_deployment = router._handle_clientside_credential(
|
|
deployment=deployment,
|
|
kwargs=kwargs,
|
|
function_name="_ageneric_api_call_with_fallbacks",
|
|
)
|
|
|
|
# Verify the result
|
|
assert isinstance(result_deployment, Deployment)
|
|
assert result_deployment.model_name == "gpt-4.1"
|
|
assert result_deployment.litellm_params.api_key == "client_side_key"
|
|
assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1"
|
|
assert result_deployment.model_info.id != "original-id-responses"
|
|
assert result_deployment.model_info.original_model_id == "original-id-responses"
|
|
|
|
# The caller-supplied credential must stay scoped to this call: it must never be
|
|
# registered as a router deployment (see LIT-7811).
|
|
assert len(router.model_list) == len(model_list)
|
|
assert router.get_deployment(model_id=result_deployment.model_info.id) is None
|
|
|
|
print(
|
|
"✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata"
|
|
)
|
|
|
|
|
|
def test_handle_clientside_credential_still_registers_custom_pricing(model_list):
|
|
"""A clientside-credential call must still price against the deployment's own
|
|
custom rate, even though the call's ephemeral deployment is never added to the
|
|
router (see LIT-7811): losing that registration would silently fall back to
|
|
public catalog pricing for every clientside-credential call on a deployment
|
|
with a custom rate configured."""
|
|
router = Router(model_list=model_list)
|
|
deployment = {
|
|
"model_name": "gpt-4.1",
|
|
"litellm_params": {
|
|
"model": "gpt-4.1",
|
|
"api_key": "test_key",
|
|
"input_cost_per_token": 0.0001234,
|
|
"output_cost_per_token": 0.0005678,
|
|
},
|
|
"model_info": {"id": "original-id-pricing"},
|
|
}
|
|
kwargs = {"api_key": "client_side_key", "metadata": {"model_group": "gpt-4.1"}}
|
|
|
|
result_deployment = router._handle_clientside_credential(
|
|
deployment=deployment, kwargs=kwargs, function_name="acompletion"
|
|
)
|
|
|
|
registered = litellm.model_cost.get(result_deployment.model_info.id)
|
|
assert registered is not None
|
|
assert registered["input_cost_per_token"] == 0.0001234
|
|
assert registered["output_cost_per_token"] == 0.0005678
|
|
|
|
|
|
def test_register_deployment_pricing_direct_call():
|
|
"""Direct-call unit test for the pricing-registration helper `_handle_clientside_credential`
|
|
relies on, so it prices a deployment that is deliberately never added to `self.model_list`."""
|
|
deployment = Deployment(
|
|
model_name="gpt-4.1",
|
|
litellm_params=LiteLLM_Params(
|
|
model="gpt-4.1",
|
|
api_key="test_key",
|
|
input_cost_per_token=0.0009999,
|
|
),
|
|
model_info=ModelInfo(id="direct-call-pricing-id"),
|
|
)
|
|
|
|
Router._register_deployment_pricing(deployment=deployment)
|
|
|
|
assert litellm.model_cost["direct-call-pricing-id"]["input_cost_per_token"] == 0.0009999
|
|
|
|
|
|
def test_get_metadata_variable_name_from_kwargs(model_list):
|
|
"""
|
|
Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content.
|
|
"""
|
|
router = Router(model_list=model_list)
|
|
|
|
# Test case 1: kwargs contains litellm_metadata - should return "litellm_metadata"
|
|
kwargs_with_litellm_metadata = {
|
|
"litellm_metadata": {"user": "test"},
|
|
"metadata": {"other": "data"},
|
|
}
|
|
result = router._get_metadata_variable_name_from_kwargs(
|
|
kwargs_with_litellm_metadata
|
|
)
|
|
assert result == "litellm_metadata"
|
|
|
|
# Test case 2: kwargs only contains metadata - should return "metadata"
|
|
kwargs_with_metadata_only = {"metadata": {"user": "test"}}
|
|
result = router._get_metadata_variable_name_from_kwargs(kwargs_with_metadata_only)
|
|
assert result == "metadata"
|
|
|
|
# Test case 3: kwargs contains neither - should return "metadata" (default)
|
|
kwargs_empty = {}
|
|
result = router._get_metadata_variable_name_from_kwargs(kwargs_empty)
|
|
assert result == "metadata"
|
|
|
|
# Test case 4: kwargs contains other keys but no metadata keys - should return "metadata"
|
|
kwargs_other = {
|
|
"model": "gpt-5.5",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}
|
|
result = router._get_metadata_variable_name_from_kwargs(kwargs_other)
|
|
assert result == "metadata"
|
|
|
|
|
|
@pytest.fixture
|
|
def search_tools():
|
|
"""Fixture for search tools configuration"""
|
|
return [
|
|
{
|
|
"search_tool_name": "test-search-tool",
|
|
"litellm_params": {
|
|
"search_provider": "perplexity",
|
|
"api_key": "test-api-key",
|
|
"api_base": "https://api.perplexity.ai",
|
|
"mode": "turbo",
|
|
},
|
|
},
|
|
{
|
|
"search_tool_name": "test-search-tool",
|
|
"litellm_params": {
|
|
"search_provider": "perplexity",
|
|
"api_key": "test-api-key-2",
|
|
"api_base": "https://api.perplexity.ai",
|
|
"mode": "turbo",
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asearch_with_fallbacks(search_tools):
|
|
"""
|
|
Test _asearch_with_fallbacks method of Router.
|
|
|
|
Tests that the _asearch_with_fallbacks method correctly:
|
|
- Accepts search parameters
|
|
- Calls async_function_with_fallbacks with correct configuration
|
|
- Returns SearchResponse
|
|
"""
|
|
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
|
|
|
|
router = Router(search_tools=search_tools)
|
|
|
|
# Create a mock search response
|
|
mock_response = SearchResponse(
|
|
object="search",
|
|
results=[
|
|
SearchResult(
|
|
title="Test Result",
|
|
url="https://example.com",
|
|
snippet="Test snippet content",
|
|
)
|
|
],
|
|
)
|
|
|
|
# Mock the async_function_with_fallbacks to return our mock response
|
|
with patch.object(
|
|
router, "async_function_with_fallbacks", new_callable=AsyncMock
|
|
) as mock_fallbacks:
|
|
mock_fallbacks.return_value = mock_response
|
|
|
|
# Mock original function
|
|
async def mock_asearch(**kwargs):
|
|
return mock_response
|
|
|
|
# Call _asearch_with_fallbacks
|
|
response = await router._asearch_with_fallbacks(
|
|
original_function=mock_asearch,
|
|
search_tool_name="test-search-tool",
|
|
query="test query",
|
|
max_results=5,
|
|
)
|
|
|
|
# Verify async_function_with_fallbacks was called
|
|
assert mock_fallbacks.called
|
|
|
|
# Verify the response
|
|
assert isinstance(response, SearchResponse)
|
|
assert response.object == "search"
|
|
assert len(response.results) == 1
|
|
assert response.results[0].title == "Test Result"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asearch_with_fallbacks_helper(search_tools):
|
|
"""
|
|
Test _asearch_with_fallbacks_helper method of Router.
|
|
|
|
Tests that the _asearch_with_fallbacks_helper method correctly:
|
|
- Selects a search tool from available options
|
|
- Calls the original search function with correct provider parameters
|
|
- Returns SearchResponse
|
|
"""
|
|
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
|
|
|
|
router = Router(search_tools=search_tools)
|
|
|
|
# Create a mock search response
|
|
mock_response = SearchResponse(
|
|
object="search",
|
|
results=[
|
|
SearchResult(
|
|
title="Helper Test Result",
|
|
url="https://example.com/helper",
|
|
snippet="Helper test snippet",
|
|
)
|
|
],
|
|
)
|
|
|
|
# Mock the original generic function
|
|
async def mock_original_function(**kwargs):
|
|
# Verify correct parameters are passed
|
|
assert "search_provider" in kwargs
|
|
assert kwargs["search_provider"] == "perplexity"
|
|
assert "api_key" in kwargs
|
|
assert kwargs["mode"] == "turbo"
|
|
assert kwargs["query"] == "helper test query"
|
|
return mock_response
|
|
|
|
# Call _asearch_with_fallbacks_helper
|
|
response = await router._asearch_with_fallbacks_helper(
|
|
model="test-search-tool",
|
|
original_generic_function=mock_original_function,
|
|
query="helper test query",
|
|
max_results=3,
|
|
)
|
|
|
|
# Verify the response
|
|
assert isinstance(response, SearchResponse)
|
|
assert response.object == "search"
|
|
assert len(response.results) == 1
|
|
assert response.results[0].title == "Helper Test Result"
|
|
assert response.results[0].url == "https://example.com/helper"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asearch_with_fallbacks_helper_missing_search_tool():
|
|
"""
|
|
Test _asearch_with_fallbacks_helper raises error when search tool not found.
|
|
|
|
Tests that the helper method raises a ValueError when the requested
|
|
search tool name doesn't exist in the router's search_tools configuration.
|
|
"""
|
|
# Create router with no search tools
|
|
router = Router(model_list=[])
|
|
|
|
async def mock_original_function(**kwargs):
|
|
return None
|
|
|
|
# Should raise ValueError for missing search tool
|
|
with pytest.raises(ValueError, match="Search tool 'nonexistent-tool' not found"):
|
|
await router._asearch_with_fallbacks_helper(
|
|
model="nonexistent-tool",
|
|
original_generic_function=mock_original_function,
|
|
query="test query",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asearch_with_fallbacks_helper_missing_search_provider():
|
|
"""
|
|
Test _asearch_with_fallbacks_helper raises error when search_provider not configured.
|
|
|
|
Tests that the helper method raises a ValueError when a search tool
|
|
is found but doesn't have search_provider in its litellm_params.
|
|
"""
|
|
# Create router with misconfigured search tool (missing search_provider)
|
|
search_tools_bad = [
|
|
{
|
|
"search_tool_name": "bad-tool",
|
|
"litellm_params": {
|
|
"api_key": "test-key"
|
|
# Missing search_provider
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(search_tools=search_tools_bad)
|
|
|
|
async def mock_original_function(**kwargs):
|
|
return None
|
|
|
|
# Should raise ValueError for missing search_provider
|
|
with pytest.raises(ValueError, match="search_provider not found in litellm_params"):
|
|
await router._asearch_with_fallbacks_helper(
|
|
model="bad-tool",
|
|
original_generic_function=mock_original_function,
|
|
query="test query",
|
|
)
|
|
|
|
|
|
def test_get_first_default_fallback():
|
|
"""Test _get_first_default_fallback method"""
|
|
# Test with default fallback ("*")
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {"model": "gpt-5-mini", "api_key": "fake-key"},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list, fallbacks=[{"*": ["gpt-5-mini"]}])
|
|
|
|
result = router._get_first_default_fallback()
|
|
assert result == "gpt-5-mini"
|
|
|
|
# Test with no fallbacks
|
|
router_no_fallbacks = Router(model_list=model_list)
|
|
result = router_no_fallbacks._get_first_default_fallback()
|
|
assert result is None
|
|
|
|
# Test with fallbacks but no default
|
|
router_no_default = Router(
|
|
model_list=model_list, fallbacks=[{"gpt-5.5": ["gpt-5-mini"]}]
|
|
)
|
|
result = router_no_default._get_first_default_fallback()
|
|
assert result is None
|
|
|
|
# Test with empty default list
|
|
router_empty_list = Router(model_list=model_list, fallbacks=[{"*": []}])
|
|
result = router_empty_list._get_first_default_fallback()
|
|
assert result is None
|
|
|
|
|
|
def test_resolve_model_name_from_model_id():
|
|
"""Test resolve_model_name_from_model_id function with various scenarios"""
|
|
|
|
# Test case 1: model_id is None
|
|
router = Router(model_list=[])
|
|
result = router.resolve_model_name_from_model_id(None)
|
|
assert result is None
|
|
|
|
# Test case 2: model_id directly matches a model_name
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
result = router.resolve_model_name_from_model_id("gpt-5-mini")
|
|
assert result == "gpt-5-mini"
|
|
|
|
# Test case 3: model_id matches litellm_params.model exactly
|
|
model_list = [
|
|
{
|
|
"model_name": "vertex-ai-sora-2",
|
|
"litellm_params": {
|
|
"model": "vertex_ai/veo-2.0-generate-001",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
result = router.resolve_model_name_from_model_id("vertex_ai/veo-2.0-generate-001")
|
|
assert result == "vertex-ai-sora-2"
|
|
|
|
# Test case 4: model_id matches when actual_model ends with /model_id
|
|
model_list = [
|
|
{
|
|
"model_name": "vertex-ai-sora-2",
|
|
"litellm_params": {
|
|
"model": "vertex_ai/veo-2.0-generate-001",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
result = router.resolve_model_name_from_model_id("veo-2.0-generate-001")
|
|
assert result == "vertex-ai-sora-2"
|
|
|
|
# Test case 5: model_id matches when actual_model ends with :model_id
|
|
# Note: We use a valid model format for router initialization, but test the function
|
|
# with a model_id that would match the pattern vertex_ai:model_id
|
|
# Since the router validates models on init, we'll test this by manually setting up
|
|
# the model_list after initialization or using a valid format
|
|
model_list = [
|
|
{
|
|
"model_name": "vertex-ai-sora-2",
|
|
"litellm_params": {
|
|
"model": "vertex_ai/veo-2.0-generate-001",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
# Test that the function can handle model_id that would match if the format was vertex_ai:model_id
|
|
# We'll test with a model_id that matches the end of the actual_model
|
|
result = router.resolve_model_name_from_model_id("veo-2.0-generate-001")
|
|
assert result == "vertex-ai-sora-2"
|
|
|
|
# Test case 6: model_id doesn't match anything
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
result = router.resolve_model_name_from_model_id("non-existent-model")
|
|
assert result is None
|
|
|
|
# Test case 7: Empty model_list
|
|
router = Router(model_list=[])
|
|
result = router.resolve_model_name_from_model_id("some-model")
|
|
assert result is None
|
|
|
|
# Test case 8: Multiple models, find the correct one
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
{
|
|
"model_name": "vertex-ai-sora-2",
|
|
"litellm_params": {
|
|
"model": "vertex_ai/veo-2.0-generate-001",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
result = router.resolve_model_name_from_model_id("veo-2.0-generate-001")
|
|
assert result == "vertex-ai-sora-2"
|
|
|
|
# Test case 9: model_id matches deployment ID (has_model_id check)
|
|
# This tests the has_model_id path in Strategy 1
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
|
|
result = router.resolve_model_name_from_model_id("gpt-5-mini")
|
|
assert result == "gpt-5-mini"
|
|
|
|
# Test case 10: model_id is a deployment ID (hash) that differs from the
|
|
# public model_name. Regression for #32580: managed batch/file IDs embed the
|
|
# deployment model_id, and it must resolve back to the public model_name so
|
|
# team model-access checks compare against the model group, not the hash.
|
|
model_list = [
|
|
{
|
|
"model_name": "bedrock-batch-model",
|
|
"litellm_params": {
|
|
"model": "bedrock/global.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
},
|
|
"model_info": {"id": "8d0eaa7e6c6f54a425dfd0062cb6b0dc"},
|
|
},
|
|
]
|
|
router = Router(model_list=model_list)
|
|
result = router.resolve_model_name_from_model_id("8d0eaa7e6c6f54a425dfd0062cb6b0dc")
|
|
assert result == "bedrock-batch-model"
|
|
|
|
|
|
def test_get_valid_args():
|
|
"""Test get_valid_args static method returns valid Router.__init__ arguments"""
|
|
# Call the static method
|
|
valid_args = Router.get_valid_args()
|
|
|
|
# Verify it returns a list
|
|
assert isinstance(valid_args, list)
|
|
assert len(valid_args) > 0
|
|
|
|
# Verify it contains expected Router.__init__ arguments
|
|
expected_args = [
|
|
"model_list",
|
|
"routing_strategy",
|
|
"cache_responses",
|
|
"num_retries",
|
|
"timeout",
|
|
"fallbacks",
|
|
]
|
|
for arg in expected_args:
|
|
assert arg in valid_args, f"Expected argument '{arg}' not found in valid_args"
|
|
|
|
# Verify "self" is not in the list (since it's removed)
|
|
assert "self" not in valid_args
|
|
|
|
# Verify it contains keyword-only arguments too
|
|
# These are common Router.__init__ parameters
|
|
assert "assistants_config" in valid_args or "search_tools" in valid_args
|
|
|
|
|
|
def test_get_router_model_info_with_deployment_object():
|
|
"""Test get_router_model_info accepts Deployment object directly and reuses LiteLLM_Params"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5.5",
|
|
"litellm_params": {"model": "gpt-5.5", "api_key": "test-key"},
|
|
"model_info": {"id": "test-id"},
|
|
}
|
|
]
|
|
)
|
|
|
|
# Get the Deployment object (not dict)
|
|
deployment = router.get_deployment(model_id="test-id")
|
|
assert deployment is not None
|
|
assert isinstance(deployment, Deployment)
|
|
assert isinstance(deployment.litellm_params, LiteLLM_Params)
|
|
|
|
# Pass Deployment directly (not .model_dump()) - this exercises the isinstance check
|
|
# that reuses the existing LiteLLM_Params instead of reconstructing it
|
|
model_info = router.get_router_model_info(
|
|
deployment=deployment,
|
|
received_model_name="gpt-5.5",
|
|
)
|
|
|
|
# Verify we got valid model info back
|
|
assert model_info is not None
|
|
assert isinstance(model_info, dict)
|
|
|
|
|
|
def test_deployment_has_budget_limits():
|
|
router = Router(model_list=[])
|
|
|
|
with_budget = Deployment(
|
|
model_name="budgeted-model",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/gpt-4o-mini",
|
|
max_budget=0.001,
|
|
budget_duration="1d",
|
|
),
|
|
model_info=ModelInfo(id="budget-deployment-id"),
|
|
)
|
|
without_budget = Deployment(
|
|
model_name="unbudgeted-model",
|
|
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
|
|
model_info=ModelInfo(id="no-budget-deployment-id"),
|
|
)
|
|
|
|
assert router._deployment_has_budget_limits(deployment=with_budget) is True
|
|
assert router._deployment_has_budget_limits(deployment=without_budget) is False
|
|
|
|
|
|
def test_sync_deployment_budget_config(monkeypatch):
|
|
import asyncio
|
|
|
|
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
|
|
|
|
router = Router(model_list=[], optional_pre_call_checks=[])
|
|
deployment = Deployment(
|
|
model_name="dynamic-budget-model",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/gpt-4o-mini",
|
|
api_key="fake-key",
|
|
max_budget=0.000000000001,
|
|
budget_duration="1d",
|
|
),
|
|
model_info=ModelInfo(id="runtime-budget-deployment"),
|
|
)
|
|
|
|
router._sync_deployment_budget_config(deployment=deployment)
|
|
|
|
budget_limiter = router._get_router_deployment_budget_limiter()
|
|
assert budget_limiter is not None
|
|
config = budget_limiter._get_budget_config_for_deployment(
|
|
"runtime-budget-deployment"
|
|
)
|
|
assert config is not None
|
|
assert config.max_budget == 0.000000000001
|
|
|
|
|
|
def test_sync_deployment_budget_config_clears_removed_limits(monkeypatch):
|
|
import asyncio
|
|
|
|
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
|
|
|
|
router = Router(model_list=[], optional_pre_call_checks=[])
|
|
model_id = "runtime-budget-deployment"
|
|
budgeted = Deployment(
|
|
model_name="dynamic-budget-model",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/gpt-4o-mini",
|
|
api_key="fake-key",
|
|
max_budget=0.000000000001,
|
|
budget_duration="1d",
|
|
),
|
|
model_info=ModelInfo(id=model_id),
|
|
)
|
|
unbudgeted = Deployment(
|
|
model_name="dynamic-budget-model",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/gpt-4o-mini",
|
|
api_key="fake-key",
|
|
),
|
|
model_info=ModelInfo(id=model_id),
|
|
)
|
|
|
|
router._sync_deployment_budget_config(deployment=budgeted)
|
|
budget_limiter = router._get_router_deployment_budget_limiter()
|
|
assert budget_limiter is not None
|
|
assert budget_limiter._get_budget_config_for_deployment(model_id) is not None
|
|
|
|
router._sync_deployment_budget_config(deployment=unbudgeted)
|
|
assert budget_limiter._get_budget_config_for_deployment(model_id) is None
|
|
|
|
|
|
def test_upsert_deployment_clears_stale_budget_config(monkeypatch):
|
|
import asyncio
|
|
|
|
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
|
|
|
|
router = Router(model_list=[], optional_pre_call_checks=[])
|
|
model_id = "upsert-budget-deployment"
|
|
budgeted = Deployment(
|
|
model_name="dynamic-budget-model",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/gpt-4o-mini",
|
|
api_key="fake-key",
|
|
max_budget=0.000000000001,
|
|
budget_duration="1d",
|
|
),
|
|
model_info=ModelInfo(id=model_id),
|
|
)
|
|
unbudgeted = Deployment(
|
|
model_name="dynamic-budget-model",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/gpt-4o-mini",
|
|
api_key="fake-key",
|
|
),
|
|
model_info=ModelInfo(id=model_id),
|
|
)
|
|
|
|
router.upsert_deployment(deployment=budgeted)
|
|
budget_limiter = router._get_router_deployment_budget_limiter()
|
|
assert budget_limiter is not None
|
|
assert budget_limiter._get_budget_config_for_deployment(model_id) is not None
|
|
|
|
router.upsert_deployment(deployment=unbudgeted)
|
|
assert budget_limiter._get_budget_config_for_deployment(model_id) is None
|