Merge branch 'main' into litellm_dev_11_23_2024

litellm/proxy/_new_secret_config.yaml

@@ -12,12 +12,3 @@ model_list:
       vertex_ai_project: "adroit-crow-413218"
       vertex_ai_location: "us-east5"  
 
-litellm_settings:
-  success_callback: ["langfuse"]
-  callbacks: ["prometheus"]
-  key_generation_settings:
-    team_key_generation:
-      allowed_team_member_roles: ["admin"]
-      required_params: ["tags"]
-    personal_key_generation: # maps to 'Default Team' on UI 
-      allowed_user_roles: ["proxy_admin"]

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -179,6 +179,77 @@ def key_generation_check(
         )
 
 
+
+def _is_team_key(data: GenerateKeyRequest):
+    return data.team_id is not None
+
+
+def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth):
+    if (
+        litellm.key_generation_settings is None
+        or litellm.key_generation_settings.get("team_key_generation") is None
+    ):
+        return True
+
+    if user_api_key_dict.team_member is None:
+        raise HTTPException(
+            status_code=400,
+            detail=f"User not assigned to team. Got team_member={user_api_key_dict.team_member}",
+        )
+
+    team_member_role = user_api_key_dict.team_member.role
+    if (
+        team_member_role
+        not in litellm.key_generation_settings["team_key_generation"][  # type: ignore
+            "allowed_team_member_roles"
+        ]
+    ):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Team member role {team_member_role} not in allowed_team_member_roles={litellm.key_generation_settings['team_key_generation']['allowed_team_member_roles']}",  # type: ignore
+        )
+    return True
+
+
+def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth):
+
+    if (
+        litellm.key_generation_settings is None
+        or litellm.key_generation_settings.get("personal_key_generation") is None
+    ):
+        return True
+
+    if (
+        user_api_key_dict.user_role
+        not in litellm.key_generation_settings["personal_key_generation"][  # type: ignore
+            "allowed_user_roles"
+        ]
+    ):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}",  # type: ignore
+        )
+    return True
+
+
+def key_generation_check(
+    user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest
+) -> bool:
+    """
+    Check if admin has restricted key creation to certain roles for teams or individuals
+    """
+    if litellm.key_generation_settings is None:
+        return True
+
+    ## check if key is for team or individual
+    is_team_key = _is_team_key(data=data)
+
+    if is_team_key:
+        return _team_key_generation_check(user_api_key_dict)
+    else:
+        return _personal_key_generation_check(user_api_key_dict=user_api_key_dict)
+
+
 router = APIRouter()