Merge 5b84fc8df9 into 98c710c411

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})
 

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
     )

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
+    )

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