From 0fbb0841d5cc6387e9c7475a34200f8adb219946 Mon Sep 17 00:00:00 2001 From: Andrey Zaytsev Date: Thu, 10 Sep 2026 13:25:35 +0200 Subject: [PATCH] feat(proxy): fall back to a shared end-user budget when exhausted --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/end_user.py | 1 + litellm/proxy/_types.py | 4 + litellm/proxy/auth/auth_checks.py | 39 ++++++- litellm/proxy/auth/user_api_key_auth.py | 59 +++++++++-- .../proxy/hooks/proxy_track_cost_callback.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 2 +- .../customer_endpoints.py | 29 ++++- litellm/proxy/schema.prisma | 2 + schema.prisma | 2 + .../proxy/auth/test_auth_checks.py | 1 + .../auth/test_custom_auth_end_user_budget.py | 10 +- .../proxy/auth/test_user_api_key_auth.py | 36 +++++++ .../hooks/test_proxy_track_cost_callback.py | 6 +- .../test_customer_endpoints.py | 56 ++++++++++ .../proxy/test_budget_reservation.py | 100 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +++ 18 files changed, 343 insertions(+), 22 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260910120000_add_end_user_fallback_budget/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910120000_add_end_user_fallback_budget/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910120000_add_end_user_fallback_budget/migration.sql new file mode 100644 index 00000000000..d064ca4da69 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910120000_add_end_user_fallback_budget/migration.sql @@ -0,0 +1,2 @@ +ALTER TABLE "LiteLLM_EndUserTable" ADD COLUMN IF NOT EXISTS "fallback_end_user_id" TEXT; +CREATE INDEX IF NOT EXISTS "LiteLLM_EndUserTable_fallback_end_user_id_idx" ON "LiteLLM_EndUserTable"("fallback_end_user_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 05c5aad9303..2e57f9f8fba 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -585,6 +585,8 @@ model LiteLLM_EndUserTable { object_permission_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + fallback_end_user_id String? + @@index([fallback_end_user_id]) blocked Boolean @default(false) } diff --git a/litellm/models/end_user.py b/litellm/models/end_user.py index 8dccf1eb5e7..f978647327d 100644 --- a/litellm/models/end_user.py +++ b/litellm/models/end_user.py @@ -22,6 +22,7 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): allowed_model_region: Literal["eu", "us"] | None = None default_model: str | None = None budget_id: str | None = None + fallback_end_user_id: str | None = None litellm_budget_table: LiteLLM_BudgetTable | None = None object_permission_id: str | None = None object_permission: LiteLLM_ObjectPermissionTable | None = None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c22bb76629d..3aaac684259 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1889,6 +1889,7 @@ class BudgetDeleteRequest(LiteLLMPydanticObjectBase): class CustomerBase(LiteLLMPydanticObjectBase): user_id: str + fallback_end_user_id: str | None = None alias: str | None = None spend: float = 0.0 allowed_model_region: AllowedModelRegion | None = None @@ -1904,6 +1905,7 @@ class NewCustomerRequest(BudgetNewRequest): """ user_id: str + fallback_end_user_id: str | None = None alias: str | None = None # human-friendly alias blocked: bool = False # allow/disallow requests for this end-user budget_id: str | None = None # give either a budget_id or max_budget @@ -1930,6 +1932,7 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): """ user_id: str + fallback_end_user_id: str | None = None alias: str | None = None # human-friendly alias blocked: bool = False # allow/disallow requests for this end-user max_budget: float | None = None @@ -3063,6 +3066,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob "user id." ), ) + billing_end_user_id: str | None = Field(default=None, exclude=True) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1efc9611fe6..1efece6d167 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1488,9 +1488,43 @@ async def _check_end_user_budget( ) +async def resolve_end_user_budget_fallback( + end_user_obj: LiteLLM_EndUserTable | None, + route: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool = False, +) -> LiteLLM_EndUserTable | None: + if end_user_obj is None or end_user_obj.fallback_end_user_id is None or skip_budget_checks: + return end_user_obj + try: + await _check_end_user_budget(end_user_obj, route) + except litellm.BudgetExceededError: + fallback: Final = await get_end_user_object( + end_user_id=end_user_obj.fallback_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + route=route, + require_exists=True, + ) + if fallback is not None and fallback.fallback_end_user_id is None: + try: + await _check_end_user_budget(fallback, route) + except litellm.BudgetExceededError: + return end_user_obj + return fallback + return end_user_obj + + #: Columns whose non-null value makes an end-user row restrict something auth enforces. ``blocked`` #: is separate: it restricts when true rather than when merely set. -_RESTRICTED_COLUMNS: Final = ("budget_id", "allowed_model_region", "default_model", "object_permission_id") +_RESTRICTED_COLUMNS: Final = ( + "budget_id", + "allowed_model_region", + "default_model", + "object_permission_id", + "fallback_end_user_id", +) def _column_is_set(column: str) -> Mapping[str, object]: @@ -1683,6 +1717,7 @@ async def get_end_user_object( parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, token_end_user_max_budget: float | None = None, + require_exists: bool = False, ) -> LiteLLM_EndUserTable | None: """ Returns end user object from database or cache. @@ -1729,7 +1764,7 @@ async def get_end_user_object( return return_obj - if await _end_user_is_known_unrestricted( + if not require_exists and await _end_user_is_known_unrestricted( end_user_id=end_user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 64c0da1c28f..48242bed5ad 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -60,6 +60,7 @@ from litellm.proxy.auth.auth_checks import ( is_valid_fallback_model, jwt_key_mapping_cache_key, resolve_and_validate_end_user_id, + resolve_end_user_budget_fallback, ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_method import AuthMethod @@ -591,7 +592,28 @@ async def user_api_key_auth_websocket(websocket: WebSocket): raise HTTPException(status_code=403, detail=str(e)) -def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: +async def update_valid_token_with_end_user_params( + valid_token: UserAPIKeyAuth, + end_user_params: dict, + end_user_obj: LiteLLM_EndUserTable | None = None, + route: str = "", + skip_budget_checks: bool = False, +) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + billing_end_user: Final = await resolve_end_user_budget_fallback( + end_user_obj, route, prisma_client, user_api_key_cache, skip_budget_checks + ) + if billing_end_user is not None and billing_end_user is not end_user_obj: + budget: Final = billing_end_user.litellm_budget_table + end_user_params.update( + billing_end_user_id=billing_end_user.user_id, + end_user_max_budget=budget.max_budget if budget is not None else None, + end_user_model_max_budget=budget.model_max_budget if budget is not None else None, + ) + valid_token.billing_end_user_id = end_user_params.get("billing_end_user_id") + if valid_token.billing_end_user_id is not None: + valid_token.end_user_max_budget = end_user_params.get("end_user_max_budget") valid_token.end_user_id = end_user_params.get("end_user_id") # Only overwrite token fields when the DB-derived value is not None. # This prevents DB lookups (where the budget table has no value set) @@ -603,8 +625,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] if end_user_params.get("allowed_model_region") is not None: valid_token.allowed_model_region = end_user_params["allowed_model_region"] - if end_user_params.get("end_user_model_max_budget") is not None: - valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] + if end_user_params.get("end_user_model_max_budget") is not None or valid_token.billing_end_user_id is not None: + valid_token.end_user_model_max_budget = end_user_params.get("end_user_model_max_budget") return valid_token @@ -1790,8 +1812,12 @@ async def _user_api_key_auth_builder( code=status.HTTP_401_UNAUTHORIZED, param=abbreviate_api_key(api_key=api_key), ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user_params + valid_token = await update_valid_token_with_end_user_params( + valid_token=valid_token, + end_user_params=end_user_params, + end_user_obj=_end_user_object, + route=route, + skip_budget_checks=_should_skip_budget_checks(request_data, route, request, llm_router), ) valid_token.parent_otel_span = parent_otel_span if _end_user_object is not None: @@ -1868,8 +1894,12 @@ async def _user_api_key_auth_builder( ) ) - _user_api_key_obj = update_valid_token_with_end_user_params( - valid_token=_user_api_key_obj, end_user_params=end_user_params + _user_api_key_obj = await update_valid_token_with_end_user_params( + valid_token=_user_api_key_obj, + end_user_params=end_user_params, + end_user_obj=_end_user_object, + route=route, + skip_budget_checks=_should_skip_budget_checks(request_data, route, request, llm_router), ) _user_api_key_obj.via_virtual_key = True @@ -1941,6 +1971,13 @@ async def _user_api_key_auth_builder( if valid_token is not None: valid_token = _update_key_budget_with_temp_budget_increase(valid_token) + valid_token = await update_valid_token_with_end_user_params( + valid_token, + end_user_params, + _end_user_object, + route, + _should_skip_budget_checks(request_data, route, request, llm_router), + ) user_obj: LiteLLM_UserTable | None = None valid_token_dict: dict = {} @@ -2180,7 +2217,7 @@ async def _user_api_key_auth_builder( ): 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_id=valid_token.billing_end_user_id or valid_token.end_user_id, end_user_model_max_budget=end_user_mmb, model=model_name, ) @@ -2471,7 +2508,7 @@ async def _run_centralized_common_checks( # resolved the end-user id and attached it here. Reuse that to avoid a # second extraction pass; fall back to extracting locally when the # function is invoked in isolation (e.g. in direct unit tests). - end_user_id = user_api_key_auth_obj.end_user_id + end_user_id = user_api_key_auth_obj.billing_end_user_id or user_api_key_auth_obj.end_user_id if end_user_id is None: raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) end_user_id = await resolve_and_validate_end_user_id( @@ -3130,7 +3167,7 @@ async def _lookup_end_user_and_apply_budget( budget_info=end_user_object.litellm_budget_table, end_user_id=valid_token.end_user_id or "", ) - valid_token = update_valid_token_with_end_user_params( + valid_token = await update_valid_token_with_end_user_params( valid_token=valid_token, end_user_params=end_user_params ) elif litellm.max_end_user_budget_id is not None: @@ -3148,7 +3185,7 @@ async def _lookup_end_user_and_apply_budget( budget_info=default_budget, end_user_id=valid_token.end_user_id or "", ) - valid_token = update_valid_token_with_end_user_params( + valid_token = await update_valid_token_with_end_user_params( valid_token=valid_token, end_user_params=end_user_params ) except Exception as e: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..ece7294608a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -207,7 +207,7 @@ class _ProxyDBLogger(CustomLogger): token=user_api_key_dict.api_key, response_cost=recovered_response_cost, user_id=user_api_key_dict.user_id, - end_user_id=user_api_key_dict.end_user_id, + end_user_id=user_api_key_dict.billing_end_user_id or user_api_key_dict.end_user_id, team_id=user_api_key_dict.team_id, kwargs=request_data, completion_response=original_exception, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index da033dc2276..d5e275eafa4 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1593,7 +1593,7 @@ class LiteLLMProxyRequestSetup: user_api_key_org_id=user_api_key_dict.org_id, user_api_key_org_alias=user_api_key_dict.organization_alias, user_api_key_team_alias=user_api_key_dict.team_alias, - user_api_key_end_user_id=user_api_key_dict.end_user_id, + user_api_key_end_user_id=user_api_key_dict.billing_end_user_id or user_api_key_dict.end_user_id, user_api_key_user_email=user_api_key_dict.user_email, user_api_key_request_route=user_api_key_dict.request_route, user_api_key_budget_reset_at=( diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index d2d87331d55..03664274539 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -89,6 +89,8 @@ if TYPE_CHECKING: data: Mapping[str, Mapping[str, object]], ) -> _RowT_co: ... + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + async def delete_many(self, where: Mapping[str, object]) -> int: ... @@ -122,6 +124,21 @@ def _end_user_cache_keys(user_ids: Sequence[str]) -> tuple[str, ...]: return (*(end_user_cache_key(user_id) for user_id in user_ids), end_user_restricted_registry_cache_key()) +async def _validate_fallback_end_user( + user_id: str, fallback_end_user_id: str | None, prisma_client: "PrismaClient" +) -> None: + if fallback_end_user_id is None: + return + if fallback_end_user_id == user_id: + raise HTTPException(status_code=422, detail="fallback_end_user_id must not equal user_id") + table: Final = _typed_table(EndUserRepository(prisma_client)) + target: Final = await table.find_first(where={"user_id": fallback_end_user_id}) + if target is None: + raise HTTPException(status_code=422, detail=f"Fallback customer {fallback_end_user_id} does not exist") + if target.fallback_end_user_id is not None: + raise HTTPException(status_code=422, detail="Fallback customers must not have another fallback (depth 1 only)") + + def _to_customer_response(record: BaseModel) -> CustomerResponse: """Validate a raw end-user DB row into the typed customer response. @@ -329,6 +346,7 @@ async def new_end_user( - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + - fallback_end_user_id: Optional[str] - Another customer whose budget is charged once this customer's own budget is exhausted. Depth is one: the fallback customer cannot itself have a fallback. - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} @@ -408,6 +426,7 @@ async def new_end_user( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) try: + await _validate_fallback_end_user(data.user_id, data.fallback_end_user_id, prisma_client) ## VALIDATION ## if data.default_model is not None: if llm_router is None: @@ -572,6 +591,7 @@ async def update_end_user( - blocked: bool = False # allow/disallow requests for this end-user - max_budget: Optional[float] = None - budget_id: Optional[str] = None # give either a budget_id or max_budget + - fallback_end_user_id: Optional[str] = None # customer charged once this customer's budget is exhausted; depth one, null clears - allowed_model_region: Optional[AllowedModelRegion] = ( None # require all user requests to use models in this specific region ) @@ -626,7 +646,9 @@ async def update_end_user( # get non default values for key non_default_values: Final = dict[str, object]() for k, v in data_json.items(): - if v is not None and ((isinstance(v, bool) and k in data.fields_set()) or v not in ([], {}, 0)): + if (v is not None or (k == "fallback_end_user_id" and k in data.fields_set())) and ( + (isinstance(v, bool) and k in data.fields_set()) or v not in ([], {}, 0) + ): non_default_values[k] = v ## Get end user table data ## @@ -642,6 +664,8 @@ async def update_end_user( param="user_id", ) + await _validate_fallback_end_user(data.user_id, data.fallback_end_user_id, prisma_client) + end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) ## Get budget table data ## @@ -783,6 +807,9 @@ async def delete_end_user( param="user_ids", ) + await _typed_table(EndUserRepository(prisma_client)).update_many( + where={"fallback_end_user_id": {"in": data.user_ids}}, data={"fallback_end_user_id": None} + ) # All users exist, proceed with deletion response: Final = await _typed_table(EndUserRepository(prisma_client)).delete_many( where={"user_id": {"in": data.user_ids}} diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 05c5aad9303..2e57f9f8fba 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -585,6 +585,8 @@ model LiteLLM_EndUserTable { object_permission_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + fallback_end_user_id String? + @@index([fallback_end_user_id]) blocked Boolean @default(false) } diff --git a/schema.prisma b/schema.prisma index 05c5aad9303..2e57f9f8fba 100644 --- a/schema.prisma +++ b/schema.prisma @@ -585,6 +585,8 @@ model LiteLLM_EndUserTable { object_permission_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + fallback_end_user_id String? + @@index([fallback_end_user_id]) blocked Boolean @default(false) } diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8777e24e209..c5374539fbe 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5813,6 +5813,7 @@ _RESTRICTED_END_USER_WHERE = { {"allowed_model_region": {"not": None}}, {"default_model": {"not": None}}, {"object_permission_id": {"not": None}}, + {"fallback_end_user_id": {"not": None}}, ] } diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index cf1f665ad21..90fc613228b 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -222,7 +222,8 @@ async def test_custom_auth_token_budget_still_loads_and_caches_unrestricted_end_ assert await cache.async_get_cache(key=end_user_cache_key("customer-1")) is not None -def test_update_valid_token_does_not_override_custom_auth_values_with_none(): +@pytest.mark.asyncio +async def test_update_valid_token_does_not_override_custom_auth_values_with_none(): """ Greptile feedback: if custom auth sets end_user_model_max_budget on the token, but the DB end_user has no model_max_budget in their budget table, the DB lookup @@ -244,7 +245,7 @@ def test_update_valid_token_does_not_override_custom_auth_values_with_none(): # No tpm_limit, rpm_limit, or model_max_budget from DB } - result = update_valid_token_with_end_user_params(valid_token, end_user_params) + result = await update_valid_token_with_end_user_params(valid_token, end_user_params) # Custom-auth-provided values should be preserved, not cleared to None assert result.end_user_tpm_limit == 100 @@ -253,7 +254,8 @@ def test_update_valid_token_does_not_override_custom_auth_values_with_none(): assert result.end_user_id == "user_1" -def test_update_valid_token_db_values_override_custom_auth_when_set(): +@pytest.mark.asyncio +async def test_update_valid_token_db_values_override_custom_auth_when_set(): """ When the DB budget table has explicit values, they should override whatever the custom auth function set (DB is source of truth). @@ -272,7 +274,7 @@ def test_update_valid_token_db_values_override_custom_auth_when_set(): "end_user_model_max_budget": db_budget, } - result = update_valid_token_with_end_user_params(valid_token, end_user_params) + result = await update_valid_token_with_end_user_params(valid_token, end_user_params) # DB values should win assert result.end_user_tpm_limit == 500 diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6cce6d0316b..8fc80b9c381 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3813,6 +3813,42 @@ async def test_auth_flow_fallback_team_object_permission_none_when_unreadable(): # --------------------------------------------------------------------------- +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback_state", ["exhausted", "missing", "nested", "available", "unlimited", "at_limit"] +) +async def test_end_user_fallback_admission(fallback_state, monkeypatch): + from litellm.proxy.auth.auth_checks import resolve_end_user_budget_fallback + import litellm.proxy.proxy_server as ps + + own = LiteLLM_EndUserTable( + user_id="person", blocked=False, fallback_end_user_id="pool", + spend=10 if fallback_state == "at_limit" else 11, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=10), + ) + pool = LiteLLM_EndUserTable( + user_id="pool", blocked=False, + spend=101 if fallback_state == "exhausted" else 1, + fallback_end_user_id="third" if fallback_state == "nested" else None, + litellm_budget_table=None if fallback_state == "unlimited" else LiteLLM_BudgetTable(max_budget=100), + ) + cache = DualCache() + db = MagicMock() + db.db.litellm_endusertable.find_unique = AsyncMock(return_value=None if fallback_state == "missing" else pool) + monkeypatch.setattr(ps, "spend_counter_cache", DualCache()) + if fallback_state == "at_limit": + assert await resolve_end_user_budget_fallback(own, "/chat/completions", db, cache) is own + db.db.litellm_endusertable.find_unique.assert_not_awaited() + return + if fallback_state in ("available", "unlimited"): + result = await resolve_end_user_budget_fallback(own, "/chat/completions", db, cache) + assert result.user_id == "pool" + assert await resolve_end_user_budget_fallback(own, "/chat/completions", db, cache) is not None + else: + assert await resolve_end_user_budget_fallback(own, "/chat/completions", db, cache) is own + assert db.db.litellm_endusertable.find_unique.await_count == 1 + + def _proxy_attrs_for_centralized_checks( user_custom_auth=None, flag=False, master_key="sk-test-master" ): diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index eff892f2d80..938f44860d0 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -16,7 +16,8 @@ from litellm.types.utils import CallTypes, Usage @pytest.mark.asyncio -async def test_async_post_call_failure_hook(): +@pytest.mark.parametrize("billing_id", [None, "pool"]) +async def test_async_post_call_failure_hook(billing_id): # Setup logger = _ProxyDBLogger() @@ -30,6 +31,7 @@ async def test_async_post_call_failure_hook(): org_id="test_org_id", team_alias="test_team_alias", end_user_id="test_end_user_id", + billing_end_user_id=billing_id, ) # Mock request data @@ -63,7 +65,7 @@ async def test_async_post_call_failure_hook(): assert call_args["token"] == "test_api_key" assert call_args["response_cost"] == 0.0 assert call_args["user_id"] == "test_user_id" - assert call_args["end_user_id"] == "test_end_user_id" + assert call_args["end_user_id"] == (billing_id or "test_end_user_id") assert call_args["team_id"] == "test_team_id" assert call_args["org_id"] == "test_org_id" assert call_args["completion_response"] == original_exception diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 1225cb80224..cde2edf8737 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -45,6 +45,7 @@ client = TestClient(app) @pytest.fixture def mock_prisma_client(): with patch("litellm.proxy.proxy_server.prisma_client") as mock: + mock.db.litellm_endusertable.update_many = AsyncMock(return_value=0) yield mock @@ -414,6 +415,7 @@ def test_update_customer_response_keeps_nested_budget_server_fields(mock_prisma_ "spend": 0.0, "allowed_model_region": None, "default_model": None, + "fallback_end_user_id": None, "budget_id": "b-1", "object_permission_id": None, "object_permission": None, @@ -690,6 +692,7 @@ _FULL_DB_ROW = { "spend": 1.5, "allowed_model_region": None, "default_model": None, + "fallback_end_user_id": None, "budget_id": "b1", "object_permission_id": "p1", "litellm_budget_table": { @@ -735,6 +738,7 @@ _EXPECTED_CUSTOMER = { "spend": 1.5, "allowed_model_region": None, "default_model": None, + "fallback_end_user_id": None, "budget_id": "b1", "litellm_budget_table": { "budget_id": "b1", @@ -975,6 +979,9 @@ def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_cl headers={"Authorization": "Bearer k"}, ) + mock_prisma_client.db.litellm_endusertable.update_many.assert_awaited_once_with( + where={"fallback_end_user_id": {"in": ["c1", "c2"]}}, data={"fallback_end_user_id": None} + ) assert response.status_code == 200, response.text assert recording_cache.deleted == [ "end_user_id:c1", @@ -986,3 +993,52 @@ def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_cl "end_user_id:c2", "end_user_restricted_registry", ] + + +@pytest.mark.parametrize("endpoint", ["new", "update"]) +@pytest.mark.parametrize("failure", ["missing", "self", "nested"]) +def test_customer_fallback_validation(mock_prisma_client, mock_user_api_key_auth, endpoint, failure): + own = LiteLLM_EndUserTable(user_id="person", blocked=False) + target = LiteLLM_EndUserTable( + user_id="pool", blocked=False, fallback_end_user_id="third" if failure == "nested" else None + ) + + table = mock_prisma_client.db.litellm_endusertable + table.find_first = AsyncMock(side_effect=[ + *([own] if endpoint == "update" else []), + None if failure == "missing" else target, + ]) + table.create = AsyncMock() + table.update = AsyncMock() + response = client.post( + f"/customer/{endpoint}", + json={"user_id": "person", "fallback_end_user_id": "person" if failure == "self" else "pool"}, + ) + assert response.status_code == 422 + table.create.assert_not_awaited() + table.update.assert_not_awaited() + + +@pytest.mark.parametrize("endpoint", ["new", "update"]) +def test_customer_fallback_round_trip(mock_prisma_client, mock_user_api_key_auth, endpoint): + own = LiteLLM_EndUserTable(user_id="person", blocked=False, fallback_end_user_id="pool") + pool = LiteLLM_EndUserTable(user_id="pool", blocked=False) + + async def find_first(where, **kwargs): + return {"person": own, "pool": pool}.get(where.get("user_id")) + + table = mock_prisma_client.db.litellm_endusertable + table.find_first = AsyncMock(side_effect=find_first) + table.create = AsyncMock(return_value=own) + table.update = AsyncMock(return_value=own) + response = client.post(f"/customer/{endpoint}", json={"user_id": "person", "fallback_end_user_id": "pool"}) + assert response.status_code == 200 + assert response.json()["fallback_end_user_id"] == "pool" + writer = table.create if endpoint == "new" else table.update + assert writer.call_args.kwargs["data"]["fallback_end_user_id"] == "pool" + cleared = own.model_copy(update={"fallback_end_user_id": None}) + table.update = AsyncMock(return_value=cleared) + response = client.post("/customer/update", json={"user_id": "person", "fallback_end_user_id": None}) + assert response.status_code == 200 + assert response.json()["fallback_end_user_id"] is None + assert table.update.call_args.kwargs["data"]["fallback_end_user_id"] is None diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 40ebc03781c..8817ce37d4c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -3348,3 +3348,103 @@ async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one assert counter_cache.in_memory_cache.get_cache( key=model_access_group_spend_counter_key("starter") ) == pytest.approx(4.2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/responses", "/v1/messages"]) +@pytest.mark.parametrize("own_spend,billing_id", [(1.0, "person"), (11.0, "pool"), (11.0, None)]) +@pytest.mark.parametrize("pool_model_cap", [0, 100]) +async def test_fallback_budget_reservation_and_spend_attribution( + spend_counter_state, monkeypatch, route, own_spend, billing_id, pool_model_cap +): + import json + from fastapi import Request + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.auth_checks import _cache_key_object + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.hooks.model_max_budget_limiter import _PROXY_VirtualKeyModelMaxBudgetLimiter + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + monkeypatch.setattr(litellm, "max_budget", 0.0) + counters, cache = spend_counter_state + own = LiteLLM_EndUserTable( + user_id="person", blocked=False, spend=own_spend, fallback_end_user_id="pool", + allowed_model_region="eu", + object_permission={"object_permission_id": "restricted", "vector_stores": ["vs-own"]}, + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=10, tpm_limit=100, rpm_limit=2, + model_max_budget={"gpt-4o-mini": {"max_budget": 0 if own_spend >= 10 else 10, "budget_duration": "1d"}}, + ), + ) + pool = LiteLLM_EndUserTable( + user_id="pool", blocked=False, spend=101.0 if billing_id is None else 3.0, + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=100, + model_max_budget={"gpt-4o-mini": {"max_budget": pool_model_cap, "budget_duration": "1d"}}, + ), + ) + cache.set_cache("end_user_id:person", own) + cache.set_cache("end_user_id:pool", pool) + counters.set_cache("spend:end_user:person", own_spend) + counters.set_cache("spend:end_user:pool", pool.spend) + prisma = MagicMock() + prisma.get_data = AsyncMock(return_value=None) + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(ps, "prisma_client", prisma) + monkeypatch.setattr(ps, "master_key", "sk-local-test") + monkeypatch.setattr(ps, "user_custom_auth", None) + monkeypatch.setattr(ps, "general_settings", {}) + router = Router(model_list=[{ + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "gpt-4o-mini", "api_key": "sk-fake"}, + "model_info": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.1}, + }]) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(ps, "model_max_budget_limiter", _PROXY_VirtualKeyModelMaxBudgetLimiter(DualCache())) + hashed_key = ps.hash_token("sk-fallback-test") + if billing_id is not None: + await _cache_key_object(hashed_key, UserAPIKeyAuth(token=hashed_key), cache, None) + request = Request({"type": "http", "path": route, "method": "POST", "headers": []}) + body = {**_request_body(), "user": "person"} + request._body = json.dumps(body).encode() + if billing_id is None: + with pytest.raises(ProxyException) as exc: + await user_api_key_auth(request=request, api_key="Bearer sk-fallback-test") + assert str(exc.value.code) == "401" + return + if billing_id == "pool" and pool_model_cap == 0: + with pytest.raises(ProxyException, match="End User: pool, exceeded budget for model") as exc: + await user_api_key_auth(request=request, api_key="Bearer sk-fallback-test") + assert int(exc.value.code) == 429 + assert counters.get_cache("spend:end_user:pool") == 3.0 + return + token = await user_api_key_auth(request=request, api_key="Bearer sk-fallback-test") + assert (token.billing_end_user_id or token.end_user_id) == billing_id + limiter = _PROXY_MaxParallelRequestsHandler_v3(ps.proxy_logging_obj.internal_usage_cache) + descriptors = limiter._create_rate_limit_descriptors(token, body, None, None, False) + end_user_descriptor = next(d for d in descriptors if d["key"] == "end_user") + assert end_user_descriptor["value"] == "person" + assert end_user_descriptor["rate_limit"]["requests_per_unit"] == 2 + assert [entry["counter_key"] for entry in token.budget_reservation["entries"]] == [f"spend:end_user:{billing_id}"] + assert counters.get_cache("spend:end_user:person") == own_spend + (1 if billing_id == "person" else 0) + assert counters.get_cache("spend:end_user:pool") == 3.0 + (1 if billing_id == "pool" else 0) + assert (token.end_user_tpm_limit, token.end_user_rpm_limit, token.allowed_model_region) == (100, 2, "eu") + assert token.end_user_object_permission == own.object_permission + assert token.end_user_max_budget == (100 if billing_id == "pool" else 10) + data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata({"metadata": {}}, token, "metadata") + payload = get_logging_payload( + {"model": "gpt-4o-mini", "call_type": "acompletion", "litellm_params": data, "response_cost": 0.25}, + {"id": "fallback-test", "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}, + datetime.now(timezone.utc), + datetime.now(timezone.utc), + ) + assert payload["end_user"] == billing_id + await ps.increment_spend_counters( + None, None, None, 0.25, end_user_id=payload["end_user"], budget_reservation=token.budget_reservation + ) + assert counters.get_cache("spend:end_user:person") == own_spend + (0.25 if billing_id == "person" else 0) + assert counters.get_cache("spend:end_user:pool") == 3.0 + (0.25 if billing_id == "pool" else 0) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d62ce2b675..ea06a010634 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -3867,6 +3867,7 @@ export interface paths { * - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. * - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. * - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + * - fallback_end_user_id: Optional[str] - Another customer whose budget is charged once this customer's own budget is exhausted. Depth is one: the fallback customer cannot itself have a fallback. * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} @@ -3972,6 +3973,7 @@ export interface paths { * - blocked: bool = False # allow/disallow requests for this end-user * - max_budget: Optional[float] = None * - budget_id: Optional[str] = None # give either a budget_id or max_budget + * - fallback_end_user_id: Optional[str] = None # customer charged once this customer's budget is exhausted; depth one, null clears * - allowed_model_region: Optional[AllowedModelRegion] = ( * None # require all user requests to use models in this specific region * ) @@ -4400,6 +4402,7 @@ export interface paths { * - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. * - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. * - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + * - fallback_end_user_id: Optional[str] - Another customer whose budget is charged once this customer's own budget is exhausted. Depth is one: the fallback customer cannot itself have a fallback. * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} @@ -4505,6 +4508,7 @@ export interface paths { * - blocked: bool = False # allow/disallow requests for this end-user * - max_budget: Optional[float] = None * - budget_id: Optional[str] = None # give either a budget_id or max_budget + * - fallback_end_user_id: Optional[str] = None # customer charged once this customer's budget is exhausted; depth one, null clears * - allowed_model_region: Optional[AllowedModelRegion] = ( * None # require all user requests to use models in this specific region * ) @@ -26779,6 +26783,8 @@ export interface components { budget_id?: string | null; /** Default Model */ default_model?: string | null; + /** Fallback End User Id */ + fallback_end_user_id?: string | null; litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTableFull"] | null; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; /** Object Permission Id */ @@ -28988,6 +28994,8 @@ export interface components { budget_id?: string | null; /** Default Model */ default_model?: string | null; + /** Fallback End User Id */ + fallback_end_user_id?: string | null; litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTable"] | null; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; /** Object Permission Id */ @@ -32125,6 +32133,8 @@ export interface components { budget_reset_at?: string | null; /** Default Model */ default_model?: string | null; + /** Fallback End User Id */ + fallback_end_user_id?: string | null; /** * Max Budget * @description Requests will fail if this budget (in USD) is exceeded. @@ -37853,6 +37863,8 @@ export interface components { budget_id?: string | null; /** Default Model */ default_model?: string | null; + /** Fallback End User Id */ + fallback_end_user_id?: string | null; /** Max Budget */ max_budget?: number | null; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null;