mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(fallback): preserve concurrent writes and serve saved configuration (#51245)
* fix(fallback): serialize fallback writes and wait for replicas on read-back POST and DELETE /fallback now read, edit and write router_settings in one writer transaction under the same per-param advisory lock the config section writer takes, so concurrent writes on any replica no longer drop each other's rules. The cache is invalidated and the router updated only after commit. The Terraform provider now polls the read after a litellm_fallback create or update until the proxy serves the written list, within new create/update timeouts, instead of clearing the ID on a stale 404 or saving a stale list. Refresh 404 still removes the resource, and error responses are not retried. Fixes #47888 * fix(fallback): bound provider requests by the apply context and inject the fallback repository The provider's fallback write and read-back requests now carry the apply context through a context-aware MakeRequest variant, so a cancelled apply sends nothing and the create/update timeout also ends a hung request. The POST and DELETE /fallback routes take a RouterFallbackRepositoryDep that returns a RouterFallbackStore. Router settings parsing moved into RouterFallbackRepository, which edits the rules through the locked ConfigRepository operation. Unit tests now check the outcome of the race and rollback through a fake lock instead of pinning the lock SQL. They also cover a racing config section write, a 404 delete inside the lock, route resolution of the repository and the no-database path. * refactor(fallback): keep atomic edits in the config repository * fix(terraform): release hung test handlers and audit MakeRequestWithContext calls The hung-request cancellation test now releases its handler through a test owned channel before closing the server, since the server does not notice the client's cancel while the POST body is unread. The endpoint audit now reads the method and path of MakeRequestWithContext calls from their shifted argument positions, and skips only MakeRequest's own forwarding call inside the helper files, so POST /fallback counts as covered again and dynamic paths elsewhere still fail closed. * test(fallback): keep stored-rule regressions with their helpers * fix(fallback): read saved rules without Terraform changes * fix(fallback): preserve valid rules when stored entries are malformed --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri
GitHub
parent
d5a1602968
commit
2c7ed0713a
7 changed files with 782 additions and 210 deletions
25
litellm/models/fallback.py
Normal file
25
litellm/models/fallback.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Final, Protocol, TypeAlias
|
||||
|
||||
from pydantic import Field, JsonValue, TypeAdapter
|
||||
|
||||
FallbackRule = dict[str, list[str]]
|
||||
StoredFallback = Annotated[FallbackRule | dict[str, object] | str, Field(union_mode="left_to_right")]
|
||||
STORED_FALLBACKS: Final = TypeAdapter(list[StoredFallback])
|
||||
SavedFallbackRule: TypeAlias = JsonValue
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FallbackRulesEdit:
|
||||
rules: Sequence[StoredFallback]
|
||||
replaced: bool
|
||||
|
||||
|
||||
FallbackRulesEditor: TypeAlias = Callable[[Sequence[StoredFallback]], FallbackRulesEdit]
|
||||
|
||||
|
||||
class RouterFallbackStore(Protocol):
|
||||
async def saved_rules(self, fallback_key: str) -> Sequence[SavedFallbackRule] | None: ...
|
||||
|
||||
async def edit_rules(self, fallback_key: str, edit: FallbackRulesEditor) -> FallbackRulesEdit: ...
|
||||
|
|
@ -3,11 +3,13 @@ from typing import Annotated, TypeAlias
|
|||
from fastapi import Depends
|
||||
|
||||
from litellm.models.deployment_warning import DeploymentWarningIgnoreStore
|
||||
from litellm.models.fallback import RouterFallbackStore
|
||||
from litellm.models.gateway_activity import GatewayActivityStore
|
||||
from litellm.models.plugin_upload import PluginUploadStore
|
||||
from litellm.models.skill_submission import SkillSubmissionStore
|
||||
from litellm.models.user_banner import UserBannerStore
|
||||
from litellm.proxy.deps.database import DatabaseDep
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
|
||||
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
|
||||
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
|
||||
|
|
@ -54,3 +56,10 @@ async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmiss
|
|||
SkillSubmissionRepositoryDep: TypeAlias = Annotated[
|
||||
SkillSubmissionStore | None, Depends(get_skill_submission_repository)
|
||||
]
|
||||
|
||||
|
||||
async def get_router_fallback_repository(database: DatabaseDep) -> RouterFallbackStore | None:
|
||||
return None if database is None else ConfigRepository(database.prisma, use_writer=True)
|
||||
|
||||
|
||||
RouterFallbackRepositoryDep: TypeAlias = Annotated[RouterFallbackStore | None, Depends(get_router_fallback_repository)]
|
||||
|
|
|
|||
|
|
@ -10,11 +10,11 @@ DELETE /fallback/{model} - Delete fallbacks for a specific model
|
|||
|
||||
# pyright: reportMissingImports=false
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import Field, JsonValue, TypeAdapter, ValidationError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -22,7 +22,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.auth.auth_utils import fallback_target_model_name
|
||||
from litellm.proxy.auth.model_checks import get_all_fallbacks
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.utils import PrismaClient, invalidate_config_param
|
||||
from litellm.proxy.utils import invalidate_config_param
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
|
|
@ -36,8 +36,16 @@ else:
|
|||
# fastapi is only required for proxy, not for SDK usage
|
||||
pass
|
||||
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.router_api import ChangeSet
|
||||
from litellm.models.fallback import (
|
||||
STORED_FALLBACKS,
|
||||
FallbackRulesEdit,
|
||||
FallbackRulesEditor,
|
||||
RouterFallbackStore,
|
||||
StoredFallback,
|
||||
)
|
||||
from litellm.proxy.deps.repositories import RouterFallbackRepositoryDep
|
||||
from litellm.repositories.config_repository import ROUTER_SETTINGS_PARAM
|
||||
from litellm.router_api import ChangeSet, fallback_chain
|
||||
from litellm.types.management_endpoints.router_settings_endpoints import (
|
||||
FallbackCreateRequest,
|
||||
FallbackDeleteResponse,
|
||||
|
|
@ -48,12 +56,10 @@ from litellm.types.router import ContentPolicyFallbackRule, ContentPolicyFallbac
|
|||
|
||||
router: Final = APIRouter()
|
||||
|
||||
ROUTER_SETTINGS_PARAM: Final = "router_settings"
|
||||
FallbackRule = dict[str, list[str]]
|
||||
StoredFallback = Annotated[FallbackRule | dict[str, object] | str, Field(union_mode="left_to_right")]
|
||||
_STORED_FALLBACKS: Final = TypeAdapter(list[StoredFallback])
|
||||
_JSON_RULES: Final = TypeAdapter(list[JsonValue])
|
||||
_STORED_ROUTER_SETTINGS: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_FALLBACK_KEY_BY_TYPE: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"general": "fallbacks", "context_window": "context_window_fallbacks", "content_policy": "content_policy_fallbacks"}
|
||||
)
|
||||
_OPTIONAL_JSON_RULES: Final[TypeAdapter[Sequence[JsonValue] | None]] = TypeAdapter(Sequence[JsonValue] | None)
|
||||
_CONTENT_POLICY_FALLBACK_RULE: Final[TypeAdapter[ContentPolicyFallbackRule]] = TypeAdapter(ContentPolicyFallbackRule)
|
||||
_FALLBACK_TARGETS: Final[TypeAdapter[FallbackTargets]] = TypeAdapter(FallbackTargets)
|
||||
|
|
@ -104,18 +110,20 @@ def _checked_content_policy_rules(
|
|||
return tuple((candidate, _checked_content_policy_rule(candidate)) for candidate in candidates)
|
||||
|
||||
|
||||
def _layered_checked_rules(
|
||||
checked_rules: Sequence[tuple[object, ContentPolicyFallbackRule | ValidationError]],
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
return _layered_over_config_rules(
|
||||
tuple(result for _, result in checked_rules if not isinstance(result, ValidationError)),
|
||||
_config_file_rules(
|
||||
litellm.content_policy_fallbacks # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped module list
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def apply_content_policy_fallbacks(llm_router: "Router", rules: object, *, log_skipped_rules: bool) -> None:
|
||||
checked_rules: Final = _checked_content_policy_rules(rules)
|
||||
runtime_rules: Final = tuple(result for _, result in checked_rules if not isinstance(result, ValidationError))
|
||||
layered_rules: Final = (
|
||||
_layered_over_config_rules(
|
||||
runtime_rules,
|
||||
_config_file_rules(
|
||||
litellm.content_policy_fallbacks # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped module list
|
||||
),
|
||||
)
|
||||
or None
|
||||
)
|
||||
layered_rules: Final = _layered_checked_rules(checked_rules) or None
|
||||
for rule, result in checked_rules:
|
||||
if log_skipped_rules and isinstance(result, ValidationError):
|
||||
_log_ignored_content_policy_rules(rule, result)
|
||||
|
|
@ -135,54 +143,52 @@ async def _reload_config_file_ownership(proxy_config: "ProxyConfig") -> None:
|
|||
await proxy_config.get_config()
|
||||
|
||||
|
||||
async def _stored_router_settings(prisma_client: PrismaClient) -> Mapping[str, JsonValue]:
|
||||
try:
|
||||
row: Final = await ConfigRepository(prisma_client, use_writer=True).get_param(ROUTER_SETTINGS_PARAM)
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
if row is None or not isinstance(row.param_value, Mapping):
|
||||
return {}
|
||||
return _STORED_ROUTER_SETTINGS.validate_python(row.param_value)
|
||||
def _with_rule(rules: Sequence[StoredFallback], model: str, fallback_models: Sequence[str]) -> FallbackRulesEdit:
|
||||
covering: Final = next((index for index, rule in enumerate(rules) if _rule_covers(rule, model)), None)
|
||||
rule: Final[StoredFallback] = {model: list(fallback_models)}
|
||||
if covering is None:
|
||||
return FallbackRulesEdit(rules=[*rules, rule], replaced=False)
|
||||
return FallbackRulesEdit(rules=[*rules[:covering], rule, *rules[covering + 1 :]], replaced=True)
|
||||
|
||||
|
||||
def _editable_rules(
|
||||
proxy_config: "ProxyConfig", stored_settings: Mapping[str, JsonValue], fallback_key: str
|
||||
) -> list[StoredFallback]:
|
||||
file_settings: Final = proxy_config.router_settings
|
||||
if not file_settings.owned_by_config(fallback_key):
|
||||
return _STORED_FALLBACKS.validate_python(stored_settings.get(fallback_key) or [])
|
||||
file_rules: Final = file_settings.get(fallback_key)
|
||||
return _STORED_FALLBACKS.validate_python(file_rules if isinstance(file_rules, list) else [])
|
||||
def _without_rule(rules: Sequence[StoredFallback], model: str, fallback_type: str) -> FallbackRulesEdit:
|
||||
kept: Final = [rule for rule in rules if not _rule_covers(rule, model)]
|
||||
if len(kept) == len(rules):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"},
|
||||
)
|
||||
return FallbackRulesEdit(rules=kept, replaced=True)
|
||||
|
||||
|
||||
def _edit_config_file_rules(
|
||||
proxy_config: "ProxyConfig",
|
||||
llm_router: "Router",
|
||||
fallback_key: str,
|
||||
edit: FallbackRulesEditor,
|
||||
) -> FallbackRulesEdit:
|
||||
file_rules: Final = proxy_config.router_settings.get(fallback_key)
|
||||
edited: Final = edit(STORED_FALLBACKS.validate_python(file_rules if isinstance(file_rules, list) else []))
|
||||
proxy_config.reject_config_owned_writes(
|
||||
section_name=ROUTER_SETTINGS_PARAM, changed_keys={fallback_key: _JSON_RULES.validate_python(edited.rules)}
|
||||
)
|
||||
llm_router.apply(ChangeSet(assign={fallback_key: edited.rules}))
|
||||
return edited
|
||||
|
||||
|
||||
async def _write_rules(
|
||||
prisma_client: PrismaClient,
|
||||
fallbacks: RouterFallbackStore,
|
||||
proxy_config: "ProxyConfig",
|
||||
llm_router: "Router",
|
||||
stored_settings: Mapping[str, JsonValue],
|
||||
fallback_key: str,
|
||||
rules: Sequence[StoredFallback],
|
||||
) -> None:
|
||||
proxy_config.reject_config_owned_writes(
|
||||
section_name=ROUTER_SETTINGS_PARAM, changed_keys={fallback_key: _JSON_RULES.validate_python(rules)}
|
||||
)
|
||||
edit: FallbackRulesEditor,
|
||||
) -> FallbackRulesEdit:
|
||||
if proxy_config.router_settings.owned_by_config(fallback_key):
|
||||
llm_router.apply(ChangeSet(assign={fallback_key: rules}))
|
||||
return
|
||||
await _persist_router_settings(prisma_client, {**stored_settings, fallback_key: rules})
|
||||
_apply_fallback_rules(llm_router, fallback_key, rules)
|
||||
|
||||
|
||||
async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None:
|
||||
router_settings_json: Final = json.dumps(router_settings)
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": ROUTER_SETTINGS_PARAM},
|
||||
data={
|
||||
"create": {"param_name": ROUTER_SETTINGS_PARAM, "param_value": router_settings_json},
|
||||
"update": {"param_value": router_settings_json},
|
||||
},
|
||||
)
|
||||
return _edit_config_file_rules(proxy_config, llm_router, fallback_key, edit)
|
||||
edited: Final = await fallbacks.edit_rules(fallback_key, edit)
|
||||
await invalidate_config_param(ROUTER_SETTINGS_PARAM)
|
||||
_apply_fallback_rules(llm_router, fallback_key, edited.rules)
|
||||
return edited
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -195,6 +201,8 @@ async def _persist_router_settings(prisma_client: PrismaClient, router_settings:
|
|||
async def create_fallback(
|
||||
data: FallbackCreateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
*,
|
||||
fallbacks: RouterFallbackRepositoryDep,
|
||||
):
|
||||
"""
|
||||
Create or update fallbacks for a specific model.
|
||||
|
|
@ -218,7 +226,6 @@ async def create_fallback(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
llm_router,
|
||||
prisma_client,
|
||||
proxy_config,
|
||||
store_model_in_db,
|
||||
)
|
||||
|
|
@ -262,7 +269,7 @@ async def create_fallback(
|
|||
)
|
||||
|
||||
# Check if we need to store in DB
|
||||
if store_model_in_db is not True or prisma_client is None:
|
||||
if store_model_in_db is not True or fallbacks is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={
|
||||
|
|
@ -271,7 +278,6 @@ async def create_fallback(
|
|||
)
|
||||
|
||||
await _reload_config_file_ownership(proxy_config)
|
||||
router_settings: Final = await _stored_router_settings(prisma_client)
|
||||
|
||||
# Get the appropriate fallback list based on type
|
||||
fallback_key = "fallbacks"
|
||||
|
|
@ -280,23 +286,13 @@ async def create_fallback(
|
|||
elif data.fallback_type == "content_policy":
|
||||
fallback_key = "content_policy_fallbacks"
|
||||
|
||||
# Get existing fallbacks
|
||||
existing_fallbacks: Final = _editable_rules(proxy_config, router_settings, fallback_key)
|
||||
|
||||
# Update or add the fallback configuration
|
||||
fallback_updated = False
|
||||
for i, rule in enumerate(existing_fallbacks):
|
||||
if _rule_covers(rule, data.model):
|
||||
# Update existing fallback
|
||||
existing_fallbacks[i] = {data.model: data.fallback_models}
|
||||
fallback_updated = True
|
||||
break
|
||||
|
||||
if not fallback_updated:
|
||||
# Add new fallback
|
||||
existing_fallbacks.append({data.model: data.fallback_models})
|
||||
|
||||
await _write_rules(prisma_client, proxy_config, llm_router, router_settings, fallback_key, existing_fallbacks)
|
||||
edited: Final = await _write_rules(
|
||||
fallbacks,
|
||||
proxy_config,
|
||||
llm_router,
|
||||
fallback_key,
|
||||
lambda rules: _with_rule(rules, data.model, data.fallback_models),
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"Fallback configured: %s -> %s (type: %s)", data.model, data.fallback_models, data.fallback_type
|
||||
|
|
@ -306,7 +302,7 @@ async def create_fallback(
|
|||
model=data.model,
|
||||
fallback_models=data.fallback_models,
|
||||
fallback_type=data.fallback_type,
|
||||
message=f"Fallback configuration {'updated' if fallback_updated else 'created'} successfully",
|
||||
message=f"Fallback configuration {'updated' if edited.replaced else 'created'} successfully",
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
|
|
@ -319,6 +315,26 @@ async def create_fallback(
|
|||
)
|
||||
|
||||
|
||||
async def _saved_rules_in_effect(
|
||||
proxy_config: "ProxyConfig", fallbacks: RouterFallbackStore, fallback_key: str
|
||||
) -> Sequence[object]:
|
||||
if proxy_config.router_settings.owned_by_config(fallback_key):
|
||||
file_rules: Final = proxy_config.router_settings.get(fallback_key)
|
||||
return file_rules if isinstance(file_rules, list) else ()
|
||||
saved: Final = await fallbacks.saved_rules(fallback_key)
|
||||
if fallback_key == "content_policy_fallbacks":
|
||||
return _layered_checked_rules(_checked_content_policy_rules(saved))
|
||||
return saved or ()
|
||||
|
||||
|
||||
def _fallback_chain_for(rules: Sequence[object], model: str) -> object:
|
||||
chain, _ = fallback_chain(
|
||||
fallbacks=rules, # pyright: ignore[reportArgumentType] # the chain also resolves the string and rich-target rules the router accepts
|
||||
model_group=model,
|
||||
)
|
||||
return chain or []
|
||||
|
||||
|
||||
@router.get(
|
||||
"/fallback/{model}",
|
||||
tags=["Fallback Management"],
|
||||
|
|
@ -329,6 +345,8 @@ async def get_fallback(
|
|||
model: str,
|
||||
fallback_type: Literal["general", "context_window", "content_policy"] = "general",
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
*,
|
||||
fallbacks: RouterFallbackRepositoryDep,
|
||||
):
|
||||
"""
|
||||
Get fallback configuration for a specific model.
|
||||
|
|
@ -342,7 +360,7 @@ async def get_fallback(
|
|||
GET /fallback/gpt-3.5-turbo?fallback_type=general
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
from litellm.proxy.proxy_server import llm_router, proxy_config, store_model_in_db
|
||||
|
||||
try:
|
||||
if llm_router is None:
|
||||
|
|
@ -351,9 +369,12 @@ async def get_fallback(
|
|||
detail={"error": "Router not initialized"},
|
||||
)
|
||||
|
||||
# Get fallbacks using the existing utility function
|
||||
fallback_targets: Final = _FALLBACK_TARGETS.validate_python(
|
||||
get_all_fallbacks(model=model, llm_router=llm_router, fallback_type=fallback_type)
|
||||
if store_model_in_db is not True or fallbacks is None
|
||||
else _fallback_chain_for(
|
||||
await _saved_rules_in_effect(proxy_config, fallbacks, _FALLBACK_KEY_BY_TYPE[fallback_type]), model
|
||||
)
|
||||
)
|
||||
fallback_models: Final = [fallback_target_model_name(target) or model for target in fallback_targets]
|
||||
|
||||
|
|
@ -389,6 +410,8 @@ async def delete_fallback(
|
|||
model: str,
|
||||
fallback_type: Literal["general", "context_window", "content_policy"] = "general",
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
*,
|
||||
fallbacks: RouterFallbackRepositoryDep,
|
||||
):
|
||||
"""
|
||||
Delete fallback configuration for a specific model.
|
||||
|
|
@ -404,7 +427,6 @@ async def delete_fallback(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
llm_router,
|
||||
prisma_client,
|
||||
proxy_config,
|
||||
store_model_in_db,
|
||||
)
|
||||
|
|
@ -416,7 +438,7 @@ async def delete_fallback(
|
|||
detail={"error": "Router not initialized"},
|
||||
)
|
||||
|
||||
if store_model_in_db is not True or prisma_client is None:
|
||||
if store_model_in_db is not True or fallbacks is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={
|
||||
|
|
@ -425,7 +447,6 @@ async def delete_fallback(
|
|||
)
|
||||
|
||||
await _reload_config_file_ownership(proxy_config)
|
||||
router_settings: Final = await _stored_router_settings(prisma_client)
|
||||
|
||||
# Get the appropriate fallback list based on type
|
||||
fallback_key = "fallbacks"
|
||||
|
|
@ -434,18 +455,13 @@ async def delete_fallback(
|
|||
elif fallback_type == "content_policy":
|
||||
fallback_key = "content_policy_fallbacks"
|
||||
|
||||
# Get existing fallbacks
|
||||
existing_fallbacks: Final = _editable_rules(proxy_config, router_settings, fallback_key)
|
||||
|
||||
updated_fallbacks: Final = [rule for rule in existing_fallbacks if not _rule_covers(rule, model)]
|
||||
|
||||
if len(updated_fallbacks) == len(existing_fallbacks):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"},
|
||||
)
|
||||
|
||||
await _write_rules(prisma_client, proxy_config, llm_router, router_settings, fallback_key, updated_fallbacks)
|
||||
await _write_rules(
|
||||
fallbacks,
|
||||
proxy_config,
|
||||
llm_router,
|
||||
fallback_key,
|
||||
lambda rules: _without_rule(rules, model, fallback_type),
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info("Fallback deleted: %s (type: %s)", model, fallback_type)
|
||||
|
||||
|
|
|
|||
|
|
@ -5748,7 +5748,6 @@ def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
|
|||
_CONFIG_PERSISTED_SECTIONS: Final = ("general_settings", "router_settings", "litellm_settings")
|
||||
_CONFIG_UNMANAGED_EXCLUSIONS: Final = frozenset(("environment_variables", "model_list"))
|
||||
_CONFIG_SECTION_VALUES: Final = TypeAdapter(Mapping[str, JsonValue])
|
||||
_CONFIG_SECTION_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtext($1))"
|
||||
|
||||
|
||||
class _ConfigParamWhere(TypedDict):
|
||||
|
|
@ -6301,7 +6300,7 @@ class ProxyConfig:
|
|||
prisma_client: PrismaClient,
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
async with prisma_client.tx() as tx:
|
||||
await tx.query_raw(_CONFIG_SECTION_LOCK_SQL, section_name)
|
||||
await tx.query_raw(CONFIG_PARAM_LOCK_SQL, section_name)
|
||||
config_table: Final = cast("TableActions[_ConfigParamRow]", tx.litellm_config)
|
||||
config_where: Final[_ConfigParamWhere] = {"param_name": section_name}
|
||||
existing_row: Final[_ConfigParamRow | None] = await config_table.find_first(where=config_where)
|
||||
|
|
@ -15287,7 +15286,7 @@ from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCoun
|
|||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
|
||||
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.config_repository import CONFIG_PARAM_LOCK_SQL, ConfigRepository
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
AccessGroupRepository,
|
||||
|
|
|
|||
|
|
@ -4,9 +4,14 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Final, Protocol, cast
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.models.fallback import STORED_FALLBACKS, FallbackRulesEdit, FallbackRulesEditor, SavedFallbackRule
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
|
@ -39,10 +44,42 @@ class _ConfigTable(Protocol):
|
|||
async def delete_many(self, *, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
class _ConfigDatabase(Protocol):
|
||||
class _ConfigTransaction(Protocol):
|
||||
@property
|
||||
def litellm_config(self) -> _ConfigTable: ...
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
|
||||
class _ConfigWriter(_ConfigTransaction, Protocol):
|
||||
def tx(self) -> AbstractAsyncContextManager[_ConfigTransaction]: ...
|
||||
|
||||
|
||||
CONFIG_PARAM_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtext($1))"
|
||||
ROUTER_SETTINGS_PARAM: Final = "router_settings"
|
||||
_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_JSON_RULES: Final = TypeAdapter(list[JsonValue])
|
||||
_ROUTER_SETTINGS: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def _stored_router_settings(param_value: object) -> Mapping[str, JsonValue]:
|
||||
try:
|
||||
decoded: Final[object] = _JSON_VALUE.validate_json(param_value) if isinstance(param_value, str) else param_value
|
||||
except ValidationError:
|
||||
return {}
|
||||
if not isinstance(decoded, Mapping):
|
||||
return {}
|
||||
return _ROUTER_SETTINGS.validate_python(decoded)
|
||||
|
||||
|
||||
def _edited_router_settings(
|
||||
param_value: object, fallback_key: str, edit: FallbackRulesEditor
|
||||
) -> tuple[Mapping[str, JsonValue], FallbackRulesEdit]:
|
||||
router_settings: Final = _stored_router_settings(param_value)
|
||||
edited: Final = edit(STORED_FALLBACKS.validate_python(router_settings.get(fallback_key) or []))
|
||||
return {**router_settings, fallback_key: _JSON_RULES.validate_python(edited.rules)}, edited
|
||||
|
||||
|
||||
class ConfigParam:
|
||||
"""Simple wrapper for config parameter from DB."""
|
||||
|
||||
|
|
@ -73,6 +110,12 @@ class ConfigRepository:
|
|||
def table(self) -> _ConfigTable:
|
||||
return self._config_table
|
||||
|
||||
@property
|
||||
def _writer(self) -> _ConfigWriter:
|
||||
return cast( # cast-ok: Prisma delegates database methods and transactions dynamically
|
||||
_ConfigWriter, self.prisma_client.writer_db
|
||||
)
|
||||
|
||||
async def get_param(self, param_name: str) -> ConfigParam | None:
|
||||
"""Get a config parameter from the database."""
|
||||
record: Final = await self._config_table.find_unique(where={"param_name": param_name})
|
||||
|
|
@ -96,9 +139,7 @@ class ConfigRepository:
|
|||
return ConfigParam(param_name=param_name, param_value=param_value)
|
||||
|
||||
async def set_param_if_revision(self, param_name: str, param_value: object, revision: int) -> bool:
|
||||
delegate: Final = self.prisma_client.writer_db
|
||||
database: Final = cast(_ConfigDatabase, delegate) # cast-ok: Prisma delegates database methods dynamically
|
||||
rows: Final = await database.query_raw(
|
||||
rows: Final = await self._writer.query_raw(
|
||||
"""INSERT INTO "LiteLLM_Config" (param_name, param_value, last_run_at)
|
||||
SELECT $1, $2::jsonb, NOW() WHERE $3::int = 0
|
||||
ON CONFLICT (param_name) DO UPDATE
|
||||
|
|
@ -115,6 +156,31 @@ class ConfigRepository:
|
|||
)
|
||||
return bool(rows)
|
||||
|
||||
async def saved_rules(self, fallback_key: str) -> Sequence[SavedFallbackRule] | None:
|
||||
row: Final = await self._config_table.find_unique(where={"param_name": ROUTER_SETTINGS_PARAM})
|
||||
if row is None:
|
||||
return None
|
||||
rules: Final = _stored_router_settings(row.param_value).get(fallback_key)
|
||||
if rules is None:
|
||||
return None
|
||||
return rules if isinstance(rules, list) else ()
|
||||
|
||||
async def edit_rules(self, fallback_key: str, edit: FallbackRulesEditor) -> FallbackRulesEdit:
|
||||
param_name: Final = ROUTER_SETTINGS_PARAM
|
||||
async with self._writer.tx() as transaction:
|
||||
await transaction.query_raw(CONFIG_PARAM_LOCK_SQL, param_name)
|
||||
row: Final = await transaction.litellm_config.find_unique(where={"param_name": param_name})
|
||||
value, result = _edited_router_settings(None if row is None else row.param_value, fallback_key, edit)
|
||||
value_json: Final = json.dumps(dict(value))
|
||||
await transaction.litellm_config.upsert(
|
||||
where={"param_name": param_name},
|
||||
data={
|
||||
"create": {"param_name": param_name, "param_value": value_json},
|
||||
"update": {"param_value": value_json},
|
||||
},
|
||||
)
|
||||
return result
|
||||
|
||||
async def delete_param(self, param_name: str) -> bool:
|
||||
"""Delete a config parameter from the database."""
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -238,6 +238,70 @@ def test_consecutive_creates_within_the_cache_window_keep_every_rule(gateway: Ga
|
|||
assert frozenset(primaries) <= covered, (primaries, covered)
|
||||
|
||||
|
||||
def _create_over_a_fresh_connection(instance: Gateway, model: str, target: str) -> httpx.Response:
|
||||
with httpx.Client(base_url=instance.client.base_url, timeout=30, trust_env=False) as client:
|
||||
return Gateway(client, instance.key, instance.upstream_url).request(
|
||||
"POST", "/fallback", _fallback_body(model, [target])
|
||||
)
|
||||
|
||||
|
||||
def _delete_over_a_fresh_connection(instance: Gateway, model: str) -> httpx.Response:
|
||||
with httpx.Client(base_url=instance.client.base_url, timeout=30, trust_env=False) as client:
|
||||
return Gateway(client, instance.key, instance.upstream_url).request(
|
||||
"DELETE", f"/fallback/{model}", params={"fallback_type": "general"}
|
||||
)
|
||||
|
||||
|
||||
def test_concurrent_writes_through_both_instances_keep_every_rule(gateway: Gateway, peer: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
target: Final = scenario.model(model=f"openai/n7t-{uuid.uuid4().hex}")
|
||||
primaries: Final = tuple(scenario.model(model=f"openai/n7{tag}-{uuid.uuid4().hex}") for tag in "abcdefgh")
|
||||
for instance in (gateway, peer):
|
||||
for model in (target, *primaries):
|
||||
_wait_until_served(instance, model)
|
||||
for primary in primaries:
|
||||
scenario.cleanups.callback(_forget_fallback, gateway, primary, "general")
|
||||
instances: Final = tuple(gateway if index % 2 == 0 else peer for index in range(len(primaries)))
|
||||
with ThreadPoolExecutor(max_workers=len(primaries)) as pool:
|
||||
created: Final = tuple(
|
||||
pool.map(partial(_create_over_a_fresh_connection, target=target), instances, primaries)
|
||||
)
|
||||
assert all(response.status_code == 200 for response in created), [response.text for response in created]
|
||||
assert frozenset(primaries) <= _covered_models(_stored_fallbacks())
|
||||
|
||||
dropped: Final = primaries[::2]
|
||||
with ThreadPoolExecutor(max_workers=len(dropped)) as pool:
|
||||
removed: Final = tuple(pool.map(_delete_over_a_fresh_connection, instances[1::2], dropped))
|
||||
assert all(response.status_code == 200 for response in removed), [response.text for response in removed]
|
||||
covered: Final = _covered_models(_stored_fallbacks()) & frozenset(primaries)
|
||||
assert covered == frozenset(primaries[1::2]), covered
|
||||
|
||||
|
||||
def _read_fallback(instance: Gateway, model: str) -> httpx.Response:
|
||||
return instance.request("GET", f"/fallback/{model}", params={"fallback_type": "general"})
|
||||
|
||||
|
||||
def test_the_other_instance_reads_each_write_immediately(gateway: Gateway, peer: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
primary: Final = scenario.model(model=f"openai/n8p-{uuid.uuid4().hex}")
|
||||
first: Final = scenario.model(model=f"openai/n8a-{uuid.uuid4().hex}")
|
||||
second: Final = scenario.model(model=f"openai/n8b-{uuid.uuid4().hex}")
|
||||
for model in (primary, first, second):
|
||||
_wait_until_served(gateway, model)
|
||||
scenario.cleanups.callback(_forget_fallback, gateway, primary, "general")
|
||||
|
||||
assert gateway.request("POST", "/fallback", _fallback_body(primary, [first])).status_code == 200
|
||||
created: Final = _read_fallback(peer, primary)
|
||||
assert gateway.request("POST", "/fallback", _fallback_body(primary, [second, first])).status_code == 200
|
||||
updated: Final = _read_fallback(peer, primary)
|
||||
assert gateway.request("DELETE", f"/fallback/{primary}", params={"fallback_type": "general"}).status_code == 200
|
||||
deleted: Final = _read_fallback(peer, primary)
|
||||
|
||||
assert created.status_code == 200 and _fallback_models(created) == [first], created.text
|
||||
assert updated.status_code == 200 and _fallback_models(updated) == [second, first], updated.text
|
||||
assert deleted.status_code == 404, deleted.text
|
||||
|
||||
|
||||
def test_unknown_model_404_lists_team_public_and_internal_names(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
team: Final = scenario.team()
|
||||
|
|
|
|||
|
|
@ -8,34 +8,44 @@ Tests:
|
|||
4. Validation tests (invalid models, duplicate fallbacks, etc.)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final, Literal, get_args
|
||||
from typing import Final, Literal, TypeAlias, get_args
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from redis.exceptions import NoPermissionError
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.config_resolvers.settings_rules import Section
|
||||
from litellm.proxy.deps.database import get_database
|
||||
from litellm.proxy.management_endpoints.fallback_management_endpoints import (
|
||||
FallbackCreateRequest,
|
||||
FallbackDeleteResponse,
|
||||
FallbackGetResponse,
|
||||
FallbackResponse,
|
||||
create_fallback,
|
||||
delete_fallback,
|
||||
get_fallback,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.proxy_server import ProxyConfig, app
|
||||
from litellm.proxy.utils import evict_config_param, get_config_param, litellm_config_cache
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.database import Database
|
||||
from litellm.router_api import ChangeSet
|
||||
from litellm.router_api import RouterSettings as RouterSettingsSnapshot
|
||||
from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict
|
||||
from tests._support.native_database import native_database
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
@ -79,9 +89,7 @@ class TestFallbackCreateRequest:
|
|||
|
||||
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"
|
||||
):
|
||||
with pytest.raises(ValueError, match="fallback_models must not contain duplicates"):
|
||||
FallbackCreateRequest(
|
||||
model="gpt-3.5-turbo",
|
||||
fallback_models=["gpt-4", "gpt-4"],
|
||||
|
|
@ -180,13 +188,12 @@ class TestCreateFallbackForTeamScopedModels:
|
|||
return _team_router()
|
||||
|
||||
@pytest.fixture
|
||||
def prisma_client(self) -> MagicMock:
|
||||
client: Final = MagicMock()
|
||||
client.writer_db = client.db
|
||||
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
client.db.litellm_config.upsert = AsyncMock()
|
||||
client.get_generic_data = AsyncMock(return_value=None)
|
||||
return client
|
||||
def stored(self) -> "_StoredRouterSettings":
|
||||
return _StoredRouterSettings(None)
|
||||
|
||||
@pytest.fixture
|
||||
def prisma_client(self, stored: "_StoredRouterSettings") -> MagicMock:
|
||||
return stored.prisma_client
|
||||
|
||||
@pytest.fixture
|
||||
def proxy_config(self) -> ProxyConfig:
|
||||
|
|
@ -197,14 +204,13 @@ class TestCreateFallbackForTeamScopedModels:
|
|||
) -> 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())
|
||||
return await create_fallback(request, MagicMock(), fallbacks=ConfigRepository(prisma_client))
|
||||
|
||||
async def test_public_names_create_a_rule_keyed_on_the_public_name(
|
||||
self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
||||
self, router: Router, stored: "_StoredRouterSettings", prisma_client: MagicMock, proxy_config: MagicMock
|
||||
) -> None:
|
||||
request: Final = FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"])
|
||||
|
||||
|
|
@ -213,10 +219,7 @@ class TestCreateFallbackForTeamScopedModels:
|
|||
assert response.model == "team-primary"
|
||||
assert response.fallback_models == ["team-fallback"]
|
||||
assert router.settings().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"]}]
|
||||
assert stored.stored_fallbacks() == [{"team-primary": ["team-fallback"]}]
|
||||
|
||||
async def test_internal_names_keep_working(
|
||||
self, router: Router, prisma_client: MagicMock, proxy_config: MagicMock
|
||||
|
|
@ -279,8 +282,45 @@ class _ConfigFileProxyConfig(ProxyConfig):
|
|||
return {"router_settings": copy.deepcopy(self.config_file)}
|
||||
|
||||
|
||||
NATIVE: Final = native_database()
|
||||
|
||||
|
||||
class _ConfigTransaction:
|
||||
def __init__(self, stored: "_StoredRouterSettings") -> None:
|
||||
self._stored: Final = stored
|
||||
self.staged: SimpleNamespace | None = None
|
||||
self.holds_lock = False
|
||||
|
||||
@property
|
||||
def litellm_config(self) -> "_ConfigTransaction":
|
||||
return self
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, int]]:
|
||||
if "router_settings" not in args or self.holds_lock:
|
||||
return ()
|
||||
if self._stored.lock.locked():
|
||||
self._stored.lock_contended.set()
|
||||
await self._stored.lock.acquire()
|
||||
self.holds_lock = True
|
||||
return ({"locked": 1},)
|
||||
|
||||
async def find_unique(self, where: Mapping[str, str]) -> SimpleNamespace | None:
|
||||
visible: Final = self.staged or copy.deepcopy(self._stored.row)
|
||||
await self._stored.after_read()
|
||||
return visible
|
||||
|
||||
async def find_first(self, where: Mapping[str, str]) -> SimpleNamespace | None:
|
||||
return await self.find_unique(where)
|
||||
|
||||
async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> None:
|
||||
if self._stored.fail_writes:
|
||||
raise RuntimeError("the database refused the write")
|
||||
self.staged = _config_row(json.loads(data["update"]["param_value"]))
|
||||
|
||||
|
||||
class _StoredRouterSettings:
|
||||
"""The LiteLLM_Config router_settings row, with the proxy objects that read it through the config cache"""
|
||||
"""The LiteLLM_Config router_settings row behind a writer with PostgreSQL's transaction semantics: a
|
||||
transaction's write lands only when it commits, and the advisory lock waits for the transaction holding it"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -288,23 +328,50 @@ class _StoredRouterSettings:
|
|||
config_file: RouterSettings | None = None,
|
||||
*,
|
||||
lagging_replica: bool = False,
|
||||
hold_first_read: bool = False,
|
||||
fail_writes: bool = False,
|
||||
) -> None:
|
||||
self.row: SimpleNamespace | None = None if router_settings is None else _config_row(router_settings)
|
||||
self.writes = 0
|
||||
self.fail_writes: Final = fail_writes
|
||||
self.lock: Final = asyncio.Lock()
|
||||
self.lock_contended: Final = asyncio.Event()
|
||||
self._holding_next_read = hold_first_read
|
||||
self.read_held: Final = asyncio.Event()
|
||||
self.held_read_released: Final = asyncio.Event()
|
||||
self.prisma_client: Final = MagicMock()
|
||||
self.prisma_client.get_generic_data = AsyncMock(
|
||||
side_effect=lambda value, **_: self.row if value == "router_settings" else None
|
||||
)
|
||||
self.prisma_client.writer_db.litellm_config.find_unique = AsyncMock(
|
||||
side_effect=lambda **_: copy.deepcopy(self.row)
|
||||
)
|
||||
self.prisma_client.writer_db.tx = self._transaction
|
||||
self.prisma_client.tx = self._transaction
|
||||
self.prisma_client.db.litellm_config.find_unique = AsyncMock(
|
||||
side_effect=lambda **_: None if lagging_replica else copy.deepcopy(self.row)
|
||||
)
|
||||
self.prisma_client.db.litellm_config.upsert = AsyncMock(side_effect=self._upsert)
|
||||
self.proxy_config: Final = _ConfigFileProxyConfig(config_file)
|
||||
self.prisma_client.writer_db.litellm_config.find_unique = AsyncMock(
|
||||
side_effect=lambda **_: copy.deepcopy(self.row)
|
||||
)
|
||||
self.fallbacks: Final = ConfigRepository(self.prisma_client, use_writer=True)
|
||||
|
||||
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"]))
|
||||
@asynccontextmanager
|
||||
async def _transaction(self) -> AsyncIterator[_ConfigTransaction]:
|
||||
transaction: Final = _ConfigTransaction(self)
|
||||
try:
|
||||
yield transaction
|
||||
if transaction.staged is not None:
|
||||
self.row = transaction.staged
|
||||
self.writes += 1
|
||||
finally:
|
||||
if transaction.holds_lock:
|
||||
self.lock.release()
|
||||
|
||||
async def after_read(self) -> None:
|
||||
if not self._holding_next_read:
|
||||
return
|
||||
self._holding_next_read = False
|
||||
self.read_held.set()
|
||||
await self.held_read_released.wait()
|
||||
|
||||
def written_by_another_instance(self, router_settings: RouterSettings) -> None:
|
||||
self.row = _config_row(router_settings)
|
||||
|
|
@ -318,6 +385,7 @@ class _StoredRouterSettings:
|
|||
return row.param_value["fallbacks"]
|
||||
|
||||
|
||||
FallbackWrite: TypeAlias = Callable[[], Awaitable[object]]
|
||||
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"}
|
||||
|
|
@ -347,20 +415,260 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
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())
|
||||
return await create_fallback(request, MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
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())
|
||||
return await delete_fallback(model, "general", MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
async def _get(
|
||||
self,
|
||||
model: str,
|
||||
stored: _StoredRouterSettings,
|
||||
fallback_type: Literal["general", "context_window", "content_policy"] = "general",
|
||||
router: Router | None = None,
|
||||
) -> FallbackGetResponse:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", router or _team_router()),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
return await get_fallback(model, fallback_type, MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
async def _read_status(self, model: str, stored: _StoredRouterSettings, router: Router | None = None) -> int:
|
||||
try:
|
||||
await self._get(model, stored, router=router)
|
||||
except HTTPException as refused:
|
||||
return refused.status_code
|
||||
return 200
|
||||
|
||||
async def test_a_replica_that_has_not_reloaded_reads_each_write_as_saved(self) -> None:
|
||||
stored: Final = _StoredRouterSettings(None, lagging_replica=True)
|
||||
stale_router: Final = _team_router()
|
||||
stale_router.apply(ChangeSet(assign={"fallbacks": [{"team-primary": ["gpt-5.4-mini"]}]}))
|
||||
|
||||
await self._create(FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]), stored)
|
||||
created: Final = await self._get("team-primary", stored, router=stale_router)
|
||||
await self._create(
|
||||
FallbackCreateRequest(model="team-primary", fallback_models=["gpt-5.4-mini", "team-fallback"]), stored
|
||||
)
|
||||
updated: Final = await self._get("team-primary", stored, router=stale_router)
|
||||
await self._delete("team-primary", stored)
|
||||
|
||||
assert created.fallback_models == ["team-fallback"]
|
||||
assert updated.fallback_models == ["gpt-5.4-mini", "team-fallback"]
|
||||
assert await self._read_status("team-primary", stored, router=stale_router) == 404
|
||||
assert stale_router.settings().fallbacks == [{"team-primary": ["gpt-5.4-mini"]}]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("fallback_type", "fallback_key"),
|
||||
[
|
||||
("general", "fallbacks"),
|
||||
("context_window", "context_window_fallbacks"),
|
||||
("content_policy", "content_policy_fallbacks"),
|
||||
],
|
||||
)
|
||||
async def test_a_read_names_every_saved_target_of_each_fallback_type(
|
||||
self, fallback_type: Literal["general", "context_window", "content_policy"], fallback_key: str
|
||||
) -> None:
|
||||
stored: Final = _StoredRouterSettings(
|
||||
{fallback_key: [{"team-primary": [{"model": "gpt-5.4-mini", "api_key": "sk-test"}, "team-fallback"]}]}
|
||||
)
|
||||
|
||||
response: Final = await self._get("team-primary", stored, fallback_type)
|
||||
|
||||
assert response.fallback_models == ["gpt-5.4-mini", "team-fallback"]
|
||||
assert response.fallback_type == fallback_type
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("fallback_type", "fallback_key"),
|
||||
[("general", "fallbacks"), ("content_policy", "content_policy_fallbacks")],
|
||||
)
|
||||
async def test_a_read_skips_a_malformed_saved_rule_and_serves_the_valid_ones(
|
||||
self, fallback_type: Literal["general", "content_policy"], fallback_key: str
|
||||
) -> None:
|
||||
stored: Final = _StoredRouterSettings(None)
|
||||
stored.row = SimpleNamespace(param_name="router_settings", param_value={fallback_key: [42, TEAM_RULE, None]})
|
||||
|
||||
response: Final = await self._get("team-primary", stored, fallback_type)
|
||||
|
||||
assert response.fallback_models == ["team-fallback"]
|
||||
|
||||
async def test_a_read_resolves_the_saved_wildcard_rule(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [GATEWAY_RULE, {"*": ["gpt-5.4-mini"]}]})
|
||||
|
||||
response: Final = await self._get("team-primary", stored)
|
||||
|
||||
assert response.fallback_models == ["gpt-5.4-mini"]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"router_settings",
|
||||
[None, {"routing_strategy": "latency-based-routing"}, {"fallbacks": None}, {"fallbacks": "team-primary"}],
|
||||
)
|
||||
async def test_a_read_does_not_answer_with_rules_only_the_routers_memory_holds(
|
||||
self, router_settings: RouterSettings | None
|
||||
) -> None:
|
||||
stored: Final = _StoredRouterSettings(router_settings)
|
||||
stale_router: Final = _team_router()
|
||||
stale_router.apply(ChangeSet(assign={"fallbacks": [TEAM_RULE]}))
|
||||
|
||||
assert await self._read_status("team-primary", stored, router=stale_router) == 404
|
||||
|
||||
async def test_a_content_policy_read_layers_the_saved_rules_over_the_litellm_settings_rules(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "content_policy_fallbacks", [TEAM_RULE, {"gpt-5.4-mini": ["team-primary"]}])
|
||||
stored: Final = _StoredRouterSettings({"content_policy_fallbacks": [GATEWAY_RULE]})
|
||||
emptied: Final = _StoredRouterSettings({"content_policy_fallbacks": []})
|
||||
|
||||
saved_rule: Final = await self._get("gpt-5.4-mini", stored, "content_policy")
|
||||
file_rule: Final = await self._get("team-primary", stored, "content_policy")
|
||||
after_delete: Final = await self._get("gpt-5.4-mini", emptied, "content_policy")
|
||||
|
||||
assert saved_rule.fallback_models == ["team-fallback"]
|
||||
assert file_rule.fallback_models == ["team-fallback"]
|
||||
assert after_delete.fallback_models == ["team-primary"]
|
||||
|
||||
async def test_a_read_outside_database_mode_answers_from_the_router(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]})
|
||||
router: Final = _team_router()
|
||||
router.apply(ChangeSet(assign={"fallbacks": [GATEWAY_RULE]}))
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", router),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
):
|
||||
response: Final = await get_fallback("gpt-5.4-mini", "general", MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
assert response.fallback_models == ["team-fallback"]
|
||||
|
||||
async def test_a_read_of_a_router_settings_row_that_fails_answers_500(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]})
|
||||
stored.prisma_client.writer_db.litellm_config.find_unique = AsyncMock(side_effect=RuntimeError("db down"))
|
||||
|
||||
with pytest.raises(HTTPException) as failed:
|
||||
await self._get("team-primary", stored)
|
||||
|
||||
assert failed.value.status_code == 500
|
||||
|
||||
async def _while_a_write_holds_the_rules(
|
||||
self, stored: _StoredRouterSettings, holder: FallbackWrite, contender: FallbackWrite
|
||||
) -> None:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", _team_router()),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
holding: Final = asyncio.create_task(holder())
|
||||
await stored.read_held.wait()
|
||||
contending: Final = asyncio.create_task(contender())
|
||||
blocked: Final = asyncio.create_task(stored.lock_contended.wait())
|
||||
await asyncio.wait((contending, blocked), return_when=asyncio.FIRST_COMPLETED)
|
||||
blocked.cancel()
|
||||
stored.held_read_released.set()
|
||||
await asyncio.gather(holding, contending)
|
||||
|
||||
async def test_concurrent_first_creates_keep_both_rules(self) -> None:
|
||||
stored: Final = _StoredRouterSettings(None, hold_first_read=True)
|
||||
|
||||
await self._while_a_write_holds_the_rules(
|
||||
stored,
|
||||
lambda: create_fallback(
|
||||
FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
),
|
||||
lambda: create_fallback(
|
||||
FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
),
|
||||
)
|
||||
|
||||
assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE]
|
||||
|
||||
async def test_a_delete_racing_a_create_keeps_the_created_rule(self) -> None:
|
||||
stored: Final = _StoredRouterSettings(
|
||||
{"fallbacks": [TEAM_RULE], "routing_strategy": "latency-based-routing"}, hold_first_read=True
|
||||
)
|
||||
|
||||
await self._while_a_write_holds_the_rules(
|
||||
stored,
|
||||
lambda: create_fallback(
|
||||
FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
),
|
||||
lambda: delete_fallback("team-primary", "general", MagicMock(), fallbacks=stored.fallbacks),
|
||||
)
|
||||
|
||||
assert stored.row is not None
|
||||
assert stored.row.param_value == {"fallbacks": [GATEWAY_RULE], "routing_strategy": "latency-based-routing"}
|
||||
|
||||
async def test_a_config_section_write_racing_a_fallback_write_keeps_both_changes(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]}, hold_first_read=True)
|
||||
|
||||
await self._while_a_write_holds_the_rules(
|
||||
stored,
|
||||
lambda: create_fallback(
|
||||
FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
),
|
||||
lambda: stored.proxy_config._upsert_changed_config_section(
|
||||
section_name="router_settings",
|
||||
changed_keys={"routing_strategy": "latency-based-routing"},
|
||||
removed_keys=frozenset(),
|
||||
prisma_client=stored.prisma_client,
|
||||
),
|
||||
)
|
||||
|
||||
assert stored.row is not None
|
||||
assert stored.row.param_value == {
|
||||
"fallbacks": [TEAM_RULE, GATEWAY_RULE],
|
||||
"routing_strategy": "latency-based-routing",
|
||||
}
|
||||
|
||||
async def test_a_failed_write_changes_neither_the_stored_rules_nor_the_router(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]}, fail_writes=True)
|
||||
router: Final = _team_router()
|
||||
router.apply(ChangeSet(assign={"fallbacks": [TEAM_RULE]}))
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", router),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
pytest.raises(HTTPException) as failed,
|
||||
):
|
||||
await create_fallback(
|
||||
FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
)
|
||||
|
||||
assert failed.value.status_code == 500
|
||||
assert stored.stored_fallbacks() == [TEAM_RULE]
|
||||
assert router.settings().fallbacks == [TEAM_RULE]
|
||||
|
||||
async def test_a_delete_of_a_missing_rule_answers_404_and_leaves_the_row_and_lock_untouched(self) -> None:
|
||||
settings: Final = {"fallbacks": [TEAM_RULE], "routing_strategy": "latency-based-routing"}
|
||||
stored: Final = _StoredRouterSettings(copy.deepcopy(settings))
|
||||
|
||||
with pytest.raises(HTTPException) as missing:
|
||||
await self._delete("gpt-5.4-mini", stored)
|
||||
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored)
|
||||
|
||||
assert missing.value.status_code == 404
|
||||
assert stored.writes == 1
|
||||
assert stored.row is not None
|
||||
assert stored.row.param_value == {**settings, "fallbacks": [TEAM_RULE, GATEWAY_RULE]}
|
||||
|
||||
async def test_a_create_reads_the_rules_from_the_writer_while_the_replica_lags(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]}, lagging_replica=True)
|
||||
|
|
@ -397,7 +705,6 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
stored: Final = _StoredRouterSettings(None)
|
||||
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),
|
||||
):
|
||||
|
|
@ -406,10 +713,15 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
model="gpt-5.4-mini", fallback_models=["team-fallback"], fallback_type="content_policy"
|
||||
),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
)
|
||||
after_create: Final = await get_fallback(
|
||||
"team-primary", "content_policy", MagicMock(), fallbacks=stored.fallbacks
|
||||
)
|
||||
await delete_fallback("gpt-5.4-mini", "content_policy", MagicMock(), fallbacks=stored.fallbacks)
|
||||
after_delete: Final = await get_fallback(
|
||||
"team-primary", "content_policy", MagicMock(), fallbacks=stored.fallbacks
|
||||
)
|
||||
after_create: Final = await get_fallback("team-primary", "content_policy", MagicMock())
|
||||
await delete_fallback("gpt-5.4-mini", "content_policy", MagicMock())
|
||||
after_delete: Final = await get_fallback("team-primary", "content_policy", MagicMock())
|
||||
|
||||
assert after_create.fallback_models == ["team-fallback"]
|
||||
assert after_delete.fallback_models == ["team-fallback"]
|
||||
|
|
@ -419,7 +731,6 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
stored: Final = _StoredRouterSettings({"content_policy_fallbacks": [MALFORMED_GATEWAY_RULE]})
|
||||
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),
|
||||
):
|
||||
|
|
@ -428,8 +739,11 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
model="team-primary", fallback_models=["team-fallback"], fallback_type="content_policy"
|
||||
),
|
||||
MagicMock(),
|
||||
fallbacks=stored.fallbacks,
|
||||
)
|
||||
after_create: Final = await get_fallback(
|
||||
"team-primary", "content_policy", MagicMock(), fallbacks=stored.fallbacks
|
||||
)
|
||||
after_create: Final = await get_fallback("team-primary", "content_policy", MagicMock())
|
||||
|
||||
assert after_create.fallback_models == ["team-fallback"]
|
||||
|
||||
|
|
@ -445,7 +759,7 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
**{router_setting: [{"team-primary": [{"model": "gpt-5.4-mini", "api_key": "sk-test"}, "team-fallback"]}]}
|
||||
)
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router):
|
||||
response: Final = await get_fallback("team-primary", fallback_type, MagicMock())
|
||||
response: Final = await get_fallback("team-primary", fallback_type, MagicMock(), fallbacks=None)
|
||||
|
||||
assert response.fallback_models == ["gpt-5.4-mini", "team-fallback"]
|
||||
|
||||
|
|
@ -459,7 +773,7 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
router: Final = _team_router()
|
||||
router.update_settings(**{router_setting: [{"team-primary": [{"temperature": 0}, "team-fallback"]}]})
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router):
|
||||
response: Final = await get_fallback("team-primary", fallback_type, MagicMock())
|
||||
response: Final = await get_fallback("team-primary", fallback_type, MagicMock(), fallbacks=None)
|
||||
|
||||
assert response.fallback_models == ["team-primary", "team-fallback"]
|
||||
|
||||
|
|
@ -467,7 +781,7 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
router: Final = _team_router()
|
||||
router.update_settings(fallbacks=[{"team-primary": "team-fallback"}])
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router), pytest.raises(HTTPException) as raised:
|
||||
await get_fallback("team-primary", "general", MagicMock())
|
||||
await get_fallback("team-primary", "general", MagicMock(), fallbacks=None)
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
|
||||
|
|
@ -530,6 +844,79 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
assert stored.row.param_value == {"fallbacks": [GATEWAY_RULE]}
|
||||
|
||||
|
||||
class TestFallbackRoutesResolveTheirRepository:
|
||||
@pytest.fixture
|
||||
def stored(self) -> Iterator[_StoredRouterSettings]:
|
||||
stored: Final = _StoredRouterSettings(None, lagging_replica=True)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", _team_router()),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
yield stored
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
app.dependency_overrides.pop(get_database, None)
|
||||
|
||||
def test_a_create_stores_the_rule_through_the_request_database(self, stored: _StoredRouterSettings) -> None:
|
||||
app.dependency_overrides[get_database] = lambda: Database(prisma=stored.prisma_client, native=NATIVE)
|
||||
|
||||
response: Final = TestClient(app).post(
|
||||
"/fallback", json={"model": "team-primary", "fallback_models": ["team-fallback"]}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert stored.stored_fallbacks() == [TEAM_RULE]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "path", "body"),
|
||||
[
|
||||
("POST", "/fallback", {"model": "team-primary", "fallback_models": ["team-fallback"]}),
|
||||
("DELETE", "/fallback/team-primary", None),
|
||||
],
|
||||
)
|
||||
def test_writes_without_a_database_are_refused(
|
||||
self, stored: _StoredRouterSettings, method: str, path: str, body: Mapping[str, JsonValue] | None
|
||||
) -> None:
|
||||
app.dependency_overrides[get_database] = lambda: None
|
||||
|
||||
response: Final = TestClient(app).request(method, path, json=body)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "Database storage not enabled" in response.text
|
||||
assert stored.row is None
|
||||
|
||||
def test_a_read_serves_the_saved_rule_the_routers_memory_lacks(self, stored: _StoredRouterSettings) -> None:
|
||||
stored.written_by_another_instance({"fallbacks": [TEAM_RULE]})
|
||||
app.dependency_overrides[get_database] = lambda: Database(prisma=stored.prisma_client, native=NATIVE)
|
||||
|
||||
response: Final = TestClient(app).get("/fallback/team-primary", params={"fallback_type": "general"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {
|
||||
"model": "team-primary",
|
||||
"fallback_models": ["team-fallback"],
|
||||
"fallback_type": "general",
|
||||
}
|
||||
|
||||
def test_a_read_without_a_database_answers_from_the_router(self, stored: _StoredRouterSettings) -> None:
|
||||
stored.written_by_another_instance({"fallbacks": [TEAM_RULE]})
|
||||
app.dependency_overrides[get_database] = lambda: None
|
||||
|
||||
response: Final = TestClient(app).get("/fallback/team-primary", params={"fallback_type": "general"})
|
||||
|
||||
assert response.status_code == 404, response.text
|
||||
|
||||
def test_an_unknown_model_is_refused_before_the_missing_database(self, stored: _StoredRouterSettings) -> None:
|
||||
app.dependency_overrides[get_database] = lambda: None
|
||||
|
||||
response: Final = TestClient(app).post(
|
||||
"/fallback", json={"model": "team-missing", "fallback_models": ["team-fallback"]}
|
||||
)
|
||||
|
||||
assert response.status_code == 404, response.text
|
||||
|
||||
|
||||
CONFIG_RULE: Final = {"gpt-5.4-mini": ["team-fallback"]}
|
||||
CONFIG_WINDOW_RULE: Final = {"team-primary": ["gpt-5.4-mini"]}
|
||||
|
||||
|
|
@ -557,20 +944,18 @@ class TestFallbackWritesHonourTheConfigFile:
|
|||
) -> 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())
|
||||
return await create_fallback(request, MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
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())
|
||||
return await delete_fallback(model, "general", MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
async def test_a_create_that_adds_a_rule_to_a_config_owned_list_is_refused(
|
||||
self, router: Router, stored: _StoredRouterSettings
|
||||
|
|
@ -584,7 +969,7 @@ class TestFallbackWritesHonourTheConfigFile:
|
|||
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 stored.writes == 0
|
||||
assert router.settings().fallbacks == [CONFIG_RULE]
|
||||
|
||||
async def test_a_create_after_the_config_file_takes_over_the_key_is_refused(self, router: Router) -> None:
|
||||
|
|
@ -610,7 +995,7 @@ class TestFallbackWritesHonourTheConfigFile:
|
|||
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 stored.writes == 0
|
||||
assert router.settings().fallbacks == [CONFIG_RULE]
|
||||
|
||||
async def test_a_delete_of_a_rule_only_the_config_file_holds_is_refused(self, router: Router) -> None:
|
||||
|
|
@ -629,10 +1014,14 @@ class TestFallbackWritesHonourTheConfigFile:
|
|||
self, router: Router, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.delenv("LITELLM_TEST_UNSET_FALLBACKS", raising=False)
|
||||
stored: Final = _StoredRouterSettings(None, config_file={"fallbacks": "os.environ/LITELLM_TEST_UNSET_FALLBACKS"})
|
||||
stored: Final = _StoredRouterSettings(
|
||||
None, config_file={"fallbacks": "os.environ/LITELLM_TEST_UNSET_FALLBACKS"}
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await self._create(FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]), router, stored)
|
||||
await self._create(
|
||||
FallbackCreateRequest(model="team-primary", fallback_models=["team-fallback"]), router, stored
|
||||
)
|
||||
with pytest.raises(HTTPException) as not_found:
|
||||
await self._delete("team-primary", router, stored)
|
||||
|
||||
|
|
@ -696,13 +1085,30 @@ class TestFallbackWritesHonourTheConfigFile:
|
|||
{"fallbacks": [TEAM_RULE]}, config_file={"routing_strategy": "latency-based-routing"}
|
||||
)
|
||||
|
||||
await self._create(FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), router, stored)
|
||||
await self._create(
|
||||
FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), router, stored
|
||||
)
|
||||
assert stored.row is not None
|
||||
assert stored.row.param_value == {"fallbacks": [TEAM_RULE, GATEWAY_RULE]}
|
||||
|
||||
await self._delete("team-primary", router, stored)
|
||||
assert stored.row.param_value == {"fallbacks": [GATEWAY_RULE]}
|
||||
|
||||
async def test_a_read_of_a_config_owned_list_answers_from_the_file_not_the_shadowed_row(self) -> None:
|
||||
stored: Final = _StoredRouterSettings({"fallbacks": [TEAM_RULE]}, config_file={"fallbacks": [CONFIG_RULE]})
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", _team_router()),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
):
|
||||
from_file: Final = await get_fallback("gpt-5.4-mini", "general", MagicMock(), fallbacks=stored.fallbacks)
|
||||
with pytest.raises(HTTPException) as shadowed:
|
||||
await get_fallback("team-primary", "general", MagicMock(), fallbacks=stored.fallbacks)
|
||||
|
||||
assert from_file.fallback_models == ["team-fallback"]
|
||||
assert shadowed.value.status_code == 404
|
||||
|
||||
async def test_a_sibling_list_the_config_file_does_not_set_stays_writable(
|
||||
self, router: Router, stored: _StoredRouterSettings
|
||||
) -> None:
|
||||
|
|
@ -737,15 +1143,12 @@ class TestCreateFallback:
|
|||
return router
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma_client(self):
|
||||
"""Create a mock prisma client"""
|
||||
client = MagicMock()
|
||||
client.writer_db = client.db
|
||||
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
client.db.litellm_config.upsert = AsyncMock()
|
||||
client.get_generic_data = AsyncMock(return_value=None)
|
||||
client.jsonify_object = lambda x: x
|
||||
return client
|
||||
def stored(self) -> _StoredRouterSettings:
|
||||
return _StoredRouterSettings(None)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma_client(self, stored: _StoredRouterSettings) -> MagicMock:
|
||||
return stored.prisma_client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_proxy_config(self):
|
||||
|
|
@ -758,7 +1161,7 @@ class TestCreateFallback:
|
|||
return MagicMock()
|
||||
|
||||
async def test_create_fallback_success(
|
||||
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
||||
self, mock_router, stored, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
||||
):
|
||||
"""Test successful fallback creation"""
|
||||
request = FallbackCreateRequest(
|
||||
|
|
@ -785,18 +1188,16 @@ class TestCreateFallback:
|
|||
True,
|
||||
),
|
||||
):
|
||||
response = await create_fallback(request, mock_user_api_key_dict)
|
||||
response = await create_fallback(
|
||||
request, mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client)
|
||||
)
|
||||
|
||||
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()
|
||||
)
|
||||
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()
|
||||
assert stored.stored_fallbacks() == [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
|
||||
|
||||
async def test_create_fallback_router_not_initialized(
|
||||
self, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
|
||||
|
|
@ -814,7 +1215,7 @@ class TestCreateFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await create_fallback(request, mock_user_api_key_dict)
|
||||
await create_fallback(request, mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client))
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "Router not initialized" in str(exc_info.value.detail)
|
||||
|
|
@ -843,7 +1244,7 @@ class TestCreateFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await create_fallback(request, mock_user_api_key_dict)
|
||||
await create_fallback(request, mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client))
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "not found in router" in str(exc_info.value.detail)
|
||||
|
|
@ -872,7 +1273,7 @@ class TestCreateFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await create_fallback(request, mock_user_api_key_dict)
|
||||
await create_fallback(request, mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client))
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Invalid fallback models" in str(exc_info.value.detail)
|
||||
|
|
@ -901,14 +1302,12 @@ class TestCreateFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await create_fallback(request, mock_user_api_key_dict)
|
||||
await create_fallback(request, mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client))
|
||||
|
||||
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
|
||||
):
|
||||
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",
|
||||
|
|
@ -926,7 +1325,7 @@ class TestCreateFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await create_fallback(request, mock_user_api_key_dict)
|
||||
await create_fallback(request, mock_user_api_key_dict, fallbacks=None)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Database storage not enabled" in str(exc_info.value.detail)
|
||||
|
|
@ -959,7 +1358,9 @@ class TestCreateFallback:
|
|||
True,
|
||||
),
|
||||
):
|
||||
response = await create_fallback(request, mock_user_api_key_dict)
|
||||
response = await create_fallback(
|
||||
request, mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client)
|
||||
)
|
||||
|
||||
assert response.fallback_type == "context_window"
|
||||
# Verify the correct setting was updated
|
||||
|
|
@ -988,25 +1389,19 @@ class TestGetFallback:
|
|||
"""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
|
||||
):
|
||||
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
|
||||
)
|
||||
response = await get_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict, fallbacks=None)
|
||||
|
||||
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
|
||||
):
|
||||
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(
|
||||
|
|
@ -1015,7 +1410,7 @@ class TestGetFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await get_fallback("gpt-4", "general", mock_user_api_key_dict)
|
||||
await get_fallback("gpt-4", "general", mock_user_api_key_dict, fallbacks=None)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "No general fallbacks configured" in str(exc_info.value.detail)
|
||||
|
|
@ -1029,7 +1424,7 @@ class TestGetFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await get_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
|
||||
await get_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict, fallbacks=None)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "Router not initialized" in str(exc_info.value.detail)
|
||||
|
|
@ -1051,17 +1446,12 @@ class TestDeleteFallback:
|
|||
return router
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma_client(self):
|
||||
"""Create a mock prisma client"""
|
||||
client = MagicMock()
|
||||
client.writer_db = client.db
|
||||
client.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_config_row({"fallbacks": [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]})
|
||||
)
|
||||
client.db.litellm_config.upsert = AsyncMock()
|
||||
client.get_generic_data = AsyncMock(return_value=None)
|
||||
client.jsonify_object = lambda x: x
|
||||
return client
|
||||
def stored(self) -> _StoredRouterSettings:
|
||||
return _StoredRouterSettings({"fallbacks": [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]})
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma_client(self, stored: _StoredRouterSettings) -> MagicMock:
|
||||
return stored.prisma_client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_proxy_config(self):
|
||||
|
|
@ -1076,6 +1466,7 @@ class TestDeleteFallback:
|
|||
async def test_delete_fallback_success(
|
||||
self,
|
||||
mock_router_with_fallbacks,
|
||||
stored,
|
||||
mock_prisma_client,
|
||||
mock_proxy_config,
|
||||
mock_user_api_key_dict,
|
||||
|
|
@ -1100,15 +1491,17 @@ class TestDeleteFallback:
|
|||
),
|
||||
):
|
||||
response = await delete_fallback(
|
||||
"gpt-3.5-turbo", "general", mock_user_api_key_dict
|
||||
"gpt-3.5-turbo",
|
||||
"general",
|
||||
mock_user_api_key_dict,
|
||||
fallbacks=ConfigRepository(mock_prisma_client),
|
||||
)
|
||||
|
||||
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()
|
||||
assert stored.stored_fallbacks() == []
|
||||
|
||||
async def test_delete_fallback_not_found(
|
||||
self,
|
||||
|
|
@ -1137,7 +1530,9 @@ class TestDeleteFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await delete_fallback("gpt-4", "general", mock_user_api_key_dict)
|
||||
await delete_fallback(
|
||||
"gpt-4", "general", mock_user_api_key_dict, fallbacks=ConfigRepository(mock_prisma_client)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "No general fallbacks configured" in str(exc_info.value.detail)
|
||||
|
|
@ -1151,14 +1546,12 @@ class TestDeleteFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
|
||||
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict, fallbacks=None)
|
||||
|
||||
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
|
||||
):
|
||||
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(
|
||||
|
|
@ -1171,7 +1564,7 @@ class TestDeleteFallback:
|
|||
),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
|
||||
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict, fallbacks=None)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Database storage not enabled" in str(exc_info.value.detail)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue