scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

feat: Per Model Budgets and Limits (#20076)

Browse source
This commit is contained in:
Harshit Jain 2026-02-01 04:09:33 +05:30 • committed by GitHub
parent c0518c35cb
commit aac9907ab8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 9250 additions and 9054 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8008
litellm/proxy/_types.py
View file

File diff suppressed because it is too large Load diff

400
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -1,192 +1,208 @@
import json
from typing import List, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import Span
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
BudgetConfig,
GenericBudgetConfigType,
StandardLoggingPayload,
)
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX = "virtual_key_spend"
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
"""
Handles budgets for model + virtual key
Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d
"""
def __init__(self, dual_cache: DualCache):
self.dual_cache = dual_cache
self.redis_increment_operation_queue = []
async def is_key_within_model_budget(
self,
user_api_key_dict: UserAPIKeyAuth,
model: str,
) -> bool:
"""
Check if the user_api_key_dict is within the model budget
Raises:
BudgetExceededError: If the user_api_key_dict has exceeded the model budget
"""
_model_max_budget = user_api_key_dict.model_max_budget
internal_model_max_budget: GenericBudgetConfigType = {}
for _model, _budget_info in _model_max_budget.items():
internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
verbose_proxy_logger.debug(
"internal_model_max_budget %s",
json.dumps(internal_model_max_budget, indent=4, default=str),
)
# check if current model is in internal_model_max_budget
_current_model_budget_info = self._get_request_model_budget_config(
model=model, internal_model_max_budget=internal_model_max_budget
)
if _current_model_budget_info is None:
verbose_proxy_logger.debug(
f"Model {model} not found in internal_model_max_budget"
)
return True
# check if current model is within budget
if (
_current_model_budget_info.max_budget
and _current_model_budget_info.max_budget > 0
):
_current_spend = await self._get_virtual_key_spend_for_model(
user_api_key_hash=user_api_key_dict.token,
model=model,
key_budget_config=_current_model_budget_info,
)
if (
_current_spend is not None
and _current_model_budget_info.max_budget is not None
and _current_spend > _current_model_budget_info.max_budget
):
raise litellm.BudgetExceededError(
message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
current_cost=_current_spend,
max_budget=_current_model_budget_info.max_budget,
)
return True
async def _get_virtual_key_spend_for_model(
self,
user_api_key_hash: Optional[str],
model: str,
key_budget_config: BudgetConfig,
) -> Optional[float]:
"""
Get the current spend for a virtual key for a model
Lookup model in this order:
1. model: directly look up `model`
2. If 1, does not exist, check if passed as {custom_llm_provider}/model
"""
# 1. model: directly look up `model`
virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}"
_current_spend = await self.dual_cache.async_get_cache(
key=virtual_key_model_spend_cache_key,
)
if _current_spend is None:
# 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
# if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview
virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}"
_current_spend = await self.dual_cache.async_get_cache(
key=virtual_key_model_spend_cache_key,
)
return _current_spend
def _get_request_model_budget_config(
self, model: str, internal_model_max_budget: GenericBudgetConfigType
) -> Optional[BudgetConfig]:
"""
Get the budget config for the request model
1. Check if `model` is in `internal_model_max_budget`
2. If not, check if `model` without custom llm provider is in `internal_model_max_budget`
"""
return internal_model_max_budget.get(
model, None
) or internal_model_max_budget.get(
self._get_model_without_custom_llm_provider(model), None
)
def _get_model_without_custom_llm_provider(self, model: str) -> str:
if "/" in model:
return model.split("/")[-1]
return model
async def async_filter_deployments(
self,
model: str,
healthy_deployments: List,
messages: Optional[List[AllMessageValues]],
request_kwargs: Optional[dict] = None,
parent_otel_span: Optional[Span] = None, # type: ignore
) -> List[dict]:
return healthy_deployments
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Track spend for virtual key + model in DualCache
Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d
"""
verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event")
standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_payload is None:
raise ValueError("standard_logging_payload is required")
_litellm_params: dict = kwargs.get("litellm_params", {}) or {}
_metadata: dict = _litellm_params.get("metadata", {}) or {}
user_api_key_model_max_budget: Optional[dict] = _metadata.get(
"user_api_key_model_max_budget", None
)
if (
user_api_key_model_max_budget is None
or len(user_api_key_model_max_budget) == 0
):
verbose_proxy_logger.debug(
"Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget is None or empty. `user_api_key_model_max_budget`=%s",
user_api_key_model_max_budget,
)
return
response_cost: float = standard_logging_payload.get("response_cost", 0)
model = standard_logging_payload.get("model")
virtual_key = standard_logging_payload.get("metadata").get("user_api_key_hash")
model = standard_logging_payload.get("model")
if virtual_key is not None:
budget_config = BudgetConfig(time_period="1d", budget_limit=0.1)
virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_config.budget_duration}"
virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}"
await self._increment_spend_for_key(
budget_config=budget_config,
spend_key=virtual_spend_key,
start_time_key=virtual_start_time_key,
response_cost=response_cost,
)
verbose_proxy_logger.debug(
"current state of in memory cache %s",
json.dumps(
self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str
),
)
import json
from typing import List, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import Span
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
BudgetConfig,
GenericBudgetConfigType,
StandardLoggingPayload,
)
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX = "virtual_key_spend"
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
"""
Handles budgets for model + virtual key
Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d
"""
def __init__(self, dual_cache: DualCache):
self.dual_cache = dual_cache
self.redis_increment_operation_queue = []
async def is_key_within_model_budget(
self,
user_api_key_dict: UserAPIKeyAuth,
model: str,
) -> bool:
"""
Check if the user_api_key_dict is within the model budget
Raises:
BudgetExceededError: If the user_api_key_dict has exceeded the model budget
"""
_model_max_budget = user_api_key_dict.model_max_budget
internal_model_max_budget: GenericBudgetConfigType = {}
for _model, _budget_info in _model_max_budget.items():
internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
verbose_proxy_logger.debug(
"internal_model_max_budget %s",
json.dumps(internal_model_max_budget, indent=4, default=str),
)
# check if current model is in internal_model_max_budget
_current_model_budget_info = self._get_request_model_budget_config(
model=model, internal_model_max_budget=internal_model_max_budget
)
if _current_model_budget_info is None:
verbose_proxy_logger.debug(
f"Model {model} not found in internal_model_max_budget"
)
return True
# check if current model is within budget
if (
_current_model_budget_info.max_budget
and _current_model_budget_info.max_budget > 0
):
_current_spend = await self._get_virtual_key_spend_for_model(
user_api_key_hash=user_api_key_dict.token,
model=model,
key_budget_config=_current_model_budget_info,
)
if (
_current_spend is not None
and _current_model_budget_info.max_budget is not None
and _current_spend > _current_model_budget_info.max_budget
):
raise litellm.BudgetExceededError(
message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
current_cost=_current_spend,
max_budget=_current_model_budget_info.max_budget,
)
return True
async def _get_virtual_key_spend_for_model(
self,
user_api_key_hash: Optional[str],
model: str,
key_budget_config: BudgetConfig,
) -> Optional[float]:
"""
Get the current spend for a virtual key for a model
Lookup model in this order:
1. model: directly look up `model`
2. If 1, does not exist, check if passed as {custom_llm_provider}/model
"""
# 1. model: directly look up `model`
virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}"
_current_spend = await self.dual_cache.async_get_cache(
key=virtual_key_model_spend_cache_key,
)
if _current_spend is None:
# 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
# if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview
virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}"
_current_spend = await self.dual_cache.async_get_cache(
key=virtual_key_model_spend_cache_key,
)
return _current_spend
def _get_request_model_budget_config(
self, model: str, internal_model_max_budget: GenericBudgetConfigType
) -> Optional[BudgetConfig]:
"""
Get the budget config for the request model
1. Check if `model` is in `internal_model_max_budget`
2. If not, check if `model` without custom llm provider is in `internal_model_max_budget`
"""
return internal_model_max_budget.get(
model, None
) or internal_model_max_budget.get(
self._get_model_without_custom_llm_provider(model), None
)
def _get_model_without_custom_llm_provider(self, model: str) -> str:
if "/" in model:
return model.split("/")[-1]
return model
async def async_filter_deployments(
self,
model: str,
healthy_deployments: List,
messages: Optional[List[AllMessageValues]],
request_kwargs: Optional[dict] = None,
parent_otel_span: Optional[Span] = None, # type: ignore
) -> List[dict]:
return healthy_deployments
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Track spend for virtual key + model in DualCache
Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d
"""
verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event")
standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_payload is None:
raise ValueError("standard_logging_payload is required")
_litellm_params: dict = kwargs.get("litellm_params", {}) or {}
_metadata: dict = _litellm_params.get("metadata", {}) or {}
user_api_key_model_max_budget: Optional[dict] = _metadata.get(
"user_api_key_model_max_budget", None
)
if (
user_api_key_model_max_budget is None
or len(user_api_key_model_max_budget) == 0
):
verbose_proxy_logger.debug(
"Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget is None or empty. `user_api_key_model_max_budget`=%s",
user_api_key_model_max_budget,
)
return
response_cost: float = standard_logging_payload.get("response_cost", 0)
model = standard_logging_payload.get("model")
virtual_key = standard_logging_payload.get("metadata", {}).get(
"user_api_key_hash"
)
if virtual_key is None or model is None:
return
# Resolve per-model budget config (same logic as is_key_within_model_budget)
internal_model_max_budget: GenericBudgetConfigType = {}
for _model, _budget_info in user_api_key_model_max_budget.items():
internal_model_max_budget[_model] = BudgetConfig(**_budget_info)
key_budget_config = self._get_request_model_budget_config(
model=model, internal_model_max_budget=internal_model_max_budget
)
if key_budget_config is None or not key_budget_config.budget_duration:
verbose_proxy_logger.debug(
"Not incrementing model spend: no budget config or budget_duration for model=%s",
model,
)
return
virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}"
virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}"
await self._increment_spend_for_key(
budget_config=key_budget_config,
spend_key=virtual_spend_key,
start_time_key=virtual_start_time_key,
response_cost=response_cost,
)
verbose_proxy_logger.debug(
"current state of in memory cache %s",
json.dumps(
self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str
),
)

673
litellm/proxy/management_endpoints/budget_management_endpoints.py
View file

@ -1,321 +1,352 @@
"""
BUDGET MANAGEMENT
All /budget management endpoints
/budget/new
/budget/info
/budget/update
/budget/delete
/budget/settings
/budget/list
"""
#### BUDGET TABLE MANAGEMENT ####
from datetime import timedelta
from fastapi import APIRouter, Depends, HTTPException
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.utils import jsonify_object
router = APIRouter()
@router.post(
"/budget/new",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def new_budget(
budget_obj: BudgetNewRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Create a new budget object. Can apply this to teams, orgs, end-users, keys.
Parameters:
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated.
- max_budget: Optional[float] - The max budget for the budget.
- soft_budget: Optional[float] - The soft budget for the budget.
- max_parallel_requests: Optional[int] - The max number of parallel requests for the budget.
- tpm_limit: Optional[int] - The tokens per minute limit for the budget.
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
# Validate budget values are not negative
if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"}
)
if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"}
)
# if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future
if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None:
budget_obj.budget_reset_at = datetime.utcnow() + timedelta(
seconds=duration_in_seconds(duration=budget_obj.budget_duration)
)
budget_obj_json = budget_obj.model_dump(exclude_none=True)
budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries
response = await prisma_client.db.litellm_budgettable.create(
data={
**budget_obj_jsonified, # type: ignore
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
} # type: ignore
)
return response
@router.post(
"/budget/update",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_budget(
budget_obj: BudgetNewRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update an existing budget object.
Parameters:
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated.
- max_budget: Optional[float] - The max budget for the budget.
- soft_budget: Optional[float] - The soft budget for the budget.
- max_parallel_requests: Optional[int] - The max number of parallel requests for the budget.
- tpm_limit: Optional[int] - The tokens per minute limit for the budget.
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if budget_obj.budget_id is None:
raise HTTPException(status_code=400, detail={"error": "budget_id is required"})
# Validate budget values are not negative
if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"}
)
if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"}
)
response = await prisma_client.db.litellm_budgettable.update(
where={"budget_id": budget_obj.budget_id},
data={
**budget_obj.model_dump(exclude_unset=True), # type: ignore
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}, # type: ignore
)
return response
@router.post(
"/budget/info",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def info_budget(data: BudgetRequest):
"""
Get the budget id specific information
Parameters:
- budgets: List[str] - The list of budget ids to get information for
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
if len(data.budgets) == 0:
raise HTTPException(
status_code=400,
detail={
"error": f"Specify list of budget id's to query. Passed in={data.budgets}"
},
)
response = await prisma_client.db.litellm_budgettable.find_many(
where={"budget_id": {"in": data.budgets}},
)
return response
@router.get(
"/budget/settings",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def budget_settings(
budget_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get list of configurable params + current value for a budget item + description of each field
Used on Admin UI.
Query Parameters:
- budget_id: str - The budget id to get information for
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
## get budget item from db
db_budget_row = await prisma_client.db.litellm_budgettable.find_first(
where={"budget_id": budget_id}
)
if db_budget_row is not None:
db_budget_row_dict = db_budget_row.model_dump(exclude_none=True)
else:
db_budget_row_dict = {}
allowed_args = {
"max_parallel_requests": {"type": "Integer"},
"tpm_limit": {"type": "Integer"},
"rpm_limit": {"type": "Integer"},
"budget_duration": {"type": "String"},
"max_budget": {"type": "Float"},
"soft_budget": {"type": "Float"},
}
return_val = []
for field_name, field_info in BudgetNewRequest.model_fields.items():
if field_name in allowed_args:
_stored_in_db = True
_response_obj = ConfigList(
field_name=field_name,
field_type=allowed_args[field_name]["type"],
field_description=field_info.description or "",
field_value=db_budget_row_dict.get(field_name, None),
stored_in_db=_stored_in_db,
field_default_value=field_info.default,
)
return_val.append(_response_obj)
return return_val
@router.get(
"/budget/list",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def list_budget(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""List all the created budgets in proxy db. Used on Admin UI."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
response = await prisma_client.db.litellm_budgettable.find_many()
return response
@router.post(
"/budget/delete",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_budget(
data: BudgetDeleteRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete budget
Parameters:
- id: str - The budget id to delete
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
response = await prisma_client.db.litellm_budgettable.delete(
where={"budget_id": data.id}
)
return response
"""
BUDGET MANAGEMENT
All /budget management endpoints
/budget/new
/budget/info
/budget/update
/budget/delete
/budget/settings
/budget/list
"""
#### BUDGET TABLE MANAGEMENT ####
from datetime import timedelta
from fastapi import APIRouter, Depends, HTTPException
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.utils import jsonify_object
router = APIRouter()
@router.post(
"/budget/new",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def new_budget(
budget_obj: BudgetNewRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Create a new budget object. Can apply this to teams, orgs, end-users, keys.
Parameters:
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated.
- max_budget: Optional[float] - The max budget for the budget.
- soft_budget: Optional[float] - The soft budget for the budget.
- max_parallel_requests: Optional[int] - The max number of parallel requests for the budget.
- tpm_limit: Optional[int] - The tokens per minute limit for the budget.
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
# Validate budget values are not negative
if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"
},
)
if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"
},
)
# Validate model_max_budget if present
if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0:
from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_model_max_budget,
)
try:
validate_model_max_budget(budget_obj.model_max_budget)
except ValueError as e:
raise HTTPException(status_code=400, detail={"error": str(e)})
# if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future
if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None:
budget_obj.budget_reset_at = datetime.utcnow() + timedelta(
seconds=duration_in_seconds(duration=budget_obj.budget_duration)
)
budget_obj_json = budget_obj.model_dump(exclude_none=True)
budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries
response = await prisma_client.db.litellm_budgettable.create(
data={
**budget_obj_jsonified, # type: ignore
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
} # type: ignore
)
return response
@router.post(
"/budget/update",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_budget(
budget_obj: BudgetNewRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update an existing budget object.
Parameters:
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated.
- max_budget: Optional[float] - The max budget for the budget.
- soft_budget: Optional[float] - The soft budget for the budget.
- max_parallel_requests: Optional[int] - The max number of parallel requests for the budget.
- tpm_limit: Optional[int] - The tokens per minute limit for the budget.
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if budget_obj.budget_id is None:
raise HTTPException(status_code=400, detail={"error": "budget_id is required"})
# Validate budget values are not negative
if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"
},
)
if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"
},
)
# Validate model_max_budget if present in update
if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0:
from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_model_max_budget,
)
try:
validate_model_max_budget(budget_obj.model_max_budget)
except ValueError as e:
raise HTTPException(status_code=400, detail={"error": str(e)})
response = await prisma_client.db.litellm_budgettable.update(
where={"budget_id": budget_obj.budget_id},
data={
**budget_obj.model_dump(exclude_unset=True), # type: ignore
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}, # type: ignore
)
return response
@router.post(
"/budget/info",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def info_budget(data: BudgetRequest):
"""
Get the budget id specific information
Parameters:
- budgets: List[str] - The list of budget ids to get information for
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
if len(data.budgets) == 0:
raise HTTPException(
status_code=400,
detail={
"error": f"Specify list of budget id's to query. Passed in={data.budgets}"
},
)
response = await prisma_client.db.litellm_budgettable.find_many(
where={"budget_id": {"in": data.budgets}},
)
return response
@router.get(
"/budget/settings",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def budget_settings(
budget_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get list of configurable params + current value for a budget item + description of each field
Used on Admin UI.
Query Parameters:
- budget_id: str - The budget id to get information for
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
## get budget item from db
db_budget_row = await prisma_client.db.litellm_budgettable.find_first(
where={"budget_id": budget_id}
)
if db_budget_row is not None:
db_budget_row_dict = db_budget_row.model_dump(exclude_none=True)
else:
db_budget_row_dict = {}
allowed_args = {
"max_parallel_requests": {"type": "Integer"},
"tpm_limit": {"type": "Integer"},
"rpm_limit": {"type": "Integer"},
"budget_duration": {"type": "String"},
"max_budget": {"type": "Float"},
"soft_budget": {"type": "Float"},
"model_max_budget": {"type": "Object"},
}
return_val = []
for field_name, field_info in BudgetNewRequest.model_fields.items():
if field_name in allowed_args:
_stored_in_db = True
_response_obj = ConfigList(
field_name=field_name,
field_type=allowed_args[field_name]["type"],
field_description=field_info.description or "",
field_value=db_budget_row_dict.get(field_name, None),
stored_in_db=_stored_in_db,
field_default_value=field_info.default,
)
return_val.append(_response_obj)
return return_val
@router.get(
"/budget/list",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def list_budget(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""List all the created budgets in proxy db. Used on Admin UI."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
response = await prisma_client.db.litellm_budgettable.find_many()
return response
@router.post(
"/budget/delete",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_budget(
data: BudgetDeleteRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete budget
Parameters:
- id: str - The budget id to delete
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
response = await prisma_client.db.litellm_budgettable.delete(
where={"budget_id": data.id}
)
return response

8961
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

File diff suppressed because it is too large Load diff

145
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -22,7 +22,6 @@ from litellm.proxy.litellm_pre_call_utils import (
_get_dynamic_logging_metadata,
add_litellm_data_to_request,
)
from litellm.types.utils import SupportedCacheControls
@pytest.fixture
@ -496,9 +495,7 @@ def test_add_litellm_data_for_backend_llm_call(
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"
)
UserAPIKeyAuth(api_key="test_api_key", user_id="test_user_id", org_id="test_org_id")
data = LiteLLMProxyRequestSetup.get_user_from_headers(
headers=headers,
@ -1059,7 +1056,7 @@ def test_update_config_fields_default_internal_user_params(monkeypatch):
},
},
}
updated_config = proxy_config._update_config_fields(**args)
proxy_config._update_config_fields(**args)
assert litellm.default_internal_user_params == {
"user_role": "proxy_admin",
@ -1320,6 +1317,61 @@ def test_litellm_verification_token_view_response_with_budget_table(
)
def test_litellm_verification_token_view_budget_does_not_override_key_model_max_budget():
"""
When key has non-empty model_max_budget, budget's model_max_budget is NOT applied.
Regression test for per-model budget: only apply budget's model_max_budget when key's is empty.
"""
from litellm.proxy._types import LiteLLM_VerificationTokenView
key_model_max_budget = {"gpt-4": {"max_budget": 50.0, "budget_duration": "1d"}}
args = {
"token": "sk-test-mock-token-303",
"key_name": "sk-...if_g",
"key_alias": None,
"soft_budget_cooldown": False,
"spend": 0.0,
"expires": None,
"models": [],
"aliases": {},
"config": {},
"user_id": None,
"team_id": "test",
"permissions": {},
"max_parallel_requests": None,
"metadata": {},
"blocked": None,
"tpm_limit": None,
"rpm_limit": None,
"max_budget": None,
"budget_duration": None,
"budget_reset_at": None,
"allowed_cache_controls": [],
"model_spend": {},
"model_max_budget": key_model_max_budget,
"budget_id": "my-test-tier",
"created_at": "2024-12-26T02:28:52.615+00:00",
"updated_at": "2024-12-26T03:01:51.159+00:00",
"team_spend": None,
"team_max_budget": None,
"team_tpm_limit": None,
"team_rpm_limit": None,
"team_models": [],
"team_metadata": {},
"team_blocked": False,
"team_alias": None,
"team_members_with_roles": [],
"team_member_spend": None,
"team_model_aliases": None,
"team_member": None,
"litellm_budget_table_model_max_budget": {
"gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"}
},
}
resp = LiteLLM_VerificationTokenView(**args)
assert resp.model_max_budget == key_model_max_budget
def test_is_allowed_to_make_key_request():
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.management_endpoints.key_management_endpoints import (
@ -1381,13 +1433,6 @@ def test_get_model_group_info():
assert len(model_list) == 1
import asyncio
import json
from unittest.mock import AsyncMock, patch
import pytest
@pytest.fixture
def mock_team_data():
return [
@ -1444,7 +1489,6 @@ async def test_get_user_info_for_proxy_admin(mock_team_data, mock_key_data):
"litellm.proxy.proxy_server.prisma_client",
MockPrismaClientDB(mock_team_data, mock_key_data),
):
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_get_user_info_for_proxy_admin,
)
@ -1558,9 +1602,6 @@ def test_update_key_budget_with_temp_budget_increase():
assert _update_key_budget_with_temp_budget_increase(valid_token).max_budget == 200
from unittest.mock import AsyncMock, MagicMock
@pytest.mark.asyncio
async def test_health_check_not_called_when_disabled(monkeypatch):
from litellm.proxy.proxy_server import ProxyStartupEvent
@ -1603,18 +1644,12 @@ async def test_health_check_not_called_when_disabled(monkeypatch):
},
)
def test_custom_openapi(mock_get_openapi_schema):
from litellm.proxy.proxy_server import app, custom_openapi
from litellm.proxy.proxy_server import custom_openapi
openapi_schema = custom_openapi()
assert openapi_schema is not None
import asyncio
from datetime import timedelta
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.utils import ProxyUpdateSpend
@ -1639,6 +1674,7 @@ async def test_end_user_transactions_reset():
async def test_spend_logs_cleanup_after_error():
# Setup test data
import asyncio
mock_client = MagicMock()
mock_client.spend_log_transactions = [
{"id": 1, "amount": 10.0},
@ -1826,7 +1862,7 @@ def test_provider_specific_header_in_request(custom_llm_provider, expected_resul
client = HTTPHandler()
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
resp = litellm.completion(
litellm.completion(
model="anthropic/claude-3-5-sonnet-v2@20241022",
messages=[{"role": "user", "content": "Hello world"}],
provider_specific_header=ProviderSpecificHeader(
@ -2063,7 +2099,7 @@ async def test_post_call_failure_hook_auth_error_key_info_route():
Test that post_call_failure_hook does NOT call _handle_logging_proxy_only_error
when we get an auth error from /key/info route (since it's not an LLM API route).
"""
from unittest.mock import AsyncMock, Mock, patch
from unittest.mock import AsyncMock, patch
from fastapi import HTTPException
@ -2117,7 +2153,7 @@ async def test_post_call_failure_hook_auth_error_llm_api_route():
Test that post_call_failure_hook DOES call _handle_logging_proxy_only_error
when we get an auth error from /v1/chat/completions route (since it is an LLM API route).
"""
from unittest.mock import AsyncMock, Mock, patch
from unittest.mock import AsyncMock, patch
from fastapi import HTTPException
@ -2182,27 +2218,27 @@ async def test_during_call_hook_parallel_execution():
cache = DualCache()
proxy_logging = ProxyLogging(user_api_key_cache=cache)
execution_order = []
class TestGuardrail(CustomGuardrail):
def __init__(self, name):
super().__init__(
guardrail_name=name,
event_hook=GuardrailEventHooks.during_call,
default_on=True
default_on=True,
)
self.name = name
async def async_moderation_hook(self, data, user_api_key_dict, call_type):
execution_order.append(f"{self.name}_start")
await asyncio.sleep(0.1)
execution_order.append(f"{self.name}_end")
return data
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [TestGuardrail(f"g{i}") for i in range(3)]
start_time = asyncio.get_event_loop().time()
result = await proxy_logging.during_call_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]},
@ -2210,14 +2246,22 @@ async def test_during_call_hook_parallel_execution():
call_type="completion",
)
execution_time = asyncio.get_event_loop().time() - start_time
# Verify parallel execution: all start before any end
first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item)
starts_before_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item)
assert starts_before_end == 3, f"Expected 3 starts before first end, got {starts_before_end}"
first_end_idx = next(
i for i, item in enumerate(execution_order) if "end" in item
)
starts_before_end = sum(
1 for item in execution_order[:first_end_idx] if "start" in item
)
assert (
starts_before_end == 3
), f"Expected 3 starts before first end, got {starts_before_end}"
# Verify timing: parallel ~0.1s vs sequential ~0.3s
assert execution_time < 0.2, f"Parallel execution took {execution_time}s, expected < 0.2s"
assert (
execution_time < 0.2
), f"Parallel execution took {execution_time}s, expected < 0.2s"
assert result["model"] == "gpt-4"
finally:
litellm.callbacks = original_callbacks
@ -2235,30 +2279,35 @@ async def test_during_call_hook_parallel_execution_with_error():
cache = DualCache()
proxy_logging = ProxyLogging(user_api_key_cache=cache)
class FailingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="failing_guardrail",
event_hook=GuardrailEventHooks.during_call,
default_on=True
default_on=True,
)
async def async_moderation_hook(self, data, user_api_key_dict, call_type):
raise ValueError("Guardrail violation detected!")
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [FailingGuardrail()]
with pytest.raises(ValueError) as exc_info:
await proxy_logging.during_call_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={
"model": "gpt-4",
"messages": [{"role": "user", "content": "test"}],
},
user_api_key_dict=UserAPIKeyAuth(
api_key="test_key", user_id="test_user"
),
call_type="completion",
)
assert "Guardrail violation detected!" in str(exc_info.value)
finally:
litellm.callbacks = original_callbacks
litellm.callbacks = original_callbacks

63
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
View file

@ -1,30 +1,20 @@
import json
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, patch
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system-path
from datetime import datetime as dt_object
import time
import pytest
import litellm
import json
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock, patch
import pytest
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy.hooks.model_max_budget_limiter import (
_PROXY_VirtualKeyModelMaxBudgetLimiter,
)
from litellm.proxy._types import UserAPIKeyAuth
import litellm
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
# Test class setup
@ -123,3 +113,48 @@ async def test_get_virtual_key_spend_for_model(budget_limiter):
key_budget_config=budget_config,
)
assert spend == 50.0
@pytest.mark.asyncio
async def test_async_log_success_event_uses_per_model_budget_duration(budget_limiter):
"""
async_log_success_event must use the per-model budget_duration for the cache key
so spend is tracked per model correctly. Regression test for per-model budget implementation.
"""
from litellm.proxy.hooks.model_max_budget_limiter import (
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
virtual_key = "test-key-hash"
model = "gpt-4"
budget_duration = "1d"
user_api_key_model_max_budget = {
model: {"budget_limit": 100.0, "time_period": budget_duration},
}
kwargs = {
"standard_logging_object": {
"response_cost": 0.05,
"model": model,
"metadata": {"user_api_key_hash": virtual_key},
},
"litellm_params": {
"metadata": {
"user_api_key_model_max_budget": user_api_key_model_max_budget
},
},
}
with patch.object(
budget_limiter,
"_increment_spend_for_key",
new_callable=AsyncMock,
) as mock_increment:
await budget_limiter.async_log_success_event(
kwargs, response_obj=None, start_time=None, end_time=None
)
mock_increment.assert_awaited_once()
call_kwargs = mock_increment.call_args.kwargs
spend_key = call_kwargs["spend_key"]
assert spend_key == (
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}"
)
assert call_kwargs["response_cost"] == 0.05

54
tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py
View file

@ -11,7 +11,6 @@ import litellm.proxy.proxy_server as ps
from litellm.proxy.proxy_server import app
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, CommonProxyErrors
import litellm.proxy.management_endpoints.budget_management_endpoints as bm
sys.path.insert(
0, os.path.abspath("../../../")
@ -22,13 +21,12 @@ sys.path.insert(
def client_and_mocks(monkeypatch):
# Setup MagicMock Prisma
mock_prisma = MagicMock()
mock_table = MagicMock()
mock_table.create = AsyncMock(side_effect=lambda *, data: data)
mock_table.update = AsyncMock(side_effect=lambda *, where, data: {**where, **data})
mock_prisma.db = types.SimpleNamespace(
litellm_budgettable = mock_table,
litellm_dailyspend = mock_table,
litellm_budgettable=mock_table,
litellm_dailyspend=mock_table,
)
# Monkeypatch Mocked Prisma client into the server module
@ -79,6 +77,7 @@ async def test_new_budget_db_not_connected(client_and_mocks, monkeypatch):
# override the prisma_client that the handler imports at runtime
import litellm.proxy.proxy_server as ps
monkeypatch.setattr(ps, "prisma_client", None)
# Call /budget/new endpoint
@ -123,6 +122,7 @@ async def test_update_budget_db_not_connected(client_and_mocks, monkeypatch):
# override the prisma_client that the handler imports at runtime
import litellm.proxy.proxy_server as ps
monkeypatch.setattr(ps, "prisma_client", None)
payload = {"budget_id": "any", "max_budget": 1.0}
@ -136,7 +136,7 @@ async def test_update_budget_db_not_connected(client_and_mocks, monkeypatch):
async def test_update_budget_allows_null_max_budget(client_and_mocks):
"""
Test that /budget/update allows setting max_budget to null.
Previously, using exclude_none=True would drop null values,
making it impossible to remove a budget limit. With exclude_unset=True,
explicitly setting max_budget to null should include it in the update.
@ -144,11 +144,11 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks):
client, _, mock_table = client_and_mocks
captured_data = {}
async def capture_update(*, where, data):
captured_data.update(data)
return {**where, **data}
mock_table.update = AsyncMock(side_effect=capture_update)
payload = {
@ -159,9 +159,11 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks):
assert resp.status_code == 200, resp.text
# Verify that max_budget=None was included in the update data
assert "max_budget" in captured_data, "max_budget should be included when explicitly set to null"
assert (
"max_budget" in captured_data
), "max_budget should be included when explicitly set to null"
assert captured_data["max_budget"] is None, "max_budget should be None"
mock_table.update.assert_awaited_once()
@ -169,7 +171,7 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks):
async def test_new_budget_negative_max_budget(client_and_mocks):
"""
Test that /budget/new rejects negative max_budget values.
This prevents the issue where negative budgets would always trigger
budget exceeded errors.
"""
@ -181,7 +183,7 @@ async def test_new_budget_negative_max_budget(client_and_mocks):
}
resp = client.post("/budget/new", json=payload)
assert resp.status_code == 400, resp.text
detail = resp.json()["detail"]
assert "max_budget cannot be negative" in str(detail)
@ -199,7 +201,7 @@ async def test_new_budget_negative_soft_budget(client_and_mocks):
}
resp = client.post("/budget/new", json=payload)
assert resp.status_code == 400, resp.text
detail = resp.json()["detail"]
assert "soft_budget cannot be negative" in str(detail)
@ -217,7 +219,7 @@ async def test_update_budget_negative_max_budget(client_and_mocks):
}
resp = client.post("/budget/update", json=payload)
assert resp.status_code == 400, resp.text
detail = resp.json()["detail"]
assert "max_budget cannot be negative" in str(detail)
@ -235,6 +237,30 @@ async def test_update_budget_negative_soft_budget(client_and_mocks):
}
resp = client.post("/budget/update", json=payload)
assert resp.status_code == 400, resp.text
detail = resp.json()["detail"]
assert "soft_budget cannot be negative" in str(detail)
@pytest.mark.asyncio
async def test_new_budget_invalid_model_max_budget(client_and_mocks, monkeypatch):
"""
Test that /budget/new validates model_max_budget and returns 400 for invalid structure.
Per-model budget implementation: validate_model_max_budget is called in new_budget.
"""
import litellm.proxy.proxy_server as ps
monkeypatch.setattr(ps, "premium_user", True)
client, _, _ = client_and_mocks
payload = {
"budget_id": "budget_invalid_mmb",
"max_budget": 10.0,
"model_max_budget": {"gpt-4": "not-a-dict"},
}
resp = client.post("/budget/new", json=payload)
# Pydantic may reject invalid structure with 422 before our validator runs
assert resp.status_code in (400, 422), resp.text
detail = resp.json()["detail"]
assert "model_max_budget" in str(detail) or "dictionary" in str(detail).lower()