scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 249

fix(proxy)!: enforce an end user's per-model budget from their linked budget (#48815)

Browse source 
* fix(proxy): enforce an end user's per-model budget from their linked budget

The shared auth checks now refuse a request once the end user's spend on
that model reaches the model_max_budget of their budget (a linked budget or
the default applied to them), for virtual keys and the master key alike.
The cap is read from the end-user row, since the master key's cached auth
object keeps whichever end user's limits it was cached with. The per-model
end-user check inside virtual-key validation is removed: it read a token
field that is never set there, so it never ran.

* fix(auth): enforce the end user's per-model budget on JWT auth and skip it on budget-exempt routes

JWT tokens now carry the end user's per-model caps so spend is charged against them and over-cap requests are refused. Routes that skip budget checks, like /spend/calculate, skip this check too. The proxy default end-user budget is covered by a test.

* fix(auth): write the end user's per-model cap onto every token and check it after the member fallback

The shared auth checks now resolve the end user's per-model cap, falling back to the key's or the proxy's default end-user budget when the end user has no row yet, and write it onto the token so the spend hook charges the cap the check enforces on every auth path. The check runs after the team-member budget fallback, on the model it picked, and the JWT builder's duplicate check is removed

* fix(auth): keep only a per-model end-user cap that custom auth itself set

A key-auth token reached through custom auth's auto hand-back now takes the current end user's cap, so a cached master-key or virtual-key object no longer carries an earlier end user's cap into another end user's request.

* test(auth): build the custom-auth token with model_copy instead of setting the flag in place

* test(auth): hold an end user to their own or their key's budget cap when the proxy has no default budget

* fix(proxy): give each request its own copy of a custom-auth token so end-user caps and ids stay per request

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
mika-lt 2026-10-28 21:24:11 -07:00 • committed by GitHub
parent 5bd47442fb
commit de42b67252
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 696 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/auth/auth_checks.py
View file

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

116
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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,

596
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

@ -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