Merge pull request #20601 from Harshit28j/litellm_fix_budget_model_v2

fix conflicts with main- (this PR is from upstream/main)

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)

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(

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 = []

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"

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
+        litellm.callbacks = original_callbacks

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

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