diff --git a/litellm/models/fallback.py b/litellm/models/fallback.py new file mode 100644 index 00000000000..b8532a7f3cb --- /dev/null +++ b/litellm/models/fallback.py @@ -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: ... diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 323cf57dfec..503314640c2 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -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)] diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 477cf44073f..f736ac1306a 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fc4d4e65660..d8320667a3a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py index 36c6ab7f711..e9973e1f674 100644 --- a/litellm/repositories/config_repository.py +++ b/litellm/repositories/config_repository.py @@ -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: diff --git a/tests/integration/management/test_fallback_management_public_names.py b/tests/integration/management/test_fallback_management_public_names.py index 6f97165fb3c..56f8309d6f4 100644 --- a/tests/integration/management/test_fallback_management_public_names.py +++ b/tests/integration/management/test_fallback_management_public_names.py @@ -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() diff --git a/tests/unit/proxy/test_fallback_management_endpoints.py b/tests/unit/proxy/test_fallback_management_endpoints.py index e3d0d072e0a..06adc0ed340 100644 --- a/tests/unit/proxy/test_fallback_management_endpoints.py +++ b/tests/unit/proxy/test_fallback_management_endpoints.py @@ -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)