litellm/tests/router_unit_tests/test_router_helper_utils.py
yuneng-jiang 14f4c34c61
fix(ci): stop stale CI reds, keep unit tests off the host env, retry CyberArk policy conflicts (#43294)
* 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
2026-09-26 09:25:13 -07:00

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