scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge 0fbb0841d5 into 0c98afa780

Browse source
This commit is contained in:
OrbitingBucket 2026-09-12 14:55:48 -04:00 • committed by GitHub
commit d5dd2f9c4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 343 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260910120000_add_end_user_fallback_budget/migration.sql Normal file
View file

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

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -586,6 +586,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)
}

1
litellm/models/end_user.py
View file

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

4
litellm/proxy/_types.py
View file

@ -1906,6 +1906,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
@ -1921,6 +1922,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
@ -1947,6 +1949,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
@ -3115,6 +3118,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)
team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)

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

@ -1502,9 +1502,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]:
@ -1697,6 +1731,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.
@ -1743,7 +1778,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,

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

@ -61,6 +61,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
@ -600,7 +601,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)
@ -612,8 +634,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
@ -1799,8 +1821,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:
@ -1877,8 +1903,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
@ -1950,6 +1980,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 = {}
@ -2192,7 +2229,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,
)
@ -2483,7 +2520,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(
@ -3176,7 +3213,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:
@ -3194,7 +3231,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:

2
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

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

2
litellm/proxy/litellm_pre_call_utils.py
View file

@ -1582,7 +1582,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=(

29
litellm/proxy/management_endpoints/customer_endpoints.py
View file

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

2
litellm/proxy/schema.prisma
View file

@ -586,6 +586,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)
}

2
schema.prisma
View file

@ -586,6 +586,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)
}

1
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -5989,6 +5989,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}},
]
}

10
tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
View file

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

36
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -3904,6 +3904,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"
):

6
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -26,7 +26,8 @@ from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, 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()
@ -40,6 +41,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
@ -73,7 +75,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

56
tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
View file

@ -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",
@ -976,6 +980,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",
@ -987,3 +994,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

100
tests/test_litellm/proxy/test_budget_reservation.py
View file

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

12
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -3905,6 +3905,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}
@ -4010,6 +4011,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
* )
@ -4439,6 +4441,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}
@ -4544,6 +4547,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
* )
@ -26778,6 +26782,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 */
@ -28987,6 +28993,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 */
@ -32130,6 +32138,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.
@ -37859,6 +37869,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;