scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

[AI Gateway] - End User Budgets - Allow pointing max_end_user budget to an id, so the default ID applies to all end users (#16456)

Browse source 
* add _apply_budget_limits_to_end_user_params

* add _apply_budget_limits_to_end_user_params

* add _apply_budget_limits_to_end_user_params

* test_default_budget_applied_to_end_user_without_budget

* docs fix

* fix config
This commit is contained in:
Ishaan Jaff 2025-11-11 08:20:13 -08:00 • committed by GitHub
parent e94186629d
commit 5c9f50d584
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 510 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
docs/my-website/docs/proxy/customers.md
View file

@ -12,7 +12,7 @@ Track spend, set budgets for your customers.
Make a /chat/completions call, pass 'user' - First call Works
```bash
```bash showLineNumbers title="Make request with customer ID"
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--header 'Authorization: Bearer sk-1234' \ # 👈 YOUR PROXY KEY
@ -39,14 +39,14 @@ If the customer_id already exists, spend will be incremented.
Call `/customer/info` to get a customer's all up spend
```bash
```bash showLineNumbers title="Get customer spend"
curl -X GET 'http://0.0.0.0:4000/customer/info?end_user_id=ishaan3' \ # 👈 CUSTOMER ID
-H 'Authorization: Bearer sk-1234' \ # 👈 YOUR PROXY KEY
```
Expected Response:
```
```json showLineNumbers title="Response"
{
"user_id": "ishaan3",
"blocked": false,
@ -67,20 +67,20 @@ E.g. if your server is `https://webhook.site` and your listening on `6ab090e8-c5
1. Add webhook url to your proxy environment:
```bash
```bash showLineNumbers title="Set webhook URL"
export WEBHOOK_URL="https://webhook.site/6ab090e8-c55f-4a23-b075-3209f5c57906"
```
2. Add 'webhook' to config.yaml
```yaml
```yaml showLineNumbers title="config.yaml"
general_settings:
alerting: ["webhook"] # 👈 KEY CHANGE
```
3. Test it!
```bash
```bash showLineNumbers title="Test webhook"
curl -X POST 'http://localhost:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
@ -99,7 +99,7 @@ curl -X POST 'http://localhost:4000/chat/completions' \
Expected Response
```json
```json showLineNumbers title="Webhook event payload"
{
"spend": 0.0011120000000000001, # 👈 SPEND
"max_budget": null,
@ -127,12 +127,51 @@ Expected Response
Set customer budgets (e.g. monthly budgets, tpm/rpm limits) on LiteLLM Proxy
### Default Budget for All Customers
Apply budget limits to all customers without explicit budgets. This is useful for rate limiting and spending controls across all end users.
**Step 1: Create a default budget**
```bash showLineNumbers title="Create default budget"
curl -X POST 'http://localhost:4000/budget/new' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"max_budget": 10,
"rpm_limit": 2,
"tpm_limit": 1000
}'
```
**Step 2: Configure the default budget ID**
```yaml showLineNumbers title="config.yaml"
litellm_settings:
max_end_user_budget_id: "budget_id_from_step_1"
```
**Step 3: Test it**
```bash showLineNumbers title="Make request with customer ID"
curl -X POST 'http://localhost:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Hello"}],
"user": "my-customer-id"
}'
```
The customer will be subject to the default budget limits (RPM, TPM, and $ budget). Customers with explicit budgets are unaffected.
### Quick Start
Create / Update a customer with budget
**Create New Customer w/ budget**
```bash
```bash showLineNumbers title="Create customer with budget"
curl -X POST 'http://0.0.0.0:4000/customer/new'
-H 'Authorization: Bearer sk-1234'
-H 'Content-Type: application/json'
@ -144,7 +183,7 @@ curl -X POST 'http://0.0.0.0:4000/customer/new'
**Test it!**
```bash
```bash showLineNumbers title="Test customer budget"
curl -X POST 'http://localhost:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
@ -180,7 +219,7 @@ Create and assign customers to pricing tiers.
Use the `/budget/new` endpoint for creating a new budget. [API Reference](https://litellm-api.up.railway.app/#/budget%20management/new_budget_budget_new_post)
```bash
```bash showLineNumbers title="Create budget via API"
curl -X POST 'http://localhost:4000/budget/new' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
@ -200,7 +239,7 @@ In your application code, assign budget when creating a new customer.
Just use the `budget_id` used when creating the budget. In our example, this is `my-free-tier`.
```bash
```bash showLineNumbers title="Assign budget to customer"
curl -X POST 'http://localhost:4000/customer/new' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
@ -215,7 +254,7 @@ curl -X POST 'http://localhost:4000/customer/new' \
<Tabs>
<TabItem value="curl" label="curl">
```bash
```bash showLineNumbers title="Test with curl"
curl -X POST 'http://localhost:4000/customer/new' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
@ -228,7 +267,7 @@ curl -X POST 'http://localhost:4000/customer/new' \
</TabItem>
<TabItem value="openai" label="OpenAI">
```python
```python showLineNumbers title="Test with OpenAI SDK"
from openai import OpenAI
client = OpenAI(
base_url="<your_proxy_base_url>",

1
litellm/__init__.py
View file

@ -369,6 +369,7 @@ max_ui_session_budget: Optional[float] = 10 # $10 USD budgets for UI Chat sessi
internal_user_budget_duration: Optional[str] = None
tag_budget_config: Optional[Dict[str, BudgetConfig]] = None
max_end_user_budget: Optional[float] = None
max_end_user_budget_id: Optional[str] = None
disable_end_user_cost_tracking: Optional[bool] = None
disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None

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

@ -29,6 +29,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
RBAC_ROLES,
CallInfo,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_JWTAuth,
@ -445,6 +446,135 @@ def get_actual_routes(allowed_routes: list) -> list:
return actual_routes
async def get_default_end_user_budget(
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
parent_otel_span: Optional[Span] = None,
) -> 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,
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)}"
)
return None
async def _apply_default_budget_to_end_user(
end_user_obj: LiteLLM_EndUserTable,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
parent_otel_span: Optional[Span] = None,
) -> 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
def _check_end_user_budget(
end_user_obj: LiteLLM_EndUserTable,
route: str,
) -> 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(
current_cost=end_user_obj.spend,
max_budget=end_user_budget,
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.spend}, Budget={end_user_budget}",
)
@log_db_metrics
async def get_end_user_object(
end_user_id: Optional[str],
@ -455,36 +585,49 @@ async def get_end_user_object(
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional[LiteLLM_EndUserTable]:
"""
Returns end user object, if in db.
Returns end user object from database or cache.
If end user exists but has no budget_id, applies the default budget
(if configured via litellm.max_end_user_budget_id).
Do a isolated check for end user in table vs. doing a combined key + team + user + end-user check, as key might come in frequently for different end-users. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (end-user).
Args:
end_user_id: The ID of the end user
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving data
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
"""
if prisma_client is None:
raise Exception("No db connected")
if end_user_id is None:
return None
_key = "end_user_id:{}".format(end_user_id)
def check_in_budget(end_user_obj: LiteLLM_EndUserTable):
if route in LiteLLMRoutes.info_routes.value: # allow calling info routes
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(
current_cost=end_user_obj.spend, max_budget=end_user_budget
)
# check if in cache
# 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)
check_in_budget(end_user_obj=return_obj)
# Apply default budget if needed
return_obj = await _apply_default_budget_to_end_user(
end_user_obj=return_obj,
prisma_client=prisma_client,
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
# else, check db
# Fetch from database
try:
response = await prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": end_user_id},
@ -494,17 +637,29 @@ async def get_end_user_object(
if response is None:
raise Exception
# save the end-user object 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()
)
# Convert to LiteLLM_EndUserTable object
_response = LiteLLM_EndUserTable(**response.dict())
check_in_budget(end_user_obj=_response)
# Apply default budget if needed
_response = await _apply_default_budget_to_end_user(
end_user_obj=_response,
prisma_client=prisma_client,
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()
)
# Check budget limits
_check_end_user_budget(end_user_obj=_response, route=route)
return _response
except Exception as e: # if end-user not in db
except Exception as e:
if isinstance(e, litellm.BudgetExceededError):
raise e
return None
@ -543,6 +698,7 @@ async def get_tag_objects_batch(
tag_objects = {}
uncached_tags = []
# Try to get all tags from cache first
for tag_name in tag_names:

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

@ -127,6 +127,33 @@ def _get_bearer_token(
return api_key
def _apply_budget_limits_to_end_user_params(
end_user_params: dict,
budget_info: LiteLLM_BudgetTable,
end_user_id: str,
) -> None:
"""
Helper function to apply budget limits to end user parameters.
Args:
end_user_params: Dictionary to update with budget parameters
budget_info: Budget table object containing limits
end_user_id: ID of the end user for logging
"""
if budget_info.tpm_limit is not None:
end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit
if budget_info.rpm_limit is not None:
end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit
if budget_info.max_budget is not None:
end_user_params["end_user_max_budget"] = budget_info.max_budget
verbose_proxy_logger.debug(
f"Applied budget limits to end user {end_user_id}"
)
async def user_api_key_auth_websocket(websocket: WebSocket):
# Accept the WebSocket connection
@ -643,19 +670,28 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
_end_user_object.allowed_model_region
)
if _end_user_object.litellm_budget_table is not None:
budget_info = _end_user_object.litellm_budget_table
if budget_info.tpm_limit is not None:
end_user_params["end_user_tpm_limit"] = (
budget_info.tpm_limit
)
if budget_info.rpm_limit is not None:
end_user_params["end_user_rpm_limit"] = (
budget_info.rpm_limit
)
if budget_info.max_budget is not None:
end_user_params["end_user_max_budget"] = (
budget_info.max_budget
)
_apply_budget_limits_to_end_user_params(
end_user_params=end_user_params,
budget_info=_end_user_object.litellm_budget_table,
end_user_id=end_user_id,
)
elif litellm.max_end_user_budget_id is not None:
# End user doesn't exist yet, but apply default budget limits if configured
from litellm.proxy.auth.auth_checks import (
get_default_end_user_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_budget_limits_to_end_user_params(
end_user_params=end_user_params,
budget_info=default_budget,
end_user_id=end_user_id,
)
except Exception as e:
if isinstance(e, litellm.BudgetExceededError):
raise e

228
tests/proxy_unit_tests/test_default_end_user_budget_simple.py Normal file
View file

@ -0,0 +1,228 @@
"""
Simplified tests for default end user budget feature.
Tests the core scenarios where litellm.max_end_user_budget_id applies
a default budget to end users without explicit budgets.
"""
import sys
import os
import uuid
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_EndUserTable
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.caching import DualCache
@pytest.mark.asyncio
async def test_default_budget_applied_to_end_user_without_budget():
"""
Core scenario: End user without explicit budget gets default budget applied.
This is the main use case - applying limits to all unbudgeted end users.
"""
end_user_id = f"test_user_{uuid.uuid4().hex}"
default_budget_id = str(uuid.uuid4())
litellm.max_end_user_budget_id = default_budget_id
default_budget = LiteLLM_BudgetTable(
budget_id=default_budget_id,
max_budget=10.0,
rpm_limit=2,
tpm_limit=10,
)
# Mock end user in DB without budget
mock_end_user_data = {
"user_id": end_user_id,
"spend": 1.0,
"litellm_budget_table": None,
"alias": None,
"allowed_model_region": None,
"default_model": None,
"blocked": False,
}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: mock_end_user_data)
)
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: default_budget.dict())
)
mock_cache = AsyncMock(spec=DualCache)
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
result = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
route="/chat/completions",
)
# Verify default budget was applied
assert result is not None
assert result.litellm_budget_table is not None
assert result.litellm_budget_table.budget_id == default_budget_id
assert result.litellm_budget_table.max_budget == 10.0
assert result.litellm_budget_table.rpm_limit == 2
assert result.litellm_budget_table.tpm_limit == 10
litellm.max_end_user_budget_id = None
@pytest.mark.asyncio
async def test_explicit_budget_not_overridden_by_default():
"""
Core scenario: End users with explicit budgets keep their budgets.
The default should not override user-specific configurations.
"""
end_user_id = f"test_user_{uuid.uuid4().hex}"
explicit_budget_id = str(uuid.uuid4())
default_budget_id = str(uuid.uuid4())
litellm.max_end_user_budget_id = default_budget_id
explicit_budget = LiteLLM_BudgetTable(
budget_id=explicit_budget_id,
max_budget=100.0,
rpm_limit=50,
)
# Mock end user with explicit budget
mock_end_user_data = {
"user_id": end_user_id,
"spend": 10.0,
"litellm_budget_table": explicit_budget.dict(),
"alias": None,
"allowed_model_region": None,
"default_model": None,
"blocked": False,
}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: mock_end_user_data)
)
mock_cache = AsyncMock(spec=DualCache)
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
result = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
route="/chat/completions",
)
# Verify explicit budget is kept (not replaced with default)
assert result is not None
assert result.litellm_budget_table.budget_id == explicit_budget_id
assert result.litellm_budget_table.max_budget == 100.0
assert result.litellm_budget_table.rpm_limit == 50
litellm.max_end_user_budget_id = None
@pytest.mark.asyncio
async def test_budget_enforcement_blocks_over_budget_users():
"""
Core scenario: Budget limits are actually enforced.
Users who exceed their budget should be blocked.
"""
end_user_id = f"test_user_{uuid.uuid4().hex}"
default_budget_id = str(uuid.uuid4())
litellm.max_end_user_budget_id = default_budget_id
default_budget = LiteLLM_BudgetTable(
budget_id=default_budget_id,
max_budget=10.0,
rpm_limit=2,
)
# Mock end user who has already spent more than budget
mock_end_user_data = {
"user_id": end_user_id,
"spend": 15.0, # Exceeds budget of 10.0
"litellm_budget_table": None,
"alias": None,
"allowed_model_region": None,
"default_model": None,
"blocked": False,
}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: mock_end_user_data)
)
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: default_budget.dict())
)
mock_cache = AsyncMock(spec=DualCache)
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
# Should raise BudgetExceededError
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await get_end_user_object(
end_user_id=end_user_id,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
route="/chat/completions",
)
assert "ExceededBudget" in str(exc_info.value)
assert end_user_id in str(exc_info.value)
litellm.max_end_user_budget_id = None
@pytest.mark.asyncio
async def test_system_works_without_default_budget_configured():
"""
Core scenario: System continues to work when no default budget is configured.
This ensures backward compatibility.
"""
end_user_id = f"test_user_{uuid.uuid4().hex}"
litellm.max_end_user_budget_id = None # Not configured
# Mock end user without budget
mock_end_user_data = {
"user_id": end_user_id,
"spend": 5.0,
"litellm_budget_table": None,
"alias": None,
"allowed_model_region": None,
"default_model": None,
"blocked": False,
}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(
return_value=MagicMock(dict=lambda: mock_end_user_data)
)
mock_cache = AsyncMock(spec=DualCache)
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
result = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
route="/chat/completions",
)
# Should work fine, just without budget limits
assert result is not None
assert result.user_id == end_user_id
assert result.litellm_budget_table is None # No budget applied