tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 01:04:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: correct budget limit validation operator (>=) for team members (#19207)

Browse Source
This commit is contained in:
Harshit Jain
2026-01-17 00:15:35 +05:30
committed by GitHub
parent b86aae0212
commit 27a2467226
1 changed files with 73 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/auth/auth_checks.py
+73 -65
View File
@@ -202,21 +202,29 @@ async def common_checks(
and general_settings["enforce_user_param"] is True
):
# Get HTTP method from request
http_method = request.method if hasattr(request, 'method') else None
http_method = request.method if hasattr(request, "method") else None
# Check if it's a POST request and if it's an OpenAI route but not MCP
is_post_method = http_method and http_method.upper() == "POST"
is_openai_route = RouteChecks.is_llm_api_route(route=route)
is_mcp_route = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
is_mcp_route = (
route in LiteLLMRoutes.mcp_routes.value
or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
)
# Enforce user param only for POST requests on OpenAI routes (excluding MCP routes)
if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body:
if (
is_post_method
and is_openai_route
and not is_mcp_route
and "user" not in request_body
):
raise Exception(
f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
)
# 6.1 [OPTIONAL] If 'reject_clientside_metadata_tags' enabled - reject request if it has client-side 'metadata.tags'
if (
general_settings.get("reject_clientside_metadata_tags", None) is not None
@@ -502,53 +510,51 @@ async def get_default_end_user_budget(
) -> Optional[LiteLLM_BudgetTable]:
"""
Fetches the default end user budget from the database if litellm.max_end_user_budget_id is configured.
This budget is applied to end users who don't have an explicit budget_id set.
Results are cached for performance.
Args:
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving budget data
parent_otel_span: Optional OpenTelemetry span for tracing
Returns:
LiteLLM_BudgetTable if configured and found, None otherwise
"""
if prisma_client is None or litellm.max_end_user_budget_id is None:
return None
cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
# Check cache first
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
if cached_budget is not None:
return LiteLLM_BudgetTable(**cached_budget)
# Fetch from database
try:
budget_record = await prisma_client.db.litellm_budgettable.find_unique(
where={"budget_id": litellm.max_end_user_budget_id}
)
if budget_record is None:
verbose_proxy_logger.warning(
f"Default end user budget not found in database: {litellm.max_end_user_budget_id}"
)
return None
# Cache the budget for 60 seconds
await user_api_key_cache.async_set_cache(
key=cache_key,
key=cache_key,
value=budget_record.dict(),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return LiteLLM_BudgetTable(**budget_record.dict())
except Exception as e:
verbose_proxy_logger.error(
f"Error fetching default end user budget: {str(e)}"
)
verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}")
return None
@@ -560,38 +566,38 @@ async def _apply_default_budget_to_end_user(
) -> LiteLLM_EndUserTable:
"""
Helper function to apply default budget to end user if they don't have a budget assigned.
Args:
end_user_obj: The end user object to potentially apply default budget to
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving data
parent_otel_span: Optional OpenTelemetry span for tracing
Returns:
Updated end user object with default budget applied if applicable
"""
# If end user already has a budget assigned, no need to apply default
if end_user_obj.litellm_budget_table is not None:
return end_user_obj
# If no default budget configured, return as-is
if litellm.max_end_user_budget_id is None:
return end_user_obj
# Fetch and apply default budget
default_budget = await get_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
if default_budget is not None:
# Apply default budget to end user object
end_user_obj.litellm_budget_table = default_budget
verbose_proxy_logger.debug(
f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}"
)
return end_user_obj
@@ -601,20 +607,20 @@ def _check_end_user_budget(
) -> None:
"""
Check if end user is within their budget limit.
Args:
end_user_obj: The end user object to check
route: The request route
Raises:
litellm.BudgetExceededError: If end user has exceeded their budget
"""
if route in LiteLLMRoutes.info_routes.value:
return
if end_user_obj.litellm_budget_table is None:
return
end_user_budget = end_user_obj.litellm_budget_table.max_budget
if end_user_budget is not None and end_user_obj.spend > end_user_budget:
raise litellm.BudgetExceededError(
@@ -635,8 +641,8 @@ async def get_end_user_object(
) -> Optional[LiteLLM_EndUserTable]:
"""
Returns end user object from database or cache.
If end user exists but has no budget_id, applies the default budget
If end user exists but has no budget_id, applies the default budget
(if configured via litellm.max_end_user_budget_id).
Args:
@@ -646,7 +652,7 @@ async def get_end_user_object(
route: The request route
parent_otel_span: Optional OpenTelemetry span for tracing
proxy_logging_obj: Optional proxy logging object
Returns:
LiteLLM_EndUserTable if found, None otherwise
"""
@@ -655,14 +661,14 @@ async def get_end_user_object(
if end_user_id is None:
return None
_key = "end_user_id:{}".format(end_user_id)
# Check cache first
cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
if cached_user_obj is not None:
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
# Apply default budget if needed
return_obj = await _apply_default_budget_to_end_user(
end_user_obj=return_obj,
@@ -670,10 +676,10 @@ async def get_end_user_object(
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
# Check budget limits
_check_end_user_budget(end_user_obj=return_obj, route=route)
return return_obj
# Fetch from database
@@ -688,7 +694,7 @@ async def get_end_user_object(
# Convert to LiteLLM_EndUserTable object
_response = LiteLLM_EndUserTable(**response.dict())
# Apply default budget if needed
_response = await _apply_default_budget_to_end_user(
end_user_obj=_response,
@@ -696,18 +702,17 @@ async def get_end_user_object(
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
# Save to cache (always store as dict for consistency)
await user_api_key_cache.async_set_cache(
key="end_user_id:{}".format(end_user_id),
value=_response.dict()
key="end_user_id:{}".format(end_user_id), value=_response.dict()
)
# Check budget limits
_check_end_user_budget(end_user_obj=_response, route=route)
return _response
except Exception as e:
if isinstance(e, litellm.BudgetExceededError):
raise e
@@ -747,7 +752,6 @@ async def get_tag_objects_batch(
tag_objects = {}
uncached_tags = []
# Try to get all tags from cache first
for tag_name in tag_names:
@@ -1138,7 +1142,6 @@ async def _cache_management_object(
user_api_key_cache: DualCache,
proxy_logging_obj: Optional[ProxyLogging],
):
await user_api_key_cache.async_set_cache(
key=key,
value=value,
@@ -1459,9 +1462,7 @@ async def get_team_object_by_alias(
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(
"Error looking up team by alias: %s", team_alias
)
verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias)
raise HTTPException(
status_code=500,
detail={
@@ -1602,11 +1603,11 @@ class ExperimentalUIJWTToken:
) -> str:
"""
Generate a JWT token for CLI authentication with 24-hour expiration.
Args:
user_info: User information from the database
team_id: Team ID for the user (optional, uses user's team if available)
Returns:
Encrypted JWT token string
"""
@@ -1800,7 +1801,7 @@ async def get_org_object(
- Check if org id in proxy Org Table
- if valid, return LiteLLM_OrganizationTable object
- if not, then raise an error
Args:
org_id: Organization ID to look up
prisma_client: Database client
@@ -1820,7 +1821,7 @@ async def get_org_object(
cache_key = "org_id:{}".format(org_id)
if include_budget_table:
cache_key = "org_id:{}:with_budget".format(org_id)
# check if in cache
cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
if cached_org_obj is not None:
@@ -1833,7 +1834,7 @@ async def get_org_object(
query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}}
if include_budget_table:
query_kwargs["include"] = {"litellm_budget_table": True}
response = await prisma_client.db.litellm_organizationtable.find_unique(
**query_kwargs
)
@@ -1844,7 +1845,9 @@ async def get_org_object(
# Cache the result
await user_api_key_cache.async_set_cache(
key=cache_key,
value=response.model_dump() if hasattr(response, "model_dump") else response,
value=response.model_dump()
if hasattr(response, "model_dump")
else response,
ttl=DEFAULT_IN_MEMORY_TTL,
)
@@ -2218,10 +2221,15 @@ async def _virtual_key_max_budget_alert_check(
and valid_token.spend is not None
and valid_token.spend > 0
):
alert_threshold = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
alert_threshold = (
valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget:
if (
valid_token.spend >= alert_threshold
and valid_token.spend < valid_token.max_budget
):
verbose_proxy_logger.debug(
"Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s",
valid_token.token,
@@ -2274,7 +2282,7 @@ async def _check_team_member_budget(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if (
team_membership is not None
and team_membership.litellm_budget_table is not None
@@ -2282,8 +2290,8 @@ async def _check_team_member_budget(
):
team_member_budget = team_membership.litellm_budget_table.max_budget
team_member_spend = team_membership.spend or 0.0
if team_member_spend > team_member_budget:
if team_member_spend >= team_member_budget:
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,
@@ -2343,11 +2351,11 @@ async def _organization_max_budget_check(
):
"""
Check if the organization is over its max budget.
This function checks the organization budget using:
1. First, tries to use valid_token.org_id (if key has organization_id set)
2. Falls back to team_object.organization_id (if key doesn't have org_id but team does)
This ensures organization budget checks work even when keys don't have organization_id
set directly, as long as their team belongs to an organization.
@@ -2364,7 +2372,7 @@ async def _organization_max_budget_check(
org_id = valid_token.org_id
elif team_object is not None and team_object.organization_id is not None:
org_id = team_object.organization_id
# If no organization_id found, skip the check
if org_id is None:
return
@@ -2655,4 +2663,4 @@ def _can_object_call_vector_stores(
code=status.HTTP_401_UNAUTHORIZED,
)
return True
return True