diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1ead470bbd2..1ed4f25dc71 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1063,6 +1063,15 @@ def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) -> return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS +class ModelMaxBudgetHolder(Protocol): + @property + def model_max_budget(self) -> Mapping[str, object] | None: ... + + +def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | None: + return holder.model_max_budget + + def route_skips_budget_checks(route: str) -> bool: return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and ( route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 79ad4824c92..292a147d200 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -40,6 +40,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_event, phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports @@ -85,6 +86,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo jwt_key_mapping_cache_key, key_access_group_grants_model, key_model_aliases_for_auth_check, + model_max_budget_of, resolve_and_validate_end_user_id, resolve_default_end_user_budget, route_skips_budget_checks, @@ -384,6 +386,16 @@ class _TeamModelBudgetLimiter(Protocol): ) -> bool: ... +class _EndUserModelBudgetLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + class _EveryModelBudgetLimiter(Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... @@ -530,6 +542,49 @@ async def _check_team_model_budget( ) +async def _end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + budget_table: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else ( + await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if end_user_id and prisma_client is not None + else None + ) + ) + return model_max_budget_of(budget_table) if budget_table is not None else None + + +async def _check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: _EndUserModelBudgetLimiter, + models: Sequence[str], + team_id: str | None, +) -> None: + if not end_user_id or not end_user_model_max_budget: + return + for model_name in models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=model_name, + team_id=team_id, + ) + + async def _is_within_every_model_budget( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _EveryModelBudgetLimiter, @@ -2094,6 +2149,10 @@ async def _resolve_object_permission_for_unresolvable_team( ) +def _custom_auth_token_for_this_request(response: object) -> UserAPIKeyAuth: + return UserAPIKeyAuth.model_validate(response).model_copy() + + async def user_api_key_auth_builder( request: Request, api_key: str, @@ -2172,7 +2231,7 @@ async def user_api_key_auth_builder( request=request, api_key=api_key, user_custom_auth=user_custom_auth ) if response is not None and isinstance(response, UserAPIKeyAuth): - validated = UserAPIKeyAuth.model_validate(response) + validated = _custom_auth_token_for_this_request(response) if getattr(litellm, "enable_post_custom_auth_checks", False): validated = await _run_post_custom_auth_checks( valid_token=validated, @@ -2188,7 +2247,7 @@ async def user_api_key_auth_builder( custom_auth_api_key = True elif user_custom_auth is not None: response = await user_custom_auth(request=request, api_key=api_key) - validated = UserAPIKeyAuth.model_validate(response) + validated = _custom_auth_token_for_this_request(response) if getattr(litellm, "enable_post_custom_auth_checks", False): validated = await _run_post_custom_auth_checks( valid_token=validated, @@ -3129,7 +3188,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ) # Recompute after a potential budget-fallback rewrite so - # the end-user check below validates the final model + # the user check below validates the final model current_model = _get_model_from_request_context( request_data=request_data, route=route, @@ -3147,23 +3206,6 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e models=current_models, ) - # Check 5b. End-user model max budget - end_user_mmb: Final = valid_token.end_user_model_max_budget - if ( - end_user_mmb is not None - and isinstance(end_user_mmb, dict) - and len(end_user_mmb) > 0 - and current_models - and valid_token.end_user_id is not None - ): - for model_name in current_models: - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_mmb, - model=model_name, - team_id=valid_token.team_id, - ) - # Check 6: Additional Common Checks across jwt + key auth if valid_token.team_id is not None: try: @@ -3733,9 +3775,7 @@ async def run_centralized_common_checks( project_object: Final[LiteLLM_ProjectTableCachedObj | None] = ( None if isinstance(project_result, BaseException) else project_result ) - end_user_object: Final[LiteLLM_EndUserTable | None] = ( - None if isinstance(end_user_result, BaseException) else end_user_result - ) + end_user_object: Final = end_user_result if isinstance(end_user_result, LiteLLM_EndUserTable) else None global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result carry_team_and_user_budget_state( valid_token=user_api_key_auth_obj, @@ -3795,6 +3835,21 @@ async def run_centralized_common_checks( parent_otel_span=parent_otel_span, keep_token_limits=user_custom_auth is not None, ) + end_user_model_max_budget: Final = await _end_user_model_max_budget( + end_user_id=end_user_id, + end_user_object=end_user_object, + key_end_user_budget_id=key_end_user_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + if end_user_id and ( + not user_api_key_auth_obj.authenticated_by_custom_auth + or user_api_key_auth_obj.end_user_model_max_budget is None + ): + user_api_key_auth_obj.end_user_model_max_budget = ( # rebind-ok: the spend hook charges from this token + dict(end_user_model_max_budget) if end_user_model_max_budget else None + ) skip_budget_checks: Final = _should_skip_budget_checks( request_data=request_data, @@ -3887,6 +3942,21 @@ async def run_centralized_common_checks( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + await _check_end_user_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=user_api_key_auth_obj.end_user_model_max_budget, + model_max_budget_limiter=model_max_budget_limiter, + team_id=user_api_key_auth_obj.team_id, + models=_get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=user_api_key_auth_obj.team_id, + ) + ), + ) if sub_call and RouteChecks.is_llm_api_route(route=route): await virtual_key_max_budget_check( valid_token=user_api_key_auth_obj, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index ae7168ed1e1..b7f3ded9291 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5,7 +5,7 @@ import os import subprocess import sys import time -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import contextmanager, nullcontext from datetime import datetime, timedelta, timezone from functools import partial @@ -37,6 +37,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( team_member_budget_entity_id, ) from litellm.proxy._types import ( + JWTAuthBuilderResult, LiteLLMRoutes, LiteLLM_JWTAuth, LiteLLM_BudgetTable, @@ -55,7 +56,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, JWTRoutingOverride, ) -from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.auth_checks import ( OrganizationNotFoundError, TeamNotFoundError, @@ -724,6 +725,44 @@ async def test_user_custom_auth_runs_post_custom_auth_checks_when_opt_in(): litellm.enable_post_custom_auth_checks = original_flag +async def _authenticate_through_custom_auth() -> UserAPIKeyAuth: + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder + + request: Final = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + return await user_api_key_auth_builder( + request=request, + api_key="Bearer sk-custom-auth-shared", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + +@pytest.mark.asyncio +async def test_user_custom_auth_returning_one_shared_token_keeps_each_requests_end_user_apart( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm.proxy.proxy_server as _proxy_server_mod + + shared_token: Final = UserAPIKeyAuth(api_key="sk-custom-auth-shared", user_id="custom-user-123") + proxy_attrs: Final = _proxy_server_attrs_for_custom_auth(user_custom_auth=AsyncMock(return_value=shared_token)) + for attr, val in proxy_attrs.items(): + monkeypatch.setattr(_proxy_server_mod, attr, val, raising=False) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", False, raising=False) + + capped_end_users_request: Final = await _authenticate_through_custom_auth() + capped_end_users_request.end_user_id = "eu-capped" + capped_end_users_request.end_user_model_max_budget = {"gpt-4o": {"max_budget": 0.00001}} + next_end_users_request: Final = await _authenticate_through_custom_auth() + + assert (next_end_users_request.end_user_id, next_end_users_request.end_user_model_max_budget) == (None, None) + + @pytest.mark.asyncio async def test_enterprise_custom_auth_skips_post_custom_auth_checks_by_default(): """ @@ -5293,6 +5332,559 @@ async def test_centralized_common_checks_enforces_team_model_max_budget_from_the assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")] +_END_USER_MODEL_CAPS: Final = {"gpt-4o": {"budget_limit": 0.00004, "time_period": "1d"}} + + +def _virtual_key_token_for_end_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", token="hashed", end_user_id="eu-1") + + +def _master_key_token_cached_with_another_end_users_cap() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + end_user_id="eu-1", + end_user_model_max_budget={"gpt-4o": {"budget_limit": 1000.0, "time_period": "1d"}}, + ) + + +def _end_user_linked_to_the_capped_budget() -> LiteLLM_EndUserTable: + return LiteLLM_EndUserTable( + user_id="eu-1", + blocked=False, + budget_id="eu-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="eu-budget", model_max_budget=_END_USER_MODEL_CAPS), + ) + + +def _end_user_on_the_proxy_default_budget() -> LiteLLM_EndUserTable: + return LiteLLM_EndUserTable(user_id="eu-1", blocked=False) + + +def _no_end_user_row_yet() -> None: + return None + + +def _jwt_token(end_user_id: str | None, end_user_object: LiteLLM_EndUserTable | None) -> UserAPIKeyAuth: + return JWTAuthManager.user_api_key_auth_from_result( + JWTAuthBuilderResult( + is_proxy_admin=False, + team_object=None, + user_object=None, + end_user_object=end_user_object, + org_object=None, + token="jwt", + team_id=None, + user_id=None, + user_email=None, + end_user_id=end_user_id, + org_id=None, + team_membership=None, + jwt_claims={}, + agent_id=None, + ) + ) + + +def _jwt_token_with_an_end_user_claim() -> UserAPIKeyAuth: + return _jwt_token(end_user_id="eu-1", end_user_object=_end_user_linked_to_the_capped_budget()) + + +def _jwt_token_naming_the_end_user_only_in_the_body() -> UserAPIKeyAuth: + return _jwt_token(end_user_id=None, end_user_object=None) + + +async def _run_centralized_checks_for_end_user( + monkeypatch: pytest.MonkeyPatch, + token: UserAPIKeyAuth, + end_user: LiteLLM_EndUserTable | None, + gpt_4o_spend: float, + model: str, + route: str, + default_budget_id: str | None, + user_custom_auth: Callable[[Request, str], Awaitable[UserAPIKeyAuth]] | None, +) -> None: + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "max_end_user_budget_id", default_budget_id) + user_api_key_cache: Final = UserApiKeyCache() + if end_user is not None: + await user_api_key_cache.async_set_cache( + key=end_user_cache_key("eu-1"), + value=end_user, + model_type=LiteLLM_EndUserTable, + ) + for budget_id in ("eu-default", "eu-key-budget"): + await user_api_key_cache.async_set_cache( + key=f"default_end_user_budget:{budget_id}", + value=LiteLLM_BudgetTable(budget_id=budget_id, model_max_budget=_END_USER_MODEL_CAPS), + model_type=LiteLLM_BudgetTable, + ) + await user_api_key_cache.async_set_cache( + key="admin", + value=LiteLLM_UserTable(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + model_type=LiteLLM_UserTable, + ) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + await limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.END_USER, + entity_id="eu-1", + budget_model="gpt-4o", + budget_duration="1d", + ), + value=gpt_4o_spend, + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + ], + ) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=user_custom_auth), + "general_settings": { + "disable_budget_reservation": True, + "custom_auth_run_common_checks": user_custom_auth is not None, + }, + "llm_model_list": router.model_list, + "llm_router": router, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + "model_max_budget_limiter": limiter, + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http"}) + request._url = URL(url=route) + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": model, "user": "eu-1", "messages": [{"role": "user", "content": "hi"}]}, + route=route, + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + + + +@pytest.mark.parametrize( + ("make_token", "make_end_user", "gpt_4o_spend", "model", "route", "must_raise"), + [ + pytest.param( + _virtual_key_token_for_end_user, + _end_user_linked_to_the_capped_budget, + 0.0001, + "gpt-4o", + "/v1/chat/completions", + True, + id="virtual-key-over-cap", + ), + pytest.param( + _master_key_token_cached_with_another_end_users_cap, + _end_user_linked_to_the_capped_budget, + 0.0001, + "gpt-4o", + "/v1/chat/completions", + True, + id="master-key-over-cap", + ), + pytest.param( + _jwt_token_with_an_end_user_claim, + _end_user_linked_to_the_capped_budget, + 0.0001, + "gpt-4o", + "/v1/chat/completions", + True, + id="jwt-claim-over-cap", + ), + pytest.param( + _jwt_token_naming_the_end_user_only_in_the_body, + _end_user_linked_to_the_capped_budget, + 0.0001, + "gpt-4o", + "/v1/chat/completions", + True, + id="jwt-body-user-over-cap", + ), + pytest.param( + _virtual_key_token_for_end_user, + _no_end_user_row_yet, + 0.0001, + "gpt-4o", + "/v1/chat/completions", + True, + id="new-end-user-on-the-proxy-default-budget-over-cap", + ), + pytest.param( + _virtual_key_token_for_end_user, + _end_user_on_the_proxy_default_budget, + 0.0001, + "gpt-4o", + "/v1/chat/completions", + True, + id="proxy-default-budget-over-cap", + ), + pytest.param( + _virtual_key_token_for_end_user, + _end_user_linked_to_the_capped_budget, + 0.00001, + "gpt-4o", + "/v1/chat/completions", + False, + id="under-cap", + ), + pytest.param( + _virtual_key_token_for_end_user, + _end_user_linked_to_the_capped_budget, + 0.0001, + "gpt-4o-mini", + "/v1/chat/completions", + False, + id="model-without-a-cap", + ), + pytest.param( + _master_key_token_cached_with_another_end_users_cap, + _end_user_linked_to_the_capped_budget, + 0.0001, + "gpt-4o", + "/spend/calculate", + False, + id="budget-exempt-route", + ), + ], +) +@pytest.mark.asyncio +async def test_centralized_common_checks_enforces_the_end_users_model_max_budget( + monkeypatch: pytest.MonkeyPatch, + make_token: Callable[[], UserAPIKeyAuth], + make_end_user: Callable[[], LiteLLM_EndUserTable | None], + gpt_4o_spend: float, + model: str, + route: str, + must_raise: bool, +) -> None: + token: Final = make_token() + with pytest.raises(litellm.BudgetExceededError) if must_raise else nullcontext(): + await _run_centralized_checks_for_end_user( + monkeypatch=monkeypatch, + token=token, + end_user=make_end_user(), + gpt_4o_spend=gpt_4o_spend, + model=model, + route=route, + default_budget_id="eu-default", + user_custom_auth=None, + ) + assert token.end_user_model_max_budget == _END_USER_MODEL_CAPS + + +def _virtual_key_token_assigning_end_users_a_budget() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", token="hashed", end_user_id="eu-1", metadata={"end_user_budget_id": "eu-key-budget"} + ) + + +@pytest.mark.parametrize( + ("make_token", "make_end_user"), + [ + pytest.param( + _virtual_key_token_for_end_user, + _end_user_linked_to_the_capped_budget, + id="end-users-own-budget", + ), + pytest.param( + _virtual_key_token_assigning_end_users_a_budget, + _end_user_on_the_proxy_default_budget, + id="keys-end-user-budget", + ), + pytest.param( + _virtual_key_token_assigning_end_users_a_budget, + _no_end_user_row_yet, + id="keys-end-user-budget-for-a-new-end-user", + ), + ], +) +@pytest.mark.asyncio +async def test_centralized_common_checks_enforces_the_end_users_model_max_budget_without_a_proxy_default( + monkeypatch: pytest.MonkeyPatch, + make_token: Callable[[], UserAPIKeyAuth], + make_end_user: Callable[[], LiteLLM_EndUserTable | None], +) -> None: + token: Final = make_token() + with pytest.raises(litellm.BudgetExceededError): + await _run_centralized_checks_for_end_user( + monkeypatch=monkeypatch, + token=token, + end_user=make_end_user(), + gpt_4o_spend=0.0001, + model="gpt-4o", + route="/v1/chat/completions", + default_budget_id=None, + user_custom_auth=None, + ) + assert token.end_user_model_max_budget == _END_USER_MODEL_CAPS + + +async def _custom_auth_that_is_never_called(request: Request, api_key: str) -> UserAPIKeyAuth: + raise AssertionError("the shared checks never call custom auth") + + +def _token_built_by_custom_auth_with_its_own_cap() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-custom", end_user_id="eu-1", end_user_model_max_budget=_END_USER_MODEL_CAPS + ).model_copy(update={"authenticated_by_custom_auth": True}) + + +def _master_key_token_custom_auth_handed_back_with_another_end_users_cap() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + end_user_id="eu-1", + end_user_model_max_budget=_END_USER_MODEL_CAPS, + ) + + +@pytest.mark.parametrize( + ("make_token", "must_raise", "expected_cap"), + [ + pytest.param( + _token_built_by_custom_auth_with_its_own_cap, + True, + _END_USER_MODEL_CAPS, + id="custom-auth-token-keeps-its-own-cap", + ), + pytest.param( + _master_key_token_custom_auth_handed_back_with_another_end_users_cap, + False, + None, + id="key-auth-token-behind-custom-auth-drops-another-end-users-cap", + ), + ], +) +@pytest.mark.asyncio +async def test_centralized_common_checks_keeps_only_a_cap_custom_auth_set( + monkeypatch: pytest.MonkeyPatch, + make_token: Callable[[], UserAPIKeyAuth], + must_raise: bool, + expected_cap: Mapping[str, object] | None, +) -> None: + token: Final = make_token() + with pytest.raises(litellm.BudgetExceededError) if must_raise else nullcontext(): + await _run_centralized_checks_for_end_user( + monkeypatch=monkeypatch, + token=token, + end_user=_end_user_on_the_proxy_default_budget(), + gpt_4o_spend=0.0001, + model="gpt-4o", + route="/v1/chat/completions", + default_budget_id=None, + user_custom_auth=_custom_auth_that_is_never_called, + ) + assert token.end_user_model_max_budget == expected_cap + + +@pytest.mark.asyncio +async def test_centralized_common_checks_holds_the_end_user_to_the_model_a_team_member_fallback_picked( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + from litellm.proxy.utils import ProxyLogging + + budget_id: Final = "member-budget-1" + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="u1", team_id="t1"), + value=LiteLLM_TeamMembership(user_id="u1", team_id="t1", budget_id=budget_id), + model_type=LiteLLM_TeamMembership, + ) + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj( + team_id="t1", + models=["gpt-4o", "gpt-4o-mini"], + metadata={"team_member_budget_id": budget_id}, + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + await user_api_key_cache.async_set_cache( + key=f"team_member_default_budget:{budget_id}", + value=LiteLLM_BudgetTable( + budget_id=budget_id, + model_max_budget={ + "gpt-4o": {"max_budget": 0.00004, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ), + model_type=LiteLLM_BudgetTable, + ) + await user_api_key_cache.async_set_cache( + key="u1", + value=LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER), + model_type=LiteLLM_UserTable, + ) + await user_api_key_cache.async_set_cache( + key=end_user_cache_key("eu-1"), + value=_end_user_linked_to_the_capped_budget(), + model_type=LiteLLM_EndUserTable, + ) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + for entity_type, entity_id in ( + (Litellm_EntityType.TEAM_MEMBER, team_member_budget_entity_id(user_id="u1", team_id="t1")), + (Litellm_EntityType.END_USER, "eu-1"), + ): + await limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model="gpt-4o", + budget_duration="1d", + ), + value=0.0001, + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + ], + fallbacks=[{"gpt-4o": ["gpt-4o-mini"]}], + ) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(), + "general_settings": {"disable_budget_reservation": True}, + "llm_model_list": router.model_list, + "llm_router": router, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + "model_max_budget_limiter": limiter, + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http", "path_params": {}, "headers": []}) + request._url = URL(url="/v1/chat/completions") + request_data: Final = {"model": "gpt-4o", "user": "eu-1", "messages": [{"role": "user", "content": "hi"}]} + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + await run_centralized_common_checks( + user_api_key_auth_obj=UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + user_id="u1", + team_id="t1", + end_user_id="eu-1", + models=["gpt-4o", "gpt-4o-mini"], + user_role=LitellmUserRoles.INTERNAL_USER, + ), + request=request, + request_data=request_data, + route="/v1/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + assert request_data["model"] == "gpt-4o-mini" + + +def _access_groups_only_team_t1_puts_gpt_4o_in(model: str, team_id: str | None) -> frozenset[str]: + return frozenset({"t1-group"}) if (model, team_id) == ("gpt-4o", "t1") else frozenset() + + +@pytest.mark.parametrize( + ("team_id", "must_raise"), + [ + pytest.param("t1", True, id="team-whose-access-group-the-end-user-capped"), + pytest.param(None, False, id="no-team-so-the-group-cap-does-not-apply"), + ], +) +@pytest.mark.asyncio +async def test_centralized_common_checks_applies_an_end_user_cap_on_the_tokens_team_access_group( + team_id: str | None, + must_raise: bool, +) -> None: + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + from litellm.proxy.utils import ProxyLogging + + group_cap: Final = {"t1-group": {"budget_limit": 0.00004, "time_period": "1d"}} + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", models=["gpt-4o"]), + model_type=LiteLLM_TeamTableCachedObj, + ) + await user_api_key_cache.async_set_cache( + key=end_user_cache_key("eu-1"), + value=LiteLLM_EndUserTable( + user_id="eu-1", + blocked=False, + budget_id="eu-group-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="eu-group-budget", model_max_budget=group_cap), + ), + model_type=LiteLLM_EndUserTable, + ) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(), model_access_groups=_access_groups_only_team_t1_puts_gpt_4o_in + ) + await limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.END_USER, + entity_id="eu-1", + budget_model="t1-group", + budget_duration="1d", + ), + value=0.0001, + ) + router: Final = litellm.Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}]) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(), + "general_settings": {"disable_budget_reservation": True}, + "llm_model_list": router.model_list, + "llm_router": router, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + "model_max_budget_limiter": limiter, + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http", "path_params": {}, "headers": []}) + request._url = URL(url="/v1/chat/completions") + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + with pytest.raises(litellm.BudgetExceededError) if must_raise else nullcontext(): + await run_centralized_common_checks( + user_api_key_auth_obj=UserAPIKeyAuth( + api_key="sk-test", token="hashed", team_id=team_id, end_user_id="eu-1", models=["gpt-4o"] + ), + request=request, + request_data={"model": "gpt-4o", "user": "eu-1", "messages": [{"role": "user", "content": "hi"}]}, + route="/v1/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + @pytest.mark.asyncio async def test_centralized_common_checks_loads_member_budget_and_exempts_discovery_routes(): from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache