diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 12d420141f1..7673c3983b4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -89,7 +89,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, ) -from litellm.proxy.common_utils.model_listing_utils import alias_map +from litellm.proxy.common_utils.model_listing_utils import alias_map, alias_target, caller_alias_maps from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, @@ -138,6 +138,7 @@ from litellm.repositories.table_repositories import ( from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime @@ -449,6 +450,10 @@ def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None: return cache if isinstance(cache, dict) else None +def _resolve_cost_model_group(model_name: str, llm_router: Router) -> str: + return resolve_model_group_alias(llm_router.model_group_alias, model_name) or model_name + + def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool: """ Check if a model has zero cost (no configured pricing). @@ -471,22 +476,23 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) for model_name in model_list: - if zero_cost_cache is not None: - cached = zero_cost_cache.get(model_name) - if cached is not None: - if cached is False: - return False - continue try: + target_group = _resolve_cost_model_group(model_name, llm_router) + if zero_cost_cache is not None: + cached = zero_cost_cache.get(target_group) + if cached is not None: + if cached is False: + return False + continue # Use router's get_model_group_info method directly for better reliability - model_group_info = llm_router.get_model_group_info(model_group=model_name) + model_group_info = llm_router.get_model_group_info(model_group=target_group) if model_group_info is None: # Model not found or no pricing info available # Conservative approach: assume it has cost verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False # Check costs for this model @@ -503,7 +509,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None output_cost, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False # If either cost is non-zero, return False @@ -512,14 +518,14 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False # Costs are 0 — verify this is from explicit configuration, # not from defaulted sparse auto-registration entries. # See: https://github.com/BerriAI/litellm/issues/24770 safe_name = str(model_name).replace("\n", "").replace("\r", "") - if not _is_cost_explicitly_configured(model_name, llm_router): + if not _is_cost_explicitly_configured(target_group, llm_router): verbose_proxy_logger.debug( "Model %s has zero cost but no explicit cost " "configuration in model_cost entry — treating as unknown " @@ -527,17 +533,17 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None safe_name, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False - if _has_ptu_flat_cost(model_name, llm_router): + if _has_ptu_flat_cost(target_group, llm_router): verbose_proxy_logger.debug( "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " "rate is not a free model (enforce budget)", safe_name, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = False + zero_cost_cache[target_group] = False return False verbose_proxy_logger.debug( @@ -547,7 +553,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None output_cost, ) if zero_cost_cache is not None: - zero_cost_cache[model_name] = True + zero_cost_cache[target_group] = True except Exception as e: # If we can't determine the cost, assume it has cost (conservative approach) @@ -558,6 +564,36 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None return True +def _dispatched_model_name( + model_name: str, valid_token: UserAPIKeyAuth, router_settings_aliases: object, router_settings_rewrite_pending: bool +) -> str: + requested: Final = ( + resolve_model_group_alias(router_settings_aliases, model_name) or model_name + if router_settings_rewrite_pending + else model_name + ) + caller_aliases: Final = caller_alias_maps( + valid_token.aliases, valid_token.team_model_aliases, valid_token.team_id, None + ) + aliased: Final = alias_target(requested, caller_aliases) or requested + return resolve_model_group_alias(router_settings_aliases, aliased) or aliased + + +def is_dispatched_model_cost_zero( + model: str | list[str] | None, + llm_router: Router | None, + valid_token: UserAPIKeyAuth, + router_settings_aliases: object, + router_settings_rewrite_pending: bool, +) -> bool: + dispatched_model: Final = ( + _dispatched_model_name(model, valid_token, router_settings_aliases, router_settings_rewrite_pending) + if isinstance(model, str) + else model + ) + return _is_model_cost_zero(model=dispatched_model, llm_router=llm_router) + + _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e3ce9bcd850..8ced8eed481 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -48,7 +48,6 @@ from litellm.proxy.auth.auth_checks import ( _check_end_user_budget, _delete_cache_key_object, _get_user_role, - _is_model_cost_zero, _is_user_proxy_admin, _team_member_max_budget_alert_check, _virtual_key_max_budget_alert_check, @@ -65,6 +64,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_membership, get_team_object, get_user_object, + is_dispatched_model_cost_zero, is_valid_fallback_model, jwt_key_mapping_cache_key, key_model_aliases_for_auth_check, @@ -1795,9 +1795,9 @@ async def _user_api_key_auth_builder( ) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero - - skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + skip_budget_checks = await _is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token, request=request + ) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) @@ -2237,9 +2237,9 @@ async def _user_api_key_auth_builder( ) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero - - skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + skip_budget_checks = await _is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token, request=request + ) if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) @@ -2960,12 +2960,12 @@ async def _run_centralized_common_checks( keep_token_limits=user_custom_auth is not None, ) - skip_budget_checks: Final = _should_skip_budget_checks( + skip_budget_checks: Final = await _should_skip_budget_checks( request_data=request_data, route=route, request=request, llm_router=llm_router, - team_id=user_api_key_auth_obj.team_id, + valid_token=user_api_key_auth_obj, ) # Pin the metadata variable name (litellm_metadata vs metadata) before @@ -3131,25 +3131,45 @@ async def _reserve_budget_after_common_checks( request.state.budget_reservation = reservation # rebind-ok: read by the release middleware -def _should_skip_budget_checks( +async def _should_skip_budget_checks( request_data: dict, route: str, request: Request | None, llm_router: Any | None, - team_id: str | None = None, + valid_token: UserAPIKeyAuth, ) -> bool: model: Final = _get_model_from_request_context( request_data=request_data, route=route, request=request, llm_router=llm_router, - team_id=team_id, + team_id=valid_token.team_id, ) if model is not None and llm_router is not None: - return _is_model_cost_zero(model=model, llm_router=llm_router) + return await _is_dispatched_model_cost_zero( + model=model, llm_router=llm_router, valid_token=valid_token, request=request + ) return False +async def _is_dispatched_model_cost_zero( + model: str | list[str], llm_router: litellm.Router, valid_token: UserAPIKeyAuth, request: Request | None +) -> bool: + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + return is_dispatched_model_cost_zero( + model=model, + llm_router=llm_router, + valid_token=valid_token, + router_settings_aliases=settings.get("model_group_alias") if isinstance(settings, Mapping) else None, + router_settings_rewrite_pending=request is not None + and request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True, + ) + + def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal: """Project the resolved identity into one per-request Principal, off the key object the builder already fetched, and stamp the request network context @@ -3715,7 +3735,9 @@ async def _run_post_custom_auth_checks( # every budget check for these; this path did not, so the same request could # be refused under custom auth and served under the other two. skip_budget_checks: Final = ( - _is_model_cost_zero(model=current_model, llm_router=llm_router) + await _is_dispatched_model_cost_zero( + model=current_model, llm_router=llm_router, valid_token=valid_token, request=request + ) if current_model is not None and llm_router is not None else False ) diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 2538556d3b5..3f1af71c06a 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -3,6 +3,7 @@ import sys, os, asyncio, time, random, uuid import traceback +from typing import Final from dotenv import load_dotenv load_dotenv() @@ -23,9 +24,12 @@ from litellm.proxy._types import ( from litellm.proxy.utils import PrismaClient from litellm.proxy.auth.auth_checks import ( can_team_access_model, + _is_model_cost_zero, + is_dispatched_model_cost_zero, _virtual_key_soft_budget_check, _team_soft_budget_check, ) +from litellm.router import Router from litellm.proxy.utils import ProxyLogging from litellm.proxy.utils import CallInfo @@ -1491,3 +1495,156 @@ async def test_key_access_group_grants_model_when_get_access_object_raises(): finally: for p in patches: p.stop() + + +def _alias_router(model_group_alias: dict[str, str | dict[str, str | bool]]) -> Router: + return Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "paid-model", + "litellm_params": { + "model": "openai/paid-model", + "api_key": "sk-fake", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + }, + }, + { + "model_name": "ptu-model", + "litellm_params": { + "model": "azure/ptu-deployment", + "api_base": "https://ptu.invalid", + "api_key": "sk-fake", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"ptu_count": 100, "cost_per_ptu_per_hour": 1.0}, + }, + ], + model_group_alias=model_group_alias, + ) + + +def _aliases_to(target: str) -> dict[str, str | dict[str, str | bool]]: + return {"visible": target, "hidden": {"model": target, "hidden": True}} + + +@pytest.mark.parametrize( + "target, model, expected", + [ + ("free-model", "visible", True), + ("free-model", "hidden", True), + ("free-model", ["visible", "hidden", "free-model"], True), + ("free-model", ["hidden", "paid-model"], False), + ("paid-model", "visible", False), + ("paid-model", "hidden", False), + ("ptu-model", "visible", False), + ("ptu-model", "hidden", False), + ("no-such-group", "visible", False), + ("no-such-group", "hidden", False), + ], +) +def test_zero_cost_check_prices_an_alias_as_its_target_group( + target: str, model: str | list[str], expected: bool +) -> None: + router: Final = _alias_router(_aliases_to(target)) + assert _is_model_cost_zero(model=model, llm_router=router) is expected + + +@pytest.mark.parametrize( + "first_target, repointed_target, expected_after_repoint", + [ + ("free-model", "paid-model", False), + ("paid-model", "free-model", True), + ], +) +@pytest.mark.parametrize("alias", ["visible", "hidden"]) +def test_zero_cost_check_follows_an_alias_repointed_at_runtime( + alias: str, first_target: str, repointed_target: str, expected_after_repoint: bool +) -> None: + router: Final = _alias_router(_aliases_to(first_target)) + assert _is_model_cost_zero(model=alias, llm_router=router) is not expected_after_repoint + router.update_settings(model_group_alias=_aliases_to(repointed_target)) + assert _is_model_cost_zero(model=alias, llm_router=router) is expected_after_repoint + + +@pytest.mark.parametrize( + "team_model_aliases, key_aliases, model, expected", + [ + (None, {"visible": "paid-model"}, "visible", False), + (None, {"free-model": "paid-model"}, "free-model", False), + ({"free-model": "paid-model"}, None, "free-model", False), + (None, {"my-free": "free-model"}, "my-free", True), + (None, {"my-free": "visible"}, "my-free", True), + ({"team-name": "key-name"}, {"key-name": "free-model"}, "team-name", True), + ({"team-name": "key-name"}, {"key-name": "paid-model"}, "team-name", False), + (None, {"other": "paid-model"}, "visible", True), + ], +) +def test_zero_cost_check_prices_the_model_the_key_and_team_aliases_dispatch_to( + team_model_aliases: dict[str, str] | None, + key_aliases: dict[str, str] | None, + model: str, + expected: bool, +) -> None: + router: Final = _alias_router(_aliases_to("free-model")) + token: Final = UserAPIKeyAuth(aliases=key_aliases or {}, team_model_aliases=team_model_aliases) + assert ( + is_dispatched_model_cost_zero( + model=model, + llm_router=router, + valid_token=token, + router_settings_aliases=None, + router_settings_rewrite_pending=False, + ) + is expected + ) + + +@pytest.mark.parametrize( + "key_aliases, router_settings_aliases, model_alias_map, rewrite_pending, model, expected", + [ + ({"my-free": "free-model"}, {"free-model": "paid-model"}, {}, True, "my-free", False), + ({"my-free": "free-model"}, {}, {"free-model": "paid-model"}, True, "my-free", False), + ({}, {"free-model": "paid-model"}, {}, True, "free-model", False), + ({}, {"router-free": "free-model"}, {}, True, "router-free", True), + ({}, {"router-free": {"model": "visible", "hidden": True}}, {}, True, "router-free", True), + ({"global-free": "free-model"}, {}, {"my-free": "global-free"}, True, "my-free", True), + ({"global-free": "paid-model"}, {}, {"my-free": "global-free"}, True, "my-free", False), + ({}, {}, {"my-free": "visible"}, True, "my-free", True), + ({"free-model": "paid-model"}, {"router-free": "free-model"}, {}, True, "router-free", False), + ({}, {"router-free": "free-model", "free-model": "paid-model"}, {}, True, "router-free", False), + ({}, {"router-free": "free-model", "free-model": "paid-model"}, {}, False, "router-free", True), + ], +) +def test_zero_cost_check_prices_the_model_after_every_alias_hop_dispatch_applies( + monkeypatch: pytest.MonkeyPatch, + key_aliases: dict[str, str], + router_settings_aliases: dict[str, str | dict[str, str | bool]], + model_alias_map: dict[str, str], + rewrite_pending: bool, + model: str, + expected: bool, +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", model_alias_map) + router: Final = _alias_router(_aliases_to("free-model")) + token: Final = UserAPIKeyAuth(aliases=key_aliases) + assert ( + is_dispatched_model_cost_zero( + model=model, + llm_router=router, + valid_token=token, + router_settings_aliases=router_settings_aliases, + router_settings_rewrite_pending=rewrite_pending, + ) + is expected + ) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 9cdac341b1f..aa09f6ae45f 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1794,3 +1794,74 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): "the mapped-key branch returns before the shared virtual-key checks, so the " "user's per-model budget is never attached and never enforced" ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, key_aliases, key_router_settings, router_settings_rewritten, expected", + [ + ("free-alias", {}, None, None, True), + ("free-alias", {"free-alias": "paid-model"}, None, None, False), + ("my-alias", {"my-alias": "free-alias"}, None, None, True), + ("my-alias", {"my-alias": "free-alias"}, {"model_group_alias": {"free-alias": "paid-model"}}, None, False), + ("rs-alias", {"free-model": "paid-model"}, {"model_group_alias": {"rs-alias": "free-model"}}, False, False), + ("rs-alias", {"free-model": "paid-model"}, {"model_group_alias": {"rs-alias": "free-model"}}, True, True), + ], +) +async def test_budget_skip_judges_the_model_the_key_aliases_dispatch_to( + model: str, + key_aliases: dict[str, str], + key_router_settings: dict[str, dict[str, str]] | None, + router_settings_rewritten: bool | None, + expected: bool, +) -> None: + from starlette.requests import Request + + from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY + from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "paid-model", + "litellm_params": { + "model": "openai/paid-model", + "api_key": "sk-fake", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + }, + }, + ], + model_group_alias={"free-alias": "free-model"}, + ) + skipped = await _should_skip_budget_checks( + request_data={"model": model}, + route="/chat/completions", + request=( + None + if router_settings_rewritten is None + else Request( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "headers": [], + "query_string": b"", + MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: router_settings_rewritten, + } + ) + ), + llm_router=router, + valid_token=UserAPIKeyAuth(aliases=key_aliases, router_settings=key_router_settings), + ) + assert skipped is expected