diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6e4131c3ce..ecefe77224 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -359,7 +359,6 @@ class LiteLLMRoutes(enum.Enum): "/v1/vector_stores/{vector_store_id}/files/{file_id}/content", "/vector_store/list", "/v1/vector_store/list", - # search "/search", "/v1/search", @@ -2232,13 +2231,22 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): last_refreshed_at: Optional[float] = None # last time joint view was pulled from db def __init__(self, **kwargs): - # Handle litellm_budget_table_* keys + # Handle litellm_budget_table_* keys (budget table overrides when key value is None or empty) for key, value in list(kwargs.items()): if key.startswith("litellm_budget_table_") and value is not None: # Extract the corresponding attribute name attr_name = key.replace("litellm_budget_table_", "") - # Check if the value is None and set the corresponding attribute - if getattr(self, attr_name, None) is None: + # Use key's value from kwargs (from DB view), not class default + current = kwargs.get(attr_name) + if current is None: + current = getattr(self, attr_name, None) + # Apply budget value when key has no value, or for model_max_budget when key has empty dict + should_apply = current is None or ( + attr_name == "model_max_budget" + and isinstance(current, dict) + and len(current) == 0 + ) + if should_apply: kwargs[attr_name] = value if key == "end_user_id" and value is not None and isinstance(value, int): kwargs[key] = str(value) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index ac02c91536..69c7e92d82 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -171,19 +171,35 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): 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" + ) - 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, + 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( diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index e43da32565..20c7f9ec41 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -59,14 +59,29 @@ async def new_budget( 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}"} + 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}"} + 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( @@ -123,14 +138,29 @@ async def update_budget( 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}"} + 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}"} + 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={ @@ -226,6 +256,7 @@ async def budget_settings( "budget_duration": {"type": "String"}, "max_budget": {"type": "Float"}, "soft_budget": {"type": "Float"}, + "model_max_budget": {"type": "Object"}, } return_val = [] diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2eb6cf6528..2e71759072 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -518,7 +518,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 ) # Handle special case where duration is "-1" (never expires) if value == "-1": - user_duration = float('inf') # Infinite duration + user_duration = float("inf") # Infinite duration else: user_duration = duration_in_seconds(duration=value) if user_duration > upperbound_duration: @@ -660,9 +660,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 request_type="key", **data_json, table_name="key" ) - response["soft_budget"] = ( - data.soft_budget - ) # include the user-input soft budget in the response + response[ + "soft_budget" + ] = data.soft_budget # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1083,12 +1083,16 @@ async def generate_key_fn( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) if user_custom_key_generate is not None: @@ -1399,8 +1403,13 @@ async def prepare_key_update_data( validate_model_max_budget(non_default_values["model_max_budget"]) # Serialize router_settings to JSON if present - if "router_settings" in non_default_values and non_default_values["router_settings"] is not None: - non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"]) + if ( + "router_settings" in non_default_values + and non_default_values["router_settings"] is not None + ): + non_default_values["router_settings"] = safe_dumps( + non_default_values["router_settings"] + ) non_default_values = prepare_metadata_fields( data=data, non_default_values=non_default_values, existing_metadata=_metadata @@ -1448,19 +1457,17 @@ def is_different_team( def _validate_max_budget(max_budget: Optional[float]) -> None: """ Validate that max_budget is not negative. - + Args: max_budget: The max_budget value to validate - + Raises: HTTPException: If max_budget is negative """ if max_budget is not None and max_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"max_budget cannot be negative. Received: {max_budget}" - }, + detail={"error": f"max_budget cannot be negative. Received: {max_budget}"}, ) @@ -1469,14 +1476,14 @@ async def _get_and_validate_existing_key( ) -> LiteLLM_VerificationToken: """ Get existing key from database and validate it exists. - + Args: token: The key token to look up prisma_client: Prisma client instance - + Returns: LiteLLM_VerificationToken: The existing key row - + Raises: HTTPException: If key is not found """ @@ -1485,19 +1492,19 @@ async def _get_and_validate_existing_key( status_code=500, detail={"error": "Database not connected"}, ) - + existing_key_row = await prisma_client.get_data( token=token, table_name="key", query_type="find_unique", ) - + if existing_key_row is None: raise HTTPException( status_code=404, detail={"error": f"Key not found: {token}"}, ) - + return existing_key_row @@ -1512,10 +1519,10 @@ async def _process_single_key_update( ) -> Dict[str, Any]: """ Process a single key update with all validations and checks. - + This function encapsulates all the logic for updating a single key, including validation, permission checks, team checks, and database updates. - + Args: key_update_item: The key update request item user_api_key_dict: The authenticated user's API key info @@ -1524,22 +1531,22 @@ async def _process_single_key_update( user_api_key_cache: User API key cache proxy_logging_obj: Proxy logging object llm_router: LLM router instance - + Returns: Dict containing the updated key information - + Raises: HTTPException: For various validation and permission errors """ # Validate max_budget _validate_max_budget(key_update_item.max_budget) - + # Get and validate existing key existing_key_row = await _get_and_validate_existing_key( token=key_update_item.key, prisma_client=prisma_client, ) - + # Check team member permissions if prisma_client is not None: await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( @@ -1549,7 +1556,7 @@ async def _process_single_key_update( existing_key_row=existing_key_row, user_api_key_cache=user_api_key_cache, ) - + # Create UpdateKeyRequest from BulkUpdateKeyRequestItem update_key_request = UpdateKeyRequest( key=key_update_item.key, @@ -1558,7 +1565,7 @@ async def _process_single_key_update( team_id=key_update_item.team_id, tags=key_update_item.tags, ) - + # Get team object and check team limits if team_id is provided team_obj: Optional[LiteLLM_TeamTableCachedObj] = None if update_key_request.team_id is not None: @@ -1568,18 +1575,16 @@ async def _process_single_key_update( user_api_key_cache=user_api_key_cache, check_db_only=True, ) - + if team_obj is not None and prisma_client is not None: await _check_team_key_limits( team_table=team_obj, data=update_key_request, prisma_client=prisma_client, ) - + # Validate team change if team is being changed - if is_different_team( - data=update_key_request, existing_key_row=existing_key_row - ): + if is_different_team(data=update_key_request, existing_key_row=existing_key_row): if llm_router is None: raise HTTPException( status_code=400, @@ -1590,9 +1595,7 @@ async def _process_single_key_update( if team_obj is None: raise HTTPException( status_code=500, - detail={ - "error": "Team object not found for team change validation" - }, + detail={"error": "Team object not found for team change validation"}, ) validate_key_team_change( key=existing_key_row, @@ -1600,31 +1603,29 @@ async def _process_single_key_update( change_initiated_by=user_api_key_dict, llm_router=llm_router, ) - + # Prepare update data non_default_values = await prepare_key_update_data( data=update_key_request, existing_key_row=existing_key_row ) - + # Update key in database if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected"}, ) - + _data = {**non_default_values, "token": key_update_item.key} - response = await prisma_client.update_data( - token=key_update_item.key, data=_data - ) - + response = await prisma_client.update_data(token=key_update_item.key, data=_data) + # Delete cache await _delete_cache_key_object( hashed_token=hash_token(key_update_item.key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - + # Trigger async hook asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -1635,19 +1636,19 @@ async def _process_single_key_update( litellm_changed_by=litellm_changed_by, ) ) - + if response is None: raise ValueError("Failed to update key got response = None") - + # Extract and format updated key info updated_key_info = response.get("data", {}) if hasattr(updated_key_info, "model_dump"): updated_key_info = updated_key_info.model_dump() elif hasattr(updated_key_info, "dict"): updated_key_info = updated_key_info.dict() - + updated_key_info.pop("token", None) - + return updated_key_info @@ -1740,7 +1741,9 @@ async def update_key_fn( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) @@ -1959,13 +1962,11 @@ async def bulk_update_keys( proxy_logging_obj, user_api_key_cache, ) - + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, - detail={ - "error": "Only proxy admins can perform bulk key updates" - }, + detail={"error": "Only proxy admins can perform bulk key updates"}, ) if prisma_client is None: @@ -2381,10 +2382,10 @@ async def info_key_fn( # if using pydantic v1 key_info = key_info.dict() key_info.pop("token") - + # Attach object_permission if object_permission_id is set key_info = await attach_object_permission_to_dict(key_info, prisma_client) - + return {"key": key, "info": key_info} except Exception as e: raise handle_exception_on_proxy(e) @@ -2509,7 +2510,9 @@ async def generate_key_helper_fn( # noqa: PLR0915 aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) - router_settings_json = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) + router_settings_json = ( + safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) + ) # Add model_rpm_limit and model_tpm_limit to metadata if model_rpm_limit is not None: @@ -2676,10 +2679,12 @@ async def generate_key_helper_fn( # noqa: PLR0915 ) key_data["created_at"] = getattr(create_key_response, "created_at", None) key_data["updated_at"] = getattr(create_key_response, "updated_at", None) - + # Deserialize router_settings from JSON string to dict for response router_settings_value = key_data.get("router_settings") - if router_settings_value is not None and isinstance(router_settings_value, str): + if router_settings_value is not None and isinstance( + router_settings_value, str + ): try: key_data["router_settings"] = yaml.safe_load(router_settings_value) except yaml.YAMLError: @@ -2762,27 +2767,27 @@ async def can_modify_verification_token( ) -> bool: """ Check if user has permission to modify (delete/regenerate) a verification token. - + Rules: - Proxy admin can modify any key - For team keys: only team admin or key owner can modify - For personal keys: only key owner can modify - + Args: key_info: The verification token to check user_api_key_cache: Cache for user API keys user_api_key_dict: The user making the request prisma_client: Prisma client for database access - + Returns: True if user can modify the key, False otherwise """ is_team_key = _is_team_key(data=key_info) - + # 1. Proxy admin can modify any key if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True - + # 2. For team keys: only team admin or key owner can modify if is_team_key and key_info.team_id is not None: # Get team object to check if user is team admin @@ -2792,34 +2797,35 @@ async def can_modify_verification_token( user_api_key_cache=user_api_key_cache, check_db_only=True, ) - + if team_table is None: return False - + # Check if user is team admin if _is_user_team_admin( user_api_key_dict=user_api_key_dict, team_obj=team_table, ): return True - + # Check if the key belongs to the user (they own it) - if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: + if ( + key_info.user_id is not None + and key_info.user_id == user_api_key_dict.user_id + ): return True - + # Not team admin and doesn't own the key return False - + # 3. For personal keys: only key owner can modify if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: return True - + # Default: deny return False - - async def delete_verification_tokens( tokens: List, user_api_key_cache: DualCache, @@ -2849,10 +2855,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} - ) + _keys_being_deleted: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} ) if len(_keys_being_deleted) == 0: @@ -2952,11 +2958,24 @@ def _transform_verification_tokens_to_deleted_records( if org_id_value is not None: record["organization_id"] = org_id_value - for json_field in ["aliases", "config", "permissions", "metadata", "model_spend", "model_max_budget", "router_settings"]: + for json_field in [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ]: if json_field in record and record[json_field] is not None: record[json_field] = json.dumps(record[json_field]) - for rel_key in ("litellm_budget_table", "litellm_organization_table", "object_permission", "id"): + for rel_key in ( + "litellm_budget_table", + "litellm_organization_table", + "object_permission", + "id", + ): record.pop(rel_key, None) records.append(record) @@ -2971,9 +2990,7 @@ async def _save_deleted_verification_token_records( """Save deleted verification token records to the database.""" if not records: return - await prisma_client.db.litellm_deletedverificationtoken.create_many( - data=records - ) + await prisma_client.db.litellm_deletedverificationtoken.create_many(data=records) async def _persist_deleted_verification_tokens( @@ -3036,9 +3053,9 @@ async def _rotate_master_key( from litellm.proxy.proxy_server import proxy_config try: - models: Optional[List] = ( - await prisma_client.db.litellm_proxymodeltable.find_many() - ) + models: Optional[ + List + ] = await prisma_client.db.litellm_proxymodeltable.find_many() except Exception: models = None # 2. process model table @@ -3115,7 +3132,9 @@ async def _rotate_master_key( updated_patch=decrypted_cred, new_encryption_key=new_master_key, ) - credential_object_jsonified = jsonify_object(encrypted_cred.model_dump()) + credential_object_jsonified = jsonify_object( + encrypted_cred.model_dump() + ) await prisma_client.db.litellm_credentialstable.update( where={"credential_name": cred.credential_name}, data={ @@ -3427,7 +3446,9 @@ def _validate_reset_spend_value( if reset_to > current_spend: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={"error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})"}, + detail={ + "error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})" + }, ) max_budget = key_in_db.max_budget @@ -3553,11 +3574,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, - ) + complete_user_info_db_obj: Optional[ + BaseModel + ] = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, ) if complete_user_info_db_obj is None: @@ -3643,10 +3664,10 @@ async def get_admin_team_ids( if complete_user_info is None: return [] # Get all teams that user is an admin of - teams: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} - ) + teams: Optional[ + List[BaseModel] + ] = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} ) if teams is None: return [] @@ -3691,8 +3712,12 @@ async def list_keys( description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), - expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), - status: Optional[str] = Query(None, description="Filter by status (e.g. 'deleted')"), + expand: Optional[List[str]] = Query( + None, description="Expand related objects (e.g. 'user')" + ), + status: Optional[str] = Query( + None, description="Filter by status (e.g. 'deleted')" + ), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. @@ -3784,7 +3809,9 @@ async def list_keys( message=getattr(e, "detail", f"error({str(e)})"), type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), - code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR), + code=getattr( + e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR + ), ) elif isinstance(e, ProxyException): raise e @@ -4617,10 +4644,16 @@ def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: for _model, _budget_info in model_max_budget.items(): assert isinstance(_model, str) + # Normalize to dict (Pydantic may already parse nested values as BudgetConfig) + _info = ( + _budget_info.model_dump() + if hasattr(_budget_info, "model_dump") + else dict(_budget_info) + ) # /CRUD endpoints can pass budget_limit as a string, so we need to convert it to a float - if "budget_limit" in _budget_info: - _budget_info["budget_limit"] = float(_budget_info["budget_limit"]) - BudgetConfig(**_budget_info) + if "budget_limit" in _info: + _info["budget_limit"] = float(_info["budget_limit"]) + BudgetConfig(**_info) except Exception as e: raise ValueError( f"Invalid model_max_budget: {str(e)}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users" diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 64f1ec2423..54b9e31a6d 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -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 \ No newline at end of file + litellm.callbacks = original_callbacks diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index fc8373a174..352db384c8 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index d5c3ecae7d..d8c505223d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -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()