mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
8270 lines
288 KiB
Python
8270 lines
288 KiB
Python
import asyncio
|
||
import copy
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
sys.path.insert(
|
||
0, os.path.abspath("../../..")
|
||
) # Adds the parent directory to the system path
|
||
|
||
|
||
import litellm
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
from litellm.integrations.custom_logger import CustomLogger
|
||
|
||
|
||
def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata():
|
||
# initialize a real Router (env‑vars can be empty)
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "azure/gpt-4.1-mini",
|
||
"api_key": os.getenv("AZURE_AI_API_KEY"),
|
||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
# override to known defaults for the test
|
||
router.default_litellm_params = {
|
||
"foo": "bar",
|
||
"metadata": {"baz": 123},
|
||
}
|
||
original = copy.deepcopy(router.default_litellm_params)
|
||
kwargs: dict = {}
|
||
|
||
# invoke the helper
|
||
router._update_kwargs_with_default_litellm_params(
|
||
kwargs=kwargs,
|
||
metadata_variable_name="litellm_metadata",
|
||
)
|
||
|
||
# 1) router.defaults must be unchanged
|
||
assert router.default_litellm_params == original
|
||
|
||
# 2) non‑metadata keys get merged
|
||
assert kwargs["foo"] == "bar"
|
||
|
||
# 3) metadata lands under "metadata"
|
||
assert kwargs["litellm_metadata"] == {"baz": 123}
|
||
|
||
|
||
def test_router_with_model_info_and_model_group():
|
||
"""
|
||
Test edge case where user specifies model_group in model_info
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
},
|
||
"model_info": {
|
||
"tpm": 1000,
|
||
"rpm": 1000,
|
||
"model_group": "gpt-3.5-turbo",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
router._set_model_group_info(
|
||
model_group="gpt-3.5-turbo",
|
||
user_facing_model_group_name="gpt-3.5-turbo",
|
||
)
|
||
|
||
|
||
def test_router_model_group_encrypted_content_affinity_callback_registration():
|
||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||
DeploymentAffinityCheck,
|
||
)
|
||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||
EncryptedContentAffinityCheck,
|
||
)
|
||
|
||
model_group = "openai.gpt-5.1-codex"
|
||
model_group_affinity_config = {
|
||
model_group: ["encrypted_content_affinity"],
|
||
}
|
||
original_callbacks = list(litellm.callbacks)
|
||
litellm.callbacks = []
|
||
router = None
|
||
|
||
try:
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": model_group,
|
||
"litellm_params": {
|
||
"model": "openai/gpt-5.1-codex",
|
||
"api_key": "mock-api-key",
|
||
},
|
||
}
|
||
],
|
||
model_group_affinity_config=model_group_affinity_config,
|
||
num_retries=0,
|
||
)
|
||
callbacks = router.optional_callbacks or []
|
||
encrypted_content_callbacks = [
|
||
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
||
]
|
||
deployment_callback = next(
|
||
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
|
||
)
|
||
assert len(encrypted_content_callbacks) == 1
|
||
assert encrypted_content_callbacks[0].enable_global_affinity is False
|
||
assert (
|
||
encrypted_content_callbacks[0].model_group_affinity_config
|
||
== model_group_affinity_config
|
||
)
|
||
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
|
||
deployment_callback
|
||
)
|
||
assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (
|
||
litellm.callbacks.index(deployment_callback)
|
||
)
|
||
|
||
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
|
||
|
||
callbacks = router.optional_callbacks or []
|
||
encrypted_content_callbacks = [
|
||
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
||
]
|
||
assert len(encrypted_content_callbacks) == 1
|
||
assert encrypted_content_callbacks[0].enable_global_affinity is True
|
||
assert encrypted_content_callbacks[0].router is router
|
||
finally:
|
||
if router is not None:
|
||
router.discard()
|
||
litellm.callbacks = original_callbacks
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_encrypted_content_affinity_model_group_config_is_additive():
|
||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||
EncryptedContentAffinityCheck,
|
||
)
|
||
|
||
model_group = "openai.gpt-5.1-codex"
|
||
target_deployment = {
|
||
"model_name": model_group,
|
||
"litellm_params": {"model": "openai/gpt-5.1-codex"},
|
||
"model_info": {"id": "deployment-b"},
|
||
}
|
||
healthy_deployments = [
|
||
{
|
||
"model_name": model_group,
|
||
"litellm_params": {"model": "openai/gpt-5.1-codex"},
|
||
"model_info": {"id": "deployment-a"},
|
||
},
|
||
target_deployment,
|
||
]
|
||
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
|
||
"deployment-b", "rs_test"
|
||
)
|
||
|
||
assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled(
|
||
{model_group: ["encrypted_content_affinity"]}
|
||
)
|
||
assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None)
|
||
|
||
per_group_check = EncryptedContentAffinityCheck(
|
||
enable_global_affinity=False,
|
||
model_group_affinity_config={
|
||
model_group: ["encrypted_content_affinity"],
|
||
},
|
||
)
|
||
request_kwargs = {
|
||
"input": [{"type": "reasoning", "id": encoded_id}],
|
||
"litellm_metadata": {},
|
||
}
|
||
filtered = await per_group_check.async_filter_deployments(
|
||
model=model_group,
|
||
healthy_deployments=healthy_deployments,
|
||
messages=None,
|
||
request_kwargs=request_kwargs,
|
||
)
|
||
|
||
assert filtered == [target_deployment]
|
||
assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"]
|
||
|
||
disabled_check = EncryptedContentAffinityCheck(
|
||
enable_global_affinity=False,
|
||
model_group_affinity_config={
|
||
"other-model-group": ["encrypted_content_affinity"],
|
||
},
|
||
)
|
||
disabled_request_kwargs = {
|
||
"input": [{"type": "reasoning", "id": encoded_id}],
|
||
"litellm_metadata": {},
|
||
}
|
||
unfiltered = await disabled_check.async_filter_deployments(
|
||
model=model_group,
|
||
healthy_deployments=healthy_deployments,
|
||
messages=None,
|
||
request_kwargs=disabled_request_kwargs,
|
||
)
|
||
|
||
assert unfiltered == healthy_deployments
|
||
assert (
|
||
"encrypted_content_affinity_enabled"
|
||
not in disabled_request_kwargs["litellm_metadata"]
|
||
)
|
||
|
||
global_check = EncryptedContentAffinityCheck(
|
||
enable_global_affinity=True,
|
||
model_group_affinity_config={
|
||
model_group: ["deployment_affinity"],
|
||
},
|
||
)
|
||
global_request_kwargs = {
|
||
"input": [{"type": "reasoning", "id": encoded_id}],
|
||
"litellm_metadata": {},
|
||
}
|
||
globally_filtered = await global_check.async_filter_deployments(
|
||
model=model_group,
|
||
healthy_deployments=healthy_deployments,
|
||
messages=None,
|
||
request_kwargs=global_request_kwargs,
|
||
)
|
||
|
||
assert globally_filtered == [target_deployment]
|
||
assert global_request_kwargs["litellm_metadata"][
|
||
"encrypted_content_affinity_enabled"
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity():
|
||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||
DeploymentAffinityCheck,
|
||
)
|
||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
||
EncryptedContentAffinityCheck,
|
||
)
|
||
|
||
model_group = "openai.gpt-5.1-codex"
|
||
user_api_key_hash = "test-user-key"
|
||
deployment_a = {
|
||
"model_name": model_group,
|
||
"litellm_params": {
|
||
"model": "openai/gpt-5.1-codex",
|
||
"api_key": "mock-api-key-a",
|
||
},
|
||
"model_info": {"id": "deployment-a"},
|
||
}
|
||
deployment_b = {
|
||
"model_name": model_group,
|
||
"litellm_params": {
|
||
"model": "openai/gpt-5.1-codex",
|
||
"api_key": "mock-api-key-b",
|
||
},
|
||
"model_info": {"id": "deployment-b"},
|
||
}
|
||
original_callbacks = list(litellm.callbacks)
|
||
litellm.callbacks = []
|
||
router = None
|
||
|
||
try:
|
||
router = litellm.Router(
|
||
model_list=[deployment_a, deployment_b],
|
||
model_group_affinity_config={
|
||
model_group: [
|
||
"deployment_affinity",
|
||
"encrypted_content_affinity",
|
||
],
|
||
},
|
||
num_retries=0,
|
||
)
|
||
callbacks = router.optional_callbacks or []
|
||
deployment_callback = next(
|
||
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
|
||
)
|
||
encrypted_content_callback = next(
|
||
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
||
)
|
||
assert callbacks.index(encrypted_content_callback) < callbacks.index(
|
||
deployment_callback
|
||
)
|
||
assert litellm.callbacks.index(encrypted_content_callback) < (
|
||
litellm.callbacks.index(deployment_callback)
|
||
)
|
||
|
||
cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
|
||
model_group=model_group,
|
||
user_key=user_api_key_hash,
|
||
)
|
||
await deployment_callback.cache.async_set_cache(
|
||
key=cache_key,
|
||
value={"model_id": "deployment-a"},
|
||
ttl=60,
|
||
)
|
||
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
|
||
"deployment-b", "rs_test"
|
||
)
|
||
request_kwargs = {
|
||
"input": [{"type": "reasoning", "id": encoded_id}],
|
||
"litellm_metadata": {"user_api_key_hash": user_api_key_hash},
|
||
}
|
||
|
||
filtered = await router.async_callback_filter_deployments(
|
||
model=model_group,
|
||
healthy_deployments=[deployment_a, deployment_b],
|
||
messages=None,
|
||
parent_otel_span=None,
|
||
request_kwargs=request_kwargs,
|
||
)
|
||
|
||
assert filtered == [deployment_b]
|
||
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
|
||
finally:
|
||
if router is not None:
|
||
router.discard()
|
||
litellm.callbacks = original_callbacks
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_arouter_with_tags_and_fallbacks():
|
||
"""
|
||
If fallback model missing tag, raise error
|
||
"""
|
||
from litellm import Router
|
||
|
||
router = Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"mock_response": "Hello, world!",
|
||
"tags": ["test"],
|
||
},
|
||
},
|
||
{
|
||
"model_name": "anthropic-claude-3-5-sonnet",
|
||
"litellm_params": {
|
||
"model": "claude-sonnet-4-5-20250929",
|
||
"mock_response": "Hello, world 2!",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[
|
||
{"gpt-3.5-turbo": ["anthropic-claude-3-5-sonnet"]},
|
||
],
|
||
enable_tag_filtering=True,
|
||
)
|
||
|
||
with pytest.raises(Exception):
|
||
response = await router.acompletion(
|
||
model="gpt-3.5-turbo",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
mock_testing_fallbacks=True,
|
||
metadata={"tags": ["test"]},
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file():
|
||
"""
|
||
Write to all deployments of a model
|
||
"""
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
},
|
||
{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-4o-mini"}},
|
||
],
|
||
)
|
||
|
||
with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file:
|
||
mock_acreate_file.return_value = MagicMock()
|
||
response = await router.acreate_file(
|
||
model="gpt-3.5-turbo",
|
||
purpose="test",
|
||
file=MagicMock(),
|
||
)
|
||
|
||
# assert that the mock_acreate_file was called twice
|
||
assert mock_acreate_file.call_count == 2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file_with_jsonl():
|
||
"""
|
||
Test router.acreate_file with both JSONL and non-JSONL files
|
||
"""
|
||
import json
|
||
from io import BytesIO
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
# Create test JSONL content
|
||
jsonl_data = [
|
||
{
|
||
"body": {
|
||
"model": "gpt-3.5-turbo-router",
|
||
"messages": [{"role": "user", "content": "test"}],
|
||
}
|
||
},
|
||
{
|
||
"body": {
|
||
"model": "gpt-3.5-turbo-router",
|
||
"messages": [{"role": "user", "content": "test2"}],
|
||
}
|
||
},
|
||
]
|
||
jsonl_content = "\n".join(json.dumps(item) for item in jsonl_data)
|
||
jsonl_file = BytesIO(jsonl_content.encode("utf-8"))
|
||
jsonl_file.name = "test.jsonl"
|
||
|
||
# Create test non-JSONL content
|
||
non_jsonl_content = "This is not a JSONL file"
|
||
non_jsonl_file = BytesIO(non_jsonl_content.encode("utf-8"))
|
||
non_jsonl_file.name = "test.txt"
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo-router",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
},
|
||
{
|
||
"model_name": "gpt-3.5-turbo-router",
|
||
"litellm_params": {"model": "gpt-4o-mini"},
|
||
},
|
||
],
|
||
)
|
||
|
||
with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file:
|
||
# Test with JSONL file
|
||
response = await router.acreate_file(
|
||
model="gpt-3.5-turbo-router",
|
||
purpose="batch",
|
||
file=jsonl_file,
|
||
)
|
||
|
||
# Verify mock was called twice (once for each deployment)
|
||
print(f"mock_acreate_file.call_count: {mock_acreate_file.call_count}")
|
||
print(f"mock_acreate_file.call_args_list: {mock_acreate_file.call_args_list}")
|
||
assert mock_acreate_file.call_count == 2
|
||
|
||
# Get the file content passed to the first call
|
||
first_call_file = mock_acreate_file.call_args_list[0][1]["file"]
|
||
first_call_content = first_call_file.read().decode("utf-8")
|
||
|
||
# Verify the model name was replaced in the JSONL content
|
||
first_line = json.loads(first_call_content.split("\n")[0])
|
||
assert first_line["body"]["model"] == "gpt-3.5-turbo"
|
||
|
||
# Reset mock for next test
|
||
mock_acreate_file.reset_mock()
|
||
|
||
# Test with non-JSONL file
|
||
response = await router.acreate_file(
|
||
model="gpt-3.5-turbo-router",
|
||
purpose="user_data",
|
||
file=non_jsonl_file,
|
||
)
|
||
|
||
# Verify mock was called twice
|
||
assert mock_acreate_file.call_count == 2
|
||
|
||
# Get the file content passed to the first call
|
||
first_call_file = mock_acreate_file.call_args_list[0][1]["file"]
|
||
first_call_content = first_call_file.read().decode("utf-8")
|
||
|
||
# Verify the non-JSONL content was not modified
|
||
assert first_call_content == non_jsonl_content
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file_does_not_fall_back_across_model_groups():
|
||
"""A file created for batches only exists under the credentials of the model group
|
||
the caller named. A cross-group fallback silently stores it with the wrong provider
|
||
and the later batch create against the named group permanently fails."""
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "azure-gpt",
|
||
"litellm_params": {
|
||
"model": "azure/my-azure-deployment",
|
||
"api_base": "http://127.0.0.1:9",
|
||
"api_key": "dummy-key",
|
||
"api_version": "2024-06-01",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "openai-gpt",
|
||
"litellm_params": {"model": "gpt-4o-mini"},
|
||
},
|
||
],
|
||
fallbacks=[{"azure-gpt": ["openai-gpt"]}],
|
||
)
|
||
|
||
def fail_azure(*args: object, **kwargs: object) -> MagicMock:
|
||
if kwargs.get("model") == "azure/my-azure-deployment":
|
||
raise litellm.APIConnectionError(
|
||
message="Connection error.",
|
||
llm_provider="azure",
|
||
model="azure/my-azure-deployment",
|
||
)
|
||
return MagicMock()
|
||
|
||
with patch("litellm.acreate_file", side_effect=fail_azure) as mock_acreate_file:
|
||
with pytest.raises(litellm.APIConnectionError):
|
||
await router.acreate_file(
|
||
model="azure-gpt",
|
||
purpose="batch",
|
||
file=MagicMock(),
|
||
)
|
||
|
||
called_models = [call.kwargs.get("model") for call in mock_acreate_file.call_args_list]
|
||
assert "azure/my-azure-deployment" in called_models
|
||
assert "gpt-4o-mini" not in called_models
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file_uses_deployment_custom_llm_provider():
|
||
"""
|
||
Ensure file routing preserves deployment custom_llm_provider instead of
|
||
inferring provider from model string alone.
|
||
"""
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "team-azure-batch",
|
||
"litellm_params": {
|
||
"model": "gpt-4.1-mini",
|
||
"custom_llm_provider": "azure",
|
||
"api_base": "https://example-resource.openai.azure.com",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file:
|
||
await router.acreate_file(
|
||
model="team-azure-batch",
|
||
purpose="batch",
|
||
file=MagicMock(),
|
||
)
|
||
|
||
assert mock_acreate_file.call_count == 1
|
||
assert mock_acreate_file.call_args.kwargs["custom_llm_provider"] == "azure"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file_forwards_target_model_names_to_litellm_proxy():
|
||
import json
|
||
from io import BytesIO
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
jsonl_file = BytesIO(
|
||
json.dumps({"body": {"model": "chained-batch", "messages": [{"role": "user", "content": "hi"}]}}).encode(
|
||
"utf-8"
|
||
)
|
||
)
|
||
jsonl_file.name = "test.jsonl"
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "chained-batch",
|
||
"litellm_params": {
|
||
"model": "litellm_proxy/gpt-4.1-batch",
|
||
"api_base": "http://localhost:4001/v1",
|
||
"api_key": "sk-proxy-b",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file:
|
||
await router.acreate_file(
|
||
model="chained-batch",
|
||
purpose="batch",
|
||
file=jsonl_file,
|
||
)
|
||
|
||
assert mock_acreate_file.call_count == 1
|
||
call_kwargs = mock_acreate_file.call_args.kwargs
|
||
assert call_kwargs["custom_llm_provider"] == "litellm_proxy"
|
||
assert call_kwargs["extra_body"] == {"target_model_names": "gpt-4.1-batch"}
|
||
uploaded_line = json.loads(call_kwargs["file"].read().decode("utf-8").split("\n")[0])
|
||
assert uploaded_line["body"]["model"] == "gpt-4.1-batch"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file_does_not_inject_target_model_names_for_other_providers():
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4.1-batch",
|
||
"litellm_params": {"model": "gpt-4.1"},
|
||
},
|
||
],
|
||
)
|
||
|
||
with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file:
|
||
await router.acreate_file(
|
||
model="gpt-4.1-batch",
|
||
purpose="batch",
|
||
file=MagicMock(),
|
||
)
|
||
|
||
assert mock_acreate_file.call_count == 1
|
||
assert mock_acreate_file.call_args.kwargs.get("extra_body") is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_acreate_file_litellm_proxy_sends_target_model_names_in_multipart_form():
|
||
import json
|
||
from io import BytesIO
|
||
|
||
import httpx
|
||
import respx
|
||
|
||
jsonl_file = BytesIO(
|
||
json.dumps({"body": {"model": "chained-batch", "messages": [{"role": "user", "content": "hi"}]}}).encode(
|
||
"utf-8"
|
||
)
|
||
)
|
||
jsonl_file.name = "test.jsonl"
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "chained-batch",
|
||
"litellm_params": {
|
||
"model": "litellm_proxy/gpt-4.1-batch",
|
||
"api_base": "http://localhost:4001/v1",
|
||
"api_key": "sk-proxy-b",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
file_object_json = {
|
||
"id": "file-abc123",
|
||
"object": "file",
|
||
"bytes": 100,
|
||
"created_at": 1700000000,
|
||
"filename": "test.jsonl",
|
||
"purpose": "batch",
|
||
"status": "processed",
|
||
}
|
||
|
||
with respx.mock(assert_all_called=True) as respx_mock:
|
||
create_route = respx_mock.post("http://localhost:4001/v1/files").mock(
|
||
return_value=httpx.Response(200, json=file_object_json)
|
||
)
|
||
response = await router.acreate_file(
|
||
model="chained-batch",
|
||
purpose="batch",
|
||
file=jsonl_file,
|
||
)
|
||
|
||
assert response.id == "file-abc123"
|
||
request_body = create_route.calls.last.request.content
|
||
assert b'name="target_model_names"' in request_body
|
||
assert b"gpt-4.1-batch" in request_body
|
||
assert b'name="purpose"' in request_body
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_router_afile_content_uses_deployment_custom_llm_provider():
|
||
"""
|
||
Regression test: Ensure afile_content preserves deployment custom_llm_provider
|
||
when model name lacks provider prefix (e.g., "gpt-4.1-mini" instead of "azure/gpt-4.1-mini").
|
||
|
||
This prevents "None is not a valid LlmProviders" errors when calling file content operations.
|
||
"""
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "team-azure-batch",
|
||
"litellm_params": {
|
||
"model": "gpt-4.1-mini", # No provider prefix
|
||
"custom_llm_provider": "azure",
|
||
"api_base": "https://example-resource.openai.azure.com",
|
||
"api_key": "test-key",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# Mock the Azure file handler's afile_content method
|
||
mock_response = MagicMock(spec=HttpxBinaryResponseContent)
|
||
mock_response.response = MagicMock()
|
||
|
||
with patch(
|
||
"litellm.llms.azure.files.handler.AzureOpenAIFilesAPI.afile_content",
|
||
return_value=mock_response,
|
||
) as mock_afile_content:
|
||
result = await router.afile_content(
|
||
model="team-azure-batch",
|
||
file_id="file-123",
|
||
)
|
||
|
||
# Verify the call was made (proves custom_llm_provider was correctly passed)
|
||
assert mock_afile_content.call_count == 1
|
||
assert result == mock_response
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_arouter_async_get_healthy_deployments():
|
||
"""
|
||
Test that afile_content returns the correct file content
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
},
|
||
],
|
||
)
|
||
|
||
result = await router.async_get_healthy_deployments(
|
||
model="gpt-3.5-turbo",
|
||
request_kwargs={},
|
||
messages=None,
|
||
input=None,
|
||
specific_deployment=False,
|
||
parent_otel_span=None,
|
||
)
|
||
|
||
assert len(result) == 1
|
||
assert result[0]["model_name"] == "gpt-3.5-turbo"
|
||
assert result[0]["litellm_params"]["model"] == "gpt-3.5-turbo"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@patch("litellm.amoderation")
|
||
async def test_arouter_amoderation_with_credential_name(mock_amoderation):
|
||
"""
|
||
Test that router.amoderation passes litellm_credential_name to the underlying litellm.amoderation call
|
||
"""
|
||
mock_amoderation.return_value = AsyncMock()
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "text-moderation-stable",
|
||
"litellm_params": {
|
||
"model": "text-moderation-stable",
|
||
"litellm_credential_name": "my-custom-auth",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
await router.amoderation(input="I love everyone!", model="text-moderation-stable")
|
||
|
||
mock_amoderation.assert_called_once()
|
||
call_kwargs = mock_amoderation.call_args[1] # Get the kwargs of the call
|
||
print(
|
||
"call kwargs for router.amoderation=",
|
||
json.dumps(call_kwargs, indent=4, default=str),
|
||
)
|
||
assert call_kwargs["litellm_credential_name"] == "my-custom-auth"
|
||
assert call_kwargs["model"] == "text-moderation-stable"
|
||
|
||
|
||
def test_arouter_test_team_model():
|
||
"""
|
||
Test that router.test_team_model returns the correct model
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
"model_info": {
|
||
"team_id": "test-team",
|
||
"team_public_model_name": "test-model",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
result = router.map_team_model(team_model_name="test-model", team_id="test-team")
|
||
assert result is not None
|
||
|
||
|
||
def test_arouter_ignore_invalid_deployments():
|
||
"""
|
||
Test that router.ignore_invalid_deployments is set to True
|
||
"""
|
||
from litellm.types.router import Deployment
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "my-bad-model"},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
|
||
assert router.ignore_invalid_deployments is True
|
||
assert router.get_model_list() == []
|
||
|
||
## check upsert deployment
|
||
router.upsert_deployment(
|
||
Deployment(
|
||
model_name="gpt-3.5-turbo",
|
||
litellm_params={"model": "my-bad-model"}, # type: ignore
|
||
model_info={"tpm": 1000, "rpm": 1000},
|
||
)
|
||
)
|
||
|
||
assert router.get_model_list() == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_arouter_aretrieve_batch():
|
||
"""
|
||
Test that router.aretrieve_batch returns the correct response
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"custom_llm_provider": "azure",
|
||
"api_key": "my-custom-key",
|
||
"api_base": "my-custom-base",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
with patch.object(
|
||
litellm, "aretrieve_batch", return_value=AsyncMock()
|
||
) as mock_aretrieve_batch:
|
||
try:
|
||
response = await router.aretrieve_batch(
|
||
model="gpt-3.5-turbo",
|
||
)
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
|
||
mock_aretrieve_batch.assert_called_once()
|
||
|
||
print(mock_aretrieve_batch.call_args.kwargs)
|
||
assert mock_aretrieve_batch.call_args.kwargs["api_key"] == "my-custom-key"
|
||
assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_arouter_aretrieve_file_content():
|
||
"""
|
||
Test that router.acreate_file with JSONL file returns the correct response
|
||
"""
|
||
|
||
with patch.object(
|
||
litellm, "afile_content", return_value=AsyncMock()
|
||
) as mock_afile_content:
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"custom_llm_provider": "azure",
|
||
"api_key": "my-custom-key",
|
||
"api_base": "my-custom-base",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
try:
|
||
response = await router.afile_content(
|
||
**{
|
||
"model": "gpt-3.5-turbo",
|
||
"file_id": "my-unique-file-id",
|
||
}
|
||
) # type: ignore
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
|
||
mock_afile_content.assert_called_once()
|
||
|
||
print(mock_afile_content.call_args.kwargs)
|
||
assert mock_afile_content.call_args.kwargs["api_key"] == "my-custom-key"
|
||
assert mock_afile_content.call_args.kwargs["api_base"] == "my-custom-base"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_arouter_filter_team_based_models():
|
||
"""
|
||
Test that router.filter_team_based_models filters out models that are not in the team
|
||
"""
|
||
from litellm.types.router import Deployment
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
"model_info": {
|
||
"team_id": "test-team",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# WORKS
|
||
result = await router.acompletion(
|
||
model="gpt-3.5-turbo",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
metadata={"user_api_key_team_id": "test-team"},
|
||
mock_response="Hello, world!",
|
||
)
|
||
|
||
assert result is not None
|
||
|
||
# FAILS
|
||
with pytest.raises(Exception) as e:
|
||
result = await router.acompletion(
|
||
model="gpt-3.5-turbo",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
metadata={"user_api_key_team_id": "test-team-2"},
|
||
mock_response="Hello, world!",
|
||
)
|
||
assert "No deployments available" in str(e.value)
|
||
|
||
## ADD A MODEL THAT IS NOT IN THE TEAM
|
||
router.add_deployment(
|
||
Deployment(
|
||
model_name="gpt-3.5-turbo",
|
||
litellm_params={"model": "gpt-3.5-turbo"}, # type: ignore
|
||
model_info={"tpm": 1000, "rpm": 1000},
|
||
)
|
||
)
|
||
|
||
result = await router.acompletion(
|
||
model="gpt-3.5-turbo",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
metadata={"user_api_key_team_id": "test-team-2"},
|
||
mock_response="Hello, world!",
|
||
)
|
||
|
||
assert result is not None
|
||
|
||
|
||
def test_arouter_should_include_deployment():
|
||
"""
|
||
Test the should_include_deployment method with various scenarios
|
||
|
||
The method logic:
|
||
1. Returns True if: team_id matches AND model_name matches team_public_model_name
|
||
2. Returns True if: model_name matches AND deployment has no team_id
|
||
3. Otherwise returns False
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
"model_info": {
|
||
"team_id": "test-team",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# Test deployment structures
|
||
deployment_with_team_and_public_name = {
|
||
"model_name": "gpt-3.5-turbo",
|
||
"model_info": {
|
||
"team_id": "test-team",
|
||
"team_public_model_name": "team-gpt-model",
|
||
},
|
||
}
|
||
|
||
deployment_with_team_no_public_name = {
|
||
"model_name": "gpt-3.5-turbo",
|
||
"model_info": {
|
||
"team_id": "test-team",
|
||
},
|
||
}
|
||
|
||
deployment_without_team = {
|
||
"model_name": "gpt-4",
|
||
"model_info": {},
|
||
}
|
||
|
||
deployment_different_team = {
|
||
"model_name": "claude-3",
|
||
"model_info": {
|
||
"team_id": "other-team",
|
||
"team_public_model_name": "team-claude-model",
|
||
},
|
||
}
|
||
|
||
# Test Case 1: Team-specific deployment - team_id and team_public_model_name match
|
||
result = router.should_include_deployment(
|
||
model_name="team-gpt-model",
|
||
model=deployment_with_team_and_public_name,
|
||
team_id="test-team",
|
||
)
|
||
assert (
|
||
result is True
|
||
), "Should return True when team_id and team_public_model_name match"
|
||
|
||
# Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name
|
||
result = router.should_include_deployment(
|
||
model_name="different-model",
|
||
model=deployment_with_team_and_public_name,
|
||
team_id="test-team",
|
||
)
|
||
assert (
|
||
result is False
|
||
), "Should return False when team_id matches but model_name doesn't match team_public_model_name"
|
||
|
||
# Test Case 3: Team-specific deployment - team_id doesn't match
|
||
result = router.should_include_deployment(
|
||
model_name="team-gpt-model",
|
||
model=deployment_with_team_and_public_name,
|
||
team_id="different-team",
|
||
)
|
||
assert result is False, "Should return False when team_id doesn't match"
|
||
|
||
# Test Case 4: Team-specific deployment with no team_public_model_name - should fail
|
||
result = router.should_include_deployment(
|
||
model_name="gpt-3.5-turbo",
|
||
model=deployment_with_team_no_public_name,
|
||
team_id="test-team",
|
||
)
|
||
assert (
|
||
result is True
|
||
), "Should return True when team deployment has no team_public_model_name to match"
|
||
|
||
# Test Case 5: Non-team deployment - model_name matches and no team_id
|
||
result = router.should_include_deployment(
|
||
model_name="gpt-4", model=deployment_without_team, team_id=None
|
||
)
|
||
assert (
|
||
result is True
|
||
), "Should return True when model_name matches and deployment has no team_id"
|
||
|
||
# Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work)
|
||
result = router.should_include_deployment(
|
||
model_name="gpt-4", model=deployment_without_team, team_id="any-team"
|
||
)
|
||
assert (
|
||
result is True
|
||
), "Should return True when model_name matches non-team deployment, regardless of team_id param"
|
||
|
||
# Test Case 7: Non-team deployment - model_name doesn't match
|
||
result = router.should_include_deployment(
|
||
model_name="different-model", model=deployment_without_team, team_id=None
|
||
)
|
||
assert result is False, "Should return False when model_name doesn't match"
|
||
|
||
# Test Case 8: Team deployment accessed without matching team_id
|
||
result = router.should_include_deployment(
|
||
model_name="gpt-3.5-turbo",
|
||
model=deployment_with_team_and_public_name,
|
||
team_id=None,
|
||
)
|
||
assert (
|
||
result is True
|
||
), "Should return True when matching model with exact model_name"
|
||
|
||
|
||
def test_arouter_responses_api_bridge():
|
||
"""
|
||
Test that router.responses_api_bridge returns the correct response
|
||
"""
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "[IP-approved] o3-pro",
|
||
"litellm_params": {
|
||
"model": "azure/responses/o_series/webinterface-o3-pro",
|
||
"api_base": "https://webhook.site/fba79dae-220a-4bb7-9a3a-8caa49604e55",
|
||
"api_key": "sk-1234567890",
|
||
"api_version": "preview",
|
||
"stream": True,
|
||
},
|
||
"model_info": {
|
||
"input_cost_per_token": 0.00002,
|
||
"output_cost_per_token": 0.00008,
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
## CONFIRM BRIDGE IS CALLED
|
||
with patch.object(litellm, "responses", return_value=AsyncMock()) as mock_responses:
|
||
result = router.completion(
|
||
model="[IP-approved] o3-pro",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
)
|
||
assert mock_responses.call_count == 1
|
||
|
||
## CONFIRM MODEL NAME IS STRIPPED
|
||
client = HTTPHandler()
|
||
|
||
mock_response = MagicMock()
|
||
mock_response.status_code = 200
|
||
mock_response.headers = {"content-type": "application/json"}
|
||
mock_response.json.return_value = {
|
||
"id": "resp_test",
|
||
"object": "response",
|
||
"status": "completed",
|
||
"output": [],
|
||
}
|
||
mock_response.text = (
|
||
'{"id": "resp_test", "object": "response", "status": "completed", "output": []}'
|
||
)
|
||
|
||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||
try:
|
||
result = router.completion(
|
||
model="[IP-approved] o3-pro",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
client=client,
|
||
num_retries=0,
|
||
)
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
|
||
assert mock_post.call_count == 1
|
||
assert (
|
||
mock_post.call_args.kwargs["url"]
|
||
== "https://webhook.site/fba79dae-220a-4bb7-9a3a-8caa49604e55/openai/v1/responses?api-version=preview"
|
||
)
|
||
assert mock_post.call_args.kwargs["json"]["model"] == "webinterface-o3-pro"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_router_v1_messages_fallbacks():
|
||
"""
|
||
Test that router.v1_messages_fallbacks returns the correct response
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "claude-sonnet-4-5-20250929",
|
||
"litellm_params": {
|
||
"model": "anthropic/claude-sonnet-4-5-20250929",
|
||
"mock_response": "litellm.InternalServerError",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "bedrock-claude",
|
||
"litellm_params": {
|
||
"model": "anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
"mock_response": "Hello, world I am a fallback!",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[
|
||
{"claude-sonnet-4-5-20250929": ["bedrock-claude"]},
|
||
],
|
||
)
|
||
|
||
result = await router.aanthropic_messages(
|
||
model="claude-sonnet-4-5-20250929",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
max_tokens=256,
|
||
)
|
||
assert result is not None
|
||
|
||
print(result)
|
||
assert result["content"][0]["text"] == "Hello, world I am a fallback!"
|
||
|
||
|
||
def test_add_invalid_provider_to_router():
|
||
"""
|
||
Test that router.add_deployment raises an error if the provider is invalid
|
||
"""
|
||
from litellm.types.router import Deployment
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
)
|
||
|
||
with pytest.raises(Exception) as e:
|
||
router.add_deployment(
|
||
Deployment(
|
||
model_name="vertex_ai/*",
|
||
litellm_params={
|
||
"model": "vertex_ai/*",
|
||
"custom_llm_provider": "vertex_ai_eu",
|
||
},
|
||
)
|
||
)
|
||
|
||
assert router.pattern_router.patterns == {}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_router_ageneric_api_call_with_fallbacks_helper():
|
||
"""
|
||
Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios
|
||
"""
|
||
from unittest.mock import patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"api_key": "test-key",
|
||
"api_base": "https://api.openai.com/v1",
|
||
},
|
||
"model_info": {
|
||
"tpm": 1000,
|
||
"rpm": 1000,
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# Test 1: Successful call
|
||
async def mock_generic_function(**kwargs):
|
||
return {"result": "success", "model": kwargs.get("model")}
|
||
|
||
with patch.object(router, "async_get_available_deployment") as mock_get_deployment:
|
||
mock_get_deployment.return_value = {
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"api_key": "test-key",
|
||
"api_base": "https://api.openai.com/v1",
|
||
},
|
||
}
|
||
|
||
with patch.object(
|
||
router, "_update_kwargs_with_deployment"
|
||
) as mock_update_kwargs:
|
||
with patch.object(
|
||
router, "async_routing_strategy_pre_call_checks"
|
||
) as mock_pre_call_checks:
|
||
with patch.object(
|
||
router, "_get_client", return_value=None
|
||
) as mock_get_client:
|
||
result = await router._ageneric_api_call_with_fallbacks_helper(
|
||
model="gpt-3.5-turbo",
|
||
original_generic_function=mock_generic_function,
|
||
messages=[{"role": "user", "content": "test"}],
|
||
)
|
||
|
||
assert result is not None
|
||
assert result["result"] == "success"
|
||
mock_get_deployment.assert_called_once()
|
||
mock_update_kwargs.assert_called_once()
|
||
mock_pre_call_checks.assert_called_once()
|
||
|
||
# Test 2: Passthrough on no deployment (success case)
|
||
async def mock_passthrough_function(**kwargs):
|
||
return {"result": "passthrough", "model": kwargs.get("model")}
|
||
|
||
with patch.object(router, "async_get_available_deployment") as mock_get_deployment:
|
||
mock_get_deployment.side_effect = Exception("No deployment available")
|
||
|
||
result = await router._ageneric_api_call_with_fallbacks_helper(
|
||
model="gpt-3.5-turbo",
|
||
original_generic_function=mock_passthrough_function,
|
||
passthrough_on_no_deployment=True,
|
||
messages=[{"role": "user", "content": "test"}],
|
||
)
|
||
|
||
assert result is not None
|
||
assert result["result"] == "passthrough"
|
||
assert result["model"] == "gpt-3.5-turbo"
|
||
|
||
# Test 3: No deployment available and passthrough=False (should raise exception)
|
||
with patch.object(router, "async_get_available_deployment") as mock_get_deployment:
|
||
mock_get_deployment.side_effect = Exception("No deployment available")
|
||
|
||
with pytest.raises(Exception) as exc_info:
|
||
await router._ageneric_api_call_with_fallbacks_helper(
|
||
model="gpt-3.5-turbo",
|
||
original_generic_function=mock_generic_function,
|
||
passthrough_on_no_deployment=False,
|
||
messages=[{"role": "user", "content": "test"}],
|
||
)
|
||
|
||
assert "No deployment available" in str(exc_info.value)
|
||
|
||
# Test 4: Test with semaphore (rate limiting)
|
||
import asyncio
|
||
|
||
async def mock_semaphore_function(**kwargs):
|
||
return {"result": "semaphore_success", "model": kwargs.get("model")}
|
||
|
||
with patch.object(router, "async_get_available_deployment") as mock_get_deployment:
|
||
mock_get_deployment.return_value = {
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"api_key": "test-key",
|
||
"api_base": "https://api.openai.com/v1",
|
||
},
|
||
}
|
||
|
||
mock_semaphore = asyncio.Semaphore(1)
|
||
|
||
with patch.object(
|
||
router, "_update_kwargs_with_deployment"
|
||
) as mock_update_kwargs:
|
||
with patch.object(
|
||
router, "_get_client", return_value=mock_semaphore
|
||
) as mock_get_client:
|
||
with patch.object(
|
||
router, "async_routing_strategy_pre_call_checks"
|
||
) as mock_pre_call_checks:
|
||
result = await router._ageneric_api_call_with_fallbacks_helper(
|
||
model="gpt-3.5-turbo",
|
||
original_generic_function=mock_semaphore_function,
|
||
messages=[{"role": "user", "content": "test"}],
|
||
)
|
||
|
||
assert result is not None
|
||
assert result["result"] == "semaphore_success"
|
||
mock_get_client.assert_called_once()
|
||
mock_pre_call_checks.assert_called_once()
|
||
|
||
# Test 5: Test call tracking (success and failure counts)
|
||
initial_success_count = router.success_calls.get("gpt-3.5-turbo", 0)
|
||
initial_fail_count = router.fail_calls.get("gpt-3.5-turbo", 0)
|
||
|
||
async def mock_failing_function(**kwargs):
|
||
raise Exception("Mock failure")
|
||
|
||
with patch.object(router, "async_get_available_deployment") as mock_get_deployment:
|
||
mock_get_deployment.return_value = {
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
"api_key": "test-key",
|
||
"api_base": "https://api.openai.com/v1",
|
||
},
|
||
}
|
||
|
||
with patch.object(
|
||
router, "_update_kwargs_with_deployment"
|
||
) as mock_update_kwargs:
|
||
with patch.object(
|
||
router, "_get_client", return_value=None
|
||
) as mock_get_client:
|
||
with patch.object(
|
||
router, "async_routing_strategy_pre_call_checks"
|
||
) as mock_pre_call_checks:
|
||
with pytest.raises(Exception) as exc_info:
|
||
await router._ageneric_api_call_with_fallbacks_helper(
|
||
model="gpt-3.5-turbo",
|
||
original_generic_function=mock_failing_function,
|
||
messages=[{"role": "user", "content": "test"}],
|
||
)
|
||
|
||
assert "Mock failure" in str(exc_info.value)
|
||
# Check that fail_calls was incremented
|
||
assert router.fail_calls["gpt-3.5-turbo"] == initial_fail_count + 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_ageneric_api_call_deployment_model_overrides_alias():
|
||
"""
|
||
Regression: when a model alias (e.g. "not-gemini-2.5-flash") maps to a deployment
|
||
with model="vertex_ai/gemini-2.5-flash", the underlying litellm function must receive
|
||
the deployment model, not the alias. Before the fix, **kwargs overwrote data["model"].
|
||
"""
|
||
from unittest.mock import patch
|
||
|
||
captured: dict = {}
|
||
|
||
async def capture_model(**kwargs):
|
||
captured["model"] = kwargs.get("model")
|
||
return {"result": "ok"}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "not-gemini-2.5-flash",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-flash",
|
||
"api_key": "fake-key",
|
||
},
|
||
}
|
||
]
|
||
)
|
||
|
||
def inject_alias_into_kwargs(deployment, kwargs, function_name=None):
|
||
# Simulate the alias leaking into kwargs (as happens when
|
||
# _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before
|
||
# calling the helper through async_function_with_fallbacks).
|
||
kwargs["model"] = "not-gemini-2.5-flash"
|
||
|
||
with (
|
||
patch.object(router, "async_get_available_deployment") as mock_dep,
|
||
patch.object(
|
||
router,
|
||
"_update_kwargs_with_deployment",
|
||
side_effect=inject_alias_into_kwargs,
|
||
),
|
||
patch.object(router, "async_routing_strategy_pre_call_checks"),
|
||
patch.object(router, "_get_client", return_value=None),
|
||
):
|
||
mock_dep.return_value = {
|
||
"model_name": "not-gemini-2.5-flash",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-flash",
|
||
"api_key": "fake-key",
|
||
},
|
||
}
|
||
|
||
await router._ageneric_api_call_with_fallbacks_helper(
|
||
model="not-gemini-2.5-flash",
|
||
original_generic_function=capture_model,
|
||
)
|
||
|
||
assert (
|
||
captured["model"] == "vertex_ai/gemini-2.5-flash"
|
||
), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'"
|
||
|
||
|
||
def test_router_get_model_access_groups_team_only_models():
|
||
"""
|
||
Test that Router.get_model_access_groups returns the correct response for team-only models
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-custom-model-name",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
"model_info": {
|
||
"team_id": "team_1",
|
||
"access_groups": ["default-models"],
|
||
"team_public_model_name": "gpt-3.5-turbo",
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
access_groups = router.get_model_access_groups(
|
||
model_name="gpt-3.5-turbo", team_id=None
|
||
)
|
||
assert len(access_groups) == 0
|
||
|
||
access_groups = router.get_model_access_groups(
|
||
model_name="gpt-3.5-turbo", team_id="team_1"
|
||
)
|
||
assert list(access_groups.keys()) == ["default-models"]
|
||
|
||
|
||
def test_cached_get_model_group_info():
|
||
"""
|
||
Test that _cached_get_model_group_info caches results and
|
||
invalidates on deployment changes.
|
||
"""
|
||
from litellm.types.router import Deployment, LiteLLM_Params
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
|
||
"model_info": {"tpm": 1000, "rpm": 100},
|
||
},
|
||
]
|
||
)
|
||
|
||
# First call should compute and cache
|
||
result1 = router._cached_get_model_group_info("gpt-4")
|
||
assert result1 is not None
|
||
assert result1.tpm == 1000
|
||
|
||
# Second call should hit cache (same object)
|
||
result2 = router._cached_get_model_group_info("gpt-4")
|
||
assert result1 is result2
|
||
|
||
# Add a deployment — cache should be invalidated
|
||
router.add_deployment(
|
||
Deployment(
|
||
model_name="gpt-4",
|
||
litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"),
|
||
model_info={"tpm": 2000, "rpm": 200},
|
||
)
|
||
)
|
||
result3 = router._cached_get_model_group_info("gpt-4")
|
||
assert result3 is not result2
|
||
assert result3 is not None
|
||
assert result3.tpm == 3000 # 1000 + 2000
|
||
|
||
# Delete a deployment — cache should be invalidated
|
||
deployment_id = router.model_list[-1]["model_info"]["id"]
|
||
router.delete_deployment(id=deployment_id)
|
||
result4 = router._cached_get_model_group_info("gpt-4")
|
||
assert result4 is not result3
|
||
assert result4 is not None
|
||
assert result4.tpm == 1000
|
||
|
||
# set_model_list — cache should be invalidated
|
||
router.set_model_list(
|
||
[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
|
||
"model_info": {"tpm": 5000},
|
||
},
|
||
]
|
||
)
|
||
result5 = router._cached_get_model_group_info("gpt-4")
|
||
assert result5 is not result4
|
||
assert result5 is not None
|
||
assert result5.tpm == 5000
|
||
|
||
# Verify cache still works after invalidation
|
||
result6 = router._cached_get_model_group_info("gpt-4")
|
||
assert result5 is result6
|
||
|
||
|
||
def test_model_group_info_cost_from_db_model_info():
|
||
"""
|
||
When get_deployment_model_info fails (model_info is None fallback),
|
||
input_cost_per_token and output_cost_per_token should be read from db model_info.
|
||
"""
|
||
from unittest.mock import patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-custom-model",
|
||
"litellm_params": {
|
||
"model": "openai/my-custom-model",
|
||
"api_key": "fake",
|
||
"api_base": "https://my-custom-endpoint.com",
|
||
},
|
||
"model_info": {
|
||
"input_cost_per_token": 0.0001,
|
||
"output_cost_per_token": 0.0002,
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
with patch.object(
|
||
router, "get_deployment_model_info", side_effect=Exception("not found")
|
||
):
|
||
result = router._cached_get_model_group_info("my-custom-model")
|
||
assert result is not None
|
||
assert result.input_cost_per_token == 0.0001
|
||
assert result.output_cost_per_token == 0.0002
|
||
|
||
|
||
def test_model_group_info_cost_none_when_db_model_info_has_no_cost():
|
||
"""
|
||
When get_deployment_model_info fails and db model_info has no cost fields,
|
||
input/output_cost_per_token should be None.
|
||
"""
|
||
from unittest.mock import patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-custom-model-no-cost",
|
||
"litellm_params": {
|
||
"model": "openai/my-custom-model-no-cost",
|
||
"api_key": "fake",
|
||
"api_base": "https://my-custom-endpoint.com",
|
||
},
|
||
"model_info": {},
|
||
},
|
||
]
|
||
)
|
||
|
||
with patch.object(
|
||
router, "get_deployment_model_info", side_effect=Exception("not found")
|
||
):
|
||
result = router._cached_get_model_group_info("my-custom-model-no-cost")
|
||
assert result is not None
|
||
assert result.input_cost_per_token is None
|
||
assert result.output_cost_per_token is None
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"value,expected",
|
||
[
|
||
("1e-05", 1e-05),
|
||
("0.00001", 1e-05),
|
||
(1e-05, 1e-05),
|
||
(5, 5.0),
|
||
(None, None),
|
||
("not-a-number", None),
|
||
],
|
||
)
|
||
def test_cost_value_as_float(value, expected):
|
||
from litellm.router import _cost_value_as_float
|
||
|
||
assert _cost_value_as_float(value) == expected
|
||
|
||
|
||
def test_model_group_info_with_stringified_cost_values():
|
||
"""
|
||
YAML 1.2 parsers emit '1e-05' (integer mantissa) as a string, so cost
|
||
values in deployment model_info can arrive as str. Aggregating the model
|
||
group must not raise TypeError('>' between str and float) and must return
|
||
float costs.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-custom-model",
|
||
"litellm_params": {
|
||
"model": "openai/my-custom-backend-1",
|
||
"api_key": "fake",
|
||
},
|
||
"model_info": {
|
||
"input_cost_per_token": "1e-05",
|
||
"output_cost_per_token": "1e-05",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "my-custom-model",
|
||
"litellm_params": {
|
||
"model": "openai/my-custom-backend-2",
|
||
"api_key": "fake",
|
||
},
|
||
"model_info": {
|
||
"input_cost_per_token": "2e-05",
|
||
"output_cost_per_token": "2e-05",
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
def _model_info_with_str_costs(model_id: str, model_name: str):
|
||
for model in router.model_list:
|
||
if model["model_info"]["id"] == model_id:
|
||
return {
|
||
"key": model_name,
|
||
"input_cost_per_token": model["model_info"]["input_cost_per_token"],
|
||
"output_cost_per_token": model["model_info"]["output_cost_per_token"],
|
||
"litellm_provider": "openai",
|
||
"mode": "chat",
|
||
}
|
||
return None
|
||
|
||
with patch.object(
|
||
router, "get_deployment_model_info", side_effect=_model_info_with_str_costs
|
||
):
|
||
result = router._set_model_group_info(
|
||
model_group="my-custom-model",
|
||
user_facing_model_group_name="my-custom-model",
|
||
)
|
||
|
||
assert result is not None
|
||
assert result.input_cost_per_token == 2e-05
|
||
assert result.output_cost_per_token == 2e-05
|
||
assert isinstance(result.input_cost_per_token, float)
|
||
assert isinstance(result.output_cost_per_token, float)
|
||
|
||
|
||
def test_model_group_info_db_fallback_with_stringified_cost_values():
|
||
"""
|
||
Fallback path: when get_deployment_model_info returns nothing, costs are
|
||
read straight from the deployment's model_info dict, which can hold
|
||
stringified floats parsed from YAML. They must be coerced to float.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-custom-model",
|
||
"litellm_params": {
|
||
"model": "openai/my-custom-backend-1",
|
||
"api_key": "fake",
|
||
},
|
||
"model_info": {
|
||
"input_cost_per_token": "1e-05",
|
||
"output_cost_per_token": "3e-05",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "my-custom-model",
|
||
"litellm_params": {
|
||
"model": "openai/my-custom-backend-2",
|
||
"api_key": "fake",
|
||
},
|
||
"model_info": {
|
||
"input_cost_per_token": "2e-05",
|
||
"output_cost_per_token": "2e-05",
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
with patch.object(
|
||
router, "get_deployment_model_info", side_effect=Exception("not found")
|
||
):
|
||
result = router._set_model_group_info(
|
||
model_group="my-custom-model",
|
||
user_facing_model_group_name="my-custom-model",
|
||
)
|
||
|
||
assert result is not None
|
||
assert result.input_cost_per_token == 2e-05
|
||
assert result.output_cost_per_token == 3e-05
|
||
assert isinstance(result.input_cost_per_token, float)
|
||
assert isinstance(result.output_cost_per_token, float)
|
||
|
||
|
||
def test_get_model_access_groups_caching():
|
||
"""
|
||
Test that get_model_access_groups caches the no-args result
|
||
and invalidates on deployment changes.
|
||
"""
|
||
from litellm.types.router import Deployment, LiteLLM_Params
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4"},
|
||
"model_info": {"access_groups": ["premium"]},
|
||
},
|
||
]
|
||
)
|
||
|
||
# First call computes and populates cache
|
||
result1 = router.get_model_access_groups()
|
||
assert "premium" in result1
|
||
|
||
# All subsequent calls should return the same cached object (including first)
|
||
result2 = router.get_model_access_groups()
|
||
assert result1 is result2
|
||
|
||
# Calls with args should bypass cache
|
||
result_with_args = router.get_model_access_groups(model_name="gpt-4")
|
||
assert result_with_args is not result2
|
||
|
||
# Add a deployment — cache should be invalidated
|
||
router.add_deployment(
|
||
Deployment(
|
||
model_name="gpt-3.5",
|
||
litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"),
|
||
model_info={"access_groups": ["default"]},
|
||
)
|
||
)
|
||
result3 = router.get_model_access_groups()
|
||
assert result3 is not result2
|
||
assert "premium" in result3
|
||
assert "default" in result3
|
||
|
||
# Delete the deployment — cache should be invalidated again
|
||
deployment_id = None
|
||
for m in router.model_list:
|
||
if m.get("model_name") == "gpt-3.5":
|
||
deployment_id = m.get("model_info", {}).get("id")
|
||
break
|
||
assert deployment_id is not None
|
||
router.delete_deployment(id=deployment_id)
|
||
result4 = router.get_model_access_groups()
|
||
assert result4 is not result3
|
||
assert "default" not in result4
|
||
assert "premium" in result4
|
||
|
||
|
||
def test_get_model_access_groups_cache_invalidation_set_model_list():
|
||
"""
|
||
Test that set_model_list invalidates the access groups cache.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4"},
|
||
"model_info": {"access_groups": ["premium"]},
|
||
},
|
||
]
|
||
)
|
||
|
||
# Populate cache
|
||
result1 = router.get_model_access_groups()
|
||
assert "premium" in result1
|
||
|
||
# set_model_list should invalidate cache
|
||
router.set_model_list(
|
||
[
|
||
{
|
||
"model_name": "claude-3",
|
||
"litellm_params": {"model": "anthropic/claude-3-opus-20240229"},
|
||
"model_info": {"access_groups": ["research"]},
|
||
},
|
||
]
|
||
)
|
||
result2 = router.get_model_access_groups()
|
||
assert result2 is not result1
|
||
assert "research" in result2
|
||
assert "premium" not in result2
|
||
|
||
|
||
def test_get_model_access_groups_cache_invalidation_upsert_deployment():
|
||
"""
|
||
Test that upsert_deployment invalidates the access groups cache.
|
||
"""
|
||
from litellm.types.router import Deployment, LiteLLM_Params
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4"},
|
||
"model_info": {"access_groups": ["premium"]},
|
||
},
|
||
]
|
||
)
|
||
|
||
# Populate cache
|
||
result1 = router.get_model_access_groups()
|
||
assert "premium" in result1
|
||
|
||
# Get the existing deployment's ID
|
||
existing_id = router.model_list[0]["model_info"]["id"]
|
||
|
||
# Upsert with the same ID but different params — triggers pop + re-add
|
||
router.upsert_deployment(
|
||
Deployment(
|
||
model_name="gpt-4-updated",
|
||
litellm_params=LiteLLM_Params(model="gpt-4-turbo"),
|
||
model_info={"id": existing_id, "access_groups": ["updated-group"]},
|
||
)
|
||
)
|
||
result2 = router.get_model_access_groups()
|
||
assert result2 is not result1
|
||
assert "updated-group" in result2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator():
|
||
"""Test _acompletion_streaming_iterator for normal streaming and fallback behavior."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
# Helper class for creating async iterators
|
||
class AsyncIterator:
|
||
def __init__(self, items, error_after=None):
|
||
self.items = items
|
||
self.index = 0
|
||
self.error_after = error_after
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
if self.error_after is not None and self.index >= self.error_after:
|
||
raise self.error_after
|
||
if self.index >= len(self.items):
|
||
raise StopAsyncIteration
|
||
item = self.items[self.index]
|
||
self.index += 1
|
||
return item
|
||
|
||
# Set up router with fallback configuration
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"},
|
||
},
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key-2"},
|
||
},
|
||
],
|
||
fallbacks=[{"gpt-4": ["gpt-3.5-turbo"]}],
|
||
set_verbose=True,
|
||
)
|
||
|
||
# Test data
|
||
messages = [{"role": "user", "content": "Hello"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True, "temperature": 0.7}
|
||
|
||
# Test 1: Successful streaming (no errors)
|
||
print("\n=== Test 1: Successful streaming ===")
|
||
|
||
# Mock successful streaming response
|
||
mock_chunks = [
|
||
MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]),
|
||
MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]),
|
||
MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]),
|
||
]
|
||
|
||
mock_response = AsyncIterator(mock_chunks)
|
||
|
||
setattr(mock_response, "model", "gpt-4")
|
||
setattr(mock_response, "custom_llm_provider", "openai")
|
||
setattr(mock_response, "logging_obj", MagicMock())
|
||
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=mock_response, messages=messages, initial_kwargs=initial_kwargs
|
||
)
|
||
|
||
# Collect streamed chunks
|
||
collected_chunks = []
|
||
async for chunk in result:
|
||
collected_chunks.append(chunk)
|
||
|
||
assert len(collected_chunks) == 3
|
||
assert all(chunk in mock_chunks for chunk in collected_chunks)
|
||
print("✓ Successfully streamed all chunks")
|
||
|
||
# Test 2: MidStreamFallbackError with generated content is re-raised, not silently continued
|
||
print("\n=== Test 2: MidStreamFallbackError re-raises when content already generated ===")
|
||
|
||
# Error with generated content and is_pre_first_chunk=False (the default):
|
||
# the router must re-raise instead of attempting a continuation-prompt fallback,
|
||
# because partial content has already been sent to the client.
|
||
error = MidStreamFallbackError(
|
||
message="Connection lost",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
generated_content="Hello",
|
||
)
|
||
|
||
class AsyncIteratorWithError:
|
||
def __init__(self, items, error_after_index):
|
||
self.items = items
|
||
self.index = 0
|
||
self.error_after_index = error_after_index
|
||
self.chunks = []
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
if self.index >= len(self.items):
|
||
raise StopAsyncIteration
|
||
if self.index == self.error_after_index:
|
||
raise error
|
||
item = self.items[self.index]
|
||
self.index += 1
|
||
return item
|
||
|
||
mock_error_response = AsyncIteratorWithError(mock_chunks, 1) # Error after first chunk
|
||
|
||
setattr(mock_error_response, "model", "gpt-4")
|
||
setattr(mock_error_response, "custom_llm_provider", "openai")
|
||
setattr(mock_error_response, "logging_obj", MagicMock())
|
||
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=mock_error_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
# Collect streamed chunks — the first chunk succeeds, then the error re-raises
|
||
collected_chunks = []
|
||
with pytest.raises(MidStreamFallbackError):
|
||
async for chunk in result:
|
||
collected_chunks.append(chunk)
|
||
|
||
assert len(collected_chunks) == 1, "one chunk yielded before the error"
|
||
print("✓ MidStreamFallbackError re-raised correctly when content was already generated")
|
||
|
||
print("\n=== All tests passed! ===")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_reraises_original_exception_when_available():
|
||
"""Async: when the mid-stream MidStreamFallbackError wraps a real provider
|
||
exception (original_exception), the router must re-raise that original
|
||
exception instead of the internal wrapper, so the client sees the
|
||
specific error type/code (e.g. RateLimitError) rather than a generic
|
||
MidStreamFallbackError."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError, RateLimitError
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
set_verbose=True,
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
original_exception = RateLimitError(
|
||
message="rate limited",
|
||
llm_provider="vertex_ai",
|
||
model="gpt-4",
|
||
)
|
||
error = MidStreamFallbackError(
|
||
message="rate limited",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
original_exception=original_exception,
|
||
generated_content="Hello",
|
||
)
|
||
|
||
mock_chunks = [
|
||
MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]),
|
||
MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]),
|
||
]
|
||
|
||
class AsyncIteratorWithError:
|
||
def __init__(self, items, error_after_index):
|
||
self.items = items
|
||
self.index = 0
|
||
self.error_after_index = error_after_index
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
if self.index >= len(self.items):
|
||
raise StopAsyncIteration
|
||
if self.index == self.error_after_index:
|
||
raise error
|
||
item = self.items[self.index]
|
||
self.index += 1
|
||
return item
|
||
|
||
mock_error_response = AsyncIteratorWithError(mock_chunks, 1)
|
||
setattr(mock_error_response, "model", "gpt-4")
|
||
setattr(mock_error_response, "custom_llm_provider", "openai")
|
||
setattr(mock_error_response, "logging_obj", MagicMock())
|
||
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=mock_error_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
with pytest.raises(RateLimitError) as exc_info:
|
||
async for _ in result:
|
||
pass
|
||
assert exc_info.value is original_exception
|
||
assert exc_info.value.type == "throttling_error"
|
||
assert exc_info.value.code == "429"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_edge_cases():
|
||
"""Test edge cases for _acompletion_streaming_iterator."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
set_verbose=True,
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
# Test: Empty generated content
|
||
empty_error = MidStreamFallbackError(
|
||
message="Error",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
generated_content="", # Empty content
|
||
)
|
||
|
||
class AsyncIteratorImmediateError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = []
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
raise empty_error
|
||
|
||
mock_response = AsyncIteratorImmediateError()
|
||
|
||
# Mock empty fallback response using AsyncIterator
|
||
class EmptyAsyncIterator:
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
raise StopAsyncIteration
|
||
|
||
mock_fallback_response = EmptyAsyncIterator()
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=mock_fallback_response,
|
||
) as mock_fallback_utils:
|
||
collected_chunks = []
|
||
iterator = await router._acompletion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
async for chunk in iterator:
|
||
collected_chunks.append(chunk)
|
||
|
||
# Should still call fallback even with empty content
|
||
assert mock_fallback_utils.called
|
||
fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"]
|
||
modified_messages = fallback_kwargs["messages"]
|
||
|
||
# Empty content → pre-first-chunk path uses original messages
|
||
# (no continuation prompt added)
|
||
assert modified_messages == messages
|
||
print("✓ Handles empty generated content correctly")
|
||
|
||
print("✓ Edge case tests passed!")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_preserves_hidden_params():
|
||
"""
|
||
Regression test: FallbackStreamWrapper must copy _hidden_params from the
|
||
original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and
|
||
other hidden params) are present in the proxy response headers for streaming.
|
||
"""
|
||
from unittest.mock import MagicMock
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
# Simulate a CustomStreamWrapper that already has timing metadata set by
|
||
# update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.)
|
||
mock_response = MagicMock()
|
||
mock_response.model = "gpt-4"
|
||
mock_response.custom_llm_provider = "openai"
|
||
mock_response.logging_obj = MagicMock()
|
||
mock_response._hidden_params = {
|
||
"litellm_overhead_time_ms": 12.34,
|
||
"_response_ms": 500.0,
|
||
"litellm_call_id": "test-call-id",
|
||
"api_base": "https://api.openai.com",
|
||
"additional_headers": {},
|
||
}
|
||
|
||
# Make the mock iterable (yields nothing — we only care about hidden_params)
|
||
async def _empty():
|
||
return
|
||
yield # make it an async generator
|
||
|
||
mock_response.__aiter__ = lambda self: _empty().__aiter__()
|
||
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=[{"role": "user", "content": "hi"}],
|
||
initial_kwargs={"model": "gpt-4", "stream": True},
|
||
)
|
||
|
||
# The returned FallbackStreamWrapper must carry the original _hidden_params
|
||
assert hasattr(result, "_hidden_params"), "result must have _hidden_params"
|
||
assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, (
|
||
"litellm_overhead_time_ms must be preserved — "
|
||
"this is what drives x-litellm-overhead-duration-ms in streaming responses"
|
||
)
|
||
assert result._hidden_params.get("litellm_call_id") == "test-call-id"
|
||
assert result._hidden_params.get("_response_ms") == 500.0
|
||
|
||
|
||
def test_completion_streaming_iterator_fallback_on_429():
|
||
"""Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback.
|
||
|
||
This is the sync counterpart of test_acompletion_streaming_iterator.
|
||
Before this fix, __next__ raised RateLimitError directly and the Router
|
||
never got a chance to fall back.
|
||
"""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
rate_limit_error = MidStreamFallbackError(
|
||
message="Resource exhausted",
|
||
model="gpt-4",
|
||
llm_provider="vertex_ai",
|
||
generated_content="",
|
||
is_pre_first_chunk=True,
|
||
)
|
||
|
||
class SyncIteratorImmediateError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = []
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __next__(self):
|
||
raise rate_limit_error
|
||
|
||
mock_response = SyncIteratorImmediateError()
|
||
|
||
# Fallback returns a simple non-streaming response (fallback may not stream)
|
||
mock_fallback_response = MagicMock()
|
||
mock_fallback_response.__iter__ = MagicMock(return_value=iter([]))
|
||
|
||
with patch.object(
|
||
router,
|
||
"function_with_fallbacks",
|
||
return_value=mock_fallback_response,
|
||
) as mock_fallback:
|
||
result = router._completion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
collected_chunks = list(result)
|
||
|
||
assert mock_fallback.called
|
||
call_kwargs = mock_fallback.call_args
|
||
# Pre-first-chunk: should use original messages, no continuation prompt
|
||
assert call_kwargs.kwargs.get("messages") == messages
|
||
# Verify original_function is _completion (sync)
|
||
assert call_kwargs.kwargs.get("original_function") == router._completion
|
||
|
||
|
||
def test_completion_streaming_iterator_preserves_hidden_params():
|
||
"""SyncFallbackStreamWrapper must copy _hidden_params from original response."""
|
||
from unittest.mock import MagicMock
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
mock_response = MagicMock()
|
||
mock_response.model = "gpt-4"
|
||
mock_response.custom_llm_provider = "openai"
|
||
mock_response.logging_obj = MagicMock()
|
||
mock_response._hidden_params = {
|
||
"litellm_overhead_time_ms": 42.0,
|
||
"litellm_call_id": "test-sync-call",
|
||
}
|
||
mock_response.__iter__ = MagicMock(return_value=iter([]))
|
||
|
||
result = router._completion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=[{"role": "user", "content": "hi"}],
|
||
initial_kwargs={"model": "gpt-4", "stream": True},
|
||
)
|
||
|
||
assert hasattr(result, "_hidden_params")
|
||
assert result._hidden_params.get("litellm_overhead_time_ms") == 42.0
|
||
assert result._hidden_params.get("litellm_call_id") == "test-sync-call"
|
||
|
||
|
||
def test_completion_streaming_iterator_reraises_mid_chunk_error():
|
||
"""Sync: MidStreamFallbackError with generated_content and is_pre_first_chunk=False
|
||
must be re-raised immediately; the router cannot recover after partial content
|
||
has already been sent to the client."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
mid_chunk_error = MidStreamFallbackError(
|
||
message="Connection reset",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
generated_content="Hello, I am",
|
||
is_pre_first_chunk=False,
|
||
)
|
||
|
||
class SyncIteratorMidChunkError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = []
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __next__(self):
|
||
raise mid_chunk_error
|
||
|
||
mock_response = SyncIteratorMidChunkError()
|
||
|
||
result = router._completion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
with pytest.raises(MidStreamFallbackError):
|
||
list(result)
|
||
|
||
|
||
def test_completion_streaming_iterator_reraises_original_exception_when_available():
|
||
"""Sync: when the mid-chunk MidStreamFallbackError wraps a real provider
|
||
exception (original_exception), the router must re-raise that original
|
||
exception instead of the internal wrapper, so the client sees the
|
||
specific error type/code (e.g. RateLimitError) rather than a generic
|
||
MidStreamFallbackError."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError, RateLimitError
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
original_exception = RateLimitError(
|
||
message="rate limited",
|
||
llm_provider="vertex_ai",
|
||
model="gpt-4",
|
||
)
|
||
mid_chunk_error = MidStreamFallbackError(
|
||
message="rate limited",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
original_exception=original_exception,
|
||
generated_content="Hello, I am",
|
||
is_pre_first_chunk=False,
|
||
)
|
||
|
||
class SyncIteratorMidChunkError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = []
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __next__(self):
|
||
raise mid_chunk_error
|
||
|
||
mock_response = SyncIteratorMidChunkError()
|
||
|
||
result = router._completion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
with pytest.raises(RateLimitError) as exc_info:
|
||
list(result)
|
||
assert exc_info.value is original_exception
|
||
assert exc_info.value.type == "throttling_error"
|
||
assert exc_info.value.code == "429"
|
||
|
||
|
||
def test_completion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content():
|
||
"""Sync: a reasoning-only chunk sets is_pre_first_chunk=False without populating
|
||
generated_content (which only tracks text deltas). The re-raise guard must still
|
||
detect this via the raw chunks on the wrapper, or the router silently retries and
|
||
the client receives duplicated/inconsistent output."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
from litellm.types.utils import Delta, StreamingChoices
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
mid_chunk_error = MidStreamFallbackError(
|
||
message="Connection reset",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
generated_content="",
|
||
is_pre_first_chunk=False,
|
||
)
|
||
|
||
reasoning_chunk = litellm.ModelResponseStream(
|
||
id="chatcmpl-partial-1",
|
||
model="gpt-4",
|
||
object="chat.completion.chunk",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
index=0,
|
||
delta=Delta(reasoning_content="Thinking about the answer", role="assistant"),
|
||
)
|
||
],
|
||
)
|
||
|
||
class SyncIteratorNoTextChunkError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = [reasoning_chunk]
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __next__(self):
|
||
raise mid_chunk_error
|
||
|
||
mock_response = SyncIteratorNoTextChunkError()
|
||
|
||
with patch.object(router, "function_with_fallbacks") as mock_fallback:
|
||
result = router._completion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
with pytest.raises(MidStreamFallbackError):
|
||
list(result)
|
||
|
||
assert not mock_fallback.called, (
|
||
"fallback must not be attempted once any content, text or non-text, has already streamed"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_pre_first_chunk_skips_continuation():
|
||
"""When MidStreamFallbackError has is_pre_first_chunk=True, use original messages."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Hello"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
pre_first_chunk_error = MidStreamFallbackError(
|
||
message="429 Resource exhausted",
|
||
model="gpt-4",
|
||
llm_provider="vertex_ai",
|
||
generated_content="",
|
||
is_pre_first_chunk=True,
|
||
)
|
||
|
||
class AsyncIteratorPreFirstChunkError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = []
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
raise pre_first_chunk_error
|
||
|
||
mock_response = AsyncIteratorPreFirstChunkError()
|
||
|
||
class EmptyAsyncIterator:
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
raise StopAsyncIteration
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=EmptyAsyncIterator(),
|
||
) as mock_fallback_utils:
|
||
iterator = await router._acompletion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
async for _ in iterator:
|
||
pass
|
||
|
||
assert mock_fallback_utils.called
|
||
fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"]
|
||
# Pre-first-chunk: should use original messages, no continuation prompt
|
||
assert fallback_kwargs["messages"] == messages
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content():
|
||
"""Async: a reasoning-only chunk sets is_pre_first_chunk=False without populating
|
||
generated_content (which only tracks text deltas). The re-raise guard must still
|
||
detect this via the raw chunks on the wrapper, or the router silently retries and
|
||
the client receives duplicated/inconsistent output."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
from litellm.types.utils import Delta, StreamingChoices
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||
}
|
||
],
|
||
)
|
||
|
||
messages = [{"role": "user", "content": "Test"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
mid_chunk_error = MidStreamFallbackError(
|
||
message="Connection reset",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
generated_content="",
|
||
is_pre_first_chunk=False,
|
||
)
|
||
|
||
reasoning_chunk = litellm.ModelResponseStream(
|
||
id="chatcmpl-partial-1",
|
||
model="gpt-4",
|
||
object="chat.completion.chunk",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
index=0,
|
||
delta=Delta(reasoning_content="Thinking about the answer", role="assistant"),
|
||
)
|
||
],
|
||
)
|
||
|
||
class AsyncIteratorNoTextChunkError:
|
||
def __init__(self):
|
||
self.model = "gpt-4"
|
||
self.custom_llm_provider = "openai"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = [reasoning_chunk]
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
raise mid_chunk_error
|
||
|
||
mock_response = AsyncIteratorNoTextChunkError()
|
||
|
||
with patch.object(router, "async_function_with_fallbacks_common_utils") as mock_fallback_utils:
|
||
iterator = await router._acompletion_streaming_iterator(
|
||
model_response=mock_response,
|
||
messages=messages,
|
||
initial_kwargs=initial_kwargs,
|
||
)
|
||
|
||
with pytest.raises(MidStreamFallbackError):
|
||
async for _ in iterator:
|
||
pass
|
||
|
||
assert not mock_fallback_utils.called, (
|
||
"fallback must not be attempted once any content, text or non-text, has already streamed"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Shared helpers for the _aresponses_streaming_iterator test suite.
|
||
# ---------------------------------------------------------------------------
|
||
def _make_responses_iterator(
|
||
*,
|
||
chunks=(),
|
||
error=None,
|
||
bridge=False,
|
||
model="gpt-4",
|
||
hidden_params=None,
|
||
chat_chunks=None,
|
||
):
|
||
"""Build a minimal mock Responses-API streaming iterator.
|
||
|
||
Bypasses BaseResponsesAPIStreamingIterator.__init__ but mirrors every
|
||
attribute production code reads. Yields *chunks*, then raises *error*
|
||
(or StopAsyncIteration). Set bridge=True to inherit from
|
||
LiteLLMCompletionStreamingIterator so the wrapper's bridge-path
|
||
isinstance check (used by usage extraction) matches.
|
||
"""
|
||
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
|
||
LiteLLMCompletionStreamingIterator,
|
||
)
|
||
from litellm.responses.streaming_iterator import (
|
||
BaseResponsesAPIStreamingIterator,
|
||
)
|
||
|
||
base = (
|
||
LiteLLMCompletionStreamingIterator
|
||
if bridge
|
||
else BaseResponsesAPIStreamingIterator
|
||
)
|
||
|
||
class _Iter(base):
|
||
def __init__(self):
|
||
self._chunks = list(chunks)
|
||
self._idx = 0
|
||
self._hidden_params = hidden_params or {}
|
||
self.model = model
|
||
self.custom_llm_provider = "anthropic"
|
||
self.logging_obj = MagicMock()
|
||
self.litellm_metadata = None
|
||
self.responses_api_provider_config = None
|
||
self.finished = False
|
||
self.completed_response = None
|
||
self.response = None
|
||
self.start_time = None
|
||
self.request_data = {}
|
||
self.call_type = None
|
||
if chat_chunks is not None:
|
||
self.collected_chat_completion_chunks = chat_chunks
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
if self._idx < len(self._chunks):
|
||
self._idx += 1
|
||
return self._chunks[self._idx - 1]
|
||
if error is not None:
|
||
raise error
|
||
raise StopAsyncIteration
|
||
|
||
return _Iter()
|
||
|
||
|
||
class _AsyncList:
|
||
"""Generic async iterator over a list — used as the fallback response."""
|
||
|
||
def __init__(self, items=()):
|
||
self._items = list(items)
|
||
self._idx = 0
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
if self._idx >= len(self._items):
|
||
raise StopAsyncIteration
|
||
item = self._items[self._idx]
|
||
self._idx += 1
|
||
return item
|
||
|
||
|
||
def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"):
|
||
return litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": primary,
|
||
"litellm_params": {"model": primary, "api_key": "k1"},
|
||
},
|
||
{
|
||
"model_name": secondary,
|
||
"litellm_params": {"model": secondary, "api_key": "k2"},
|
||
},
|
||
],
|
||
fallbacks=[{primary: [secondary]}],
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_streaming_iterator_fallback():
|
||
"""Catches MidStreamFallbackError, re-enters the fallback chain via
|
||
async_function_with_fallbacks_common_utils with the per-attempt helper
|
||
and original_generic_function preserved. Mirrors
|
||
test_acompletion_streaming_iterator for the aresponses path."""
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
from litellm.responses.streaming_iterator import (
|
||
BaseResponsesAPIStreamingIterator,
|
||
)
|
||
|
||
router = _make_router_with_fallback(
|
||
"anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6"
|
||
)
|
||
src = _make_responses_iterator(
|
||
chunks=[MagicMock(type="response.created")],
|
||
error=MidStreamFallbackError(
|
||
message="anthropic socket timeout",
|
||
model="anthropic/claude-sonnet-4-6",
|
||
llm_provider="anthropic",
|
||
is_pre_first_chunk=False,
|
||
generated_content="",
|
||
),
|
||
model="anthropic/claude-sonnet-4-6",
|
||
hidden_params={"model_id": "src-deployment-1"},
|
||
)
|
||
fallback_chunks = [
|
||
MagicMock(type="response.output_text.delta"),
|
||
MagicMock(type="response.completed"),
|
||
]
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=_AsyncList(fallback_chunks),
|
||
) as mock_fallback_utils:
|
||
wrapped = await router._aresponses_streaming_iterator(
|
||
response=src,
|
||
initial_kwargs={
|
||
"model": "anthropic/claude-sonnet-4-6",
|
||
"stream": True,
|
||
"input": "Hi",
|
||
"original_generic_function": litellm.aresponses,
|
||
},
|
||
)
|
||
assert isinstance(wrapped, BaseResponsesAPIStreamingIterator)
|
||
assert wrapped._hidden_params.get("model_id") == "src-deployment-1"
|
||
collected = [c async for c in wrapped]
|
||
|
||
assert len(collected) == 3 # 1 primary chunk + 2 fallback chunks
|
||
call_kwargs = mock_fallback_utils.call_args.kwargs
|
||
fbk = call_kwargs["kwargs"]
|
||
# Bound methods compare equal when they share the same instance + __func__.
|
||
assert fbk["original_function"] == router._ageneric_api_call_with_fallbacks_helper
|
||
assert fbk["original_generic_function"] is litellm.aresponses
|
||
assert call_kwargs["model_group"] == "anthropic/claude-sonnet-4-6"
|
||
assert call_kwargs["disable_fallbacks"] is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback():
|
||
"""Regression: model_group must land under "litellm_metadata" (the key
|
||
litellm.aresponses reads), not the default "metadata"."""
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = _make_router_with_fallback()
|
||
src = _make_responses_iterator(
|
||
error=MidStreamFallbackError(
|
||
message="boom",
|
||
model="gpt-4",
|
||
llm_provider="anthropic",
|
||
is_pre_first_chunk=True,
|
||
generated_content="",
|
||
)
|
||
)
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=_AsyncList(),
|
||
) as mock_fallback_utils:
|
||
wrapped = await router._aresponses_streaming_iterator(
|
||
response=src,
|
||
initial_kwargs={
|
||
"model": "gpt-4",
|
||
"stream": True,
|
||
"input": "Hello",
|
||
"original_generic_function": litellm.aresponses,
|
||
},
|
||
)
|
||
async for _ in wrapped:
|
||
pass
|
||
|
||
fbk = mock_fallback_utils.call_args.kwargs["kwargs"]
|
||
assert "litellm_metadata" in fbk, "wrong metadata_variable_name"
|
||
assert fbk["litellm_metadata"]["model_group"] == "gpt-4"
|
||
assert "model_group" not in fbk.get(
|
||
"metadata", {}
|
||
), "model_group leaked into 'metadata' instead of 'litellm_metadata'"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_streaming_iterator_pre_first_chunk_skips_continuation():
|
||
"""Pre-first-chunk error: original input is preserved unchanged."""
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = _make_router_with_fallback()
|
||
src = _make_responses_iterator(
|
||
error=MidStreamFallbackError(
|
||
message="socket timeout before first chunk",
|
||
model="gpt-4",
|
||
llm_provider="anthropic",
|
||
is_pre_first_chunk=True,
|
||
generated_content="",
|
||
)
|
||
)
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=_AsyncList(),
|
||
) as mock_fallback_utils:
|
||
wrapped = await router._aresponses_streaming_iterator(
|
||
response=src,
|
||
initial_kwargs={
|
||
"model": "gpt-4",
|
||
"stream": True,
|
||
"input": "Hello",
|
||
"original_generic_function": litellm.aresponses,
|
||
},
|
||
)
|
||
async for _ in wrapped:
|
||
pass
|
||
|
||
fbk = mock_fallback_utils.call_args.kwargs["kwargs"]
|
||
assert fbk["input"] == "Hello" # original input, no continuation messages
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_streaming_iterator_partial_content_injects_continuation():
|
||
"""Mid-stream error: input is rewritten to include user prompt +
|
||
developer instruction + prior assistant message with partial output."""
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
|
||
router = _make_router_with_fallback()
|
||
src = _make_responses_iterator(
|
||
chunks=[MagicMock(type="response.output_text.delta")],
|
||
error=MidStreamFallbackError(
|
||
message="socket reset mid-stream",
|
||
model="gpt-4",
|
||
llm_provider="anthropic",
|
||
is_pre_first_chunk=False,
|
||
generated_content="The capital of France is",
|
||
),
|
||
)
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=_AsyncList(),
|
||
) as mock_fallback_utils:
|
||
wrapped = await router._aresponses_streaming_iterator(
|
||
response=src,
|
||
initial_kwargs={
|
||
"model": "gpt-4",
|
||
"stream": True,
|
||
"input": "What's the capital of France?",
|
||
"original_generic_function": litellm.aresponses,
|
||
},
|
||
)
|
||
async for _ in wrapped:
|
||
pass
|
||
|
||
new_input = mock_fallback_utils.call_args.kwargs["kwargs"]["input"]
|
||
assert isinstance(new_input, list)
|
||
assert new_input[0]["role"] == "user"
|
||
assert new_input[0]["content"][0]["text"] == "What's the capital of France?"
|
||
assert new_input[1]["role"] == "developer"
|
||
assert "do not repeat" in new_input[1]["content"][0]["text"].lower()
|
||
assert new_input[2]["role"] == "assistant"
|
||
assert new_input[2]["content"][0]["type"] == "output_text"
|
||
assert new_input[2]["content"][0]["text"] == "The capital of France is"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_streaming_iterator_combines_partial_usage():
|
||
"""Partial usage from the bridge path is normalized to ResponseAPIUsage
|
||
and summed onto the fallback's response.completed event — no token-name
|
||
split, clean ResponseAPIUsage on output."""
|
||
from types import SimpleNamespace
|
||
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
from litellm.types.llms.openai import (
|
||
ResponseAPIUsage,
|
||
ResponseCompletedEvent,
|
||
ResponsesAPIResponse,
|
||
ResponsesAPIStreamEvents,
|
||
)
|
||
|
||
router = _make_router_with_fallback()
|
||
src = _make_responses_iterator(
|
||
bridge=True,
|
||
chat_chunks=[MagicMock()],
|
||
chunks=[MagicMock(type="response.output_text.delta")],
|
||
error=MidStreamFallbackError(
|
||
message="boom",
|
||
model="gpt-4",
|
||
llm_provider="anthropic",
|
||
is_pre_first_chunk=False,
|
||
generated_content="hello",
|
||
),
|
||
)
|
||
|
||
fallback_response_object = ResponsesAPIResponse(
|
||
id="resp_test", created_at=0, model="gpt-4", object="response", output=[]
|
||
)
|
||
fallback_response_object.usage = ResponseAPIUsage(
|
||
input_tokens=20, output_tokens=15, total_tokens=35
|
||
)
|
||
fallback_event = ResponseCompletedEvent(
|
||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||
response=fallback_response_object,
|
||
)
|
||
|
||
with (
|
||
patch(
|
||
"litellm.main.stream_chunk_builder",
|
||
return_value=SimpleNamespace(
|
||
usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)
|
||
),
|
||
),
|
||
patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
return_value=_AsyncList([fallback_event]),
|
||
),
|
||
):
|
||
wrapped = await router._aresponses_streaming_iterator(
|
||
response=src,
|
||
initial_kwargs={
|
||
"model": "gpt-4",
|
||
"stream": True,
|
||
"input": "hi",
|
||
"original_generic_function": litellm.aresponses,
|
||
},
|
||
)
|
||
async for _ in wrapped:
|
||
pass
|
||
|
||
merged = fallback_response_object.usage
|
||
assert isinstance(merged, ResponseAPIUsage)
|
||
assert merged.input_tokens == 30 # 10 (translated from prompt_tokens) + 20
|
||
assert merged.output_tokens == 19 # 4 (translated from completion_tokens) + 15
|
||
assert merged.total_tokens == 49
|
||
|
||
|
||
def _midstream_rate_limit_error():
|
||
rate_limit_error = litellm.RateLimitError(
|
||
message="vertex_ai_betaException - Resource exhausted.",
|
||
model="gemini",
|
||
llm_provider="vertex_ai_beta",
|
||
)
|
||
midstream_error = MidStreamFallbackError(
|
||
message=str(rate_limit_error),
|
||
model="gemini",
|
||
llm_provider="vertex_ai_beta",
|
||
original_exception=rate_limit_error,
|
||
is_pre_first_chunk=True,
|
||
)
|
||
return rate_limit_error, midstream_error
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks():
|
||
"""Regression for #26015: a mid-stream 429 with no fallbacks configured must
|
||
surface a clean RateLimitError, not leak the internal MidStreamFallbackError
|
||
wrapper to the client, and must terminate instead of hanging."""
|
||
rate_limit_error, midstream_error = _midstream_rate_limit_error()
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gemini",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.0-flash",
|
||
"api_key": "fake-key",
|
||
},
|
||
},
|
||
],
|
||
num_retries=0,
|
||
)
|
||
|
||
class _RaisingStream:
|
||
def __init__(self):
|
||
self.chunks = []
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
raise midstream_error
|
||
|
||
stream = _RaisingStream()
|
||
setattr(stream, "model", "gemini")
|
||
setattr(stream, "custom_llm_provider", "vertex_ai_beta")
|
||
setattr(stream, "logging_obj", MagicMock())
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
new=AsyncMock(side_effect=midstream_error),
|
||
):
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=stream,
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
initial_kwargs={"model": "gemini", "stream": True},
|
||
)
|
||
|
||
async def _consume():
|
||
async for _ in result:
|
||
pass
|
||
|
||
with pytest.raises(litellm.RateLimitError) as exc_info:
|
||
await asyncio.wait_for(_consume(), timeout=10)
|
||
|
||
assert not isinstance(exc_info.value, MidStreamFallbackError)
|
||
assert exc_info.value.status_code == 429
|
||
assert exc_info.value is rate_limit_error
|
||
|
||
|
||
def test_completion_streaming_iterator_surfaces_rate_limit_without_fallbacks():
|
||
"""Sync counterpart of
|
||
test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks."""
|
||
rate_limit_error, midstream_error = _midstream_rate_limit_error()
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gemini",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.0-flash",
|
||
"api_key": "fake-key",
|
||
},
|
||
},
|
||
],
|
||
num_retries=0,
|
||
)
|
||
|
||
class _RaisingSyncStream:
|
||
def __init__(self):
|
||
self.model = "gemini"
|
||
self.custom_llm_provider = "vertex_ai_beta"
|
||
self.logging_obj = MagicMock()
|
||
self.chunks = []
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __next__(self):
|
||
raise midstream_error
|
||
|
||
with patch.object(
|
||
router,
|
||
"function_with_fallbacks",
|
||
side_effect=midstream_error,
|
||
):
|
||
result = router._completion_streaming_iterator(
|
||
model_response=_RaisingSyncStream(),
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
initial_kwargs={"model": "gemini", "stream": True},
|
||
)
|
||
|
||
with pytest.raises(litellm.RateLimitError) as exc_info:
|
||
list(result)
|
||
|
||
assert not isinstance(exc_info.value, MidStreamFallbackError)
|
||
assert exc_info.value.status_code == 429
|
||
assert exc_info.value is rate_limit_error
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_streaming_iterator_surfaces_rate_limit_without_fallbacks():
|
||
"""Responses-API counterpart of
|
||
test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks."""
|
||
rate_limit_error, midstream_error = _midstream_rate_limit_error()
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gemini",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.0-flash",
|
||
"api_key": "fake-key",
|
||
},
|
||
},
|
||
],
|
||
num_retries=0,
|
||
)
|
||
src = _make_responses_iterator(error=midstream_error, model="gemini")
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
new=AsyncMock(side_effect=midstream_error),
|
||
):
|
||
wrapped = await router._aresponses_streaming_iterator(
|
||
response=src,
|
||
initial_kwargs={
|
||
"model": "gemini",
|
||
"stream": True,
|
||
"input": "Hello",
|
||
"original_generic_function": litellm.aresponses,
|
||
},
|
||
)
|
||
|
||
async def _consume():
|
||
async for _ in wrapped:
|
||
pass
|
||
|
||
with pytest.raises(litellm.RateLimitError) as exc_info:
|
||
await asyncio.wait_for(_consume(), timeout=10)
|
||
|
||
assert not isinstance(exc_info.value, MidStreamFallbackError)
|
||
assert exc_info.value.status_code == 429
|
||
assert exc_info.value is rate_limit_error
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_function_with_fallbacks_common_utils():
|
||
"""Test the async_function_with_fallbacks_common_utils method"""
|
||
# Create a basic router for testing
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {
|
||
"model": "gpt-3.5-turbo",
|
||
},
|
||
}
|
||
],
|
||
max_fallbacks=5,
|
||
)
|
||
|
||
# Test case 1: disable_fallbacks=True should raise original exception
|
||
test_exception = Exception("Test error")
|
||
with pytest.raises(Exception, match="Test error"):
|
||
await router.async_function_with_fallbacks_common_utils(
|
||
e=test_exception,
|
||
disable_fallbacks=True,
|
||
fallbacks=None,
|
||
context_window_fallbacks=None,
|
||
content_policy_fallbacks=None,
|
||
model_group="gpt-3.5-turbo",
|
||
args=(),
|
||
kwargs=MagicMock(),
|
||
)
|
||
|
||
# Test case 2: original_model_group=None should raise original exception
|
||
with pytest.raises(Exception, match="Test error"):
|
||
await router.async_function_with_fallbacks_common_utils(
|
||
e=test_exception,
|
||
disable_fallbacks=False,
|
||
fallbacks=None,
|
||
context_window_fallbacks=None,
|
||
content_policy_fallbacks=None,
|
||
model_group="gpt-3.5-turbo",
|
||
args=(),
|
||
kwargs={}, # No model key
|
||
)
|
||
|
||
|
||
def test_should_include_deployment():
|
||
"""Test that Router.should_include_deployment returns the correct response"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "model_name_a28a12f9-3e44-4861-bd4f-325f2d309ce8_cd5dc6fb-b046-4e05-ae1d-32ba4d936266",
|
||
"litellm_params": {"model": "openai/*"},
|
||
"model_info": {
|
||
"team_id": "a28a12f9-3e44-4861-bd4f-325f2d309ce8",
|
||
"team_public_model_name": "openai/*",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
model = {
|
||
"model_name": "model_name_a28a12f9-3e44-4861-bd4f-325f2d309ce8_cd5dc6fb-b046-4e05-ae1d-32ba4d936266",
|
||
"litellm_params": {
|
||
"api_key": "sk-proj-1234567890",
|
||
"custom_llm_provider": "openai",
|
||
"use_in_pass_through": False,
|
||
"use_litellm_proxy": False,
|
||
"merge_reasoning_content_in_choices": False,
|
||
"model": "openai/*",
|
||
},
|
||
"model_info": {
|
||
"id": "95f58039-d54a-4d1c-b700-5e32e99a1120",
|
||
"db_model": True,
|
||
"updated_by": "64a2f787-0863-4d76-9516-2dc49c1598e8",
|
||
"created_by": "64a2f787-0863-4d76-9516-2dc49c1598e8",
|
||
"team_id": "a28a12f9-3e44-4861-bd4f-325f2d309ce8",
|
||
"team_public_model_name": "openai/*",
|
||
"mode": "completion",
|
||
"access_groups": ["restricted-models-openai"],
|
||
},
|
||
}
|
||
model_name = "openai/o4-mini-deep-research"
|
||
team_id = "a28a12f9-3e44-4861-bd4f-325f2d309ce8"
|
||
assert router.get_model_list(
|
||
model_name=model_name,
|
||
team_id=team_id,
|
||
)
|
||
|
||
|
||
def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch):
|
||
"""
|
||
tiktoken token counting is the dominant on-loop cost for large prompts. When no
|
||
deployment in the group declares max_input_tokens, the count is never consumed, so
|
||
_pre_call_checks must not run it at all.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {})
|
||
|
||
calls = []
|
||
monkeypatch.setattr(
|
||
litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000
|
||
)
|
||
|
||
deployments = [
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}},
|
||
]
|
||
result = router._pre_call_checks(
|
||
model="m",
|
||
healthy_deployments=deployments,
|
||
messages=[{"role": "user", "content": "hi"}],
|
||
)
|
||
|
||
assert calls == []
|
||
assert len(result) == 2
|
||
|
||
|
||
def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch):
|
||
"""
|
||
When a deployment declares max_input_tokens the count must still run, be performed
|
||
at most once across the group (memoized), and filter deployments whose limit is
|
||
exceeded.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
monkeypatch.setattr(
|
||
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}
|
||
)
|
||
|
||
calls = []
|
||
monkeypatch.setattr(
|
||
litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000
|
||
)
|
||
|
||
deployments = [
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}},
|
||
]
|
||
with pytest.raises(litellm.ContextWindowExceededError):
|
||
router._pre_call_checks(
|
||
model="m",
|
||
healthy_deployments=deployments,
|
||
messages=[{"role": "user", "content": "hi"}],
|
||
)
|
||
|
||
assert calls == [1]
|
||
|
||
|
||
def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch):
|
||
"""
|
||
Responses API calls pass `input` (str) instead of `messages`. Context-window
|
||
checks must count tokens from `input` and filter deployments over the limit. Uses
|
||
the real token_counter so the transform + counting path is a true regression guard.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
monkeypatch.setattr(
|
||
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}
|
||
)
|
||
|
||
deployments = [
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
|
||
]
|
||
with pytest.raises(litellm.ContextWindowExceededError):
|
||
router._pre_call_checks(
|
||
model="m",
|
||
healthy_deployments=deployments,
|
||
input="a very long prompt that exceeds the tiny context window",
|
||
)
|
||
|
||
|
||
def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch):
|
||
"""
|
||
Responses API `input` can be a list of input items. It must be normalized to
|
||
chat messages and counted so oversized requests are filtered out. Uses the real
|
||
token_counter (no mock) so the transform + counting path is a true regression guard.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
monkeypatch.setattr(
|
||
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}
|
||
)
|
||
|
||
deployments = [
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
|
||
]
|
||
with pytest.raises(litellm.ContextWindowExceededError):
|
||
router._pre_call_checks(
|
||
model="m",
|
||
healthy_deployments=deployments,
|
||
input=[
|
||
{"role": "user", "content": "count these tokens against the one token limit please"},
|
||
],
|
||
)
|
||
|
||
|
||
def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch):
|
||
"""
|
||
Responses API `instructions` become a system message the model receives, so their
|
||
tokens must be counted too. A request whose `input` alone fits under the limit but
|
||
whose `input` + `instructions` exceeds it must be filtered (regression for the
|
||
context-window check under-filtering when instructions were ignored).
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
|
||
deployments = [
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
|
||
]
|
||
|
||
short_input = "hi"
|
||
long_instructions = "you are a helpful assistant. " * 20
|
||
|
||
input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input)
|
||
with_instructions_tokens = router._count_pre_call_check_tokens(
|
||
messages=None, input=short_input, instructions=long_instructions
|
||
)
|
||
assert with_instructions_tokens > input_only_tokens
|
||
|
||
monkeypatch.setattr(
|
||
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}
|
||
)
|
||
with pytest.raises(litellm.ContextWindowExceededError):
|
||
router._pre_call_checks(
|
||
model="m",
|
||
healthy_deployments=deployments,
|
||
input=short_input,
|
||
request_kwargs={"instructions": long_instructions},
|
||
)
|
||
|
||
|
||
def test_count_pre_call_check_tokens_across_api_surfaces():
|
||
"""
|
||
_count_pre_call_check_tokens must count tokens from chat `messages`, a Responses
|
||
API string `input`, and a Responses API list `input`, and raise when given neither.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
)
|
||
|
||
messages_tokens = router._count_pre_call_check_tokens(
|
||
messages=[{"role": "user", "content": "hello world"}], input=None
|
||
)
|
||
string_input_tokens = router._count_pre_call_check_tokens(messages=None, input="hello world")
|
||
list_input_tokens = router._count_pre_call_check_tokens(
|
||
messages=None, input=[{"role": "user", "content": "hello world"}]
|
||
)
|
||
|
||
assert messages_tokens > 0
|
||
assert string_input_tokens > 0
|
||
assert list_input_tokens > 0
|
||
|
||
with pytest.raises(ValueError):
|
||
router._count_pre_call_check_tokens(messages=None, input=None)
|
||
|
||
|
||
def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch):
|
||
"""
|
||
When neither messages nor input is provided (e.g. endpoints without prompt text),
|
||
token counting is skipped gracefully and all deployments are returned.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
monkeypatch.setattr(
|
||
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}
|
||
)
|
||
|
||
counted: list[dict] = []
|
||
original = router._count_pre_call_check_tokens
|
||
monkeypatch.setattr(
|
||
router,
|
||
"_count_pre_call_check_tokens",
|
||
lambda **kwargs: counted.append(kwargs) or original(**kwargs),
|
||
)
|
||
|
||
deployments = [
|
||
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
|
||
]
|
||
result = router._pre_call_checks(model="m", healthy_deployments=deployments)
|
||
assert len(result) == 1
|
||
assert counted == [] # token counting skipped entirely, so no misleading error is logged
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aresponses_enforces_context_window_pre_call_check():
|
||
"""
|
||
End-to-end router regression: a Responses API call whose `input` exceeds the
|
||
deployment's max_input_tokens must be filtered by the pre-call check, raising
|
||
ContextWindowExceededError instead of being silently routed. This guards the
|
||
wiring that forwards `input` from the generic-call path into deployment selection
|
||
(the deployment uses mock_response, so the check must trip before any real call).
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "small-ctx",
|
||
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
|
||
"model_info": {"max_input_tokens": 5},
|
||
}
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
with pytest.raises(litellm.ContextWindowExceededError):
|
||
await router.aresponses(
|
||
model="small-ctx",
|
||
input="this responses input is definitely much longer than five tokens for sure",
|
||
)
|
||
|
||
|
||
def test_get_deployment_model_info_base_model_flow():
|
||
"""Test that get_deployment_model_info correctly handles the base model flow"""
|
||
from unittest.mock import patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "test-model",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
)
|
||
|
||
# Mock data for the test
|
||
mock_custom_model_info = {
|
||
"base_model": "gpt-3.5-turbo",
|
||
"input_cost_per_token": 0.001,
|
||
"output_cost_per_token": 0.002,
|
||
"custom_field": "custom_value",
|
||
}
|
||
|
||
mock_base_model_info = {
|
||
"key": "gpt-3.5-turbo",
|
||
"max_tokens": 4096,
|
||
"max_input_tokens": 4096,
|
||
"max_output_tokens": 4096,
|
||
"input_cost_per_token": 0.0015, # This should be overridden by custom model info
|
||
"output_cost_per_token": 0.002,
|
||
"litellm_provider": "openai",
|
||
"mode": "chat",
|
||
"supported_openai_params": ["temperature", "max_tokens"],
|
||
}
|
||
|
||
mock_litellm_model_name_info = {
|
||
"key": "test-model",
|
||
"max_tokens": 2048,
|
||
"max_input_tokens": 2048,
|
||
"max_output_tokens": 2048,
|
||
"input_cost_per_token": 0.0005,
|
||
"output_cost_per_token": 0.001,
|
||
"litellm_provider": "test_provider",
|
||
"mode": "completion",
|
||
"supported_openai_params": ["temperature"],
|
||
}
|
||
|
||
# Test Case 1: Base model flow with custom model info that has base_model
|
||
with patch.object(
|
||
litellm, "model_cost", {"test-custom-model": mock_custom_model_info}
|
||
):
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
# Configure mock returns
|
||
mock_get_model_info.side_effect = lambda model: {
|
||
"gpt-3.5-turbo": mock_base_model_info,
|
||
"test-model": mock_litellm_model_name_info,
|
||
}.get(model)
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="test-custom-model", model_name="test-model"
|
||
)
|
||
|
||
# Verify that get_model_info was called for both base model and model name
|
||
assert mock_get_model_info.call_count == 2
|
||
mock_get_model_info.assert_any_call(
|
||
model="gpt-3.5-turbo"
|
||
) # base model call
|
||
mock_get_model_info.assert_any_call(model="test-model") # model name call
|
||
|
||
# Verify the result contains merged information
|
||
assert result is not None
|
||
|
||
# Test the correct merging behavior after fix:
|
||
# 1. base_model_info provides defaults, custom_model_info overrides (correct priority)
|
||
# 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm)
|
||
|
||
# Fields from custom model (should override base model values)
|
||
assert (
|
||
result["input_cost_per_token"] == 0.001
|
||
) # From custom model (overrides base 0.0015)
|
||
assert (
|
||
result["output_cost_per_token"] == 0.002
|
||
) # From custom model (same as base)
|
||
assert result["custom_field"] == "custom_value" # From custom model
|
||
|
||
# Fields from base model that weren't overridden by custom
|
||
assert result["max_tokens"] == 4096 # From base model
|
||
assert result["litellm_provider"] == "openai" # From base model
|
||
assert (
|
||
result["mode"] == "chat"
|
||
) # From base model (overrides litellm "completion")
|
||
|
||
# The key field comes from base model since both base and litellm have it
|
||
# and base model info overrides litellm model name info in final merge
|
||
assert (
|
||
result["key"] == "gpt-3.5-turbo"
|
||
) # From base model (overrides litellm key)
|
||
|
||
# Test Case 2: Custom model info without base_model
|
||
mock_custom_model_info_no_base = {
|
||
"input_cost_per_token": 0.001,
|
||
"output_cost_per_token": 0.002,
|
||
"custom_field": "custom_value",
|
||
}
|
||
|
||
with patch.object(
|
||
litellm,
|
||
"model_cost",
|
||
{"test-custom-model-no-base": mock_custom_model_info_no_base},
|
||
):
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
mock_get_model_info.side_effect = lambda model: {
|
||
"test-model": mock_litellm_model_name_info,
|
||
}.get(model)
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="test-custom-model-no-base", model_name="test-model"
|
||
)
|
||
|
||
# Should only call get_model_info once for model name (no base model)
|
||
assert mock_get_model_info.call_count == 1
|
||
mock_get_model_info.assert_called_with(model="test-model")
|
||
|
||
# Verify the result contains merged information
|
||
assert result is not None
|
||
assert result["input_cost_per_token"] == 0.001 # From custom model
|
||
assert result["max_tokens"] == 2048 # From litellm model name info
|
||
assert result["custom_field"] == "custom_value" # From custom model
|
||
assert result["mode"] == "completion" # From litellm model name info
|
||
|
||
# Test Case 3: No custom model info, only litellm model name info
|
||
with patch.object(litellm, "model_cost", {}): # Empty model cost
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
mock_get_model_info.side_effect = lambda model: {
|
||
"test-model": mock_litellm_model_name_info,
|
||
}.get(model)
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="non-existent-model", model_name="test-model"
|
||
)
|
||
|
||
# Should only call get_model_info once for model name
|
||
assert mock_get_model_info.call_count == 1
|
||
mock_get_model_info.assert_called_with(model="test-model")
|
||
|
||
# Result should be just the litellm model name info
|
||
assert result is not None
|
||
assert result == mock_litellm_model_name_info
|
||
|
||
# Test Case 4: Base model info retrieval fails (exception handling)
|
||
mock_custom_model_info_invalid_base = {
|
||
"base_model": "invalid-base-model",
|
||
"input_cost_per_token": 0.001,
|
||
"output_cost_per_token": 0.002,
|
||
}
|
||
|
||
with patch.object(
|
||
litellm,
|
||
"model_cost",
|
||
{"test-custom-model-invalid": mock_custom_model_info_invalid_base},
|
||
):
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
# Mock get_model_info to raise exception for invalid base model
|
||
def mock_get_model_info_side_effect(model):
|
||
if model == "invalid-base-model":
|
||
raise Exception("Model not found")
|
||
elif model == "test-model":
|
||
return mock_litellm_model_name_info
|
||
return None
|
||
|
||
mock_get_model_info.side_effect = mock_get_model_info_side_effect
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="test-custom-model-invalid", model_name="test-model"
|
||
)
|
||
|
||
# Should handle exception gracefully and still return merged result
|
||
assert result is not None
|
||
assert result["input_cost_per_token"] == 0.001 # From custom model
|
||
assert result["mode"] == "completion" # From litellm model name info
|
||
|
||
# Test Case 5: Both model_cost.get() and get_model_info() return None
|
||
with patch.object(litellm, "model_cost", {}):
|
||
with patch.object(
|
||
litellm, "get_model_info", side_effect=Exception("Not found")
|
||
):
|
||
result = router.get_deployment_model_info(
|
||
model_id="non-existent", model_name="non-existent"
|
||
)
|
||
|
||
# Should return None when no model info is found
|
||
assert result is None
|
||
|
||
# Test Case 6: custom_model_info present but litellm_model_name_model_info is None
|
||
# (model has custom pricing in config but is not in built-in model_prices_and_context_window.json)
|
||
mock_custom_pricing_only = {
|
||
"input_cost_per_token": 1.74e-06,
|
||
"output_cost_per_token": 3.48e-06,
|
||
"cache_read_input_token_cost": 1.45e-08,
|
||
"mode": "chat",
|
||
}
|
||
|
||
with patch.object(
|
||
litellm,
|
||
"model_cost",
|
||
{"custom-model-id": mock_custom_pricing_only},
|
||
):
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
# Model NOT in built-in cost map — raise exception
|
||
mock_get_model_info.side_effect = Exception("Model not in cost map")
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="custom-model-id", model_name="unknown-model"
|
||
)
|
||
|
||
# Should return custom_model_info even when litellm_model_name_model_info is None
|
||
assert result is not None
|
||
assert result["input_cost_per_token"] == 1.74e-06
|
||
assert result["output_cost_per_token"] == 3.48e-06
|
||
assert result["cache_read_input_token_cost"] == 1.45e-08
|
||
assert result["mode"] == "chat"
|
||
|
||
# Test Case 7: custom_model_info with base_model but litellm_model_name_model_info None
|
||
mock_custom_with_base = {
|
||
"base_model": "some-base-model",
|
||
"input_cost_per_token": 0.01,
|
||
"output_cost_per_token": 0.02,
|
||
}
|
||
mock_base_info = {
|
||
"key": "some-base-model",
|
||
"max_tokens": 8192,
|
||
"mode": "chat",
|
||
"litellm_provider": "openai",
|
||
}
|
||
|
||
with patch.object(
|
||
litellm,
|
||
"model_cost",
|
||
{"custom-with-base": mock_custom_with_base},
|
||
):
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
|
||
def get_info_side_effect(model):
|
||
if model == "some-base-model":
|
||
return mock_base_info
|
||
raise Exception("Model not in cost map")
|
||
|
||
mock_get_model_info.side_effect = get_info_side_effect
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="custom-with-base", model_name="unknown-model"
|
||
)
|
||
|
||
# Should return custom_model_info merged with base model info
|
||
assert result is not None
|
||
assert (
|
||
result["input_cost_per_token"] == 0.01
|
||
) # From custom (overrides base)
|
||
assert result["max_tokens"] == 8192 # From base model
|
||
assert result["litellm_provider"] == "openai" # From base model
|
||
|
||
print("✓ All base model flow test cases passed!")
|
||
|
||
|
||
@patch("litellm.model_cost", {})
|
||
def test_get_deployment_model_info_base_model_merge_priority():
|
||
"""Test that base model info merging respects the correct priority order"""
|
||
from unittest.mock import patch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "test-model",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
)
|
||
|
||
# Test data with overlapping fields to test merge priority
|
||
mock_custom_model_info = {
|
||
"base_model": "gpt-4",
|
||
"input_cost_per_token": 0.01, # Should override base model value
|
||
"max_tokens": 8000, # Should override base model value
|
||
"custom_only_field": "custom_value",
|
||
}
|
||
|
||
mock_base_model_info = {
|
||
"key": "gpt-4",
|
||
"max_tokens": 4096, # Should be overridden by custom model
|
||
"input_cost_per_token": 0.03, # Should be overridden by custom model
|
||
"output_cost_per_token": 0.06, # Should be preserved (not in custom)
|
||
"litellm_provider": "openai",
|
||
"base_only_field": "base_value",
|
||
}
|
||
|
||
mock_litellm_model_name_info = {
|
||
"key": "test-model",
|
||
"max_tokens": 2048, # Should be overridden by final custom model info
|
||
"input_cost_per_token": 0.005, # Should be overridden by final custom model info
|
||
"output_cost_per_token": 0.01, # Should be overridden by final custom model info
|
||
"mode": "completion",
|
||
"litellm_only_field": "litellm_value",
|
||
}
|
||
|
||
with patch.object(
|
||
litellm, "model_cost", {"custom-model-id": mock_custom_model_info}
|
||
):
|
||
with patch.object(litellm, "get_model_info") as mock_get_model_info:
|
||
mock_get_model_info.side_effect = lambda model: {
|
||
"gpt-4": mock_base_model_info,
|
||
"test-model": mock_litellm_model_name_info,
|
||
}.get(model)
|
||
|
||
result = router.get_deployment_model_info(
|
||
model_id="custom-model-id", model_name="test-model"
|
||
)
|
||
|
||
assert result is not None
|
||
|
||
# Test correct merge priority after fix:
|
||
# 1. base_model_info provides defaults
|
||
# 2. custom_model_info overrides base_model_info
|
||
# 3. Result from steps 1-2 overrides litellm_model_name_info
|
||
|
||
# Fields that should come from custom model info (highest priority)
|
||
assert (
|
||
result["input_cost_per_token"] == 0.01
|
||
) # From custom model (overrides base 0.03)
|
||
assert (
|
||
result["max_tokens"] == 8000
|
||
) # From custom model (overrides base 4096)
|
||
assert result["custom_only_field"] == "custom_value" # From custom model
|
||
|
||
# Fields that should come from base model (not overridden by custom)
|
||
assert (
|
||
result["output_cost_per_token"] == 0.06
|
||
) # From base model (not in custom)
|
||
assert (
|
||
result["litellm_provider"] == "openai"
|
||
) # From base model (not in custom)
|
||
assert (
|
||
result["base_only_field"] == "base_value"
|
||
) # From base model (not in custom)
|
||
|
||
# Fields that should come from litellm model name info (not overridden by custom+base)
|
||
assert (
|
||
result["mode"] == "completion"
|
||
) # From litellm model name info (not in custom or base)
|
||
assert (
|
||
result["litellm_only_field"] == "litellm_value"
|
||
) # From litellm model name info (not in custom or base)
|
||
|
||
# Key comes from base model since both base and litellm have key fields
|
||
# and the merged custom+base overrides litellm in the final merge
|
||
assert result["key"] == "gpt-4"
|
||
|
||
print("✓ Base model merge priority test passed!")
|
||
|
||
|
||
def test_add_deployment_model_to_endpoint_for_llm_passthrough_route():
|
||
"""
|
||
Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "special-bedrock-model",
|
||
"litellm_params": {
|
||
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
# Test Case 1: Bedrock model with provider prefix - should strip "bedrock/" prefix
|
||
kwargs = {
|
||
"endpoint": "/model/special-bedrock-model/invoke",
|
||
"custom_llm_provider": "bedrock",
|
||
}
|
||
result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route(
|
||
kwargs=kwargs,
|
||
model="special-bedrock-model",
|
||
model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
)
|
||
assert (
|
||
result["endpoint"]
|
||
== "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke"
|
||
), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'"
|
||
|
||
# Test Case 2: Bedrock invoke-with-response-stream endpoint
|
||
kwargs = {
|
||
"endpoint": "/model/special-bedrock-model/invoke-with-response-stream",
|
||
"custom_llm_provider": "bedrock",
|
||
}
|
||
result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route(
|
||
kwargs=kwargs,
|
||
model="special-bedrock-model",
|
||
model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
)
|
||
assert (
|
||
result["endpoint"]
|
||
== "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream"
|
||
), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'"
|
||
|
||
# Test Case 3: Bedrock converse endpoint
|
||
kwargs = {
|
||
"endpoint": "/model/bedrock-model/converse",
|
||
"custom_llm_provider": "bedrock",
|
||
}
|
||
result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route(
|
||
kwargs=kwargs,
|
||
model="bedrock-model",
|
||
model_name="bedrock/us.meta.llama3-8b-instruct-v1:0",
|
||
)
|
||
assert (
|
||
result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse"
|
||
), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'"
|
||
|
||
# Test Case 4: Bedrock provider prefix auto-detected from model_name
|
||
kwargs = {
|
||
"endpoint": "/model/router-model/invoke",
|
||
}
|
||
result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route(
|
||
kwargs=kwargs,
|
||
model="router-model",
|
||
model_name="bedrock/us.meta.llama3-8b-instruct-v1:0",
|
||
)
|
||
assert (
|
||
result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke"
|
||
), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'"
|
||
|
||
|
||
def test_update_kwargs_with_deployment_uses_pass_through_request_timeout():
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-bedrock-model",
|
||
"litellm_params": {
|
||
"model": "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
deployment = router.model_list[0]
|
||
kwargs: dict = {}
|
||
|
||
with patch(
|
||
"litellm.proxy.proxy_server.general_settings",
|
||
{"pass_through_request_timeout": 6},
|
||
):
|
||
router._update_kwargs_with_deployment(
|
||
deployment=deployment,
|
||
kwargs=kwargs,
|
||
function_name="_ageneric_api_call_with_fallbacks",
|
||
)
|
||
|
||
assert kwargs["timeout"] == 6.0
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_router_acompletion_with_unknown_model_and_default_fallback():
|
||
"""
|
||
Test that the router successfully uses a default fallback when a completely
|
||
unknown model is requested. It should not raise a BadRequestError.
|
||
This test verifies the fix for issue #15114.
|
||
"""
|
||
model_list = [
|
||
{
|
||
"model_name": "gpt-4o", # This is the fallback model
|
||
"litellm_params": {
|
||
"model": "azure/gpt-4o-real", # The actual underlying model name
|
||
"api_key": "fake-key",
|
||
"api_base": "https://fake-endpoint.openai.azure.com/",
|
||
"mock_response": "this is the fallback response", # Mocked response to prevent real API calls
|
||
},
|
||
}
|
||
]
|
||
|
||
# Initialize the router with a default fallback
|
||
router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"])
|
||
|
||
messages = [
|
||
{"role": "user", "content": "This call should succeed by falling back."}
|
||
]
|
||
|
||
# Call completion with a model name that is NOT in the model_list
|
||
response = await router.acompletion(
|
||
model="completely-unknown-model", messages=messages
|
||
)
|
||
|
||
# Check that the call did not fail and we received a valid response object.
|
||
assert response is not None
|
||
|
||
# Check that the content of the response is from the MOCKED fallback model.
|
||
assert response.choices[0].message.content == "this is the fallback response"
|
||
|
||
# Check that the response object reports the model that was *actually* called.
|
||
assert response.model == "gpt-4o-real"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_router_acompletion_with_unknown_model_and_no_fallback():
|
||
"""
|
||
Test that the router still raises a BadRequestError for an unknown model
|
||
when no default fallbacks are configured. This ensures we don't break
|
||
the original behavior.
|
||
"""
|
||
model_list = [
|
||
{
|
||
"model_name": "gpt-4o",
|
||
"litellm_params": {
|
||
"model": "azure/gpt-4o-real",
|
||
"api_key": "fake-key",
|
||
"mock_response": "this should not be called",
|
||
},
|
||
}
|
||
]
|
||
|
||
# Initialize the router WITHOUT any default fallbacks
|
||
router = litellm.Router(model_list=model_list)
|
||
|
||
messages = [{"role": "user", "content": "This call should fail."}]
|
||
|
||
# Use pytest.raises to assert that a BadRequestError is thrown.
|
||
with pytest.raises(litellm.BadRequestError) as excinfo:
|
||
await router.acompletion(model="completely-unknown-model", messages=messages)
|
||
|
||
# Check that the error message is correct.
|
||
# The router returns 'no healthy deployments' because get_model_list returns [] not None.
|
||
assert "no healthy deployments for this model" in str(excinfo.value)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_router_unknown_model_error_message_renders_model_name_literally():
|
||
"""
|
||
The unknown-model error message renders the caller-supplied model name
|
||
verbatim. A name containing Python format-field syntax must be treated as
|
||
literal text, not re-interpreted as a format template, which would distort
|
||
the message and balloon its length.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4o",
|
||
"litellm_params": {"model": "azure/gpt-4o-real", "api_key": "fake-key"},
|
||
}
|
||
]
|
||
)
|
||
|
||
weird_model = "ghost{:>200}model"
|
||
messages = [{"role": "user", "content": "hi"}]
|
||
|
||
with pytest.raises(litellm.BadRequestError) as excinfo:
|
||
await router.acompletion(model=weird_model, messages=messages)
|
||
|
||
message = str(excinfo.value)
|
||
assert weird_model in message
|
||
assert " " not in message # no padding run from an expanded format field
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint():
|
||
"""
|
||
Test that get_deployment_credentials_with_provider correctly copies
|
||
aws_bedrock_runtime_endpoint from deployment litellm_params to credentials.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "bedrock-claude-model",
|
||
"litellm_params": {
|
||
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
"aws_access_key_id": "test-access-key",
|
||
"aws_secret_access_key": "test-secret-key",
|
||
"aws_region_name": "us-east-1",
|
||
"aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-east-1.amazonaws.com",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="bedrock-claude-model"
|
||
)
|
||
|
||
assert credentials is not None
|
||
assert (
|
||
credentials["aws_bedrock_runtime_endpoint"]
|
||
== "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||
)
|
||
assert credentials["aws_access_key_id"] == "test-access-key"
|
||
assert credentials["aws_secret_access_key"] == "test-secret-key"
|
||
assert credentials["aws_region_name"] == "us-east-1"
|
||
assert credentials["custom_llm_provider"] == "bedrock"
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_includes_bucket_name():
|
||
"""
|
||
Regression: bucket_name must survive the CredentialLiteLLMParams filter so
|
||
managed-files batch retrieval can resolve the GCS/S3 bucket. Previously it was
|
||
dropped, causing "GCS bucket_name is required" when fetching batch output files.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "vertex-gemini",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-3.5-flash",
|
||
"vertex_project": "my-project",
|
||
"vertex_location": "global",
|
||
"gcs_bucket_name": "my-batch-bucket",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="vertex-gemini"
|
||
)
|
||
|
||
assert credentials is not None
|
||
assert credentials["gcs_bucket_name"] == "my-batch-bucket"
|
||
assert credentials["vertex_project"] == "my-project"
|
||
assert credentials["custom_llm_provider"] == "vertex_ai"
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_resolves_credential_name():
|
||
"""
|
||
Test that get_deployment_credentials_with_provider correctly resolves
|
||
litellm_credential_name to actual credential values (for UI-created models).
|
||
"""
|
||
from litellm.types.utils import CredentialItem
|
||
|
||
# Setup credential list with a test credential
|
||
litellm.credential_list = [
|
||
CredentialItem(
|
||
credential_name="test-azure-cred",
|
||
credential_info={"custom_llm_provider": "azure"},
|
||
credential_values={
|
||
"api_key": "resolved-api-key",
|
||
"api_base": "https://resolved.openai.azure.com",
|
||
"api_version": "2024-02-01",
|
||
},
|
||
)
|
||
]
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "azure-gpt-4",
|
||
"litellm_params": {
|
||
"model": "azure/gpt-4",
|
||
"litellm_credential_name": "test-azure-cred",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="azure-gpt-4"
|
||
)
|
||
|
||
assert credentials is not None
|
||
assert credentials["api_key"] == "resolved-api-key"
|
||
assert credentials["api_base"] == "https://resolved.openai.azure.com"
|
||
assert credentials["api_version"] == "2024-02-01"
|
||
assert credentials["custom_llm_provider"] == "azure"
|
||
# Ensure credential name is removed after resolution
|
||
assert "litellm_credential_name" not in credentials
|
||
|
||
# Cleanup
|
||
litellm.credential_list = []
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_bedrock_batch_fields():
|
||
"""
|
||
Test that get_deployment_credentials_with_provider returns the deployment's
|
||
model and the Bedrock batch/S3 fields (s3_region_name, s3_encryption_key_id,
|
||
aws_batch_role_arn) instead of silently dropping them (#25104).
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "bedrock-batch-model",
|
||
"litellm_params": {
|
||
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
"aws_region_name": "us-west-2",
|
||
"s3_bucket_name": "my-batch-bucket",
|
||
"s3_region_name": "us-east-1",
|
||
"s3_encryption_key_id": "arn:aws:kms:us-west-2:123:key/abc",
|
||
"aws_batch_role_arn": "arn:aws:iam::123:role/batch-role",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="bedrock-batch-model"
|
||
)
|
||
|
||
assert credentials is not None
|
||
assert credentials["custom_llm_provider"] == "bedrock"
|
||
assert credentials["model"] == "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||
assert credentials["aws_region_name"] == "us-west-2"
|
||
assert credentials["s3_bucket_name"] == "my-batch-bucket"
|
||
assert credentials["s3_region_name"] == "us-east-1"
|
||
assert credentials["s3_encryption_key_id"] == "arn:aws:kms:us-west-2:123:key/abc"
|
||
assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role"
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_preserves_aws_auth_params():
|
||
"""
|
||
Test that get_deployment_credentials_with_provider preserves every AWS auth
|
||
selector (session token, assume-role, web identity, profile) so bedrock
|
||
files/batches deployments using temporary or role-based credentials do not
|
||
silently fall back to the server's ambient identity (#36155).
|
||
"""
|
||
aws_auth_params = {
|
||
"aws_access_key_id": "deployment-access-key",
|
||
"aws_secret_access_key": "deployment-secret",
|
||
"aws_session_token": "deployment-session-token",
|
||
"aws_region_name": "us-west-2",
|
||
"aws_session_name": "deployment-session",
|
||
"aws_profile_name": "deployment-profile",
|
||
"aws_role_name": "arn:aws:iam::123:role/deployment-role",
|
||
"aws_web_identity_token": "deployment-web-identity",
|
||
"aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com",
|
||
"aws_external_id": "deployment-external-id",
|
||
}
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "bedrock-batch-model",
|
||
"litellm_params": {
|
||
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||
**aws_auth_params,
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="bedrock-batch-model"
|
||
)
|
||
|
||
assert credentials is not None
|
||
for key, value in aws_auth_params.items():
|
||
assert credentials.get(key) == value, key
|
||
|
||
|
||
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
|
||
return {
|
||
"model_name": f"model_name_team-1_{model_id}",
|
||
"litellm_params": {"model": "openai/*", "api_key": api_key},
|
||
"model_info": {
|
||
"id": model_id,
|
||
"team_id": "team-1",
|
||
"team_public_model_name": "openai/*",
|
||
},
|
||
}
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_team_wildcard_priority():
|
||
"""
|
||
Regression: a global wildcard pattern (e.g. "openai/*") must not shadow a
|
||
team's own wildcard entry. When team_id is provided, the team wildcard
|
||
deployment's credentials win; without team_id the global one is used.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "openai/*",
|
||
"litellm_params": {"model": "openai/*", "api_key": "global-key"},
|
||
},
|
||
_team_wildcard_model(api_key="team-key"),
|
||
],
|
||
)
|
||
|
||
team_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-1"
|
||
)
|
||
assert team_credentials is not None
|
||
assert team_credentials["api_key"] == "team-key"
|
||
|
||
global_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2"
|
||
)
|
||
assert global_credentials is not None
|
||
assert global_credentials["api_key"] == "global-key"
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_skips_other_team_deployment():
|
||
"""
|
||
Regression: a team-scoped deployment sharing a model_name with a global
|
||
deployment must never resolve for another team's (or an unscoped) caller,
|
||
even when it is indexed first; the shared global deployment wins instead.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gemini-2.5-pro",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-pro",
|
||
"vertex_project": "team-b-project",
|
||
},
|
||
"model_info": {
|
||
"id": "team-b-vertex",
|
||
"team_id": "team-b",
|
||
"team_public_model_name": "gemini-2.5-pro",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "gemini-2.5-pro",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-pro",
|
||
"vertex_project": "shared-project",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
other_team_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="gemini-2.5-pro", team_id="team-a"
|
||
)
|
||
assert other_team_credentials is not None
|
||
assert other_team_credentials["vertex_project"] == "shared-project"
|
||
|
||
unscoped_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="gemini-2.5-pro"
|
||
)
|
||
assert unscoped_credentials is not None
|
||
assert unscoped_credentials["vertex_project"] == "shared-project"
|
||
|
||
owner_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="gemini-2.5-pro", team_id="team-b"
|
||
)
|
||
assert owner_credentials is not None
|
||
assert owner_credentials["vertex_project"] == "team-b-project"
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only_name():
|
||
"""
|
||
When the only deployments under a model name belong to another team, other
|
||
callers must get None (env fallback) instead of that team's credentials.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gemini-2.5-pro",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-pro",
|
||
"vertex_project": "team-b-project",
|
||
},
|
||
"model_info": {
|
||
"id": "team-b-vertex",
|
||
"team_id": "team-b",
|
||
"team_public_model_name": "gemini-2.5-pro",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
assert (
|
||
router.get_deployment_credentials_with_provider(
|
||
model_id="gemini-2.5-pro", team_id="team-a"
|
||
)
|
||
is None
|
||
)
|
||
assert (
|
||
router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro")
|
||
is None
|
||
)
|
||
|
||
|
||
def test_deployment_usable_by_team_helpers():
|
||
"""
|
||
Direct coverage of the team-ownership filter: a team-scoped deployment is
|
||
usable only by its owning team, shared deployments by anyone, and the
|
||
model-group picker returns the first usable deployment or None.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gemini-2.5-pro",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-pro",
|
||
"vertex_project": "team-b-project",
|
||
},
|
||
"model_info": {
|
||
"id": "team-b-vertex",
|
||
"team_id": "team-b",
|
||
"team_public_model_name": "gemini-2.5-pro",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "gemini-2.5-pro",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.5-pro",
|
||
"vertex_project": "shared-project",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
team_owned, shared = router.model_list
|
||
assert router._deployment_usable_by_team(team_owned, "team-b") is True
|
||
assert router._deployment_usable_by_team(team_owned, "team-a") is False
|
||
assert router._deployment_usable_by_team(team_owned, None) is False
|
||
assert router._deployment_usable_by_team(shared, "team-a") is True
|
||
assert router._deployment_usable_by_team(shared, None) is True
|
||
|
||
picked = router._get_model_group_deployment_usable_by_team(
|
||
model_group_name="gemini-2.5-pro", team_id="team-a"
|
||
)
|
||
assert picked is not None
|
||
assert picked.litellm_params.vertex_project == "shared-project"
|
||
|
||
owner_picked = router._get_model_group_deployment_usable_by_team(
|
||
model_group_name="gemini-2.5-pro", team_id="team-b"
|
||
)
|
||
assert owner_picked is not None
|
||
assert owner_picked.litellm_params.vertex_project == "team-b-project"
|
||
|
||
assert (
|
||
router._get_model_group_deployment_usable_by_team(
|
||
model_group_name="unknown-model", team_id="team-a"
|
||
)
|
||
is None
|
||
)
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_skips_other_team_wildcard():
|
||
"""
|
||
Global wildcard resolution must skip a team-scoped wildcard deployment for
|
||
callers outside that team, falling through to the shared wildcard entry.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "openai/*",
|
||
"litellm_params": {"model": "openai/*", "api_key": "team-b-key"},
|
||
"model_info": {
|
||
"id": "team-b-wildcard",
|
||
"team_id": "team-b",
|
||
"team_public_model_name": "openai/*",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "openai/*",
|
||
"litellm_params": {"model": "openai/*", "api_key": "global-key"},
|
||
},
|
||
],
|
||
)
|
||
|
||
other_team_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-a"
|
||
)
|
||
assert other_team_credentials is not None
|
||
assert other_team_credentials["api_key"] == "global-key"
|
||
|
||
owner_credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-b"
|
||
)
|
||
assert owner_credentials is not None
|
||
assert owner_credentials["api_key"] == "team-b-key"
|
||
|
||
|
||
def test_team_wildcard_credentials_not_usable_after_delete_deployment():
|
||
"""
|
||
Regression: team_pattern_routers retained deleted deployments, so a team
|
||
user could keep resolving credentials of a deleted wildcard deployment.
|
||
"""
|
||
router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")])
|
||
|
||
assert (
|
||
router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-1"
|
||
)
|
||
is not None
|
||
)
|
||
|
||
router.delete_deployment(id="team-wildcard-id")
|
||
|
||
assert (
|
||
router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-1"
|
||
)
|
||
is None
|
||
)
|
||
|
||
|
||
def test_pattern_match_router_remove_deployment():
|
||
"""
|
||
remove_deployment must drop only the deployment with the given model id and
|
||
delete patterns whose deployment list becomes empty.
|
||
"""
|
||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
||
|
||
pattern_router = PatternMatchRouter()
|
||
pattern_router.add_pattern(
|
||
"openai/*",
|
||
{"litellm_params": {"model": "openai/*", "api_key": "key-a"}, "model_info": {"id": "dep-a"}},
|
||
)
|
||
pattern_router.add_pattern(
|
||
"openai/*",
|
||
{"litellm_params": {"model": "openai/*", "api_key": "key-b"}, "model_info": {"id": "dep-b"}},
|
||
)
|
||
|
||
pattern_router.remove_deployment(model_id="dep-a")
|
||
matches = pattern_router.route("openai/gpt-5.2")
|
||
assert matches is not None
|
||
assert [m["model_info"]["id"] for m in matches] == ["dep-b"]
|
||
|
||
pattern_router.remove_deployment(model_id="dep-b")
|
||
assert pattern_router.patterns == {}
|
||
assert pattern_router.route("openai/gpt-5.2") is None
|
||
|
||
|
||
def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list():
|
||
"""
|
||
Regression: replacing a team wildcard deployment (upsert or model list
|
||
reload) must serve the new credentials, not the stale cached ones.
|
||
"""
|
||
from litellm.types.router import Deployment
|
||
|
||
router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")])
|
||
|
||
router.upsert_deployment(
|
||
deployment=Deployment(**_team_wildcard_model(api_key="new-key"))
|
||
)
|
||
credentials = router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-1"
|
||
)
|
||
assert credentials is not None
|
||
assert credentials["api_key"] == "new-key"
|
||
|
||
router.set_model_list(model_list=[])
|
||
assert (
|
||
router.get_deployment_credentials_with_provider(
|
||
model_id="openai/gpt-5.2", team_id="team-1"
|
||
)
|
||
is None
|
||
)
|
||
|
||
|
||
def test_get_available_guardrail_single_deployment():
|
||
"""
|
||
Test get_available_guardrail returns the single guardrail when only one exists.
|
||
"""
|
||
guardrail_config = {
|
||
"guardrail_name": "content-filter",
|
||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||
"id": "guardrail-1",
|
||
}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
guardrail_list=[guardrail_config],
|
||
)
|
||
|
||
result = router.get_available_guardrail(guardrail_name="content-filter")
|
||
assert result == guardrail_config
|
||
|
||
|
||
def test_get_available_guardrail_multiple_deployments():
|
||
"""
|
||
Test get_available_guardrail load balances across multiple guardrails.
|
||
"""
|
||
guardrail_1 = {
|
||
"guardrail_name": "content-filter",
|
||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||
"id": "guardrail-1",
|
||
}
|
||
guardrail_2 = {
|
||
"guardrail_name": "content-filter",
|
||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||
"id": "guardrail-2",
|
||
}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
guardrail_list=[guardrail_1, guardrail_2],
|
||
)
|
||
|
||
# Call multiple times to verify load balancing
|
||
results = set()
|
||
for _ in range(20):
|
||
result = router.get_available_guardrail(guardrail_name="content-filter")
|
||
results.add(result["id"])
|
||
|
||
# Both guardrails should be selected at least once
|
||
assert "guardrail-1" in results or "guardrail-2" in results
|
||
|
||
|
||
def test_get_available_guardrail_not_found():
|
||
"""
|
||
Test get_available_guardrail raises ValueError when guardrail not found.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
guardrail_list=[],
|
||
)
|
||
|
||
with pytest.raises(ValueError, match="No guardrail found with name"):
|
||
router.get_available_guardrail(guardrail_name="non-existent")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aguardrail_helper():
|
||
"""
|
||
Test _aguardrail_helper selects a guardrail and executes the original function.
|
||
"""
|
||
guardrail_config = {
|
||
"guardrail_name": "content-filter",
|
||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||
"id": "guardrail-1",
|
||
}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
guardrail_list=[guardrail_config],
|
||
)
|
||
|
||
# Mock the original function
|
||
async def mock_original_function(**kwargs):
|
||
return {
|
||
"result": "success",
|
||
"selected_guardrail": kwargs.get("selected_guardrail"),
|
||
}
|
||
|
||
result = await router._aguardrail_helper(
|
||
model="content-filter",
|
||
original_generic_function=mock_original_function,
|
||
)
|
||
|
||
assert result["result"] == "success"
|
||
assert result["selected_guardrail"] == guardrail_config
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_aguardrail():
|
||
"""
|
||
Test aguardrail executes a guardrail with load balancing and fallbacks.
|
||
"""
|
||
guardrail_config = {
|
||
"guardrail_name": "content-filter",
|
||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||
"id": "guardrail-1",
|
||
}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-3.5-turbo",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
guardrail_list=[guardrail_config],
|
||
)
|
||
|
||
# Mock the original function
|
||
async def mock_original_function(**kwargs):
|
||
return {
|
||
"result": "success",
|
||
"selected_guardrail": kwargs.get("selected_guardrail"),
|
||
}
|
||
|
||
result = await router.aguardrail(
|
||
guardrail_name="content-filter",
|
||
original_function=mock_original_function,
|
||
)
|
||
|
||
assert result["result"] == "success"
|
||
assert result["selected_guardrail"]["id"] == "guardrail-1"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_anthropic_messages_call_type_is_cached():
|
||
"""
|
||
Regression test: Verify that anthropic_messages call type is allowed
|
||
in PromptCachingDeploymentCheck.async_log_success_event.
|
||
"""
|
||
import asyncio
|
||
|
||
from litellm.caching.dual_cache import DualCache
|
||
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
|
||
PromptCachingDeploymentCheck,
|
||
)
|
||
from litellm.router_utils.prompt_caching_cache import PromptCachingCache
|
||
from litellm.types.utils import (
|
||
CallTypes,
|
||
StandardLoggingHiddenParams,
|
||
StandardLoggingMetadata,
|
||
StandardLoggingModelInformation,
|
||
StandardLoggingPayload,
|
||
)
|
||
|
||
# Create mock standard logging payload inline
|
||
def create_standard_logging_payload() -> StandardLoggingPayload:
|
||
return StandardLoggingPayload(
|
||
id="test_id",
|
||
call_type="completion",
|
||
response_cost=0.1,
|
||
response_cost_failure_debug_info=None,
|
||
status="success",
|
||
total_tokens=30,
|
||
prompt_tokens=20,
|
||
completion_tokens=10,
|
||
startTime=1234567890.0,
|
||
endTime=1234567891.0,
|
||
completionStartTime=1234567890.5,
|
||
model_map_information=StandardLoggingModelInformation(
|
||
model_map_key="gpt-3.5-turbo", model_map_value=None
|
||
),
|
||
model="gpt-3.5-turbo",
|
||
model_id="model-123",
|
||
model_group="openai-gpt",
|
||
api_base="https://api.openai.com",
|
||
metadata=StandardLoggingMetadata(
|
||
user_api_key_hash="test_hash",
|
||
user_api_key_org_id=None,
|
||
user_api_key_alias="test_alias",
|
||
user_api_key_team_id="test_team",
|
||
user_api_key_user_id="test_user",
|
||
user_api_key_team_alias="test_team_alias",
|
||
spend_logs_metadata=None,
|
||
requester_ip_address="127.0.0.1",
|
||
requester_metadata=None,
|
||
),
|
||
cache_hit=False,
|
||
cache_key=None,
|
||
saved_cache_cost=0.0,
|
||
request_tags=[],
|
||
end_user=None,
|
||
requester_ip_address="127.0.0.1",
|
||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||
response={"choices": [{"message": {"content": "Hi there!"}}]},
|
||
error_str=None,
|
||
model_parameters={"stream": True},
|
||
hidden_params=StandardLoggingHiddenParams(
|
||
model_id="model-123",
|
||
cache_key=None,
|
||
api_base="https://api.openai.com",
|
||
response_cost="0.1",
|
||
additional_headers=None,
|
||
),
|
||
)
|
||
|
||
cache = DualCache()
|
||
deployment_check = PromptCachingDeploymentCheck(cache=cache)
|
||
prompt_cache = PromptCachingCache(cache=cache)
|
||
|
||
# Create messages with enough tokens to pass the caching threshold
|
||
test_messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "test long message here" * 1024,
|
||
"cache_control": {"type": "ephemeral", "ttl": "5m"},
|
||
}
|
||
],
|
||
}
|
||
]
|
||
test_model_id = "test-model-id-123"
|
||
|
||
# Create a payload with anthropic_messages call type
|
||
payload = create_standard_logging_payload()
|
||
payload["call_type"] = CallTypes.anthropic_messages.value
|
||
payload["messages"] = test_messages
|
||
payload["model"] = "anthropic/claude-3-5-sonnet-20240620"
|
||
payload["model_id"] = test_model_id
|
||
|
||
# Log the success event (should cache the model_id)
|
||
await deployment_check.async_log_success_event(
|
||
kwargs={"standard_logging_object": payload},
|
||
response_obj={},
|
||
start_time=1234567890.0,
|
||
end_time=1234567891.0,
|
||
)
|
||
|
||
# Small delay to ensure cache write completes
|
||
await asyncio.sleep(0.1)
|
||
|
||
# Verify that the model_id was actually cached
|
||
cached_result = await prompt_cache.async_get_model_id(
|
||
messages=test_messages,
|
||
tools=None,
|
||
)
|
||
|
||
# This assertion will FAIL if anthropic_messages is filtered out
|
||
assert (
|
||
cached_result is not None
|
||
), "Model ID should be cached for anthropic_messages call type"
|
||
assert (
|
||
cached_result["model_id"] == test_model_id
|
||
), f"Expected {test_model_id}, got {cached_result['model_id']}"
|
||
|
||
|
||
def test_update_kwargs_with_deployment_propagates_model_tags():
|
||
"""
|
||
Test that deployment-level tags from litellm_params are merged into
|
||
kwargs metadata when _update_kwargs_with_deployment is called.
|
||
|
||
This ensures model-level tags defined in config.yaml appear in SpendLogs.
|
||
See: https://github.com/BerriAI/litellm/issues/XXXX
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4o-mini",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "fake-key",
|
||
"tags": ["openai-account", "production"],
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {"metadata": {}}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="gpt-4o-mini"
|
||
)
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
# Deployment tags should be propagated to kwargs metadata
|
||
assert "tags" in kwargs["metadata"]
|
||
assert "openai-account" in kwargs["metadata"]["tags"]
|
||
assert "production" in kwargs["metadata"]["tags"]
|
||
|
||
|
||
def test_update_kwargs_with_deployment_merges_tags_without_duplicates():
|
||
"""
|
||
Test that when both request-level and deployment-level tags exist,
|
||
they are merged without duplicates.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4o-mini",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "fake-key",
|
||
"tags": ["openai-account", "shared-tag"],
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# Simulate request that already has tags (from request body or key/team level)
|
||
kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="gpt-4o-mini"
|
||
)
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
# Both sources should be merged, no duplicates
|
||
assert "user-tag" in kwargs["metadata"]["tags"]
|
||
assert "openai-account" in kwargs["metadata"]["tags"]
|
||
assert "shared-tag" in kwargs["metadata"]["tags"]
|
||
assert kwargs["metadata"]["tags"].count("shared-tag") == 1
|
||
|
||
|
||
def test_update_kwargs_with_deployment_no_tags():
|
||
"""
|
||
Test that when deployment has no tags, kwargs metadata is not affected.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4o-mini",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "fake-key",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {"metadata": {}}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="gpt-4o-mini"
|
||
)
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
# No tags key should be added if deployment has no tags
|
||
assert "tags" not in kwargs["metadata"]
|
||
|
||
|
||
def test_update_kwargs_with_deployment_merges_tools():
|
||
"""
|
||
Test that when both deployment litellm_params and request have tools,
|
||
they are merged (deployment tools first, then request tools).
|
||
|
||
Supports proxy-configured tools (e.g. for o3 deep research) merged with
|
||
client-provided tools.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "o3-deep-research",
|
||
"litellm_params": {
|
||
"model": "openai/o3-deep-research",
|
||
"api_key": "fake-key",
|
||
"tools": [{"type": "web_search"}],
|
||
"tool_choice": "auto",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {
|
||
"metadata": {},
|
||
"tools": [
|
||
{
|
||
"type": "function",
|
||
"function": {"name": "get_weather", "description": "Get weather"},
|
||
},
|
||
],
|
||
}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="o3-deep-research"
|
||
)
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
# Tools should be merged: deployment first, then request
|
||
assert "tools" in kwargs
|
||
assert len(kwargs["tools"]) == 2
|
||
assert kwargs["tools"][0] == {"type": "web_search"}
|
||
assert kwargs["tools"][1]["function"]["name"] == "get_weather"
|
||
# tool_choice from request (none) - deployment's should be used
|
||
assert kwargs["tool_choice"] == "auto"
|
||
|
||
|
||
def test_update_kwargs_with_deployment_merge_tools_deployment_only():
|
||
"""
|
||
Test that when only deployment has tools, they are applied to kwargs.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "o3-deep-research",
|
||
"litellm_params": {
|
||
"model": "openai/o3-deep-research",
|
||
"api_key": "fake-key",
|
||
"tools": [{"type": "web_search"}],
|
||
"tool_choice": "required",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {"metadata": {}}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="o3-deep-research"
|
||
)
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
assert kwargs["tools"] == [{"type": "web_search"}]
|
||
assert kwargs["tool_choice"] == "required"
|
||
|
||
|
||
def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice():
|
||
"""
|
||
Test that when request has tool_choice, it overrides deployment's.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "o3-deep-research",
|
||
"litellm_params": {
|
||
"model": "openai/o3-deep-research",
|
||
"api_key": "fake-key",
|
||
"tools": [{"type": "web_search"}],
|
||
"tool_choice": "auto",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {
|
||
"metadata": {},
|
||
"tool_choice": "none",
|
||
}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="o3-deep-research"
|
||
)
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
# Request tool_choice should be preserved (merged tools still applied)
|
||
assert kwargs["tool_choice"] == "none"
|
||
|
||
|
||
def test_credential_name_injected_as_tag():
|
||
"""
|
||
Test that litellm_credential_name from deployment litellm_params
|
||
is injected as a tag into metadata during _update_kwargs_with_deployment.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "xai-model",
|
||
"litellm_params": {
|
||
"model": "xai/grok-4-1-fast",
|
||
"litellm_credential_name": "xAI",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {"metadata": {"tags": ["A.101"]}}
|
||
deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model")
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
assert "Credential: xAI" in kwargs["metadata"]["tags"]
|
||
assert "A.101" in kwargs["metadata"]["tags"]
|
||
|
||
|
||
def test_credential_name_not_duplicated_in_tags():
|
||
"""
|
||
Test that if the credential tag already exists in the tags list,
|
||
it is not duplicated.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "xai-model",
|
||
"litellm_params": {
|
||
"model": "xai/grok-4-1-fast",
|
||
"litellm_credential_name": "xAI",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}}
|
||
deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model")
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1
|
||
|
||
|
||
def test_credential_name_not_injected_when_absent():
|
||
"""
|
||
Test that when no litellm_credential_name is set, tags are unchanged.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-model",
|
||
"litellm_params": {
|
||
"model": "gpt-4o",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {"metadata": {"tags": ["A.101"]}}
|
||
deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model")
|
||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||
|
||
assert kwargs["metadata"]["tags"] == ["A.101"]
|
||
|
||
|
||
def test_update_kwargs_with_deployment_model_info_in_litellm_metadata():
|
||
"""For generic_api_call, model_info with pricing must go to litellm_metadata.
|
||
|
||
Routes like /messages and /responses use generic_api_call which stores
|
||
model_info under litellm_metadata. Regression test for #23185.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "claude-sonnet-4",
|
||
"litellm_params": {
|
||
"model": "anthropic/claude-sonnet-4-20250514",
|
||
"api_key": "fake-key",
|
||
},
|
||
"model_info": {
|
||
"id": "custom-pricing-id",
|
||
"input_cost_per_token": 0.0003,
|
||
"output_cost_per_token": 0.0015,
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="claude-sonnet-4"
|
||
)
|
||
router._update_kwargs_with_deployment(
|
||
deployment=deployment, kwargs=kwargs, function_name="generic_api_call"
|
||
)
|
||
|
||
assert "litellm_metadata" in kwargs
|
||
model_info = kwargs["litellm_metadata"]["model_info"]
|
||
assert model_info["id"] == "custom-pricing-id"
|
||
assert model_info["input_cost_per_token"] == 0.0003
|
||
assert model_info["output_cost_per_token"] == 0.0015
|
||
|
||
|
||
def test_update_kwargs_with_deployment_model_info_in_metadata():
|
||
"""For acompletion (function_name=None), model_info goes to metadata.
|
||
|
||
/chat/completions uses acompletion which stores model_info under metadata.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "claude-sonnet-4",
|
||
"litellm_params": {
|
||
"model": "anthropic/claude-sonnet-4-20250514",
|
||
"api_key": "fake-key",
|
||
},
|
||
"model_info": {
|
||
"id": "custom-pricing-id",
|
||
"input_cost_per_token": 0.0003,
|
||
"output_cost_per_token": 0.0015,
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
kwargs: dict = {}
|
||
deployment = router.get_deployment_by_model_group_name(
|
||
model_group_name="claude-sonnet-4"
|
||
)
|
||
router._update_kwargs_with_deployment(
|
||
deployment=deployment, kwargs=kwargs, function_name=None
|
||
)
|
||
|
||
assert "metadata" in kwargs
|
||
model_info = kwargs["metadata"]["model_info"]
|
||
assert model_info["id"] == "custom-pricing-id"
|
||
assert model_info["input_cost_per_token"] == 0.0003
|
||
assert model_info["output_cost_per_token"] == 0.0015
|
||
|
||
|
||
def test_combine_fallback_usage():
|
||
"""Test that _combine_fallback_usage merges partial and fallback usage."""
|
||
from litellm.router import Router
|
||
from litellm.types.utils import Usage
|
||
|
||
# Create a stream chunk with usage
|
||
chunk = litellm.ModelResponseStream(
|
||
id="test",
|
||
model="gpt-4o",
|
||
choices=[],
|
||
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||
)
|
||
|
||
# Call _combine_fallback_usage with no extra usage
|
||
Router._combine_fallback_usage(chunk, None)
|
||
assert chunk.usage is not None
|
||
assert chunk.usage.prompt_tokens == 10
|
||
assert chunk.usage.completion_tokens == 5
|
||
assert chunk.usage.total_tokens == 15
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure():
|
||
"""A mid-stream failure with no successful fallback raises and is logged as
|
||
a failure, so the router must never dispatch it as a success. Partial-spend
|
||
recovery for the failure row happens in the streaming handler, not here, so
|
||
this guards only against reintroducing a success log for a failed stream.
|
||
"""
|
||
from litellm.exceptions import MidStreamFallbackError
|
||
from litellm.types.utils import Delta, StreamingChoices, Usage
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"},
|
||
},
|
||
],
|
||
set_verbose=True,
|
||
)
|
||
|
||
error = MidStreamFallbackError(
|
||
message="Connection lost",
|
||
model="gpt-4",
|
||
llm_provider="openai",
|
||
generated_content="The Roman Empire began when",
|
||
)
|
||
|
||
def _make_interrupted_model_response():
|
||
partial_chunk = litellm.ModelResponseStream(
|
||
id="chatcmpl-partial-1",
|
||
created=1742056047,
|
||
model="gpt-4",
|
||
object="chat.completion.chunk",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
index=0,
|
||
delta=Delta(content="The Roman Empire began when", role="assistant"),
|
||
)
|
||
],
|
||
usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26),
|
||
)
|
||
|
||
class _RaisingStream:
|
||
def __init__(self):
|
||
self.index = 0
|
||
self.chunks = [partial_chunk]
|
||
|
||
def __aiter__(self):
|
||
return self
|
||
|
||
async def __anext__(self):
|
||
if self.index == 0:
|
||
self.index += 1
|
||
return partial_chunk
|
||
raise error
|
||
|
||
stream = _RaisingStream()
|
||
logging_obj = MagicMock()
|
||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||
logging_obj.model_call_details = {}
|
||
setattr(stream, "model", "gpt-4")
|
||
setattr(stream, "custom_llm_provider", "openai")
|
||
setattr(stream, "logging_obj", logging_obj)
|
||
return stream, logging_obj
|
||
|
||
messages = [{"role": "user", "content": "Hello"}]
|
||
initial_kwargs = {"model": "gpt-4", "stream": True}
|
||
|
||
# Terminal path: no successful fallback -> the error propagates and the
|
||
# router never dispatches a success for the failed stream.
|
||
model_response, logging_obj = _make_interrupted_model_response()
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
new=AsyncMock(side_effect=error),
|
||
):
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=model_response,
|
||
messages=messages,
|
||
initial_kwargs=dict(initial_kwargs),
|
||
)
|
||
collected = []
|
||
with pytest.raises(MidStreamFallbackError):
|
||
async for chunk in result:
|
||
collected.append(chunk)
|
||
|
||
assert len(collected) == 1
|
||
logging_obj.dispatch_success_handlers.assert_not_called()
|
||
|
||
# Mid-stream errors with generated content are now re-raised immediately;
|
||
# no continuation-prompt fallback is attempted. Success handlers must
|
||
# still not be dispatched in this path.
|
||
model_response, logging_obj = _make_interrupted_model_response()
|
||
|
||
with patch.object(
|
||
router,
|
||
"async_function_with_fallbacks_common_utils",
|
||
new=AsyncMock(),
|
||
) as mock_fallback:
|
||
result = await router._acompletion_streaming_iterator(
|
||
model_response=model_response,
|
||
messages=messages,
|
||
initial_kwargs=dict(initial_kwargs),
|
||
)
|
||
collected = []
|
||
with pytest.raises(MidStreamFallbackError):
|
||
async for chunk in result:
|
||
collected.append(chunk)
|
||
|
||
assert len(collected) == 1, "only the partial chunk before the error"
|
||
mock_fallback.assert_not_called()
|
||
logging_obj.dispatch_success_handlers.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_team_scoped_model_fallback():
|
||
"""
|
||
Test that fallback works correctly for team-scoped models.
|
||
|
||
When a team-scoped model fails and the fallback model is also team-scoped,
|
||
the router should find the fallback deployment by matching team_public_model_name.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "team-a-primary-internal",
|
||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"},
|
||
"model_info": {
|
||
"team_id": "team-a",
|
||
"team_public_model_name": "primary-model",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "team-a-fallback-internal",
|
||
"litellm_params": {
|
||
"model": "gpt-4",
|
||
"api_key": "fake",
|
||
"mock_response": "fallback success from team-a",
|
||
},
|
||
"model_info": {
|
||
"team_id": "team-a",
|
||
"team_public_model_name": "fallback-model",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"primary-model": ["fallback-model"]}],
|
||
)
|
||
|
||
response = await router.acompletion(
|
||
model="primary-model",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
metadata={"user_api_key_team_id": "team-a"},
|
||
mock_testing_fallbacks=True,
|
||
)
|
||
assert response is not None
|
||
assert response.choices[0].message.content == "fallback success from team-a"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_team_scoped_model_fallback_to_global():
|
||
"""
|
||
Test that a team-scoped model can fall back to a global (non-team) model.
|
||
|
||
Global models (no team_id on deployment) should be accessible as fallback
|
||
targets for team-scoped requests.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "team-a-primary-internal",
|
||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"},
|
||
"model_info": {
|
||
"team_id": "team-a",
|
||
"team_public_model_name": "primary-model",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "global-fallback",
|
||
"litellm_params": {
|
||
"model": "gpt-4",
|
||
"api_key": "fake",
|
||
"mock_response": "global fallback success",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"primary-model": ["global-fallback"]}],
|
||
)
|
||
|
||
response = await router.acompletion(
|
||
model="primary-model",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
metadata={"user_api_key_team_id": "team-a"},
|
||
mock_testing_fallbacks=True,
|
||
)
|
||
assert response is not None
|
||
assert response.choices[0].message.content == "global fallback success"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_team_scoped_model_fallback_cross_team_blocked():
|
||
"""
|
||
Test that cross-team fallback is correctly blocked.
|
||
|
||
When team-a's model fails and the fallback target is scoped to team-b,
|
||
the router should NOT use it (team isolation).
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "team-a-primary-internal",
|
||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"},
|
||
"model_info": {
|
||
"team_id": "team-a",
|
||
"team_public_model_name": "primary-model",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "team-b-fallback-internal",
|
||
"litellm_params": {
|
||
"model": "gpt-4",
|
||
"api_key": "fake",
|
||
"mock_response": "team-b response - should not reach here",
|
||
},
|
||
"model_info": {
|
||
"team_id": "team-b",
|
||
"team_public_model_name": "fallback-model",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"primary-model": ["fallback-model"]}],
|
||
)
|
||
|
||
with pytest.raises(Exception):
|
||
await router.acompletion(
|
||
model="primary-model",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
metadata={"user_api_key_team_id": "team-a"},
|
||
mock_testing_fallbacks=True,
|
||
)
|
||
|
||
|
||
def test_get_all_deployments_with_team_id():
|
||
"""
|
||
Test that _get_all_deployments with team_id can find deployments
|
||
by team_public_model_name when the model_name is not in the index.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "internal-team-deployment",
|
||
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
|
||
"model_info": {
|
||
"team_id": "team-x",
|
||
"team_public_model_name": "gpt-4",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# Without team_id: "gpt-4" is not in the model_name index (internal name is different)
|
||
deployments = router._get_all_deployments(model_name="gpt-4")
|
||
assert len(deployments) == 0
|
||
|
||
# With correct team_id: should find via O(n) scan matching team_public_model_name
|
||
deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-x")
|
||
assert len(deployments) == 1
|
||
assert deployments[0]["model_name"] == "internal-team-deployment"
|
||
|
||
# With wrong team_id: should find nothing
|
||
deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-y")
|
||
assert len(deployments) == 0
|
||
|
||
|
||
def test_multiregion_team_deployments_unique_model_names():
|
||
"""
|
||
Simulates athenahealth's exact setup: unique model_names per deployment,
|
||
same team_public_model_name, multiple regions.
|
||
|
||
Verifies that _get_all_deployments returns ALL regional deployments
|
||
for a team when queried by team_public_model_name.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "metis-claude-us-east-1",
|
||
"litellm_params": {
|
||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||
"aws_region_name": "us-east-1",
|
||
"api_key": "fake",
|
||
},
|
||
"model_info": {
|
||
"team_id": "metis-team",
|
||
"team_public_model_name": "claude-sonnet",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "metis-claude-us-west-2",
|
||
"litellm_params": {
|
||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||
"aws_region_name": "us-west-2",
|
||
"api_key": "fake",
|
||
},
|
||
"model_info": {
|
||
"team_id": "metis-team",
|
||
"team_public_model_name": "claude-sonnet",
|
||
},
|
||
},
|
||
],
|
||
)
|
||
|
||
# "claude-sonnet" is NOT in the model_name index
|
||
assert "claude-sonnet" not in router.model_names
|
||
|
||
# Without team_id: returns nothing (no model_name="claude-sonnet" in index, no O(n) scan)
|
||
deployments = router._get_all_deployments(model_name="claude-sonnet")
|
||
assert len(deployments) == 0
|
||
|
||
# With team_id: O(n) scan finds BOTH regional deployments
|
||
deployments = router._get_all_deployments(
|
||
model_name="claude-sonnet", team_id="metis-team"
|
||
)
|
||
assert len(deployments) == 2
|
||
deployment_names = {d["model_name"] for d in deployments}
|
||
assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"}
|
||
|
||
# Each deployment has a unique ID (critical for cooldown/retry to work)
|
||
deployment_ids = {d["model_info"]["id"] for d in deployments}
|
||
assert (
|
||
len(deployment_ids) == 2
|
||
), "Each deployment must have a unique ID for cooldown tracking"
|
||
|
||
# Wrong team: returns nothing
|
||
deployments = router._get_all_deployments(
|
||
model_name="claude-sonnet", team_id="other-team"
|
||
)
|
||
assert len(deployments) == 0
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_multiregion_team_failover_between_regions():
|
||
"""
|
||
Simulates athenahealth's multiregion failover scenario:
|
||
- Two Bedrock deployments (us-east-1 and us-west-2) with unique model_names
|
||
- Same team_public_model_name ("claude-sonnet")
|
||
- Primary region fails → router should failover to second region
|
||
|
||
This is the exact scenario Sean Glover from athenahealth will demonstrate.
|
||
"""
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "metis-claude-us-east-1",
|
||
"litellm_params": {
|
||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||
"api_key": "fake",
|
||
"mock_response": "response from us-east-1",
|
||
},
|
||
"model_info": {
|
||
"team_id": "metis-team",
|
||
"team_public_model_name": "claude-sonnet",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "metis-claude-us-west-2",
|
||
"litellm_params": {
|
||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||
"api_key": "fake",
|
||
"mock_response": "response from us-west-2",
|
||
},
|
||
"model_info": {
|
||
"team_id": "metis-team",
|
||
"team_public_model_name": "claude-sonnet",
|
||
},
|
||
},
|
||
],
|
||
num_retries=1,
|
||
)
|
||
|
||
# Verify the router finds both deployments for the team
|
||
deployments = router._get_all_deployments(
|
||
model_name="claude-sonnet", team_id="metis-team"
|
||
)
|
||
assert (
|
||
len(deployments) == 2
|
||
), "Router must find both regional deployments by team_public_model_name"
|
||
|
||
# Make a normal request — should succeed from one of the regions
|
||
response = await router.acompletion(
|
||
model="claude-sonnet",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
metadata={"user_api_key_team_id": "metis-team"},
|
||
)
|
||
assert response is not None
|
||
assert response.choices[0].message.content in [
|
||
"response from us-east-1",
|
||
"response from us-west-2",
|
||
]
|
||
|
||
|
||
def test_access_group_scoped_key_filters_deployments_with_same_public_model():
|
||
"""
|
||
If a key can access a model only via access group membership,
|
||
router candidate deployments for that public model should be constrained
|
||
to deployments in the allowed access group.
|
||
"""
|
||
from litellm.proxy._types import UserAPIKeyAuth
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-5.1",
|
||
"api_key": "key1",
|
||
"mock_response": "response-via-AG1",
|
||
},
|
||
"model_info": {"access_groups": ["AG1"]},
|
||
},
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o",
|
||
"api_key": "key2",
|
||
"mock_response": "response-via-AG2",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
]
|
||
)
|
||
|
||
scoped_key = UserAPIKeyAuth(
|
||
api_key="hashed-key",
|
||
team_id="team2",
|
||
models=["AG2"],
|
||
team_models=["AG2"],
|
||
)
|
||
|
||
_model, deployments = router._common_checks_available_deployment(
|
||
model="gpt-5",
|
||
request_kwargs={
|
||
"metadata": {
|
||
"user_api_key_team_id": "team2",
|
||
"user_api_key_auth": scoped_key,
|
||
}
|
||
},
|
||
)
|
||
|
||
assert len(deployments) == 1
|
||
assert deployments[0].get("model_info", {}).get("access_groups") == ["AG2"]
|
||
|
||
seen = set()
|
||
for _ in range(20):
|
||
response = router.completion(
|
||
model="gpt-5",
|
||
messages=[{"role": "user", "content": "hello"}],
|
||
metadata={"user_api_key_team_id": "team2", "user_api_key_auth": scoped_key},
|
||
)
|
||
seen.add(response.choices[0].message.content)
|
||
|
||
assert seen == {"response-via-AG2"}
|
||
|
||
|
||
def test_explicit_model_access_does_not_force_access_group_filtering():
|
||
"""
|
||
If a key has explicit model access in addition to access group entries,
|
||
do not force access-group-only filtering for deployment selection.
|
||
"""
|
||
from litellm.proxy._types import UserAPIKeyAuth
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-5.1",
|
||
"api_key": "key1",
|
||
"mock_response": "response-via-AG1",
|
||
},
|
||
"model_info": {"access_groups": ["AG1"]},
|
||
},
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o",
|
||
"api_key": "key2",
|
||
"mock_response": "response-via-AG2",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
]
|
||
)
|
||
|
||
explicit_key = UserAPIKeyAuth(
|
||
api_key="hashed-key",
|
||
team_id="team2",
|
||
models=["AG2", "gpt-5"],
|
||
team_models=["AG2", "gpt-5"],
|
||
)
|
||
|
||
_model, deployments = router._common_checks_available_deployment(
|
||
model="gpt-5",
|
||
request_kwargs={
|
||
"metadata": {
|
||
"user_api_key_team_id": "team2",
|
||
"user_api_key_auth": explicit_key,
|
||
}
|
||
},
|
||
)
|
||
|
||
deployment_groups = [
|
||
d.get("model_info", {}).get("access_groups") for d in deployments
|
||
]
|
||
assert ["AG1"] in deployment_groups
|
||
assert ["AG2"] in deployment_groups
|
||
|
||
|
||
def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
):
|
||
"""
|
||
When access-group filtering removes all candidates, _get_deployment_by_litellm_model
|
||
must not run: it does not re-apply access groups and could return blocked deployments
|
||
that share the same litellm_params.model as the request model string.
|
||
|
||
``get_model_access_groups`` is patched to expose AG1 for the public model (so the
|
||
access-group filter runs with a non-empty allowed set) while every deployment
|
||
returned for that name is AG2-only — filtered to empty. Without the guard, the
|
||
litellm-model fallback would return both rows because ``litellm_params.model`` matches.
|
||
"""
|
||
from litellm.proxy._types import UserAPIKeyAuth
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "gpt-5",
|
||
"api_key": "key1",
|
||
"mock_response": "blocked-dep-1",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "gpt-5",
|
||
"api_key": "key2",
|
||
"mock_response": "blocked-dep-2",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
]
|
||
)
|
||
|
||
orig_groups = router.get_model_access_groups
|
||
|
||
def fake_get_model_access_groups(
|
||
model_name=None, model_access_group=None, team_id=None
|
||
):
|
||
if model_name == "gpt-5" and model_access_group is None:
|
||
return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]}
|
||
return orig_groups(
|
||
model_name=model_name,
|
||
model_access_group=model_access_group,
|
||
team_id=team_id,
|
||
)
|
||
|
||
monkeypatch.setattr(router, "get_model_access_groups", fake_get_model_access_groups)
|
||
|
||
scoped_key = UserAPIKeyAuth(
|
||
api_key="hashed-key",
|
||
team_id="team2",
|
||
models=["AG1"],
|
||
team_models=["AG1"],
|
||
)
|
||
|
||
with pytest.raises(litellm.BadRequestError):
|
||
router._common_checks_available_deployment(
|
||
model="gpt-5",
|
||
request_kwargs={
|
||
"metadata": {
|
||
"user_api_key_team_id": "team2",
|
||
"user_api_key_auth": scoped_key,
|
||
}
|
||
},
|
||
)
|
||
|
||
|
||
def test_access_group_block_does_not_silently_use_default_fallback_model(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
):
|
||
"""
|
||
When access-group filtering empties candidates for model X, the router must not use
|
||
``fallbacks`` default ``*`` routing to model Y: Y may have no ``access_groups``, so
|
||
``_filter_deployments_by_model_access_groups`` would not constrain Y and the caller
|
||
would be served despite being blocked from X.
|
||
"""
|
||
from litellm.proxy._types import UserAPIKeyAuth
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "gpt-5",
|
||
"api_key": "key1",
|
||
"mock_response": "blocked-dep-1",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "gpt-5",
|
||
"api_key": "key2",
|
||
"mock_response": "blocked-dep-2",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
{
|
||
"model_name": "gpt-4-fallback",
|
||
"litellm_params": {
|
||
"model": "gpt-4",
|
||
"api_key": "fallback-key",
|
||
"mock_response": "should-not-reach",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"*": ["gpt-4-fallback"]}],
|
||
)
|
||
|
||
orig_groups = router.get_model_access_groups
|
||
|
||
def fake_get_model_access_groups(
|
||
model_name=None, model_access_group=None, team_id=None
|
||
):
|
||
if model_name == "gpt-5" and model_access_group is None:
|
||
return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]}
|
||
return orig_groups(
|
||
model_name=model_name,
|
||
model_access_group=model_access_group,
|
||
team_id=team_id,
|
||
)
|
||
|
||
monkeypatch.setattr(router, "get_model_access_groups", fake_get_model_access_groups)
|
||
|
||
scoped_key = UserAPIKeyAuth(
|
||
api_key="hashed-key",
|
||
team_id="team2",
|
||
models=["AG1"],
|
||
team_models=["AG1"],
|
||
)
|
||
|
||
with pytest.raises(litellm.BadRequestError):
|
||
router._common_checks_available_deployment(
|
||
model="gpt-5",
|
||
request_kwargs={
|
||
"metadata": {
|
||
"user_api_key_team_id": "team2",
|
||
"user_api_key_auth": scoped_key,
|
||
}
|
||
},
|
||
)
|
||
|
||
|
||
def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallback(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
):
|
||
"""
|
||
When the by-name lookup returns no deployments and the litellm-model fallback
|
||
branch finds candidates that access-group filtering then empties, the router
|
||
must not fall through to default ``fallbacks`` routing — the default fallback
|
||
model may have no ``access_groups`` and would short-circuit the filter,
|
||
silently serving a caller blocked by access-group restrictions.
|
||
"""
|
||
from litellm.proxy._types import UserAPIKeyAuth
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-5-alias",
|
||
"litellm_params": {
|
||
"model": "gpt-5",
|
||
"api_key": "key1",
|
||
"mock_response": "blocked-dep-1",
|
||
},
|
||
"model_info": {"access_groups": ["AG2"]},
|
||
},
|
||
{
|
||
"model_name": "gpt-4-fallback",
|
||
"litellm_params": {
|
||
"model": "gpt-4",
|
||
"api_key": "fallback-key",
|
||
"mock_response": "should-not-reach",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"*": ["gpt-4-fallback"]}],
|
||
)
|
||
|
||
orig_groups = router.get_model_access_groups
|
||
|
||
def fake_get_model_access_groups(
|
||
model_name=None, model_access_group=None, team_id=None
|
||
):
|
||
if model_name == "gpt-5" and model_access_group is None:
|
||
return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]}
|
||
return orig_groups(
|
||
model_name=model_name,
|
||
model_access_group=model_access_group,
|
||
team_id=team_id,
|
||
)
|
||
|
||
monkeypatch.setattr(router, "get_model_access_groups", fake_get_model_access_groups)
|
||
|
||
scoped_key = UserAPIKeyAuth(
|
||
api_key="hashed-key",
|
||
team_id="team2",
|
||
models=["AG1"],
|
||
team_models=["AG1"],
|
||
)
|
||
|
||
with pytest.raises(litellm.BadRequestError):
|
||
router._common_checks_available_deployment(
|
||
model="gpt-5",
|
||
request_kwargs={
|
||
"metadata": {
|
||
"user_api_key_team_id": "team2",
|
||
"user_api_key_auth": scoped_key,
|
||
}
|
||
},
|
||
)
|
||
|
||
|
||
def test_try_early_resolve_deployments_for_model_not_in_names():
|
||
"""
|
||
Direct coverage for ``_try_early_resolve_deployments_for_model_not_in_names``:
|
||
|
||
- Returns ``None`` when the requested model is already in ``self.model_names``
|
||
(the by-name lookup path will handle it).
|
||
- Returns ``None`` when there are no team deployments, no pattern matches, and
|
||
no default deployment to fall back to.
|
||
- Returns the pattern-router match when the model matches a wildcard route.
|
||
- Returns the default deployment with the request model substituted in when one
|
||
is configured, without mutating the stored default.
|
||
"""
|
||
router_in_names = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-5",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-5",
|
||
"api_key": "key1",
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
assert (
|
||
router_in_names._try_early_resolve_deployments_for_model_not_in_names(
|
||
model="gpt-5", request_team_id=None
|
||
)
|
||
is None
|
||
)
|
||
assert (
|
||
router_in_names._try_early_resolve_deployments_for_model_not_in_names(
|
||
model="some-unknown-model", request_team_id=None
|
||
)
|
||
is None
|
||
)
|
||
|
||
pattern_router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "openai/*",
|
||
"litellm_params": {
|
||
"model": "openai/*",
|
||
"api_key": "key-pattern",
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
pattern_result = (
|
||
pattern_router._try_early_resolve_deployments_for_model_not_in_names(
|
||
model="openai/gpt-4o-mini", request_team_id=None
|
||
)
|
||
)
|
||
assert pattern_result is not None
|
||
resolved_model, pattern_deployments = pattern_result
|
||
assert resolved_model == "openai/gpt-4o-mini"
|
||
assert isinstance(pattern_deployments, list) and len(pattern_deployments) == 1
|
||
|
||
default_router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "named-model",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o",
|
||
"api_key": "key-named",
|
||
},
|
||
},
|
||
]
|
||
)
|
||
default_router.default_deployment = {
|
||
"model_name": "default",
|
||
"litellm_params": {
|
||
"model": "openai/will-be-overridden",
|
||
"api_key": "key-default",
|
||
},
|
||
}
|
||
|
||
default_result = (
|
||
default_router._try_early_resolve_deployments_for_model_not_in_names(
|
||
model="brand-new-model", request_team_id=None
|
||
)
|
||
)
|
||
assert default_result is not None
|
||
resolved_model, default_deployment = default_result
|
||
assert resolved_model == "brand-new-model"
|
||
assert isinstance(default_deployment, dict)
|
||
assert default_deployment["litellm_params"]["model"] == "brand-new-model"
|
||
# The original default_deployment must not be mutated.
|
||
assert (
|
||
default_router.default_deployment["litellm_params"]["model"]
|
||
== "openai/will-be-overridden"
|
||
)
|
||
|
||
|
||
def _router_with_two_deployments(blocked_flags):
|
||
import litellm
|
||
|
||
model_list = []
|
||
for idx, blocked in enumerate(blocked_flags):
|
||
model_list.append(
|
||
{
|
||
"model_name": "gpt-4o",
|
||
"litellm_params": {"model": f"openai/gpt-4o-{idx}"},
|
||
"model_info": {"id": f"dep-{idx}", "blocked": blocked},
|
||
}
|
||
)
|
||
return litellm.Router(model_list=model_list)
|
||
|
||
|
||
def test_get_fully_blocked_model_names_marks_name_when_all_deployments_blocked():
|
||
router = _router_with_two_deployments([True, True])
|
||
assert router.get_fully_blocked_model_names() == {"gpt-4o"}
|
||
|
||
|
||
def test_get_fully_blocked_model_names_keeps_name_when_partial_blocked():
|
||
router = _router_with_two_deployments([True, False])
|
||
assert router.get_fully_blocked_model_names() == set()
|
||
|
||
|
||
def test_get_fully_blocked_model_names_treats_missing_key_as_unblocked():
|
||
import litellm
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4o",
|
||
"litellm_params": {"model": "openai/gpt-4o"},
|
||
"model_info": {"id": "dep-0"},
|
||
}
|
||
]
|
||
)
|
||
assert router.get_fully_blocked_model_names() == set()
|
||
|
||
|
||
def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None):
|
||
import time
|
||
|
||
ts = timestamp if timestamp is not None else time.time()
|
||
router.health_state_cache.set_deployment_health_states(
|
||
{
|
||
uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"}
|
||
for uid in unhealthy_ids
|
||
}
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_fully_unhealthy_model_names_marks_name_when_all_unhealthy():
|
||
router = _router_with_two_deployments([False, False])
|
||
_seed_unhealthy_states(router, {"dep-0", "dep-1"})
|
||
assert await router.async_get_fully_unhealthy_model_names() == {"gpt-4o"}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_fully_unhealthy_model_names_keeps_name_when_partial():
|
||
router = _router_with_two_deployments([False, False])
|
||
_seed_unhealthy_states(router, {"dep-0"})
|
||
assert await router.async_get_fully_unhealthy_model_names() == set()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_fully_unhealthy_model_names_empty_without_health_state():
|
||
router = _router_with_two_deployments([False, False])
|
||
assert await router.async_get_fully_unhealthy_model_names() == set()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_fully_unhealthy_model_names_ignores_stale_state():
|
||
import time
|
||
|
||
router = _router_with_two_deployments([False, False])
|
||
stale_ts = time.time() - (router.health_state_cache.staleness_threshold + 10)
|
||
_seed_unhealthy_states(router, {"dep-0", "dep-1"}, timestamp=stale_ts)
|
||
assert await router.async_get_fully_unhealthy_model_names() == set()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_fully_unhealthy_model_names_includes_team_alias():
|
||
import litellm
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "gpt-4o",
|
||
"litellm_params": {"model": "openai/gpt-4o"},
|
||
"model_info": {
|
||
"id": "dep-0",
|
||
"team_id": "team-1",
|
||
"team_public_model_name": "team-gpt",
|
||
},
|
||
}
|
||
]
|
||
)
|
||
_seed_unhealthy_states(router, {"dep-0"})
|
||
assert await router.async_get_fully_unhealthy_model_names() == {
|
||
"gpt-4o",
|
||
"team-gpt",
|
||
}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_policy():
|
||
from litellm.types.router import AllowedFailsPolicy
|
||
|
||
router = _router_with_two_deployments([False, False])
|
||
router.allowed_fails_policy = AllowedFailsPolicy(BadRequestErrorAllowedFails=1)
|
||
_seed_unhealthy_states(router, {"dep-0", "dep-1"})
|
||
assert await router.async_get_fully_unhealthy_model_names() == set()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_async_get_healthy_deployments_skips_blocked_deployment():
|
||
router = _router_with_two_deployments([True, False])
|
||
healthy, all_dep = await router._async_get_healthy_deployments(
|
||
model="gpt-4o", parent_otel_span=None
|
||
)
|
||
healthy_ids = [d["model_info"]["id"] for d in healthy]
|
||
assert "dep-0" not in healthy_ids
|
||
assert "dep-1" in healthy_ids
|
||
assert len(all_dep) == 2
|
||
|
||
|
||
def test_get_healthy_deployments_sync_skips_blocked_deployment():
|
||
router = _router_with_two_deployments([False, True])
|
||
healthy, all_dep = router._get_healthy_deployments(
|
||
model="gpt-4o", parent_otel_span=None
|
||
)
|
||
healthy_ids = [d["model_info"]["id"] for d in healthy]
|
||
assert "dep-0" in healthy_ids
|
||
assert "dep-1" not in healthy_ids
|
||
assert len(all_dep) == 2
|
||
|
||
|
||
def test_filter_blocked_deployments_drops_blocked_keeps_unblocked():
|
||
router = _router_with_two_deployments([True, False])
|
||
filtered = router._filter_blocked_deployments(router.get_model_list() or [])
|
||
ids = [d["model_info"]["id"] for d in filtered]
|
||
assert ids == ["dep-1"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path():
|
||
router = _router_with_two_deployments([True, False])
|
||
deployments = await router.async_get_healthy_deployments(
|
||
model="gpt-4o", request_kwargs={}
|
||
)
|
||
assert isinstance(deployments, list)
|
||
ids = [d["model_info"]["id"] for d in deployments]
|
||
assert "dep-0" not in ids
|
||
assert "dep-1" in ids
|
||
|
||
|
||
def test_public_get_available_deployment_skips_blocked_on_primary_path():
|
||
router = _router_with_two_deployments([True, False])
|
||
deployment = router.get_available_deployment(model="gpt-4o", request_kwargs={})
|
||
assert deployment["model_info"]["id"] == "dep-1"
|
||
|
||
|
||
def test_get_available_deployment_raises_when_addressed_dict_is_blocked():
|
||
import litellm
|
||
|
||
router = _router_with_two_deployments([True, True])
|
||
with pytest.raises(litellm.ServiceUnavailableError):
|
||
router.get_available_deployment(model="dep-0", request_kwargs={})
|
||
|
||
|
||
def _router_with_two_pass_through_deployments(blocked_flags):
|
||
import litellm
|
||
|
||
model_list = []
|
||
for idx, blocked in enumerate(blocked_flags):
|
||
model_list.append(
|
||
{
|
||
"model_name": "gpt-4o",
|
||
"litellm_params": {
|
||
"model": f"openai/gpt-4o-{idx}",
|
||
"api_key": "sk-fake-for-tests",
|
||
"use_in_pass_through": True,
|
||
},
|
||
"model_info": {"id": f"pt-{idx}", "blocked": blocked},
|
||
}
|
||
)
|
||
return litellm.Router(model_list=model_list)
|
||
|
||
|
||
def test_get_available_deployment_for_pass_through_skips_blocked():
|
||
router = _router_with_two_pass_through_deployments([True, False])
|
||
deployment = router.get_available_deployment_for_pass_through(
|
||
model="gpt-4o", request_kwargs={}
|
||
)
|
||
assert deployment["model_info"]["id"] == "pt-1"
|
||
|
||
|
||
def test_get_available_deployment_for_pass_through_raises_when_dict_blocked():
|
||
import litellm
|
||
|
||
router = _router_with_two_pass_through_deployments([True, True])
|
||
with pytest.raises(litellm.ServiceUnavailableError):
|
||
router.get_available_deployment_for_pass_through(
|
||
model="pt-0", request_kwargs={}
|
||
)
|
||
|
||
|
||
def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment():
|
||
"""
|
||
Bedrock deployments using IAM/OIDC auth have no api_key; pass-through
|
||
init must not raise and drop them from routing (#27728).
|
||
"""
|
||
import litellm
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "bedrock-claude",
|
||
"litellm_params": {
|
||
"model": "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||
"aws_role_name": "arn:aws:iam::123456789012:role/my-role",
|
||
"aws_session_name": "my-session",
|
||
"use_in_pass_through": True,
|
||
},
|
||
"model_info": {"id": "bedrock-iam-pt"},
|
||
}
|
||
]
|
||
)
|
||
assert [m["model_info"]["id"] for m in router.get_model_list()] == [
|
||
"bedrock-iam-pt"
|
||
]
|
||
|
||
|
||
def test_pass_through_deployment_api_key_resolves_via_get_credentials():
|
||
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
|
||
PassthroughEndpointRouter,
|
||
)
|
||
|
||
router = _router_with_two_pass_through_deployments([False, False])
|
||
passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router)
|
||
assert len(router.get_model_list()) == 2
|
||
assert (
|
||
passthrough_router.get_credentials(
|
||
custom_llm_provider="openai", region_name=None
|
||
)
|
||
== "sk-fake-for-tests"
|
||
)
|
||
|
||
|
||
def test_get_deployment_credentials_returns_none_for_blocked_deployment():
|
||
router = _router_with_two_deployments([True, False])
|
||
assert router.get_deployment_credentials(model_id="dep-0") is None
|
||
assert router.get_deployment_credentials(model_id="dep-1") is not None
|
||
|
||
|
||
def test_get_deployment_credentials_with_provider_returns_none_for_blocked_deployment():
|
||
router = _router_with_two_deployments([True, False])
|
||
assert router.get_deployment_credentials_with_provider(model_id="dep-0") is None
|
||
assert router.get_deployment_credentials_with_provider(model_id="dep-1") is not None
|
||
|
||
|
||
def test_is_deployment_blocked_static_helper_reflects_blocked_flag():
|
||
"""
|
||
Exercises Router._is_deployment_blocked so router_code_coverage.py (AST call graph)
|
||
marks the helper as covered by router-named tests.
|
||
"""
|
||
import types
|
||
|
||
import litellm
|
||
|
||
router = _router_with_two_deployments([True, False])
|
||
blocked_dep = router.get_deployment("dep-0")
|
||
unblocked_dep = router.get_deployment("dep-1")
|
||
assert blocked_dep is not None and unblocked_dep is not None
|
||
assert litellm.Router._is_deployment_blocked(blocked_dep) is True
|
||
assert litellm.Router._is_deployment_blocked(unblocked_dep) is False
|
||
|
||
# No model_info on deployment object → treated as not blocked
|
||
assert litellm.Router._is_deployment_blocked(object()) is False
|
||
missing_blocked = types.SimpleNamespace()
|
||
assert (
|
||
litellm.Router._is_deployment_blocked(
|
||
types.SimpleNamespace(model_info=missing_blocked)
|
||
)
|
||
is False
|
||
)
|
||
assert (
|
||
litellm.Router._is_deployment_blocked(
|
||
types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))
|
||
)
|
||
is True
|
||
)
|
||
|
||
|
||
class TestRouterRequestTimeoutPropagation:
|
||
"""litellm_settings.request_timeout must act as an independent per-attempt timeout.
|
||
|
||
Regression for LIT-2369: request_timeout was shadowed by router_settings.timeout,
|
||
so Bedrock (and other provider) calls fell back to the hardcoded 600s httpx
|
||
default instead of the configured value.
|
||
"""
|
||
|
||
def _make_router(self, timeout=None, stream_timeout=None):
|
||
return litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "test-model",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4",
|
||
"api_key": "sk-test",
|
||
},
|
||
}
|
||
],
|
||
timeout=timeout,
|
||
stream_timeout=stream_timeout,
|
||
)
|
||
|
||
@pytest.fixture
|
||
def explicit_request_timeout(self):
|
||
original_value = litellm.request_timeout
|
||
original_flag = litellm.request_timeout_explicitly_set
|
||
litellm.request_timeout = 300
|
||
litellm.request_timeout_explicitly_set = True
|
||
try:
|
||
yield 300
|
||
finally:
|
||
litellm.request_timeout = original_value
|
||
litellm.request_timeout_explicitly_set = original_flag
|
||
|
||
def test_request_timeout_stored_independently_when_both_set(
|
||
self, explicit_request_timeout
|
||
):
|
||
router = self._make_router(timeout=330)
|
||
assert router.timeout == 330
|
||
assert router.request_timeout == 300
|
||
|
||
def test_request_timeout_none_when_not_explicitly_configured(self):
|
||
original_value = litellm.request_timeout
|
||
original_flag = litellm.request_timeout_explicitly_set
|
||
litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||
litellm.request_timeout_explicitly_set = False
|
||
try:
|
||
router = self._make_router(timeout=330)
|
||
assert router.timeout == 330
|
||
assert router.request_timeout is None
|
||
finally:
|
||
litellm.request_timeout = original_value
|
||
litellm.request_timeout_explicitly_set = original_flag
|
||
|
||
def test_non_stream_prefers_request_timeout_over_router_timeout(
|
||
self, explicit_request_timeout
|
||
):
|
||
router = self._make_router(timeout=330)
|
||
assert router._get_non_stream_timeout(kwargs={}, data={}) == 300
|
||
|
||
def test_stream_prefers_request_timeout_over_router_timeout(
|
||
self, explicit_request_timeout
|
||
):
|
||
router = self._make_router(timeout=330)
|
||
# stream=True resolves through _get_stream_timeout; request_timeout must win.
|
||
assert router._get_timeout(kwargs={"stream": True}, data={}) == 300
|
||
|
||
def test_explicit_stream_timeout_still_wins_over_request_timeout(
|
||
self, explicit_request_timeout
|
||
):
|
||
router = self._make_router(timeout=330, stream_timeout=45)
|
||
assert router._get_stream_timeout(kwargs={}, data={}) == 45
|
||
|
||
def test_non_stream_falls_through_to_router_timeout_without_request_timeout(self):
|
||
original_value = litellm.request_timeout
|
||
original_flag = litellm.request_timeout_explicitly_set
|
||
litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||
litellm.request_timeout_explicitly_set = False
|
||
try:
|
||
router = self._make_router(timeout=330)
|
||
assert router._get_non_stream_timeout(kwargs={}, data={}) == 330
|
||
finally:
|
||
litellm.request_timeout = original_value
|
||
litellm.request_timeout_explicitly_set = original_flag
|
||
|
||
def test_per_deployment_timeout_overrides_request_timeout(
|
||
self, explicit_request_timeout
|
||
):
|
||
router = self._make_router(timeout=330)
|
||
assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120
|
||
|
||
def test_per_request_timeout_overrides_request_timeout(
|
||
self, explicit_request_timeout
|
||
):
|
||
router = self._make_router(timeout=330)
|
||
assert (
|
||
router._get_non_stream_timeout(
|
||
kwargs={"timeout": 60}, data={"timeout": 120}
|
||
)
|
||
== 60
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Deferred-stream eager-fetch tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_deferred_stream_wrapper(make_call_fn):
|
||
"""Return a CustomStreamWrapper with completion_stream=None and the given make_call."""
|
||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||
|
||
logging_obj = MagicMock()
|
||
logging_obj.model_call_details = {"litellm_params": {}}
|
||
return CustomStreamWrapper(
|
||
completion_stream=None,
|
||
model="vertex_ai/gemini-2.0-flash",
|
||
logging_obj=logging_obj,
|
||
custom_llm_provider="vertex_ai_beta",
|
||
make_call=make_call_fn,
|
||
)
|
||
|
||
|
||
def _make_router_with_vertex_and_fallback():
|
||
return litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-gemini",
|
||
"litellm_params": {
|
||
"model": "vertex_ai/gemini-2.0-flash",
|
||
"vertex_project": "test-project",
|
||
"vertex_location": "us-central1",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "my-fallback",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "sk-fake",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"my-gemini": ["my-fallback"]}],
|
||
num_retries=0,
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_deferred_stream_error_propagates_through_acompletion():
|
||
"""Regression: a deferred-stream CustomStreamWrapper whose make_call raises a 429
|
||
must propagate the exception from within _acompletion's except block so that
|
||
fail_calls is incremented (i.e., deployment cooldown fires) and the standard
|
||
router fallback chain can handle it.
|
||
|
||
Before the fix, the HTTP call happened inside __anext__ (outside the except block),
|
||
so fail_calls was never incremented.
|
||
"""
|
||
import litellm as _litellm
|
||
|
||
rate_limit_err = _litellm.RateLimitError(
|
||
message="Resource exhausted",
|
||
llm_provider="vertex_ai",
|
||
model="gemini-2.0-flash",
|
||
)
|
||
|
||
async def failing_make_call(**kwargs):
|
||
raise rate_limit_err
|
||
|
||
router = _make_router_with_vertex_and_fallback()
|
||
deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call)
|
||
|
||
with patch(
|
||
"litellm.acompletion",
|
||
new_callable=AsyncMock,
|
||
return_value=deferred_wrapper,
|
||
):
|
||
with pytest.raises(_litellm.RateLimitError):
|
||
await router._acompletion(
|
||
model="vertex_ai/gemini-2.0-flash",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
stream=True,
|
||
specific_deployment=router.model_list[0],
|
||
)
|
||
|
||
model_name = router.model_list[0]["litellm_params"]["model"]
|
||
assert router.fail_calls[model_name] == 1, (
|
||
"fail_calls must be incremented when the deferred HTTP call fails; "
|
||
"without the eager fetch_stream() fix this stays at 0"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_deferred_stream_preserves_original_headers_on_error():
|
||
"""Router is used both by the proxy and directly as an SDK. HTTP-framing headers
|
||
(Content-Length, Transfer-Encoding, ...) must NOT be stripped at this layer, or
|
||
direct SDK callers lose legitimate provider metadata (e.g. content-type,
|
||
proxy-authenticate) that only the proxy's own response construction needs to
|
||
worry about. Stripping happens in the proxy layer instead
|
||
(_handle_llm_api_exception)."""
|
||
import litellm as _litellm
|
||
|
||
err = _litellm.RateLimitError(
|
||
message="Resource exhausted",
|
||
llm_provider="vertex_ai",
|
||
model="gemini-2.0-flash",
|
||
)
|
||
err.headers = {
|
||
"content-length": "42",
|
||
"transfer-encoding": "chunked",
|
||
"content-encoding": "gzip",
|
||
"content-type": "application/json",
|
||
"x-request-id": "abc-123",
|
||
}
|
||
|
||
async def failing_make_call(**kwargs):
|
||
raise err
|
||
|
||
router = _make_router_with_vertex_and_fallback()
|
||
deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call)
|
||
|
||
with patch(
|
||
"litellm.acompletion",
|
||
new_callable=AsyncMock,
|
||
return_value=deferred_wrapper,
|
||
):
|
||
with pytest.raises(_litellm.RateLimitError) as exc_info:
|
||
await router._acompletion(
|
||
model="vertex_ai/gemini-2.0-flash",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
stream=True,
|
||
specific_deployment=router.model_list[0],
|
||
)
|
||
|
||
raised = exc_info.value
|
||
headers = getattr(raised, "headers", {})
|
||
assert headers.get("content-length") == "42"
|
||
assert headers.get("transfer-encoding") == "chunked"
|
||
assert headers.get("content-encoding") == "gzip"
|
||
assert headers.get("content-type") == "application/json"
|
||
assert headers.get("x-request-id") == "abc-123"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_deferred_stream_skipped_when_stream_already_set():
|
||
"""When completion_stream is already populated (non-deferred provider), the eager
|
||
fetch_stream() call must be skipped entirely; no exception should be raised even
|
||
if make_call would fail.
|
||
"""
|
||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||
|
||
async def would_fail(**kwargs):
|
||
raise RuntimeError("should not be called")
|
||
|
||
logging_obj = MagicMock()
|
||
logging_obj.model_call_details = {"litellm_params": {}}
|
||
|
||
async def noop_aiter():
|
||
return
|
||
yield
|
||
|
||
noop_stream = noop_aiter()
|
||
already_set_wrapper = CustomStreamWrapper(
|
||
completion_stream=noop_stream,
|
||
model="openai/gpt-4o",
|
||
logging_obj=logging_obj,
|
||
custom_llm_provider="openai",
|
||
make_call=would_fail,
|
||
)
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-model",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o",
|
||
"api_key": "sk-fake",
|
||
},
|
||
}
|
||
],
|
||
)
|
||
|
||
with patch(
|
||
"litellm.acompletion",
|
||
new_callable=AsyncMock,
|
||
return_value=already_set_wrapper,
|
||
):
|
||
result = await router._acompletion(
|
||
model="openai/gpt-4o",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
stream=True,
|
||
specific_deployment=router.model_list[0],
|
||
)
|
||
|
||
assert result is not None, "should return a streaming wrapper without errors"
|
||
assert already_set_wrapper.completion_stream is noop_stream, "completion_stream must not be re-fetched"
|
||
await noop_stream.aclose()
|
||
|
||
|
||
def test_completion_deferred_stream_error_propagates_through_completion():
|
||
"""Regression: the sync router path needs the same eager fetch as the async one.
|
||
|
||
A deferred-stream CustomStreamWrapper hands back a wrapper whose HTTP call has
|
||
not happened yet, so without fetch_sync_stream() the provider error surfaces on
|
||
first iteration, outside _completion's except block. The deployment is then never
|
||
marked failed and function_with_fallbacks never sees the error.
|
||
"""
|
||
import litellm as _litellm
|
||
|
||
rate_limit_err = _litellm.RateLimitError(
|
||
message="Resource exhausted",
|
||
llm_provider="vertex_ai",
|
||
model="gemini-2.0-flash",
|
||
)
|
||
make_call_invocations = []
|
||
|
||
def failing_make_call(**kwargs):
|
||
make_call_invocations.append(kwargs)
|
||
raise rate_limit_err
|
||
|
||
router = _make_router_with_vertex_and_fallback()
|
||
deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call)
|
||
|
||
with patch("litellm.completion", return_value=deferred_wrapper):
|
||
with pytest.raises(_litellm.RateLimitError):
|
||
router._completion(
|
||
model="vertex_ai/gemini-2.0-flash",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
stream=True,
|
||
specific_deployment=router.model_list[0],
|
||
)
|
||
|
||
assert len(make_call_invocations) == 1, (
|
||
"the deferred HTTP call must run inside _completion's try block; "
|
||
"without the eager fetch_sync_stream() fix it is deferred to first iteration"
|
||
)
|
||
|
||
|
||
def test_completion_deferred_stream_skipped_when_stream_already_set():
|
||
"""A non-deferred sync provider already has completion_stream populated, so the
|
||
eager fetch must be skipped and make_call left untouched.
|
||
"""
|
||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||
|
||
def would_fail(**kwargs):
|
||
raise RuntimeError("should not be called")
|
||
|
||
logging_obj = MagicMock()
|
||
logging_obj.model_call_details = {"litellm_params": {}}
|
||
already_set_stream = iter([])
|
||
|
||
already_set_wrapper = CustomStreamWrapper(
|
||
completion_stream=already_set_stream,
|
||
model="openai/gpt-4o",
|
||
logging_obj=logging_obj,
|
||
custom_llm_provider="openai",
|
||
make_call=would_fail,
|
||
)
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-model",
|
||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"},
|
||
}
|
||
],
|
||
)
|
||
|
||
with patch("litellm.completion", return_value=already_set_wrapper):
|
||
result = router._completion(
|
||
model="openai/gpt-4o",
|
||
messages=[{"role": "user", "content": "Hello"}],
|
||
stream=True,
|
||
specific_deployment=router.model_list[0],
|
||
)
|
||
|
||
assert result is not None, "should return a streaming wrapper without errors"
|
||
assert already_set_wrapper.completion_stream is already_set_stream, "completion_stream must not be re-fetched"
|
||
|
||
|
||
class TestAdvisorSubCallCooldown:
|
||
"""Regression for LIT-4565: an advisor orchestration failure must not cool
|
||
down the selected (healthy) deployment, which would reject unrelated
|
||
callers to the same model group."""
|
||
|
||
def _router(self):
|
||
return litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "claude-sonnet-5",
|
||
"litellm_params": {"model": "bedrock/us.anthropic.claude-opus-4-8"},
|
||
"model_info": {"id": "dep-1"},
|
||
}
|
||
],
|
||
)
|
||
|
||
def _kwargs(self, exception):
|
||
return {
|
||
"exception": exception,
|
||
"litellm_params": {"model_info": {"id": "dep-1"}, "metadata": {}},
|
||
}
|
||
|
||
def _auth_error(self):
|
||
return litellm.AuthenticationError(
|
||
message="x-api-key header is required",
|
||
llm_provider="anthropic",
|
||
model="claude-opus-4-8",
|
||
)
|
||
|
||
def _cooled_down_ids(self, router):
|
||
active = router.cooldown_cache.get_active_cooldowns(
|
||
model_ids=["dep-1"], parent_otel_span=None
|
||
)
|
||
return [entry[0] for entry in active]
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_untagged_auth_error_cools_down_deployment(self):
|
||
from datetime import datetime
|
||
|
||
router = self._router()
|
||
now = datetime.now()
|
||
assert (
|
||
router.deployment_callback_on_failure(
|
||
self._kwargs(self._auth_error()), None, now, now
|
||
)
|
||
is True
|
||
)
|
||
assert "dep-1" in self._cooled_down_ids(router)
|
||
|
||
def test_advisor_orchestration_failure_does_not_cool_down_deployment(self):
|
||
from datetime import datetime
|
||
|
||
from litellm.router_utils.cooldown_handlers import (
|
||
mark_advisor_orchestration_failure,
|
||
)
|
||
|
||
router = self._router()
|
||
exception = self._auth_error()
|
||
mark_advisor_orchestration_failure(exception)
|
||
|
||
now = datetime.now()
|
||
assert (
|
||
router.deployment_callback_on_failure(
|
||
self._kwargs(exception), None, now, now
|
||
)
|
||
is False
|
||
)
|
||
assert "dep-1" not in self._cooled_down_ids(router)
|
||
|
||
|
||
def test_stream_chunks_have_generated_content_detects_text_and_non_text():
|
||
from litellm.router import _stream_chunks_have_generated_content
|
||
from litellm.types.utils import (
|
||
ChatCompletionDeltaToolCall,
|
||
Delta,
|
||
Function,
|
||
StreamingChoices,
|
||
)
|
||
|
||
def _chunk(delta):
|
||
return litellm.ModelResponseStream(
|
||
id="chatcmpl-1",
|
||
model="gpt-4",
|
||
object="chat.completion.chunk",
|
||
choices=[StreamingChoices(finish_reason=None, index=0, delta=delta)],
|
||
)
|
||
|
||
assert _stream_chunks_have_generated_content([]) is False
|
||
|
||
empty_chunk = _chunk(Delta(role="assistant"))
|
||
assert _stream_chunks_have_generated_content([empty_chunk]) is False
|
||
|
||
text_chunk = _chunk(Delta(content="Hello"))
|
||
assert _stream_chunks_have_generated_content([text_chunk]) is True
|
||
|
||
reasoning_chunk = _chunk(Delta(reasoning_content="Thinking"))
|
||
assert _stream_chunks_have_generated_content([reasoning_chunk]) is True
|
||
|
||
tool_call_delta = Delta(
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
id="call_1",
|
||
function=Function(name="get_weather", arguments="{}"),
|
||
type="function",
|
||
index=0,
|
||
)
|
||
]
|
||
)
|
||
tool_call_chunk = _chunk(tool_call_delta)
|
||
assert _stream_chunks_have_generated_content([tool_call_chunk]) is True
|
||
|
||
thinking_delta = Delta(thinking_blocks=[{"type": "thinking", "thinking": "Let me think..."}])
|
||
thinking_chunk = _chunk(thinking_delta)
|
||
assert _stream_chunks_have_generated_content([thinking_chunk]) is True
|
||
|
||
reasoning_items_delta = Delta(reasoning_items=[{"type": "reasoning", "id": "rs_1"}])
|
||
reasoning_items_chunk = _chunk(reasoning_items_delta)
|
||
assert _stream_chunks_have_generated_content([reasoning_items_chunk]) is True
|
||
|
||
audio_delta = Delta(audio={"data": "abc123", "expires_at": 1234567890, "transcript": "hello"})
|
||
audio_chunk = _chunk(audio_delta)
|
||
assert _stream_chunks_have_generated_content([audio_chunk]) is True
|
||
|
||
images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}])
|
||
images_chunk = _chunk(images_delta)
|
||
assert _stream_chunks_have_generated_content([images_chunk]) is True
|
||
|
||
annotations_delta = Delta(
|
||
annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]
|
||
)
|
||
annotations_chunk = _chunk(annotations_delta)
|
||
assert _stream_chunks_have_generated_content([annotations_chunk]) is True
|
||
|
||
|
||
def test_get_configured_token_limits_reads_deployment_model_info():
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "my-custom-model",
|
||
"litellm_params": {"model": "openai/some-unmapped-model"},
|
||
"model_info": {"max_input_tokens": 32000, "max_output_tokens": 8000},
|
||
}
|
||
]
|
||
)
|
||
|
||
assert router.get_configured_token_limits("my-custom-model") == (32000, 8000)
|
||
|
||
|
||
def test_get_configured_token_limits_returns_none_for_unset_or_unknown():
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "no-limits-model",
|
||
"litellm_params": {"model": "openai/some-unmapped-model"},
|
||
}
|
||
]
|
||
)
|
||
|
||
assert router.get_configured_token_limits("no-limits-model") == (None, None)
|
||
assert router.get_configured_token_limits("not-a-real-model") == (None, None)
|
||
|
||
|
||
def test_get_configured_token_limits_skips_wildcard_pattern_matching():
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "bedrock/*",
|
||
"litellm_params": {"model": "bedrock/*"},
|
||
"model_info": {"max_input_tokens": 12345},
|
||
}
|
||
]
|
||
)
|
||
|
||
with patch.object(
|
||
router.pattern_router, "route", side_effect=AssertionError("pattern route called")
|
||
):
|
||
assert router.get_configured_token_limits(
|
||
"bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
|
||
) == (None, None)
|
||
|
||
|
||
def test_get_configured_token_limits_treats_malformed_values_as_absent():
|
||
malformed = ["", "unlimited", "128,000", [128000], {"max": 128000}, True]
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": f"bad-limit-{i}",
|
||
"litellm_params": {"model": "openai/some-unmapped-model"},
|
||
"model_info": {"max_input_tokens": bad, "max_output_tokens": bad},
|
||
}
|
||
for i, bad in enumerate(malformed)
|
||
]
|
||
)
|
||
|
||
for i in range(len(malformed)):
|
||
assert router.get_configured_token_limits(f"bad-limit-{i}") == (None, None)
|
||
|
||
|
||
def test_get_configured_token_limits_coerces_numeric_strings():
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "quoted-limits-model",
|
||
"litellm_params": {"model": "openai/some-unmapped-model"},
|
||
"model_info": {"max_input_tokens": "32000", "max_output_tokens": "8000"},
|
||
}
|
||
]
|
||
)
|
||
|
||
assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error():
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "owning-model",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "sk-owning",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "fallback-model",
|
||
"litellm_params": {
|
||
"model": "azure/gpt-4o-mini",
|
||
"api_key": "sk-fallback",
|
||
"api_base": "https://fallback.openai.azure.com",
|
||
"api_version": "2024-08-01-preview",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"owning-model": ["fallback-model"]}],
|
||
num_retries=0,
|
||
)
|
||
owning_provider_error = litellm.BadRequestError(
|
||
message="completion_window must be one of: 24h",
|
||
model="openai/gpt-4o-mini",
|
||
llm_provider="openai",
|
||
)
|
||
mock_create = AsyncMock(side_effect=owning_provider_error)
|
||
|
||
with patch.object(router, "_acreate_batch", mock_create):
|
||
with pytest.raises(litellm.BadRequestError, match="24h"):
|
||
await router.acreate_batch(
|
||
model="owning-model",
|
||
input_file_id="file-owned-by-openai",
|
||
endpoint="/v1/chat/completions",
|
||
completion_window="5m",
|
||
disable_fallbacks=True,
|
||
)
|
||
|
||
mock_create.assert_awaited_once()
|
||
assert mock_create.call_args.kwargs["model"] == "owning-model"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acreate_batch_surfaces_owning_provider_error_without_disable_fallbacks():
|
||
"""The router itself has to keep a batch inside the group that owns the input file:
|
||
the proxy only sets disable_fallbacks on the managed-files route, so the caller
|
||
otherwise gets the fallback provider's error for a file it never received."""
|
||
from litellm.types.utils import LiteLLMBatch
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "owning-model",
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "sk-owning",
|
||
},
|
||
},
|
||
{
|
||
"model_name": "fallback-model",
|
||
"litellm_params": {
|
||
"model": "azure/gpt-4o-mini",
|
||
"api_key": "sk-fallback",
|
||
"api_base": "https://fallback.openai.azure.com",
|
||
"api_version": "2024-08-01-preview",
|
||
},
|
||
},
|
||
],
|
||
fallbacks=[{"owning-model": ["fallback-model"]}],
|
||
num_retries=0,
|
||
)
|
||
attempted_models = []
|
||
|
||
async def _acreate_batch(model, **kwargs):
|
||
attempted_models.append(model)
|
||
if model == "owning-model":
|
||
raise litellm.APIConnectionError(
|
||
message="Connection error - openai is unreachable",
|
||
model="openai/gpt-4o-mini",
|
||
llm_provider="openai",
|
||
)
|
||
return LiteLLMBatch(
|
||
id="batch-created-on-the-wrong-provider",
|
||
completion_window="24h",
|
||
created_at=0,
|
||
endpoint="/v1/chat/completions",
|
||
input_file_id="file-owned-by-openai",
|
||
object="batch",
|
||
status="validating",
|
||
)
|
||
|
||
with patch.object(router, "_acreate_batch", _acreate_batch):
|
||
with pytest.raises(litellm.APIConnectionError, match="openai is unreachable"):
|
||
await router.acreate_batch(
|
||
model="owning-model",
|
||
input_file_id="file-owned-by-openai",
|
||
endpoint="/v1/chat/completions",
|
||
completion_window="24h",
|
||
metadata={"team": "batch-jobs"},
|
||
)
|
||
|
||
assert attempted_models == ["owning-model"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acreate_batch_request_bedrock_tags_override_deployment_tags():
|
||
import httpx
|
||
|
||
from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils
|
||
|
||
deployment_tags = [{"key": "application", "value": "config-level"}]
|
||
request_tags = [{"key": "application", "value": "request-level"}]
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "bedrock-batch-model",
|
||
"litellm_params": {
|
||
"model": "bedrock/us.anthropic.claude-sonnet-5",
|
||
"aws_batch_role_arn": "arn:aws:iam::123:role/batch-role",
|
||
"aws_region_name": "us-west-2",
|
||
"bedrock_tags": deployment_tags,
|
||
},
|
||
}
|
||
]
|
||
)
|
||
|
||
def fake_response():
|
||
return httpx.Response(
|
||
status_code=200,
|
||
json={
|
||
"jobArn": "arn:aws:bedrock:us-west-2:123:model-invocation-job/abc1234567",
|
||
"status": "Submitted",
|
||
},
|
||
)
|
||
|
||
mock_client = MagicMock()
|
||
mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response())
|
||
|
||
with patch.object(
|
||
CommonBatchFilesUtils,
|
||
"sign_aws_request",
|
||
return_value=({"Authorization": "signed"}, b"{}"),
|
||
) as mock_sign, patch(
|
||
"litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client",
|
||
return_value=mock_client,
|
||
):
|
||
await router.acreate_batch(
|
||
model="bedrock-batch-model",
|
||
input_file_id="s3://bucket/input.jsonl",
|
||
endpoint="/v1/chat/completions",
|
||
completion_window="24h",
|
||
)
|
||
assert mock_sign.call_args.kwargs["data"]["tags"] == deployment_tags
|
||
|
||
await router.acreate_batch(
|
||
model="bedrock-batch-model",
|
||
input_file_id="s3://bucket/input.jsonl",
|
||
endpoint="/v1/chat/completions",
|
||
completion_window="24h",
|
||
bedrock_tags=request_tags,
|
||
)
|
||
assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags
|
||
|
||
|
||
class TestPreRoutingStrategyRegistryLifecycle:
|
||
"""
|
||
Regression tests: a deployment leaving the model_list must release the
|
||
pre-routing strategy slot it holds in `auto_routers` / `complexity_routers` /
|
||
`adaptive_routers` / `quality_routers`.
|
||
|
||
Before this fix, editing an auto-router-family model (a UI save, which reaches
|
||
every other pod as an `upsert_deployment` from the periodic DB reload) popped
|
||
the deployment out of the model_list and then failed to re-add it: registration
|
||
raised "already exists" against the stale registry entry, and
|
||
`ignore_invalid_deployments=True` swallowed the error. The router vanished from
|
||
the Models page and stayed gone until a proxy restart, while the DB row and the
|
||
"saved successfully" response both looked fine.
|
||
"""
|
||
|
||
@staticmethod
|
||
def _complexity_router_params(default_model: str, tags=None) -> dict:
|
||
return {
|
||
"model": "auto_router/complexity_router",
|
||
"complexity_router_config": {
|
||
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}
|
||
},
|
||
"complexity_router_default_model": default_model,
|
||
**({"tags": tags} if tags else {}),
|
||
}
|
||
|
||
@classmethod
|
||
def _router_with_complexity_router(cls, default_model: str = "gpt-4o") -> "litellm.Router":
|
||
return litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
|
||
{
|
||
"model_name": "smart-router",
|
||
"litellm_params": cls._complexity_router_params(default_model),
|
||
"model_info": {"id": "router-1", "db_model": True},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
|
||
@staticmethod
|
||
def _model_names(router: "litellm.Router") -> list:
|
||
return [model["model_name"] for model in router.model_list]
|
||
|
||
def test_upsert_of_edited_router_keeps_it_routable(self):
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = self._router_with_complexity_router()
|
||
|
||
router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="smart-router",
|
||
litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")),
|
||
model_info=ModelInfo(id="router-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "smart-router" in self._model_names(router)
|
||
registered = router.complexity_routers["smart-router"]
|
||
assert len(registered) == 1
|
||
# the surviving strategy is the edited one, not the pre-edit leftover
|
||
assert registered[0].strategy.config.default_model == "gpt-4o-mini"
|
||
|
||
def test_unchanged_upsert_leaves_router_untouched(self):
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = self._router_with_complexity_router()
|
||
strategy_before = router.complexity_routers["smart-router"][0].strategy
|
||
|
||
for _ in range(3):
|
||
router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="smart-router",
|
||
litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o")),
|
||
model_info=ModelInfo(id="router-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "smart-router" in self._model_names(router)
|
||
assert router.complexity_routers["smart-router"][0].strategy is strategy_before
|
||
|
||
def test_delete_frees_the_name_for_a_new_router(self):
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = self._router_with_complexity_router()
|
||
|
||
router.delete_deployment(id="router-1")
|
||
assert "smart-router" not in router.complexity_routers
|
||
|
||
router.add_deployment(
|
||
deployment=Deployment(
|
||
model_name="smart-router",
|
||
litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")),
|
||
model_info=ModelInfo(id="router-2", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "smart-router" in self._model_names(router)
|
||
assert router.complexity_routers["smart-router"][0].strategy.config.default_model == "gpt-4o-mini"
|
||
|
||
def test_delete_only_frees_the_matching_tag_slot(self):
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
|
||
{
|
||
"model_name": "shared-router",
|
||
"litellm_params": self._complexity_router_params("gpt-4o", tags=["team-a"]),
|
||
"model_info": {"id": "router-a"},
|
||
},
|
||
{
|
||
"model_name": "shared-router",
|
||
"litellm_params": self._complexity_router_params("gpt-4o-mini", tags=["team-b"]),
|
||
"model_info": {"id": "router-b"},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
assert len(router.complexity_routers["shared-router"]) == 2
|
||
|
||
router.delete_deployment(id="router-a")
|
||
|
||
remaining = router.complexity_routers["shared-router"]
|
||
assert len(remaining) == 1
|
||
assert remaining[0].tags == ("team-b",)
|
||
|
||
def test_delete_of_regular_model_preserves_router_sharing_its_name(self):
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
|
||
{
|
||
"model_name": "shared-name",
|
||
"litellm_params": self._complexity_router_params("gpt-4o"),
|
||
"model_info": {"id": "router-1"},
|
||
},
|
||
{
|
||
"model_name": "shared-name",
|
||
"litellm_params": {"model": "openai/gpt-4o"},
|
||
"model_info": {"id": "regular-1"},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
strategy = router.complexity_routers["shared-name"][0].strategy
|
||
|
||
router.delete_deployment(id="regular-1")
|
||
|
||
assert router.complexity_routers["shared-name"][0].strategy is strategy
|
||
|
||
def test_upsert_of_edited_adaptive_router_rebuilds_it(self):
|
||
"""Adaptive routers are built by set_model_list()'s deferred pass, not by
|
||
add_deployment(), so releasing the slot on edit must be paired with a rebuild -
|
||
otherwise the edit silently turns adaptive routing off."""
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
def adaptive_params(available_models: list) -> dict:
|
||
return {
|
||
"model": "auto_router/adaptive_router",
|
||
"adaptive_router_config": {"available_models": available_models},
|
||
}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
|
||
{
|
||
"model_name": "adaptive-router",
|
||
"litellm_params": adaptive_params(["gpt-4o-mini"]),
|
||
"model_info": {"id": "router-1", "db_model": True},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
assert "adaptive-router" in router.adaptive_routers
|
||
|
||
router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="adaptive-router",
|
||
litellm_params=LiteLLM_Params(**adaptive_params(["gpt-4o", "gpt-4o-mini"])),
|
||
model_info=ModelInfo(id="router-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "adaptive-router" in self._model_names(router)
|
||
registered = router.adaptive_routers["adaptive-router"]
|
||
assert len(registered) == 1
|
||
assert set(registered[0].strategy.config.available_models) == {"gpt-4o", "gpt-4o-mini"}
|
||
|
||
def test_delete_repairs_indices_even_when_strategy_release_fails(self):
|
||
"""Structural removal and strategy release are not equally critical. Once the entry
|
||
leaves model_list the index maps must be repaired no matter what, so releasing the
|
||
registry slot runs after that repair and cannot abandon the router half-updated."""
|
||
router = self._router_with_complexity_router()
|
||
idx = router.model_id_to_deployment_index_map["router-1"]
|
||
router.model_list[idx] = {"model_name": "smart-router", "litellm_params": None}
|
||
|
||
returned = router.delete_deployment(id="router-1")
|
||
|
||
assert returned is not None
|
||
assert "router-1" not in router.model_id_to_deployment_index_map
|
||
assert all(entry.get("model_info", {}).get("id") != "router-1" for entry in router.model_list)
|
||
assert router.get_deployment(model_id="router-1") is None
|
||
assert "gpt-4o" in self._model_names(router)
|
||
|
||
def test_delete_of_adaptive_enabled_complexity_router_frees_both_registries(self):
|
||
"""A complexity router with adaptive set is registered in BOTH complexity_routers
|
||
and adaptive_routers under the same (model_name, tags). Releasing only the first
|
||
match leaves the adaptive strategy live, so a deleted alias stays routable and its
|
||
post-call hook keeps recording."""
|
||
import litellm as litellm_module
|
||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||
|
||
params = {
|
||
"model": "auto_router/complexity_router",
|
||
"complexity_router_config": {
|
||
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
|
||
"adaptive": True,
|
||
},
|
||
"complexity_router_default_model": "gpt-4o",
|
||
}
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
|
||
{
|
||
"model_name": "hybrid-router",
|
||
"litellm_params": params,
|
||
"model_info": {"id": "router-1", "db_model": True},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
assert "hybrid-router" in router.complexity_routers
|
||
assert "hybrid-router" in router.adaptive_routers
|
||
hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
|
||
assert len(hooks) == 1
|
||
|
||
router.delete_deployment(id="router-1")
|
||
|
||
assert "hybrid-router" not in router.complexity_routers
|
||
assert "hybrid-router" not in router.adaptive_routers
|
||
remaining_hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(
|
||
AdaptiveRouterPostCallHook
|
||
)
|
||
assert remaining_hooks == []
|
||
|
||
def test_upsert_of_edited_quality_router_keeps_it_routable(self):
|
||
"""_unregister_pre_routing_strategy_for_deployment dispatches on four prefixes;
|
||
quality_router is one of them and would otherwise go unexercised."""
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
def quality_params(default_model: str) -> dict:
|
||
return {
|
||
"model": "auto_router/quality_router",
|
||
"quality_router_default_model": default_model,
|
||
}
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
|
||
{
|
||
"model_name": "quality-router",
|
||
"litellm_params": quality_params("gpt-4o"),
|
||
"model_info": {"id": "router-1", "db_model": True},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
assert "quality-router" in router.quality_routers
|
||
|
||
router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="quality-router",
|
||
litellm_params=LiteLLM_Params(**quality_params("gpt-4o-mini")),
|
||
model_info=ModelInfo(id="router-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "quality-router" in self._model_names(router)
|
||
registered = router.quality_routers["quality-router"]
|
||
assert len(registered) == 1
|
||
assert registered[0].strategy.config.default_model == "gpt-4o-mini"
|
||
|
||
@staticmethod
|
||
def _hybrid_router_params(tiers: dict) -> dict:
|
||
return {
|
||
"model": "auto_router/complexity_router",
|
||
"complexity_router_config": {"tiers": tiers, "adaptive": True},
|
||
"complexity_router_default_model": "gpt-4o",
|
||
}
|
||
|
||
@classmethod
|
||
def _router_with_hybrid_router(cls) -> "litellm.Router":
|
||
return litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
|
||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
|
||
{
|
||
"model_name": "hybrid-router",
|
||
"litellm_params": cls._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}),
|
||
"model_info": {"id": "router-1", "db_model": True},
|
||
},
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
|
||
def test_upsert_of_edited_hybrid_complexity_router_relinks_adaptive(self):
|
||
"""Editing an adaptive-enabled complexity router releases its adaptive companion
|
||
along with the complexity slot; the finalize re-run must fire for it (not just for
|
||
`auto_router/adaptive_router` deployments) or the rebuilt complexity router keeps
|
||
routing while bandit recording, DB persistence and /adaptive_router/state all
|
||
silently stop until the next full reload."""
|
||
import litellm as litellm_module
|
||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = self._router_with_hybrid_router()
|
||
assert "hybrid-router" in router.adaptive_routers
|
||
|
||
router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="hybrid-router",
|
||
litellm_params=LiteLLM_Params(
|
||
**self._hybrid_router_params(
|
||
{"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}
|
||
)
|
||
),
|
||
model_info=ModelInfo(id="router-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "hybrid-router" in self._model_names(router)
|
||
assert "hybrid-router" in router.complexity_routers
|
||
assert "hybrid-router" in router.adaptive_routers
|
||
rebuilt = router.complexity_routers["hybrid-router"][0].strategy
|
||
assert router.adaptive_routers["hybrid-router"][0].strategy is rebuilt.adaptive_router
|
||
hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
|
||
assert len(hooks) == 1
|
||
|
||
def test_upsert_turning_adaptive_on_builds_the_companion(self):
|
||
"""An edit that flips `adaptive: true` on an existing complexity router must
|
||
register the companion immediately; neither side of the old prefix-only gate
|
||
matches a complexity deployment, so the flip was a silent no-op until restart."""
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = self._router_with_complexity_router()
|
||
assert "smart-router" not in router.adaptive_routers
|
||
|
||
params = self._complexity_router_params("gpt-4o")
|
||
params["complexity_router_config"] = {**params["complexity_router_config"], "adaptive": True}
|
||
router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="smart-router",
|
||
litellm_params=LiteLLM_Params(**params),
|
||
model_info=ModelInfo(id="router-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "smart-router" in router.adaptive_routers
|
||
|
||
def test_unregister_pre_routing_strategy_scopes_the_drop_by_tags(self):
|
||
"""The bool return drives the hook re-sync; a tag mismatch must report False and
|
||
leave the registry untouched, and dropping the last entry must free the key."""
|
||
from litellm.types.router import TaggedPreRoutingStrategy
|
||
|
||
registry = {
|
||
"m": [
|
||
TaggedPreRoutingStrategy(tags=("team-a",), strategy=object()),
|
||
TaggedPreRoutingStrategy(tags=(), strategy=object()),
|
||
]
|
||
}
|
||
|
||
assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-b",)) is False
|
||
assert len(registry["m"]) == 2
|
||
|
||
assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-a",)) is True
|
||
assert [entry.tags for entry in registry["m"]] == [()]
|
||
|
||
assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ()) is True
|
||
assert "m" not in registry
|
||
|
||
def test_unregister_for_deployment_ignores_non_router_deployments(self):
|
||
"""Direct twin of the endpoint-level test: a regular deployment that shares a
|
||
router's model_name must not evict the router's registry slot."""
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = self._router_with_complexity_router()
|
||
|
||
router._unregister_pre_routing_strategy_for_deployment(
|
||
deployment=Deployment(
|
||
model_name="smart-router",
|
||
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
|
||
model_info=ModelInfo(id="plain-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert "smart-router" in router.complexity_routers
|
||
|
||
def test_sync_adaptive_router_hooks_keeps_one_hook_per_registered_router(self):
|
||
"""Re-syncing must replace, not accumulate: a duplicated hook double-fires
|
||
bandit signal recording for every request."""
|
||
import litellm as litellm_module
|
||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||
|
||
router = self._router_with_hybrid_router()
|
||
|
||
router._sync_adaptive_router_hooks()
|
||
router._sync_adaptive_router_hooks()
|
||
|
||
hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
|
||
assert len(hooks) == 1
|
||
|
||
def test_deployment_participates_in_adaptive_routing_matrix(self):
|
||
"""The upsert finalize re-run keys off this predicate for both the incoming and
|
||
outgoing deployment; a false negative silently strands the adaptive companion."""
|
||
from litellm.types.router import LiteLLM_Params
|
||
|
||
router = self._router_with_complexity_router()
|
||
|
||
cases = [
|
||
({"model": "auto_router/adaptive_router", "adaptive_router_config": {}}, True),
|
||
(self._hybrid_router_params({"SIMPLE": "gpt-4o-mini"}), True),
|
||
(self._complexity_router_params("gpt-4o"), False),
|
||
(
|
||
{
|
||
"model": "auto_router/complexity_router",
|
||
"complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "adaptive": False},
|
||
"complexity_router_default_model": "gpt-4o",
|
||
},
|
||
False,
|
||
),
|
||
({"model": "openai/gpt-4o"}, False),
|
||
]
|
||
for params, expected in cases:
|
||
actual = router._deployment_participates_in_adaptive_routing(
|
||
litellm_params=LiteLLM_Params(**params)
|
||
)
|
||
assert actual is expected, params["model"]
|
||
|
||
|
||
def test_model_info_is_active_for_environment_matrix(monkeypatch):
|
||
"""The model-write endpoints consult this predicate to tell a deliberately
|
||
environment-inactive model from one dropped by a failed reload; the Router's own
|
||
deployment gate delegates to it, so the two can never diverge."""
|
||
from litellm.router import model_info_is_active_for_environment
|
||
|
||
assert model_info_is_active_for_environment(model_info=None) is True
|
||
assert model_info_is_active_for_environment(model_info={"id": "m1"}) is True
|
||
assert model_info_is_active_for_environment(model_info={"supported_environments": None}) is True
|
||
|
||
monkeypatch.setenv("LITELLM_ENVIRONMENT", "development")
|
||
assert model_info_is_active_for_environment(model_info={"supported_environments": ["development"]}) is True
|
||
assert model_info_is_active_for_environment(model_info={"supported_environments": ["production"]}) is False
|
||
|
||
monkeypatch.delenv("LITELLM_ENVIRONMENT")
|
||
with pytest.raises(ValueError, match="LITELLM_ENVIRONMENT"):
|
||
model_info_is_active_for_environment(model_info={"supported_environments": ["production"]})
|
||
|
||
|
||
def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(monkeypatch):
|
||
"""
|
||
The supported-params check must run against the deployment's own
|
||
provider-qualified model. Resolving the per-deployment model only after the
|
||
model-info lookup leaves it unset whenever that lookup raises (an
|
||
unregistered custom model), so the check falls back to the bare model group
|
||
name and the request dies with 'LLM Provider NOT provided'.
|
||
"""
|
||
monkeypatch.setattr(litellm, "drop_params", False)
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "custom-alias",
|
||
"litellm_params": {"model": "hosted_vllm/not-in-the-catalog"},
|
||
}
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
|
||
def _raise_unmapped(**kwargs):
|
||
raise ValueError("This model isn't mapped yet")
|
||
|
||
monkeypatch.setattr(router, "get_router_model_info", _raise_unmapped)
|
||
|
||
seen: list[tuple] = []
|
||
original_get_supported_openai_params = litellm.get_supported_openai_params
|
||
|
||
def _record(model, custom_llm_provider=None, **kwargs):
|
||
seen.append((model, custom_llm_provider))
|
||
return original_get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider, **kwargs)
|
||
|
||
monkeypatch.setattr(litellm, "get_supported_openai_params", _record)
|
||
|
||
deployments = [
|
||
{
|
||
"litellm_params": {"model": "hosted_vllm/not-in-the-catalog"},
|
||
"model_info": {"id": "d1"},
|
||
}
|
||
]
|
||
result = router._pre_call_checks(
|
||
model="custom-alias",
|
||
healthy_deployments=deployments,
|
||
messages=[{"role": "user", "content": "hi"}],
|
||
request_kwargs={},
|
||
)
|
||
|
||
assert len(result) == 1
|
||
assert seen == [("not-in-the-catalog", "hosted_vllm")]
|
||
|
||
|
||
def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypatch):
|
||
"""
|
||
Pre-call checks filter deployments; they must never be the thing that fails
|
||
a request. A deployment whose provider cannot be resolved simply skips the
|
||
supported-params check instead of raising out of deployment selection.
|
||
"""
|
||
monkeypatch.setattr(litellm, "drop_params", False)
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "custom-alias",
|
||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||
}
|
||
],
|
||
enable_pre_call_checks=True,
|
||
)
|
||
|
||
def _raise_no_provider(**kwargs):
|
||
raise litellm.BadRequestError(
|
||
message="LLM Provider NOT provided.",
|
||
model="custom-alias",
|
||
llm_provider="",
|
||
)
|
||
|
||
monkeypatch.setattr(litellm, "get_llm_provider", _raise_no_provider)
|
||
|
||
deployments = [
|
||
{
|
||
"litellm_params": {"model": "some-unresolvable-model"},
|
||
"model_info": {"id": "d1"},
|
||
}
|
||
]
|
||
result = router._pre_call_checks(
|
||
model="custom-alias",
|
||
healthy_deployments=deployments,
|
||
messages=[{"role": "user", "content": "hi"}],
|
||
request_kwargs={},
|
||
)
|
||
|
||
assert len(result) == 1
|
||
|
||
|
||
class TestUpsertDeploymentRollback:
|
||
"""
|
||
Regression tests: `upsert_deployment` pops the previous deployment before
|
||
re-adding the edited one. When the re-add raises under
|
||
`ignore_invalid_deployments=True`, the pop must be rolled back so this pod
|
||
keeps serving the previous configuration instead of silently dropping a live
|
||
deployment (the "Error upserting deployment" drop behind the access-group
|
||
reload 500 in the 2-replica e2e suite).
|
||
"""
|
||
|
||
def test_failed_upsert_keeps_previous_deployment_serving(self):
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "prod-model",
|
||
"litellm_params": {"model": "gpt-4o", "api_key": "sk-old"},
|
||
"model_info": {"id": "prod-1", "db_model": True},
|
||
}
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
|
||
result = router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="prod-model",
|
||
litellm_params=LiteLLM_Params(model="auto_router/broken"),
|
||
model_info=ModelInfo(id="prod-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert result is None
|
||
restored = router.get_deployment(model_id="prod-1")
|
||
assert restored is not None
|
||
assert restored.litellm_params.model == "gpt-4o"
|
||
assert [model["model_name"] for model in router.model_list] == ["prod-model"]
|
||
|
||
def test_failed_fresh_add_returns_none_without_restore(self):
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
router = litellm.Router(model_list=[], ignore_invalid_deployments=True)
|
||
|
||
result = router.upsert_deployment(
|
||
deployment=Deployment(
|
||
model_name="fresh-router",
|
||
litellm_params=LiteLLM_Params(model="auto_router/broken"),
|
||
model_info=ModelInfo(id="fresh-1", db_model=True),
|
||
)
|
||
)
|
||
|
||
assert result is None
|
||
assert router.get_deployment(model_id="fresh-1") is None
|
||
assert router.model_list == []
|
||
|
||
def test_restore_re_adds_popped_deployment(self):
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": "prod-model",
|
||
"litellm_params": {"model": "gpt-4o", "api_key": "sk-old"},
|
||
"model_info": {"id": "prod-1", "db_model": True},
|
||
}
|
||
],
|
||
ignore_invalid_deployments=True,
|
||
)
|
||
previous = router.get_deployment(model_id="prod-1")
|
||
router.delete_deployment(id="prod-1")
|
||
assert router.has_model_id("prod-1") is False
|
||
|
||
router._restore_deployment_after_failed_upsert(
|
||
previous_deployment=previous, model_id="prod-1"
|
||
)
|
||
|
||
restored = router.get_deployment(model_id="prod-1")
|
||
assert restored is not None
|
||
assert restored.litellm_params.model == "gpt-4o"
|
||
|
||
router._restore_deployment_after_failed_upsert(
|
||
previous_deployment=previous, model_id="prod-1"
|
||
)
|
||
assert len(router.model_list) == 1
|
||
|
||
router._restore_deployment_after_failed_upsert(
|
||
previous_deployment=None, model_id="prod-1"
|
||
)
|
||
assert len(router.model_list) == 1
|
||
|
||
|
||
class TestConsumedRequestTagsStamp:
|
||
"""Issue #36621: when a request's tags select a tagged pre-routing strategy, those
|
||
tags are consumed by the selection; the hook must stamp the rewritten model group so
|
||
tag filtering skips request-body tags there, and must clear the stamp on every
|
||
re-entry (fallbacks reuse the same request_kwargs) so it cannot leak elsewhere."""
|
||
|
||
class _RewriteStrategy:
|
||
def __init__(self, rewrite_to: str):
|
||
self.rewrite_to = rewrite_to
|
||
|
||
async def async_pre_routing_hook(
|
||
self, model, request_kwargs, messages=None, input=None, specific_deployment=False
|
||
):
|
||
from litellm.types.router import PreRoutingHookResponse
|
||
|
||
return PreRoutingHookResponse(model=self.rewrite_to, messages=messages)
|
||
|
||
@classmethod
|
||
def _router(cls, marker_tags=("route",)) -> "litellm.Router":
|
||
from litellm.types.router import TaggedPreRoutingStrategy
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt4o", "litellm_params": {"model": "openai/gpt-4o"}},
|
||
{"model_name": "gemini-flash", "litellm_params": {"model": "gemini/gemini-3.6-flash"}},
|
||
],
|
||
enable_tag_filtering=True,
|
||
)
|
||
router.auto_routers = {
|
||
"gpt4o": [TaggedPreRoutingStrategy(tags=marker_tags, strategy=cls._RewriteStrategy("gemini-flash"))]
|
||
}
|
||
return router
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_stamps_the_rewritten_group_when_request_tags_selected_the_router(self):
|
||
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
|
||
from litellm.types.router import ConsumedRequestTagsStamp
|
||
|
||
router = self._router()
|
||
request_kwargs = {"metadata": {"tags": ["route"]}}
|
||
|
||
await router.async_pre_routing_hook(model="gpt4o", request_kwargs=request_kwargs)
|
||
|
||
assert request_kwargs["metadata"][CONSUMED_REQUEST_TAGS_METADATA_KEY] == ConsumedRequestTagsStamp(
|
||
model_group="gemini-flash", tags=("route",)
|
||
)
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_stamps_into_litellm_metadata_when_the_request_uses_that_bucket(self):
|
||
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
|
||
from litellm.types.router import ConsumedRequestTagsStamp
|
||
|
||
router = self._router()
|
||
request_kwargs = {"litellm_metadata": {"tags": ["route"]}}
|
||
|
||
await router.async_pre_routing_hook(model="gpt4o", request_kwargs=request_kwargs)
|
||
|
||
assert request_kwargs["litellm_metadata"][CONSUMED_REQUEST_TAGS_METADATA_KEY] == ConsumedRequestTagsStamp(
|
||
model_group="gemini-flash", tags=("route",)
|
||
)
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fallback_reentry_with_a_plain_group_clears_the_stale_stamp(self):
|
||
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
|
||
|
||
router = self._router()
|
||
request_kwargs = {"metadata": {"tags": ["route"]}}
|
||
|
||
await router.async_pre_routing_hook(model="gpt4o", request_kwargs=request_kwargs)
|
||
await router.async_pre_routing_hook(model="gemini-flash", request_kwargs=request_kwargs)
|
||
|
||
assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"]
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_no_stamp_when_the_request_is_untagged(self):
|
||
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
|
||
|
||
router = self._router()
|
||
request_kwargs = {"metadata": {}}
|
||
|
||
await router.async_pre_routing_hook(model="gpt4o", request_kwargs=request_kwargs)
|
||
|
||
assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"]
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_no_stamp_when_the_selected_strategy_carries_no_tags(self):
|
||
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
|
||
|
||
router = self._router(marker_tags=())
|
||
request_kwargs = {"metadata": {"tags": ["route"]}}
|
||
|
||
await router.async_pre_routing_hook(model="gpt4o", request_kwargs=request_kwargs)
|
||
|
||
assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"]
|
||
|
||
|
||
class TestAutoRouterMaxInputCharsWiring:
|
||
"""`auto_router_max_input_chars` on the deployment has to reach the AutoRouter that embeds prompts.
|
||
|
||
Without it the cap silently reverts to the default, so an operator whose embedding model has a
|
||
512-token window cannot lower it and every long prompt falls back to the default model instead
|
||
of being routed.
|
||
"""
|
||
|
||
@staticmethod
|
||
def _router(**extra_params) -> "litellm.Router":
|
||
pytest.importorskip("semantic_router", reason="auto-router needs the semantic-router extra")
|
||
return litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
|
||
{
|
||
"model_name": "my-auto-router",
|
||
"litellm_params": {
|
||
"model": "auto_router/my-auto-router",
|
||
"auto_router_config": json.dumps(
|
||
{"routes": [{"name": "gpt-4o", "utterances": ["write me code"]}]}
|
||
),
|
||
"auto_router_default_model": "gpt-4o",
|
||
"auto_router_embedding_model": "text-embedding-3-small",
|
||
**extra_params,
|
||
},
|
||
},
|
||
]
|
||
)
|
||
|
||
@staticmethod
|
||
def _registered_auto_router(router: "litellm.Router"):
|
||
return router.auto_routers["my-auto-router"][0].strategy
|
||
|
||
def test_should_pass_the_configured_cap_to_the_auto_router(self):
|
||
router = self._router(auto_router_max_input_chars=512)
|
||
|
||
assert self._registered_auto_router(router).max_input_chars == 512
|
||
|
||
def test_should_fall_back_to_the_shared_default_when_the_deployment_omits_it(self):
|
||
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
|
||
|
||
router = self._router()
|
||
|
||
assert self._registered_auto_router(router).max_input_chars == DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
|
||
|
||
|
||
class TestTaggedAutoRouterOnSharedModelName:
|
||
"""A tagged auto-router marker sharing its model_name with a plain deployment must not
|
||
capture requests whose tags don't match it when tag filtering is enabled (#36620)."""
|
||
|
||
class _FixedRouteLayer:
|
||
def __call__(self, text: str):
|
||
from semantic_router.schema import RouteChoice
|
||
|
||
return RouteChoice(name="gemini-flash")
|
||
|
||
@classmethod
|
||
def _router(cls, marker_tags, include_plain_sibling: bool, enable_tag_filtering: bool) -> "litellm.Router":
|
||
pytest.importorskip("semantic_router", reason="auto-router needs the semantic-router extra")
|
||
marker = {
|
||
"model_name": "gpt4o",
|
||
"litellm_params": {
|
||
"model": "auto_router/gpt4o-router",
|
||
"auto_router_config": json.dumps(
|
||
{"routes": [{"name": "gemini-flash", "utterances": ["capital city questions"]}]}
|
||
),
|
||
"auto_router_default_model": "gemini-flash",
|
||
"auto_router_embedding_model": "text-embedding-3-small",
|
||
**({"tags": marker_tags} if marker_tags else {}),
|
||
},
|
||
}
|
||
plain = {"model_name": "gpt4o", "litellm_params": {"model": "openai/gpt-4o"}}
|
||
tier = {"model_name": "gemini-flash", "litellm_params": {"model": "gemini/gemini-3.6-flash"}}
|
||
router = litellm.Router(
|
||
model_list=[plain, marker, tier] if include_plain_sibling else [marker, tier],
|
||
enable_tag_filtering=enable_tag_filtering,
|
||
)
|
||
router.auto_routers["gpt4o"][0].strategy.routelayer = cls._FixedRouteLayer()
|
||
return router
|
||
|
||
@staticmethod
|
||
async def _hook_response(router: "litellm.Router", request_kwargs: dict):
|
||
return await router.async_pre_routing_hook(
|
||
model="gpt4o",
|
||
request_kwargs=request_kwargs,
|
||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||
)
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_untagged_request_bypasses_the_tagged_marker_when_a_plain_deployment_shares_the_name(self):
|
||
router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True)
|
||
|
||
assert await self._hook_response(router, {}) is None
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_request_tagged_for_the_marker_is_still_semantically_routed(self):
|
||
router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True)
|
||
|
||
response = await self._hook_response(router, {"metadata": {"tags": ["route"]}})
|
||
|
||
assert response is not None
|
||
assert response.model == "gemini-flash"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_request_level_tag_filtering_from_key_settings_bypasses_the_marker(self):
|
||
router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=False)
|
||
|
||
assert await self._hook_response(router, {"enable_tag_filtering": True}) is None
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_globally_disabled_filtering_still_lets_the_marker_capture_untagged_requests(self):
|
||
router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=False)
|
||
|
||
response = await self._hook_response(router, {})
|
||
|
||
assert response is not None
|
||
assert response.model == "gemini-flash"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_marker_only_alias_still_captures_untagged_requests(self):
|
||
router = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True)
|
||
|
||
response = await self._hook_response(router, {})
|
||
|
||
assert response is not None
|
||
assert response.model == "gemini-flash"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_untagged_marker_sharing_the_name_still_captures_untagged_requests(self):
|
||
router = self._router(marker_tags=None, include_plain_sibling=True, enable_tag_filtering=True)
|
||
|
||
response = await self._hook_response(router, {})
|
||
|
||
assert response is not None
|
||
assert response.model == "gemini-flash"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_untagged_selection_never_lands_on_the_marker_deployment(self):
|
||
router = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True)
|
||
|
||
for _ in range(20):
|
||
deployment = await router.async_get_available_deployment(
|
||
model="gpt4o",
|
||
request_kwargs={},
|
||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||
)
|
||
assert deployment["litellm_params"]["model"] == "openai/gpt-4o"
|
||
|
||
def test_deployment_without_litellm_params_mapping_is_not_a_marker(self):
|
||
assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False
|
||
|
||
def test_model_name_has_plain_deployments_reflects_the_pool(self):
|
||
mixed = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True)
|
||
marker_only = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True)
|
||
|
||
assert mixed._model_name_has_plain_deployments("gpt4o") is True
|
||
assert marker_only._model_name_has_plain_deployments("gpt4o") is False
|
||
|
||
|
||
class TestGetAllowedFailsFromPolicy:
|
||
def _make_router(self, **policy_kwargs) -> litellm.Router:
|
||
from litellm.types.router import AllowedFailsPolicy
|
||
|
||
return litellm.Router(
|
||
model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "fake"}}],
|
||
allowed_fails_policy=AllowedFailsPolicy(**policy_kwargs),
|
||
)
|
||
|
||
def test_no_policy_returns_none(self):
|
||
router = litellm.Router(
|
||
model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "fake"}}],
|
||
)
|
||
assert router.get_allowed_fails_from_policy(litellm.RateLimitError("429", "openai", "gpt-4")) is None
|
||
|
||
def test_internal_server_error_allowed_fails(self):
|
||
router = self._make_router(InternalServerErrorAllowedFails=7)
|
||
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
||
assert router.get_allowed_fails_from_policy(exc) == 7
|
||
|
||
def test_service_unavailable_error_allowed_fails(self):
|
||
router = self._make_router(ServiceUnavailableErrorAllowedFails=4)
|
||
exc = litellm.ServiceUnavailableError("503", "openai", "gpt-4")
|
||
assert router.get_allowed_fails_from_policy(exc) == 4
|
||
|
||
def test_bad_gateway_error_allowed_fails(self):
|
||
router = self._make_router(BadGatewayErrorAllowedFails=2)
|
||
exc = litellm.BadGatewayError("502", "openai", "gpt-4")
|
||
assert router.get_allowed_fails_from_policy(exc) == 2
|
||
|
||
def test_not_found_error_allowed_fails(self):
|
||
router = self._make_router(NotFoundErrorAllowedFails=1)
|
||
exc = litellm.NotFoundError("404", "openai", "gpt-4")
|
||
assert router.get_allowed_fails_from_policy(exc) == 1
|
||
|
||
def test_unmatched_exception_returns_none(self):
|
||
router = self._make_router(InternalServerErrorAllowedFails=5)
|
||
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
||
assert router.get_allowed_fails_from_policy(exc) is None
|
||
|
||
|
||
class _LogCapture(logging.Handler):
|
||
def __init__(self, level):
|
||
super().__init__(level=level)
|
||
self._level = level
|
||
self.messages = []
|
||
|
||
def emit(self, record):
|
||
if record.levelno == self._level:
|
||
self.messages.append(record.getMessage())
|
||
|
||
|
||
class _FallbackAttemptRecorder(CustomLogger):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.failed_targets = []
|
||
|
||
async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception):
|
||
self.failed_targets.append(kwargs.get("model"))
|
||
|
||
|
||
def _cyclic_fallback_router(num_retries=0):
|
||
groups = ["group-a", "group-b", "group-c", "group-d"]
|
||
return litellm.Router(
|
||
model_list=[
|
||
{
|
||
"model_name": group,
|
||
"litellm_params": {
|
||
"model": "openai/gpt-4o-mini",
|
||
"api_key": "sk-fake",
|
||
"mock_response": "litellm.InternalServerError",
|
||
},
|
||
}
|
||
for group in groups
|
||
],
|
||
fallbacks=[
|
||
{"group-a": ["group-b", "group-c"]},
|
||
{"group-b": ["group-a", "group-c"]},
|
||
{"group-c": ["group-d"]},
|
||
{"group-d": ["group-b", "group-a"]},
|
||
],
|
||
num_retries=num_retries,
|
||
)
|
||
|
||
|
||
async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwargs):
|
||
router_logger = logging.getLogger("LiteLLM Router")
|
||
previous_level = router_logger.level
|
||
router_logger.setLevel(capture.level)
|
||
router_logger.addHandler(capture)
|
||
if recorder is not None:
|
||
litellm.callbacks.append(recorder)
|
||
try:
|
||
with pytest.raises(litellm.InternalServerError):
|
||
await router.acompletion(
|
||
model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs
|
||
)
|
||
finally:
|
||
router_logger.removeHandler(capture)
|
||
router_logger.setLevel(previous_level)
|
||
if recorder is not None:
|
||
litellm.callbacks.remove(recorder)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_cyclic_fallback_graph_does_not_amplify_one_request():
|
||
"""A fallback graph whose entries loop back on each other is easy to build by accident,
|
||
and every group in the loop fails identically on a deterministic error, so the walk must
|
||
not revisit a group and must not re-emit a growing chained traceback at each level. Left
|
||
unbounded, one request blocks the event loop long enough for health probes to fail."""
|
||
recorder = _FallbackAttemptRecorder()
|
||
capture = _LogCapture(logging.ERROR)
|
||
|
||
await _drive_cyclic_fallback(_cyclic_fallback_router(), capture, recorder)
|
||
|
||
assert sorted(set(recorder.failed_targets)) == ["group-b", "group-c", "group-d"]
|
||
assert len(recorder.failed_targets) == len(set(recorder.failed_targets))
|
||
assert not any("Traceback (most recent call last)" in message for message in capture.messages)
|
||
assert sum(len(message) for message in capture.messages) < 5_000
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retry_breadcrumbs_do_not_carry_the_walk_state():
|
||
"""log_retry copies every kwarg into previous_models, which reaches spend logs and
|
||
logging callbacks. The set of already-attempted groups is router-internal walk state
|
||
with no diagnostic value there, and it is the one entry that is not a plain scalar.
|
||
A retry has to be configured for the walk state to reach log_retry at all."""
|
||
router = _cyclic_fallback_router(num_retries=1)
|
||
capture = _LogCapture(logging.ERROR)
|
||
|
||
await _drive_cyclic_fallback(router, capture)
|
||
|
||
assert router.previous_models, "no retry breadcrumbs were recorded"
|
||
assert any(
|
||
"fallback_depth" in breadcrumb for breadcrumb in router.previous_models
|
||
), "no breadcrumb carried router walk state, so this test cannot see the leak"
|
||
for breadcrumb in router.previous_models:
|
||
assert "attempted_targets" not in breadcrumb
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fallback_traceback_stays_available_at_debug_level():
|
||
"""Dropping the stack from the ERROR line is only safe because the fallback path still
|
||
emits it once per level at DEBUG, which is what an operator needs to diagnose why every
|
||
fallback failed. This pins that remaining debug traceback."""
|
||
capture = _LogCapture(logging.DEBUG)
|
||
|
||
await _drive_cyclic_fallback(_cyclic_fallback_router(), capture)
|
||
|
||
assert any("Traceback (most recent call last)" in message for message in capture.messages)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fallback_failure_detail_from_upstream_is_bounded():
|
||
"""The detail each level records about the level below it is attacker-influenced, since
|
||
it carries whatever the upstream error said. It has to be bounded on its own, so a walk
|
||
over several groups cannot compound one large message into the log or into the message
|
||
handed back to the caller."""
|
||
huge_message = "z" * 50_000
|
||
capture = _LogCapture(logging.ERROR)
|
||
|
||
await _drive_cyclic_fallback(
|
||
_cyclic_fallback_router(),
|
||
capture,
|
||
mock_response=litellm.InternalServerError(
|
||
message=huge_message, llm_provider="openai", model="group-a"
|
||
),
|
||
)
|
||
|
||
assert capture.messages, "the fallback failure path did not log at ERROR"
|
||
assert huge_message not in "".join(capture.messages)
|
||
assert max(len(message) for message in capture.messages) < 5_000
|
||
|
||
|
||
def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets():
|
||
request_kwargs = {"metadata": {}}
|
||
litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7)
|
||
assert request_kwargs["metadata"]["probe"] == 7
|
||
|
||
stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}}
|
||
litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None)
|
||
assert "probe" not in stale_kwargs["metadata"]
|
||
assert "probe" not in stale_kwargs["litellm_metadata"]
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"complexity_router_config,expect_callback",
|
||
[
|
||
({"tiers": {"SIMPLE": "gpt-4o"}}, True),
|
||
({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False}, False),
|
||
({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False, "session_affinity": True}, True),
|
||
],
|
||
)
|
||
def test_complexity_router_registers_affinity_callback_for_deployment_pin(complexity_router_config, expect_callback):
|
||
"""The marker the complexity router stamps is inert unless a DeploymentAffinityCheck is
|
||
registered to read it, so deployment_affinity has to pull the callback in, and its default-on
|
||
means a bare config registers one. Opting out must skip the callback entirely rather than
|
||
register a filter that can never fire, including when session_affinity is on, since the two
|
||
pins are independent."""
|
||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||
DeploymentAffinityCheck,
|
||
)
|
||
|
||
router = litellm.Router(
|
||
model_list=[
|
||
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
|
||
{
|
||
"model_name": "my-complexity-router",
|
||
"litellm_params": {
|
||
"model": "auto_router/complexity_router",
|
||
"complexity_router_config": complexity_router_config,
|
||
},
|
||
},
|
||
]
|
||
)
|
||
try:
|
||
registered = any(isinstance(cb, DeploymentAffinityCheck) for cb in router.optional_callbacks or [])
|
||
assert registered is expect_callback
|
||
finally:
|
||
for cb in router.optional_callbacks or []:
|
||
litellm.logging_callback_manager.remove_callback_from_all_lists(cb)
|
||
|
||
|
||
def test_ensure_deployment_affinity_callback_is_idempotent():
|
||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||
DeploymentAffinityCheck,
|
||
)
|
||
|
||
router = litellm.Router(model_list=[])
|
||
try:
|
||
router._ensure_deployment_affinity_callback()
|
||
router._ensure_deployment_affinity_callback()
|
||
affinity_callbacks = [
|
||
cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)
|
||
]
|
||
assert len(affinity_callbacks) == 1
|
||
finally:
|
||
for cb in router.optional_callbacks or []:
|
||
litellm.logging_callback_manager.remove_callback_from_all_lists(cb)
|
||
|
||
|
||
def test_get_router_model_info_does_not_wipe_cached_pricing():
|
||
"""A Deployment's model_info declares the mirrored pricing fields with None defaults;
|
||
merging it must not write those Nones into the lru_cache'd dict get_model_info() owns,
|
||
or /model/info loses built-in prices for every model a worker serves."""
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
litellm.get_model_info.cache_clear()
|
||
expected = copy.deepcopy(litellm.get_model_info(model="anthropic/claude-sonnet-4-5"))
|
||
|
||
router = litellm.Router(model_list=[])
|
||
merged = router.get_router_model_info(
|
||
deployment=Deployment(
|
||
model_name="sonnet",
|
||
litellm_params=LiteLLM_Params(model="claude-sonnet-4-5", custom_llm_provider="anthropic"),
|
||
model_info=ModelInfo(id="sonnet-1"),
|
||
),
|
||
received_model_name="sonnet",
|
||
)
|
||
|
||
assert litellm.get_model_info(model="anthropic/claude-sonnet-4-5") == expected
|
||
for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"):
|
||
assert merged[field] == expected[field]
|
||
|
||
|
||
def test_get_router_model_info_keeps_explicit_pricing_overrides():
|
||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||
|
||
litellm.get_model_info.cache_clear()
|
||
router = litellm.Router(model_list=[])
|
||
merged = router.get_router_model_info(
|
||
deployment=Deployment(
|
||
model_name="sonnet",
|
||
litellm_params=LiteLLM_Params(model="claude-sonnet-4-5", custom_llm_provider="anthropic"),
|
||
model_info=ModelInfo(id="sonnet-1", input_cost_per_token=1e-08),
|
||
),
|
||
received_model_name="sonnet",
|
||
)
|
||
|
||
assert merged["input_cost_per_token"] == 1e-08
|
||
assert litellm.get_model_info(model="anthropic/claude-sonnet-4-5")["input_cost_per_token"] != 1e-08
|