mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
929 lines
35 KiB
Python
929 lines
35 KiB
Python
"""
|
|
Tests for fallback management endpoints
|
|
|
|
Tests:
|
|
1. Create fallback configuration
|
|
2. Get fallback configuration
|
|
3. Delete fallback configuration
|
|
4. Validation tests (invalid models, duplicate fallbacks, etc.)
|
|
"""
|
|
|
|
import copy
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from types import SimpleNamespace
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from pydantic import JsonValue, TypeAdapter
|
|
|
|
from litellm import Router
|
|
from litellm.proxy.management_endpoints.fallback_management_endpoints import (
|
|
FallbackCreateRequest,
|
|
FallbackDeleteResponse,
|
|
FallbackResponse,
|
|
create_fallback,
|
|
delete_fallback,
|
|
get_fallback,
|
|
)
|
|
from litellm.proxy.proxy_server import ProxyConfig
|
|
from litellm.proxy.utils import evict_config_param, get_config_param
|
|
from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict
|
|
|
|
|
|
class TestFallbackCreateRequest:
|
|
"""Test the FallbackCreateRequest validation"""
|
|
|
|
def test_valid_request(self):
|
|
"""Test valid fallback request"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4", "claude-3-haiku"],
|
|
fallback_type="general",
|
|
)
|
|
assert request.model == "gpt-3.5-turbo"
|
|
assert request.fallback_models == ["gpt-4", "claude-3-haiku"]
|
|
assert request.fallback_type == "general"
|
|
|
|
def test_default_fallback_type(self):
|
|
"""Test default fallback type is 'general'"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
assert request.fallback_type == "general"
|
|
|
|
def test_empty_fallback_models(self):
|
|
"""Test that empty fallback_models raises validation error"""
|
|
with pytest.raises(ValueError, match="at least 1 item"):
|
|
FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=[],
|
|
)
|
|
|
|
def test_duplicate_fallback_models(self):
|
|
"""Test that duplicate fallback models raise validation error"""
|
|
with pytest.raises(
|
|
ValueError, match="fallback_models must not contain duplicates"
|
|
):
|
|
FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4", "gpt-4"],
|
|
)
|
|
|
|
def test_empty_model_name(self):
|
|
"""Test that empty model name raises validation error"""
|
|
with pytest.raises(ValueError, match="model must be a non-empty string"):
|
|
FallbackCreateRequest(
|
|
model="",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
|
|
def test_whitespace_model_name(self):
|
|
"""Test that whitespace-only model name raises validation error"""
|
|
with pytest.raises(ValueError, match="model must be a non-empty string"):
|
|
FallbackCreateRequest(
|
|
model=" ",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
|
|
def test_model_name_trimmed(self):
|
|
"""Test that model name is trimmed"""
|
|
request = FallbackCreateRequest(
|
|
model=" gpt-3.5-turbo ",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
assert request.model == "gpt-3.5-turbo"
|
|
|
|
def test_context_window_fallback_type(self):
|
|
"""Test context_window fallback type"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4-32k"],
|
|
fallback_type="context_window",
|
|
)
|
|
assert request.fallback_type == "context_window"
|
|
|
|
def test_content_policy_fallback_type(self):
|
|
"""Test content_policy fallback type"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4"],
|
|
fallback_type="content_policy",
|
|
)
|
|
assert request.fallback_type == "content_policy"
|
|
|
|
|
|
TEAM_ID: Final = "team-a"
|
|
PRIMARY_INTERNAL_NAME: Final = f"model_name_{TEAM_ID}_1a6437cb-4cab-432c-8099-1d7411731a8b"
|
|
FALLBACK_INTERNAL_NAME: Final = f"model_name_{TEAM_ID}_83151607-5556-4bbf-ac65-c474dbc64eba"
|
|
PRIMARY_DEPLOYMENT_ID: Final = "team-primary-id"
|
|
FALLBACK_DEPLOYMENT_ID: Final = "team-fallback-id"
|
|
|
|
FallbackRules = list[dict[str, list[str]] | dict[str, object] | str]
|
|
RouterSettings = dict[str, FallbackRules | None]
|
|
_ROUTER_SETTINGS_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
|
|
|
|
|
def _litellm_params(mock_response: str | None) -> LiteLLMParamsTypedDict:
|
|
if mock_response is None:
|
|
return {"model": "openai/gpt-5.4-mini", "api_key": "fake"}
|
|
return {"model": "openai/gpt-5.4-mini", "api_key": "fake", "mock_response": mock_response}
|
|
|
|
|
|
def _team_scoped_deployment(
|
|
internal_name: str, public_name: str, deployment_id: str, mock_response: str | None = None
|
|
) -> DeploymentTypedDict:
|
|
return {
|
|
"model_name": internal_name,
|
|
"litellm_params": _litellm_params(mock_response),
|
|
"model_info": {"id": deployment_id, "team_id": TEAM_ID, "team_public_model_name": public_name},
|
|
}
|
|
|
|
|
|
def _team_router(primary_mock_response: str | None = None, fallback_mock_response: str | None = None) -> Router:
|
|
return Router(
|
|
model_list=[
|
|
{"model_name": "gpt-5.4-mini", "litellm_params": _litellm_params(None)},
|
|
_team_scoped_deployment(
|
|
PRIMARY_INTERNAL_NAME, "team-primary", PRIMARY_DEPLOYMENT_ID, primary_mock_response
|
|
),
|
|
_team_scoped_deployment(
|
|
FALLBACK_INTERNAL_NAME, "team-fallback", FALLBACK_DEPLOYMENT_ID, fallback_mock_response
|
|
),
|
|
],
|
|
num_retries=0,
|
|
)
|
|
|
|
|
|
class TestCreateFallbackForTeamScopedModels:
|
|
"""POST /fallback takes the public name a caller invokes a team-scoped model by, not only the stored internal one"""
|
|
|
|
@pytest.fixture
|
|
def router(self) -> Router:
|
|
return _team_router()
|
|
|
|
@pytest.fixture
|
|
def prisma_client(self) -> MagicMock:
|
|
client: Final = MagicMock()
|
|
client.db.litellm_config.upsert = AsyncMock()
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def proxy_config(self) -> MagicMock:
|
|
config: Final = MagicMock()
|
|
config.get_config = AsyncMock(return_value={"router_settings": {}})
|
|
return config
|
|
|
|
async def _create(
|
|
self, request: FallbackCreateRequest, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
|
) -> FallbackResponse:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.llm_router", router),
|
|
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
|
|
patch("litellm.proxy.proxy_server.proxy_config", proxy_config),
|
|
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
|
):
|
|
return await create_fallback(request, MagicMock())
|
|
|
|
async def test_public_names_create_a_rule_keyed_on_the_public_name(
|
|
self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
|
) -> None:
|
|
request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"])
|
|
|
|
response: Final = await self._create(request, router, prisma_client, proxy_config)
|
|
|
|
assert response.model == "team-primary"
|
|
assert response.fallback_models == ["team-fallback"]
|
|
assert router.fallbacks == [{"team-primary": ["team-fallback"]}]
|
|
persisted: Final = json.loads(
|
|
prisma_client.db.litellm_config.upsert.call_args.kwargs["data"]["create"]["param_value"]
|
|
)
|
|
assert persisted["fallbacks"] == [{"team-primary": ["team-fallback"]}]
|
|
|
|
async def test_internal_names_keep_working(
|
|
self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
|
) -> None:
|
|
request: Final = FallbackCreateRequest(model=PRIMARY_INTERNAL_NAME, fallback_models=[FALLBACK_INTERNAL_NAME])
|
|
|
|
response: Final = await self._create(request, router, prisma_client, proxy_config)
|
|
|
|
assert router.fallbacks == [{PRIMARY_INTERNAL_NAME: [FALLBACK_INTERNAL_NAME]}]
|
|
assert response.model == PRIMARY_INTERNAL_NAME
|
|
|
|
async def test_public_fallback_target_behind_a_gateway_primary(
|
|
self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
|
) -> None:
|
|
request: Final = FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"])
|
|
|
|
await self._create(request, router, prisma_client, proxy_config)
|
|
|
|
assert router.fallbacks == [{"gpt-5.4-mini": ["team-fallback"]}]
|
|
|
|
async def test_unknown_name_is_rejected_and_the_error_names_the_public_names(
|
|
self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
|
) -> None:
|
|
request: Final = FallbackCreateRequest(model="team-missing", fallback_models=["team-fallback"])
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await self._create(request, router, prisma_client, proxy_config)
|
|
|
|
assert exc_info.value.status_code == 404
|
|
assert {"team-primary", "team-fallback", "gpt-5.4-mini"} <= set(exc_info.value.detail["available_models"])
|
|
|
|
async def test_a_team_request_fails_over_to_the_public_name_fallback(
|
|
self, prisma_client: MagicMock, proxy_config: MagicMock
|
|
) -> None:
|
|
router: Final = _team_router(primary_mock_response="litellm.RateLimitError", fallback_mock_response="pong")
|
|
request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"])
|
|
await self._create(request, router, prisma_client, proxy_config)
|
|
|
|
response: Final = await router.acompletion(
|
|
model="team-primary",
|
|
messages=[{"role": "user", "content": "ping"}],
|
|
metadata={"user_api_key_team_id": TEAM_ID},
|
|
)
|
|
|
|
assert response.choices[0].message.content == "pong"
|
|
assert response._hidden_params["model_id"] == FALLBACK_DEPLOYMENT_ID
|
|
|
|
|
|
def _config_row(router_settings: RouterSettings) -> SimpleNamespace:
|
|
return SimpleNamespace(param_name="router_settings", param_value=router_settings)
|
|
|
|
|
|
class _StoredRouterSettings:
|
|
"""The LiteLLM_Config router_settings row, with the proxy objects that read it through the config cache"""
|
|
|
|
def __init__(self, router_settings: RouterSettings | None, config_file: RouterSettings | None = None) -> None:
|
|
self.row: SimpleNamespace | None = None if router_settings is None else _config_row(router_settings)
|
|
self.prisma_client: Final = MagicMock()
|
|
self.prisma_client.get_generic_data = AsyncMock(side_effect=lambda **_: self.row)
|
|
self.prisma_client.db.litellm_config.upsert = AsyncMock(side_effect=self._upsert)
|
|
self.proxy_config: Final = MagicMock() if config_file is None else ProxyConfig()
|
|
if config_file is not None:
|
|
self.proxy_config.router_settings.load_yaml(_ROUTER_SETTINGS_JSON.validate_python(config_file))
|
|
self.proxy_config.get_config = AsyncMock(side_effect=self._get_config)
|
|
|
|
async def _upsert(self, where: dict[str, str], data: dict[str, dict[str, str]]) -> None:
|
|
self.row = _config_row(json.loads(data["update"]["param_value"]))
|
|
|
|
async def _get_config(self) -> dict[str, RouterSettings]:
|
|
row: Final = await get_config_param(self.prisma_client, "router_settings")
|
|
return {"router_settings": copy.deepcopy(row.param_value) if row is not None else {}}
|
|
|
|
def written_by_another_instance(self, router_settings: RouterSettings) -> None:
|
|
self.row = _config_row(router_settings)
|
|
|
|
def stored_fallbacks(self) -> FallbackRules:
|
|
assert self.row is not None
|
|
return self.row.param_value["fallbacks"]
|
|
|
|
async def cached_fallbacks(self) -> FallbackRules:
|
|
row: Final = await get_config_param(self.prisma_client, "router_settings")
|
|
return row.param_value["fallbacks"]
|
|
|
|
|
|
TEAM_RULE: Final = {"team-primary": ["team-fallback"]}
|
|
GATEWAY_RULE: Final = {"gpt-5.4-mini": ["team-fallback"]}
|
|
MALFORMED_GATEWAY_RULE: Final = {"gpt-5.4-mini": "team-fallback"}
|
|
NON_STANDARD_RULES: Final = [
|
|
"claude-3-haiku",
|
|
{"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "retry"}]},
|
|
]
|
|
|
|
|
|
class TestFallbackWritesSeeTheLatestStoredRules:
|
|
"""A write reads the rules the database holds now and leaves no stale copy in the config cache behind"""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
async def clean_config_cache(self) -> AsyncIterator[None]:
|
|
await evict_config_param("router_settings")
|
|
yield
|
|
await evict_config_param("router_settings")
|
|
|
|
async def _create(self, request: FallbackCreateRequest, stored: _StoredRouterSettings) -> FallbackResponse:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.llm_router", _team_router()),
|
|
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client),
|
|
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
|
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
|
):
|
|
return await create_fallback(request, MagicMock())
|
|
|
|
async def _delete(self, model: str, stored: _StoredRouterSettings) -> FallbackDeleteResponse:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.llm_router", _team_router()),
|
|
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client),
|
|
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
|
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
|
):
|
|
return await delete_fallback(model, "general", MagicMock())
|
|
|
|
async def test_a_second_create_keeps_the_first_rule(self) -> None:
|
|
stored: Final = _StoredRouterSettings(None)
|
|
|
|
await self._create(FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]), stored)
|
|
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
|
|
|
assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE]
|
|
assert await stored.cached_fallbacks() == [TEAM_RULE, GATEWAY_RULE]
|
|
|
|
async def test_a_create_keeps_a_rule_another_instance_stored_since_this_one_last_read(self) -> None:
|
|
stored: Final = _StoredRouterSettings({})
|
|
await get_config_param(stored.prisma_client, "router_settings")
|
|
stored.written_by_another_instance({"fallbacks": [TEAM_RULE]})
|
|
|
|
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
|
|
|
assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE]
|
|
|
|
async def test_a_delete_keeps_a_rule_another_instance_stored_since_this_one_last_read(self) -> None:
|
|
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]})
|
|
await get_config_param(stored.prisma_client, "router_settings")
|
|
stored.written_by_another_instance({"fallbacks": [TEAM_RULE, GATEWAY_RULE]})
|
|
|
|
await self._delete("team-primary", stored)
|
|
|
|
assert stored.stored_fallbacks() == [GATEWAY_RULE]
|
|
assert await stored.cached_fallbacks() == [GATEWAY_RULE]
|
|
|
|
async def test_writes_keep_the_non_standard_rules_the_router_accepts(self) -> None:
|
|
stored: Final = _StoredRouterSettings({"fallbacks": [*NON_STANDARD_RULES, TEAM_RULE]})
|
|
|
|
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
|
assert stored.stored_fallbacks() == [*NON_STANDARD_RULES, TEAM_RULE, GATEWAY_RULE]
|
|
|
|
await self._delete("team-primary", stored)
|
|
assert stored.stored_fallbacks() == [*NON_STANDARD_RULES, GATEWAY_RULE]
|
|
|
|
async def test_a_create_replaces_a_same_key_rule_whatever_its_targets_shape(self) -> None:
|
|
stored: Final = _StoredRouterSettings({"fallbacks": [MALFORMED_GATEWAY_RULE, TEAM_RULE]})
|
|
|
|
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
|
|
|
assert stored.stored_fallbacks() == [GATEWAY_RULE, TEAM_RULE]
|
|
|
|
async def test_a_delete_removes_a_same_key_rule_whatever_its_targets_shape(self) -> None:
|
|
stored: Final = _StoredRouterSettings({"fallbacks": [MALFORMED_GATEWAY_RULE, TEAM_RULE]})
|
|
|
|
await self._delete("gpt-5.4-mini", stored)
|
|
|
|
assert stored.stored_fallbacks() == [TEAM_RULE]
|
|
|
|
async def test_a_create_treats_a_null_rule_list_as_empty(self) -> None:
|
|
stored: Final = _StoredRouterSettings({"fallbacks": None})
|
|
|
|
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
|
|
|
assert stored.stored_fallbacks() == [GATEWAY_RULE]
|
|
|
|
|
|
CONFIG_RULE: Final = {"gpt-5.4-mini": ["team-fallback"]}
|
|
CONFIG_WINDOW_RULE: Final = {"team-primary": ["gpt-5.4-mini"]}
|
|
|
|
|
|
class TestFallbackWritesHonourTheConfigFile:
|
|
"""A fallbacks list set in config.yaml owns the key: /fallback writes that would change it are refused and
|
|
the router keeps serving the file's rules"""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
async def clean_config_cache(self) -> AsyncIterator[None]:
|
|
await evict_config_param("router_settings")
|
|
yield
|
|
await evict_config_param("router_settings")
|
|
|
|
@pytest.fixture
|
|
def router(self) -> Router:
|
|
return _team_router()
|
|
|
|
@pytest.fixture
|
|
def stored(self) -> _StoredRouterSettings:
|
|
return _StoredRouterSettings({"fallbacks": [CONFIG_RULE]}, config_file={"fallbacks": [CONFIG_RULE]})
|
|
|
|
async def _create(
|
|
self, request: FallbackCreateRequest, router: Router, stored: _StoredRouterSettings
|
|
) -> FallbackResponse:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.llm_router", router),
|
|
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client),
|
|
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
|
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
|
):
|
|
return await create_fallback(request, MagicMock())
|
|
|
|
async def _delete(self, model: str, router: Router, stored: _StoredRouterSettings) -> FallbackDeleteResponse:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.llm_router", router),
|
|
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client),
|
|
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
|
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
|
):
|
|
return await delete_fallback(model, "general", MagicMock())
|
|
|
|
async def test_a_create_that_adds_a_rule_to_a_config_owned_list_is_refused(
|
|
self, router: Router, stored: _StoredRouterSettings
|
|
) -> None:
|
|
router.fallbacks = [CONFIG_RULE]
|
|
request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"])
|
|
|
|
with pytest.raises(HTTPException) as refused:
|
|
await self._create(request, router, stored)
|
|
|
|
assert refused.value.status_code == 400
|
|
assert "key 'fallbacks' is set in the config file" in str(refused.value.detail)
|
|
assert stored.stored_fallbacks() == [CONFIG_RULE]
|
|
stored.prisma_client.db.litellm_config.upsert.assert_not_awaited()
|
|
assert router.fallbacks == [CONFIG_RULE]
|
|
|
|
async def test_a_delete_of_a_config_owned_rule_is_refused(
|
|
self, router: Router, stored: _StoredRouterSettings
|
|
) -> None:
|
|
router.fallbacks = [CONFIG_RULE]
|
|
|
|
with pytest.raises(HTTPException) as refused:
|
|
await self._delete("gpt-5.4-mini", router, stored)
|
|
|
|
assert refused.value.status_code == 400
|
|
assert "key 'fallbacks' is set in the config file" in str(refused.value.detail)
|
|
assert stored.stored_fallbacks() == [CONFIG_RULE]
|
|
stored.prisma_client.db.litellm_config.upsert.assert_not_awaited()
|
|
assert router.fallbacks == [CONFIG_RULE]
|
|
|
|
async def test_a_create_that_restates_the_config_rule_is_accepted(
|
|
self, router: Router, stored: _StoredRouterSettings
|
|
) -> None:
|
|
request: Final = FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"])
|
|
|
|
response: Final = await self._create(request, router, stored)
|
|
|
|
assert response.fallback_models == ["team-fallback"]
|
|
assert stored.stored_fallbacks() == [CONFIG_RULE]
|
|
|
|
async def test_a_sibling_list_the_config_file_does_not_set_stays_writable(
|
|
self, router: Router, stored: _StoredRouterSettings
|
|
) -> None:
|
|
request: Final = FallbackCreateRequest(
|
|
model="team-primary", fallback_models=["gpt-5.4-mini"], fallback_type="context_window"
|
|
)
|
|
|
|
response: Final = await self._create(request, router, stored)
|
|
|
|
assert response.fallback_type == "context_window"
|
|
assert stored.row is not None
|
|
assert stored.row.param_value["context_window_fallbacks"] == [CONFIG_WINDOW_RULE]
|
|
assert router.context_window_fallbacks == [CONFIG_WINDOW_RULE]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestCreateFallback:
|
|
"""Test the create_fallback endpoint"""
|
|
|
|
@pytest.fixture
|
|
def mock_router(self):
|
|
"""Create a mock router"""
|
|
router = MagicMock()
|
|
router.model_names = {"gpt-3.5-turbo", "gpt-4", "claude-3-haiku"}
|
|
router.team_public_model_names = frozenset()
|
|
router.fallbacks = []
|
|
router.context_window_fallbacks = []
|
|
router.content_policy_fallbacks = []
|
|
return router
|
|
|
|
@pytest.fixture
|
|
def mock_prisma_client(self):
|
|
"""Create a mock prisma client"""
|
|
client = MagicMock()
|
|
client.db.litellm_config.upsert = AsyncMock()
|
|
client.jsonify_object = lambda x: x
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def mock_proxy_config(self):
|
|
"""Create a mock proxy config"""
|
|
config = MagicMock()
|
|
config.get_config = AsyncMock(return_value={"router_settings": {}})
|
|
return config
|
|
|
|
@pytest.fixture
|
|
def mock_user_api_key_dict(self):
|
|
"""Create a mock user API key dict"""
|
|
return MagicMock()
|
|
|
|
async def test_create_fallback_success(
|
|
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
|
):
|
|
"""Test successful fallback creation"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4", "claude-3-haiku"],
|
|
fallback_type="general",
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.proxy_config",
|
|
mock_proxy_config,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
):
|
|
response = await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert response.model == "gpt-3.5-turbo"
|
|
assert response.fallback_models == ["gpt-4", "claude-3-haiku"]
|
|
assert response.fallback_type == "general"
|
|
assert (
|
|
"created" in response.message.lower()
|
|
or "updated" in response.message.lower()
|
|
)
|
|
|
|
# Verify database was updated
|
|
mock_prisma_client.db.litellm_config.upsert.assert_called_once()
|
|
|
|
async def test_create_fallback_router_not_initialized(
|
|
self, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
|
):
|
|
"""Test error when router is not initialized"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
None,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 500
|
|
assert "Router not initialized" in str(exc_info.value.detail)
|
|
|
|
async def test_create_fallback_model_not_found(
|
|
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
|
):
|
|
"""Test error when model is not found in router"""
|
|
request = FallbackCreateRequest(
|
|
model="invalid-model",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 404
|
|
assert "not found in router" in str(exc_info.value.detail)
|
|
|
|
async def test_create_fallback_invalid_fallback_model(
|
|
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
|
):
|
|
"""Test error when fallback model is not found in router"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["invalid-fallback-model"],
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert "Invalid fallback models" in str(exc_info.value.detail)
|
|
|
|
async def test_create_fallback_model_is_own_fallback(
|
|
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
|
):
|
|
"""Test error when model is its own fallback"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-3.5-turbo", "gpt-4"],
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert "cannot be its own fallback" in str(exc_info.value.detail)
|
|
|
|
async def test_create_fallback_db_not_enabled(
|
|
self, mock_router, mock_user_api_key_dict
|
|
):
|
|
"""Test error when database storage is not enabled"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4"],
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
False,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert "Database storage not enabled" in str(exc_info.value.detail)
|
|
|
|
async def test_create_fallback_context_window_type(
|
|
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
|
):
|
|
"""Test creating context_window fallback"""
|
|
request = FallbackCreateRequest(
|
|
model="gpt-3.5-turbo",
|
|
fallback_models=["gpt-4"],
|
|
fallback_type="context_window",
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.proxy_config",
|
|
mock_proxy_config,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
):
|
|
response = await create_fallback(request, mock_user_api_key_dict)
|
|
|
|
assert response.fallback_type == "context_window"
|
|
# Verify the correct attribute was updated
|
|
assert hasattr(mock_router, "context_window_fallbacks")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestGetFallback:
|
|
"""Test the get_fallback endpoint"""
|
|
|
|
@pytest.fixture
|
|
def mock_router_with_fallbacks(self):
|
|
"""Create a mock router with fallbacks configured"""
|
|
router = MagicMock()
|
|
router.fallbacks = [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
|
|
router.context_window_fallbacks = []
|
|
router.content_policy_fallbacks = []
|
|
return router
|
|
|
|
@pytest.fixture
|
|
def mock_user_api_key_dict(self):
|
|
"""Create a mock user API key dict"""
|
|
return MagicMock()
|
|
|
|
async def test_get_fallback_success(
|
|
self, mock_router_with_fallbacks, mock_user_api_key_dict
|
|
):
|
|
"""Test successful fallback retrieval"""
|
|
with patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router_with_fallbacks,
|
|
):
|
|
response = await get_fallback(
|
|
"gpt-3.5-turbo", "general", mock_user_api_key_dict
|
|
)
|
|
|
|
assert response.model == "gpt-3.5-turbo"
|
|
assert response.fallback_models == ["gpt-4", "claude-3-haiku"]
|
|
assert response.fallback_type == "general"
|
|
|
|
async def test_get_fallback_not_found(
|
|
self, mock_router_with_fallbacks, mock_user_api_key_dict
|
|
):
|
|
"""Test error when fallback is not found"""
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router_with_fallbacks,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await get_fallback("gpt-4", "general", mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 404
|
|
assert "No general fallbacks configured" in str(exc_info.value.detail)
|
|
|
|
async def test_get_fallback_router_not_initialized(self, mock_user_api_key_dict):
|
|
"""Test error when router is not initialized"""
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
None,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await get_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 500
|
|
assert "Router not initialized" in str(exc_info.value.detail)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestDeleteFallback:
|
|
"""Test the delete_fallback endpoint"""
|
|
|
|
@pytest.fixture
|
|
def mock_router_with_fallbacks(self):
|
|
"""Create a mock router with fallbacks configured"""
|
|
router = MagicMock()
|
|
router.fallbacks = [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
|
|
router.context_window_fallbacks = []
|
|
router.content_policy_fallbacks = []
|
|
return router
|
|
|
|
@pytest.fixture
|
|
def mock_prisma_client(self):
|
|
"""Create a mock prisma client"""
|
|
client = MagicMock()
|
|
client.db.litellm_config.upsert = AsyncMock()
|
|
client.jsonify_object = lambda x: x
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def mock_proxy_config(self):
|
|
"""Create a mock proxy config"""
|
|
config = MagicMock()
|
|
config.get_config = AsyncMock(
|
|
return_value={
|
|
"router_settings": {
|
|
"fallbacks": [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
|
|
}
|
|
}
|
|
)
|
|
return config
|
|
|
|
@pytest.fixture
|
|
def mock_user_api_key_dict(self):
|
|
"""Create a mock user API key dict"""
|
|
return MagicMock()
|
|
|
|
async def test_delete_fallback_success(
|
|
self,
|
|
mock_router_with_fallbacks,
|
|
mock_prisma_client,
|
|
mock_proxy_config,
|
|
mock_user_api_key_dict,
|
|
):
|
|
"""Test successful fallback deletion"""
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router_with_fallbacks,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.proxy_config",
|
|
mock_proxy_config,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
):
|
|
response = await delete_fallback(
|
|
"gpt-3.5-turbo", "general", mock_user_api_key_dict
|
|
)
|
|
|
|
assert response.model == "gpt-3.5-turbo"
|
|
assert response.fallback_type == "general"
|
|
assert "deleted" in response.message.lower()
|
|
|
|
# Verify database was updated
|
|
mock_prisma_client.db.litellm_config.upsert.assert_called_once()
|
|
|
|
async def test_delete_fallback_not_found(
|
|
self,
|
|
mock_router_with_fallbacks,
|
|
mock_prisma_client,
|
|
mock_proxy_config,
|
|
mock_user_api_key_dict,
|
|
):
|
|
"""Test error when fallback to delete is not found"""
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router_with_fallbacks,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.prisma_client",
|
|
mock_prisma_client,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.proxy_config",
|
|
mock_proxy_config,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
True,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await delete_fallback("gpt-4", "general", mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 404
|
|
assert "No general fallbacks configured" in str(exc_info.value.detail)
|
|
|
|
async def test_delete_fallback_router_not_initialized(self, mock_user_api_key_dict):
|
|
"""Test error when router is not initialized"""
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
None,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 500
|
|
assert "Router not initialized" in str(exc_info.value.detail)
|
|
|
|
async def test_delete_fallback_db_not_enabled(
|
|
self, mock_router_with_fallbacks, mock_user_api_key_dict
|
|
):
|
|
"""Test error when database storage is not enabled"""
|
|
with (
|
|
patch(
|
|
"litellm.proxy.proxy_server.llm_router",
|
|
mock_router_with_fallbacks,
|
|
),
|
|
patch(
|
|
"litellm.proxy.proxy_server.store_model_in_db",
|
|
False,
|
|
),
|
|
pytest.raises(HTTPException) as exc_info,
|
|
):
|
|
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert "Database storage not enabled" in str(exc_info.value.detail)
|