fix(key_management_endpoints.py): correctly raise an error when tags set on /key/update by non-premium user

Closes https://github.com/BerriAI/litellm/issues/14366

litellm/proxy/management_endpoints/common_utils.py

@@ -93,7 +93,7 @@ async def _upsert_budget_and_membership(
         create_data["tpm_limit"] = tpm_limit
     if rpm_limit is not None:
         create_data["rpm_limit"] = rpm_limit
-    
+
     new_budget = await tx.litellm_budgettable.create(
         data=create_data,
         include={"team_membership": True},

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -346,6 +346,7 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
         data_json["allowed_routes"] = ["info_routes"]
     return data_json
 
+
 async def validate_team_id_used_in_service_account_request(
     team_id: Optional[str],
     prisma_client: Optional[PrismaClient],
@@ -358,13 +359,13 @@ async def validate_team_id_used_in_service_account_request(
             status_code=400,
             detail="team_id is required for service account keys. Please specify `team_id` in the request body.",
         )
-    
+
     if prisma_client is None:
         raise HTTPException(
             status_code=400,
             detail="prisma_client is required for service account keys. Please specify `prisma_client` in the request body.",
         )
-    
+
     # check if team_id exists in the database
     team = await prisma_client.db.litellm_teamtable.find_unique(
         where={"team_id": team_id},
@@ -376,6 +377,7 @@ async def validate_team_id_used_in_service_account_request(
         )
     return True
 
+
 async def _common_key_generation_helper(  # noqa: PLR0915
     data: GenerateKeyRequest,
     user_api_key_dict: UserAPIKeyAuth,
@@ -557,7 +559,7 @@ async def _common_key_generation_helper(  # noqa: PLR0915
             status_code=400,
             detail={
                 "error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {data.key}"
-            }
+            },
         )
 
     response = await generate_key_helper_fn(
@@ -923,6 +925,15 @@ async def prepare_key_update_data(
             detail="team_id is required for service account keys. Please specify `team_id` in the request body.",
         )
     non_default_values = {}
+    # ADD METADATA FIELDS
+    # Set Management Endpoint Metadata Fields
+    for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
+        if getattr(data, field, None) is not None:
+            _set_object_metadata_field(
+                object_data=data,
+                field_name=field,
+                value=getattr(data, field),
+            )
     for k, v in data_json.items():
         if (
             k in LiteLLM_ManagementEndpoint_MetadataFields
@@ -1135,6 +1146,9 @@ async def update_key_fn(
                 change_initiated_by=user_api_key_dict,
                 llm_router=llm_router,
             )
+
+            # Set Management Endpoint Metadata Fields
+
         non_default_values = await prepare_key_update_data(
             data=data, existing_key_row=existing_key_row
         )