tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-12 15:05:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "feat(guardrails): implement team-based isolation guardrails mgmnt (#1…" (#20393)

Browse Source 
This reverts commit 76a399ba69.
This commit is contained in:
Krish Dholakia
2026-02-03 20:55:46 -08:00
committed by GitHub
parent 49c678589e
commit ca34b4ee13
12 changed files with 618 additions and 1293 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql
-8
View File
@@ -1,8 +0,0 @@
-- AlterTable
ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN "team_id" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
-- AlterTable
ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
litellm/proxy/_types.py
+1 -9
View File
@@ -359,6 +359,7 @@ 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",
@@ -630,9 +631,6 @@ class LiteLLMRoutes(enum.Enum):
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
"/guardrails",
"/guardrails/{guardrail_id}",
"/v2/guardrails/list",
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
@@ -1484,9 +1482,6 @@ class TeamBase(LiteLLMPydanticObjectBase):
members: list = []
members_with_roles: List[Member] = []
team_member_permissions: Optional[List[str]] = None
allow_team_guardrail_config: Optional[
bool
] = None # if True, team admin can configure guardrails for this team
metadata: Optional[dict] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
@@ -1584,9 +1579,6 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
model_tpm_limit: Optional[Dict[str, int]] = None
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
router_settings: Optional[dict] = None
allow_team_guardrail_config: Optional[
bool
] = None # if True, team admin can configure guardrails for this team
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
litellm/proxy/auth/login_utils.py
+2 -58
View File
@@ -34,58 +34,6 @@ from litellm.secret_managers.main import get_secret_bool
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
async def expire_previous_ui_session_tokens(
user_id: str, prisma_client: Optional[PrismaClient]
) -> None:
"""
Expire (block) all other valid UI session tokens for a user.
This prevents accumulation of multiple valid UI session tokens that
are supposed to be short-lived test keys. Only affects keys with
team_id = "litellm-dashboard" and that haven't expired yet.
Args:
user_id: The user ID whose previous UI session tokens should be expired
prisma_client: Database client for performing the update
"""
if prisma_client is None:
return
try:
from datetime import datetime, timezone
current_time = datetime.now(timezone.utc)
# Find all unblocked AND non-expired UI session tokens for this user
ui_session_tokens = await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": user_id,
"team_id": "litellm-dashboard",
"OR": [
{"blocked": None}, # Tokens that have never been blocked (null)
{"blocked": False}, # Tokens explicitly set to not blocked
],
"expires": {"gt": current_time}, # Only get tokens that haven't expired
}
)
if not ui_session_tokens:
return
# Block all the found tokens
tokens_to_block = [token.token for token in ui_session_tokens if token.token]
if tokens_to_block:
await prisma_client.db.litellm_verificationtoken.update_many(
where={"token": {"in": tokens_to_block}}, data={"blocked": True}
)
except Exception:
# Silently fail - don't block login if cleanup fails
# This is a best-effort operation
pass
def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]:
"""
Get UI username and password from environment variables or master key.
@@ -297,11 +245,6 @@ async def authenticate_user( # noqa: PLR0915
)
user_email = getattr(_user_row, "user_email", "unknown")
_password = getattr(_user_row, "password", "unknown")
user_team_id = getattr(_user_row, "team_id", "litellm-dashboard")
# if user_team_id is None, set it to "litellm-dashboard"
if user_team_id is None:
user_team_id = "litellm-dashboard"
if _password is None:
raise ProxyException(
@@ -328,7 +271,7 @@ async def authenticate_user( # noqa: PLR0915
"config": {},
"spend": 0,
"user_id": user_id,
"team_id": user_team_id,
"team_id": "litellm-dashboard",
},
)
else:
@@ -397,3 +340,4 @@ def create_ui_token_object(
disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
server_root_path=get_server_root_path(),
)
litellm/proxy/guardrails/guardrail_endpoints.py
+86 -194
View File
@@ -11,7 +11,7 @@ from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
from litellm.types.guardrails import (
@@ -32,8 +32,6 @@ from litellm.types.guardrails import (
PresidioPresidioConfigModelUserInterface,
SupportedGuardrailIntegrations,
ToolPermissionGuardrailConfigModel,
CreateGuardrailRequest,
UpdateGuardrailRequest,
)
#### GUARDRAILS ENDPOINTS ####
@@ -42,37 +40,6 @@ router = APIRouter()
GUARDRAIL_REGISTRY = GuardrailRegistry()
async def _check_team_can_configure_guardrails(
user_api_key_dict: UserAPIKeyAuth,
team_id: Optional[str],
prisma_client: Any,
user_api_key_cache: Any,
) -> None:
"""
If the user is not proxy admin and team_id is set, verify the team has
allow_team_guardrail_config enabled. Raise HTTPException 403 otherwise.
"""
if team_id is None or team_id == "litellm-dashboard":
return
user_role = getattr(user_api_key_dict, "user_role", None)
if user_role == LitellmUserRoles.PROXY_ADMIN:
return
from litellm.proxy.auth.auth_checks import get_team_object
team_table = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
allow_config = getattr(team_table, "allow_team_guardrail_config", False)
if not allow_config:
raise HTTPException(
status_code=403,
detail="Guardrail configuration is not enabled for this team. Contact your administrator to enable it.",
)
def _get_guardrails_list_response(
guardrails_config: List[Dict],
) -> ListGuardrailsResponse:
@@ -97,19 +64,9 @@ def _get_guardrails_list_response(
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
@router.get(
"/v2/guardrails/list",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
async def list_guardrails(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
async def list_guardrails():
"""
List the guardrails that are available in the database using GuardrailRegistry
Now supports both DB-persisted and In-Memory (Config) guardrails.
List the guardrails that are available on the proxy server
👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
@@ -118,6 +75,59 @@ async def list_guardrails(
curl -X GET "http://localhost:4000/guardrails/list" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
"guardrails": [
{
"guardrail_name": "bedrock-pre-guard",
"guardrail_info": {
"params": [
{
"name": "toxicity_score",
"type": "float",
"description": "Score between 0-1 indicating content toxicity level"
},
{
"name": "pii_detection",
"type": "boolean"
}
]
}
}
]
}
```
"""
from litellm.proxy.proxy_server import proxy_config
config = proxy_config.config
_guardrails_config = cast(Optional[list[dict]], config.get("guardrails"))
if _guardrails_config is None:
return _get_guardrails_list_response([])
return _get_guardrails_list_response(_guardrails_config)
@router.get(
"/v2/guardrails/list",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
async def list_guardrails_v2():
"""
List the guardrails that are available in the database using GuardrailRegistry
👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
Example Request:
```bash
curl -X GET "http://localhost:4000/v2/guardrails/list" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
@@ -129,14 +139,12 @@ async def list_guardrails(
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "ff6ujrregl1q",
"guardrailVersion": "1.0",
"guardrailVersion": "DRAFT",
"default_on": true
},
"guardrail_info": {
"description": "Updated Bedrock content moderation guardrail"
},
"created_at": "2023-11-09T12:34:56.789Z",
"updated_at": "2023-11-09T13:45:12.345Z"
"description": "Bedrock content moderation guardrail"
}
}
]
}
@@ -145,29 +153,12 @@ async def list_guardrails(
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
# Check if DB is connected
if prisma_client is None:
# Fallback to config only if DB is missing (though prisma_client usually exists)
from litellm.proxy.proxy_server import proxy_config
config = proxy_config.config
_guardrails_config = cast(Optional[list[dict]], config.get("guardrails"))
if _guardrails_config is None:
return _get_guardrails_list_response([])
return _get_guardrails_list_response(_guardrails_config)
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
team_id = getattr(user_api_key_dict, "team_id", None)
user_role = getattr(user_api_key_dict, "user_role", None)
filter_team_id = None
if user_role != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
filter_team_id = team_id
guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db(
prisma_client=prisma_client, team_id=filter_team_id
prisma_client=prisma_client
)
guardrail_configs: List[GuardrailInfoResponse] = []
@@ -182,7 +173,6 @@ async def list_guardrails(
created_at=guardrail.get("created_at"),
updated_at=guardrail.get("updated_at"),
guardrail_definition_location="db",
team_id=guardrail.get("team_id"),
)
)
seen_guardrail_ids.add(guardrail.get("guardrail_id"))
@@ -190,15 +180,6 @@ async def list_guardrails(
# get guardrails initialized on litellm config.yaml
in_memory_guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails()
for guardrail in in_memory_guardrails:
# Check access for in-memory guardrails too
if filter_team_id:
g_team = guardrail.get("team_id")
# If guardrail has a team_id, it must match.
# If guardrail has NO team_id, it is likely a global config guardrail, so we usually allow it.
if g_team and g_team != filter_team_id:
continue
# only add guardrails that are not in DB guardrail list already
if guardrail.get("guardrail_id") not in seen_guardrail_ids:
guardrail_configs.append(
@@ -208,7 +189,6 @@ async def list_guardrails(
litellm_params=dict(guardrail.get("litellm_params") or {}),
guardrail_info=dict(guardrail.get("guardrail_info") or {}),
guardrail_definition_location="config",
team_id=guardrail.get("team_id"),
)
)
seen_guardrail_ids.add(guardrail.get("guardrail_id"))
@@ -219,15 +199,16 @@ async def list_guardrails(
raise HTTPException(status_code=500, detail=str(e))
class CreateGuardrailRequest(BaseModel):
guardrail: Guardrail
@router.post(
"/guardrails",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_guardrail(
request: CreateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
async def create_guardrail(request: CreateGuardrailRequest):
"""
Create a new guardrail
@@ -276,21 +257,14 @@ async def create_guardrail(
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
team_id = getattr(user_api_key_dict, "team_id", None)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
result = await GUARDRAIL_REGISTRY.add_guardrail_to_db(
guardrail=request.guardrail, prisma_client=prisma_client, team_id=team_id
guardrail=request.guardrail, prisma_client=prisma_client
)
guardrail_name = result.get("guardrail_name", "Unknown")
@@ -309,23 +283,21 @@ async def create_guardrail(
)
return result
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error adding guardrail to db: {e}")
raise HTTPException(status_code=500, detail=str(e))
class UpdateGuardrailRequest(BaseModel):
guardrail: Guardrail
@router.put(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_guardrail(
guardrail_id: str,
request: UpdateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
"""
Update an existing guardrail
@@ -374,7 +346,7 @@ async def update_guardrail(
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@@ -390,25 +362,6 @@ async def update_guardrail(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if existing_guardrail.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if existing_guardrail.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=existing_guardrail.get("team_id"),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
result = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=request.guardrail,
@@ -441,9 +394,7 @@ async def update_guardrail(
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_guardrail(
guardrail_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)
):
async def delete_guardrail(guardrail_id: str):
"""
Delete a guardrail
@@ -463,7 +414,7 @@ async def delete_guardrail(
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@@ -479,25 +430,6 @@ async def delete_guardrail(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if existing_guardrail.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if existing_guardrail.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=existing_guardrail.get("team_id"),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db(
guardrail_id=guardrail_id, prisma_client=prisma_client
)
@@ -528,11 +460,7 @@ async def delete_guardrail(
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def patch_guardrail(
guardrail_id: str,
request: PatchGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
"""
Partially update an existing guardrail
@@ -579,7 +507,7 @@ async def patch_guardrail(
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@@ -595,25 +523,6 @@ async def patch_guardrail(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if existing_guardrail.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if existing_guardrail.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=existing_guardrail.get("team_id"),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
# Create updated guardrail object
guardrail_name = (
request.guardrail_name
@@ -685,10 +594,7 @@ async def patch_guardrail(
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_guardrail_info(
guardrail_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
async def get_guardrail_info(guardrail_id: str):
"""
Get detailed information about a specific guardrail by ID
@@ -747,19 +653,6 @@ async def get_guardrail_info(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if result.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if result.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
litellm_params: Optional[Union[LitellmParams, dict]] = result.get(
"litellm_params"
)
@@ -782,7 +675,6 @@ async def get_guardrail_info(
created_at=result.get("created_at"),
updated_at=result.get("updated_at"),
guardrail_definition_location=guardrail_definition_location,
team_id=result.get("team_id"),
)
except HTTPException as e:
raise e
@@ -1317,9 +1209,9 @@ async def get_provider_specific_params():
lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel)
tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel)
tool_permission_fields[
"ui_friendly_name"
] = ToolPermissionGuardrailConfigModel.ui_friendly_name()
tool_permission_fields["ui_friendly_name"] = (
ToolPermissionGuardrailConfigModel.ui_friendly_name()
)
# Return the provider-specific parameters
provider_params = {
@@ -1627,10 +1519,10 @@ async def apply_guardrail(
from litellm.proxy.utils import handle_exception_on_proxy
try:
active_guardrail: Optional[
CustomGuardrail
] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
guardrail_name=request.guardrail_name
active_guardrail: Optional[CustomGuardrail] = (
GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
guardrail_name=request.guardrail_name
)
)
if active_guardrail is None:
raise HTTPException(
litellm/proxy/guardrails/guardrail_registry.py
+9 -90
View File
@@ -233,10 +233,7 @@ class GuardrailRegistry:
########### DB management helpers for guardrails ###########
############################################################
async def add_guardrail_to_db(
self,
guardrail: Guardrail,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
self, guardrail: Guardrail, prisma_client: PrismaClient
):
"""
Add a guardrail to the database
@@ -251,8 +248,6 @@ class GuardrailRegistry:
litellm_params_dict = (
dict(litellm_params_obj) if litellm_params_obj else {}
)
# Use safe_dumps to store as string, as Prisma client seems to prefer this for now
litellm_params: str = safe_dumps(litellm_params_dict)
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
@@ -262,7 +257,6 @@ class GuardrailRegistry:
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
"guardrail_info": guardrail_info,
"team_id": team_id,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
@@ -271,32 +265,18 @@ class GuardrailRegistry:
# Add guardrail_id to the returned guardrail object
guardrail_dict = dict(guardrail)
guardrail_dict["guardrail_id"] = created_guardrail.guardrail_id
guardrail_dict["team_id"] = team_id
return guardrail_dict
except Exception as e:
raise Exception(f"Error adding guardrail to DB: {str(e)}")
async def delete_guardrail_from_db(
self,
guardrail_id: str,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
self, guardrail_id: str, prisma_client: PrismaClient
):
"""
Delete a guardrail from the database
"""
try:
# Check ownership if team_id is provided
if team_id:
existing_guardrail = (
await prisma_client.db.litellm_guardrailstable.find_unique(
where={"guardrail_id": guardrail_id}
)
)
if not existing_guardrail or existing_guardrail.team_id != team_id:
raise Exception("Guardrail not found or access denied")
# Delete from DB
await prisma_client.db.litellm_guardrailstable.delete(
where={"guardrail_id": guardrail_id}
@@ -307,26 +287,12 @@ class GuardrailRegistry:
raise Exception(f"Error deleting guardrail from DB: {str(e)}")
async def update_guardrail_in_db(
self,
guardrail_id: str,
guardrail: Guardrail,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient
):
"""
Update a guardrail in the database
"""
try:
# Check ownership if team_id is provided
if team_id:
existing_guardrail = (
await prisma_client.db.litellm_guardrailstable.find_unique(
where={"guardrail_id": guardrail_id}
)
)
if not existing_guardrail or existing_guardrail.team_id != team_id:
raise Exception("Guardrail not found or access denied")
guardrail_name = guardrail.get("guardrail_name")
# Properly serialize LitellmParams Pydantic model to dict
litellm_params_obj: Any = guardrail.get("litellm_params", {})
@@ -336,7 +302,6 @@ class GuardrailRegistry:
litellm_params_dict = (
dict(litellm_params_obj) if litellm_params_obj else {}
)
litellm_params: str = safe_dumps(litellm_params_dict)
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
@@ -359,67 +324,27 @@ class GuardrailRegistry:
@staticmethod
async def get_all_guardrails_from_db(
prisma_client: PrismaClient,
team_id: Optional[str] = None,
) -> List[Guardrail]:
"""
Get all guardrails from the database
"""
try:
# Normal ORM Fetch
all_guardrails = await prisma_client.db.litellm_guardrailstable.find_many(
order={"created_at": "desc"},
guardrails_from_db = (
await prisma_client.db.litellm_guardrailstable.find_many(
order={"created_at": "desc"},
)
)
# Filter in Python
if team_id:
guardrails_from_db = [g for g in all_guardrails if g.team_id == team_id]
else:
guardrails_from_db = all_guardrails
guardrails: List[Guardrail] = []
for guardrail in guardrails_from_db:
try:
# Deep copy to avoid mutating cache/original
g_dict = (
guardrail.dict()
if hasattr(guardrail, "dict")
else dict(guardrail)
)
# Handle litellm_params
params = g_dict.get("litellm_params")
if isinstance(params, str):
import json
try:
g_dict["litellm_params"] = json.loads(params)
except Exception:
g_dict["litellm_params"] = {}
# Handle guardrail_info
info = g_dict.get("guardrail_info")
if isinstance(info, str):
import json
try:
g_dict["guardrail_info"] = json.loads(info)
except Exception:
g_dict["guardrail_info"] = {}
# Construct
guardrails.append(Guardrail(**g_dict)) # type: ignore
except Exception:
continue
guardrails.append(Guardrail(**(dict(guardrail)))) # type: ignore
return guardrails
except Exception as e:
raise Exception(f"Error getting guardrails from DB: {str(e)}")
async def get_guardrail_by_id_from_db(
self,
guardrail_id: str,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
self, guardrail_id: str, prisma_client: PrismaClient
) -> Optional[Guardrail]:
"""
Get a guardrail by its ID from the database
@@ -432,11 +357,6 @@ class GuardrailRegistry:
if not guardrail:
return None
if team_id and guardrail.team_id != team_id:
# Return None if not found or not owned by team
# Alternatively could raise exception, but returning None resembles "not found"
return None
return Guardrail(**(dict(guardrail))) # type: ignore
except Exception as e:
raise Exception(f"Error getting guardrail from DB: {str(e)}")
@@ -553,7 +473,6 @@ class InMemoryGuardrailHandler:
guardrail_id=guardrail.get("guardrail_id"),
guardrail_name=guardrail["guardrail_name"],
litellm_params=litellm_params,
team_id=guardrail.get("team_id"),
)
# store references to the guardrail in memory
litellm/proxy/management_endpoints/team_endpoints.py
+175 -104
View File
@@ -753,16 +753,12 @@ async def new_team( # noqa: PLR0915
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.team_member_budget is not None and data.team_member_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"
},
detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
)
# Check if license is over limit
@@ -922,16 +918,12 @@ async def new_team( # noqa: PLR0915
complete_team_data.members_with_roles = []
complete_team_data_dict = complete_team_data.model_dump(exclude_none=True)
# Serialize router_settings to JSON (matching key creation pattern)
router_settings_value = getattr(data, "router_settings", None)
router_settings_json = (
safe_dumps(router_settings_value)
if router_settings_value is not None
else safe_dumps({})
)
router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({})
complete_team_data_dict["router_settings"] = router_settings_json
complete_team_data_dict = prisma_client.jsonify_team_object(
db_data=complete_team_data_dict
)
@@ -1107,9 +1099,7 @@ async def fetch_and_validate_organization(
validate_team_org_change(
team=LiteLLM_TeamTable(**existing_team_row.model_dump()),
organization=LiteLLM_OrganizationTableWithMembers(
**organization_row.model_dump()
),
organization=LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()),
llm_router=llm_router,
)
@@ -1117,9 +1107,7 @@ async def fetch_and_validate_organization(
def validate_team_org_change(
team: LiteLLM_TeamTable,
organization: LiteLLM_OrganizationTableWithMembers,
llm_router: Router,
team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router
) -> bool:
"""
Validate that a team can be moved to an organization.
@@ -1170,9 +1158,7 @@ def validate_team_org_change(
# Check if the team's user_id is a member of the org
team_members = [m.user_id for m in team.members_with_roles]
org_members = (
[m.user_id for m in organization.members] if organization.members else []
)
org_members = [m.user_id for m in organization.members] if organization.members else []
not_in_org = [
m
for m in team_members
@@ -1218,7 +1204,7 @@ def validate_team_org_change(
"/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
)
@management_endpoint_wrapper
async def update_team( # noqa: PLR0915
async def update_team( # noqa: PLR0915
data: UpdateTeamRequest,
http_request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -1302,25 +1288,19 @@ async def update_team( # noqa: PLR0915
)
if data.team_id is None:
raise HTTPException(
status_code=400, detail={"error": "No team id passed in"}
)
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
verbose_proxy_logger.debug("/team/update - %s", data)
# Validate budget values are not negative
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.team_member_budget is not None and data.team_member_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"
},
detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
)
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
@@ -1387,22 +1367,6 @@ async def update_team( # noqa: PLR0915
updated_kv = data.json(exclude_unset=True)
# Only proxy admin can change allow_team_guardrail_config
if (
"allow_team_guardrail_config" in updated_kv
and user_api_key_dict.user_role
not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN.value,
)
):
raise HTTPException(
status_code=403,
detail={
"error": "Only proxy admin can change 'Allow team to configure guardrails'. Contact your administrator."
},
)
# Check budget_duration and budget_reset_at
_set_budget_reset_at(data, updated_kv)
@@ -1447,19 +1411,16 @@ async def update_team( # noqa: PLR0915
updated_kv["model_id"] = _model_id
# Serialize router_settings to JSON if present (matching key update pattern)
if (
"router_settings" in updated_kv
and updated_kv["router_settings"] is not None
):
if "router_settings" in updated_kv and updated_kv["router_settings"] is not None:
updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"])
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
team_row: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data=updated_kv,
include={"litellm_model_table": True}, # type: ignore
team_row: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data=updated_kv,
include={"litellm_model_table": True}, # type: ignore
)
)
if team_row is None or team_row.team_id is None:
@@ -1468,9 +1429,7 @@ async def update_team( # noqa: PLR0915
detail={"error": "Team doesn't exist. Got={}".format(team_row)},
)
verbose_proxy_logger.info(
"Successfully updated team - %s, info", team_row.team_id
)
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
await _cache_team_object(
team_id=team_row.team_id,
team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()),
@@ -1812,6 +1771,113 @@ async def _add_team_members_to_team(
return updated_team, updated_users, updated_team_memberships
async def _validate_and_populate_member_user_info(
member: Member,
prisma_client: PrismaClient,
) -> Member:
"""
Validate and populate user_email/user_id for a member.
Logic:
1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth)
2. If only user_email is provided, populate user_id from DB
3. If only user_id is provided, populate user_email from DB (if user exists)
4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later)
5. If user_email and user_id mismatch, throw error
Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist).
"""
if member.user_email is None and member.user_id is None:
raise HTTPException(
status_code=400,
detail={"error": "Either user_id or user_email must be provided"},
)
# Case 1: Both user_email and user_id provided - verify they match
if member.user_email is not None and member.user_id is not None:
# Use user_email as source of truth
# Check for multiple users with same email first
users_by_email = await prisma_client.get_data(
key_val={"user_email": member.user_email},
table_name="user",
query_type="find_all",
)
if users_by_email is None or (
isinstance(users_by_email, list) and len(users_by_email) == 0
):
# User doesn't exist yet - this is fine, will be created later
return member
if isinstance(users_by_email, list) and len(users_by_email) > 1:
raise HTTPException(
status_code=400,
detail={
"error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead."
},
)
# Get the single user
user_by_email = users_by_email[0]
# Verify the user_id matches
if user_by_email.user_id != member.user_id:
raise HTTPException(
status_code=400,
detail={
"error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user."
},
)
# Both match, return as is
return member
# Case 2: Only user_email provided - populate user_id from DB
if member.user_email is not None and member.user_id is None:
user_by_email = await prisma_client.db.litellm_usertable.find_first(
where={"user_email": {"equals": member.user_email, "mode": "insensitive"}}
)
if user_by_email is None:
# User doesn't exist yet - this is fine, will be created later
return member
# Check for multiple users with same email
users_by_email = await prisma_client.get_data(
key_val={"user_email": member.user_email},
table_name="user",
query_type="find_all",
)
if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1:
raise HTTPException(
status_code=400,
detail={
"error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead."
},
)
# Populate user_id
member.user_id = user_by_email.user_id
return member
# Case 3: Only user_id provided - populate user_email from DB if user exists
if member.user_id is not None and member.user_email is None:
user_by_id = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": member.user_id}
)
if user_by_id is None:
# User doesn't exist yet - allow it to pass with user_email as None
# Will be upserted later with just user_id and null email
return member
# Populate user_email
member.user_email = user_by_id.user_email
return member
return member
@router.post(
"/team/member_add",
tags=["team management"],
@@ -1887,16 +1953,27 @@ async def team_member_add(
complete_team_data=complete_team_data,
)
(
updated_team,
updated_users,
updated_team_memberships,
) = await _add_team_members_to_team(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
# Validate and populate user_email/user_id for members before processing
if isinstance(data.member, Member):
await _validate_and_populate_member_user_info(
member=data.member,
prisma_client=prisma_client,
)
elif isinstance(data.member, List):
for m in data.member:
await _validate_and_populate_member_user_info(
member=m,
prisma_client=prisma_client,
)
updated_team, updated_users, updated_team_memberships = (
await _add_team_members_to_team(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
)
# Check if updated_team is None
@@ -2075,15 +2152,15 @@ async def team_member_delete(
)
# Fetch keys before deletion to persist them
keys_to_delete: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"team_id": data.team_id,
}
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"team_id": data.team_id,
}
)
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
keys=keys_to_delete,
@@ -2462,10 +2539,10 @@ async def delete_team(
team_rows: List[LiteLLM_TeamTable] = []
for team_id in data.team_ids:
try:
team_row_base: Optional[
BaseModel
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
team_row_base: Optional[BaseModel] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
)
if team_row_base is None:
raise Exception
@@ -2524,10 +2601,10 @@ async def delete_team(
_persist_deleted_verification_tokens,
)
keys_to_delete: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": {"in": data.team_ids}}
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": {"in": data.team_ids}}
)
)
if keys_to_delete:
@@ -2566,6 +2643,7 @@ async def delete_team(
return deleted_teams
def _transform_teams_to_deleted_records(
teams: List[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
@@ -2588,13 +2666,7 @@ def _transform_teams_to_deleted_records(
)
record = deleted_record.model_dump()
for json_field in [
"members_with_roles",
"metadata",
"model_spend",
"model_max_budget",
"router_settings",
]:
for json_field in ["members_with_roles", "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])
@@ -2613,7 +2685,9 @@ async def _save_deleted_team_records(
"""Save deleted team records to the database."""
if not records:
return
await prisma_client.db.litellm_deletedteamtable.create_many(data=records)
await prisma_client.db.litellm_deletedteamtable.create_many(
data=records
)
async def _persist_deleted_team_records(
@@ -2633,7 +2707,6 @@ async def _persist_deleted_team_records(
prisma_client=prisma_client,
)
def validate_membership(
user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable
):
@@ -2757,11 +2830,11 @@ async def team_info(
)
try:
team_info: Optional[
BaseModel
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id},
include={"object_permission": True},
team_info: Optional[BaseModel] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id},
include={"object_permission": True},
)
)
if team_info is None:
raise Exception
@@ -3224,9 +3297,7 @@ async def list_team_v2(
order=order_by if order_by else {"created_at": "desc"}, # Default sort
)
# Get total count for pagination
total_count = await prisma_client.db.litellm_teamtable.count(
where=where_conditions
)
total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions)
# Calculate total pages
total_pages = -(-total_count // page_size) # Ceiling division
litellm/proxy/schema.prisma
+1 -4
View File
@@ -128,7 +128,6 @@ model LiteLLM_TeamTable {
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
policies String[] @default([])
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
@@ -160,9 +159,8 @@ model LiteLLM_DeletedTeamTable {
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
policies String[] @default([])
allow_team_guardrail_config Boolean @default(false)
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
@@ -776,7 +774,6 @@ model LiteLLM_GuardrailsTable {
guardrail_name String @unique
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
litellm/types/guardrails.py
-10
View File
@@ -731,7 +731,6 @@ class Guardrail(TypedDict, total=False):
guardrail_name: Required[str]
litellm_params: Required[LitellmParams]
guardrail_info: Optional[Dict]
team_id: Optional[str]
created_at: Optional[datetime]
updated_at: Optional[datetime]
@@ -763,7 +762,6 @@ class GuardrailInfoResponse(BaseModel):
guardrail_name: str
litellm_params: Optional[BaseLitellmParams] = None
guardrail_info: Optional[Dict] = None
team_id: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = (
@@ -810,11 +808,3 @@ class PatchGuardrailRequest(BaseModel):
guardrail_name: Optional[str] = None
litellm_params: Optional[BaseLitellmParams] = None
guardrail_info: Optional[Dict[str, Any]] = None
class CreateGuardrailRequest(BaseModel):
guardrail: Guardrail
class UpdateGuardrailRequest(BaseModel):
guardrail: Guardrail
schema.prisma
+1 -4
View File
@@ -129,7 +129,6 @@ model LiteLLM_TeamTable {
team_member_permissions String[] @default([])
policies String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
@@ -161,8 +160,7 @@ model LiteLLM_DeletedTeamTable {
team_member_permissions String[] @default([])
policies String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false)
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
@@ -776,7 +774,6 @@ model LiteLLM_GuardrailsTable {
guardrail_name String @unique
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
+327 -446
View File
File diff suppressed because it is too large Load Diff
tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py
-295
View File
@@ -1,295 +0,0 @@
import pytest
import sys
import os
from unittest.mock import AsyncMock
from fastapi import HTTPException
# Add repo root to path
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.proxy.guardrails.guardrail_endpoints import (
create_guardrail,
update_guardrail,
list_guardrails,
)
from litellm.types.guardrails import (
CreateGuardrailRequest,
UpdateGuardrailRequest,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
# Fixtures
@pytest.fixture
def mock_team_user_auth():
return UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, team_id="team-123")
@pytest.fixture
def mock_other_team_user_auth():
return UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, team_id="team-456")
@pytest.fixture
def mock_proxy_admin_auth():
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, team_id="team-123")
@pytest.fixture
def mock_prisma_client(mocker):
return mocker.Mock()
@pytest.fixture
def mock_in_memory_handler(mocker):
mock = mocker.Mock()
mock.get_guardrail_by_id.return_value = None
mock.list_in_memory_guardrails.return_value = []
return mock
def _make_mock_team(mocker, allow_team_guardrail_config: bool):
"""Create a mock team object for get_team_object. Used by guardrail permission checks."""
team = mocker.Mock()
team.allow_team_guardrail_config = allow_team_guardrail_config
return team
@pytest.mark.asyncio
async def test_team_list_guardrails_v2_isolation(
mocker,
mock_prisma_client,
mock_in_memory_handler,
mock_team_user_auth,
mock_other_team_user_auth,
):
"""Test that teams only see their own guardrails in list endpoint"""
# Setup mocks
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
async def mock_get_all_guardrails(prisma_client=None, team_id=None):
all_guardrails = [
{
"guardrail_id": "g1",
"guardrail_name": "G1",
"team_id": "team-123",
"guardrail_config": {},
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
},
{
"guardrail_id": "g2",
"guardrail_name": "G2",
"team_id": "team-456",
"guardrail_config": {},
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
},
{
"guardrail_id": "g3",
"guardrail_name": "G3",
"team_id": None,
"guardrail_config": {},
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
},
]
if team_id:
return [g for g in all_guardrails if g["team_id"] == team_id]
return all_guardrails
mock_registry = mocker.Mock()
mock_registry.get_all_guardrails_from_db = AsyncMock(
side_effect=mock_get_all_guardrails
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
# Test for Team 123
response_123 = await list_guardrails(user_api_key_dict=mock_team_user_auth)
assert len(response_123.guardrails) == 1
assert response_123.guardrails[0].guardrail_id == "g1"
# Test for Team 456
response_456 = await list_guardrails(user_api_key_dict=mock_other_team_user_auth)
assert len(response_456.guardrails) == 1
assert response_456.guardrails[0].guardrail_id == "g2"
@pytest.mark.asyncio
async def test_team_create_guardrail_sets_team_id(
mocker, mock_prisma_client, mock_team_user_auth
):
"""Test that creating a guardrail as a team user sets the team_id when team has allow_team_guardrail_config."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Team user must have allow_team_guardrail_config=True to create guardrails
mock_team = _make_mock_team(mocker, allow_team_guardrail_config=True)
mocker.patch(
"litellm.proxy.auth.auth_checks.get_team_object",
AsyncMock(return_value=mock_team),
)
mock_registry = mocker.Mock()
mock_registry.get_guardrail_by_name_from_db = AsyncMock(return_value=None)
async def mock_add_guardrail(guardrail, team_id=None, **kwargs):
return {
"guardrail_id": "new-g",
"guardrail_name": guardrail["guardrail_name"],
"team_id": team_id,
"created_at": "2024-01-01T00:00:00.000Z",
"updated_at": "2024-01-01T00:00:00.000Z",
}
mock_registry.add_guardrail_to_db = AsyncMock(side_effect=mock_add_guardrail)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
# Mock initialize_guardrail which is called after success
# NOTE: The endpoint calls initialize_guardrail on IN_MEMORY_GUARDRAIL_HANDLER imported from guardrail_registry
mock_in_memory = mocker.Mock()
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory,
)
request = CreateGuardrailRequest(
guardrail={
"guardrail_name": "New Team Guardrail",
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
"guardrail_info": {},
}
)
response = await create_guardrail(
request=request, user_api_key_dict=mock_team_user_auth
)
mock_registry.add_guardrail_to_db.assert_called_once()
call_kwargs = mock_registry.add_guardrail_to_db.call_args[1]
assert call_kwargs["team_id"] == "team-123"
assert response["team_id"] == "team-123"
@pytest.mark.asyncio
async def test_team_update_other_team_guardrail_fails(
mocker, mock_prisma_client, mock_team_user_auth
):
"""Test that a team user cannot update another team's guardrail"""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_registry = mocker.Mock()
async def mock_get_by_id(guardrail_id, **kwargs):
g = {
"guardrail_id": "g2",
"guardrail_name": "G2",
"team_id": "team-456",
"guardrail_config": {},
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
}
if g["guardrail_id"] == guardrail_id:
return g
return None
mock_registry.get_guardrail_by_id_from_db = AsyncMock(side_effect=mock_get_by_id)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
request = UpdateGuardrailRequest(
guardrail={
"guardrail_name": "Updated Name",
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
}
)
# Team 123 tries to update G2 (owned by 456)
with pytest.raises(HTTPException) as excinfo:
await update_guardrail(
guardrail_id="g2", request=request, user_api_key_dict=mock_team_user_auth
)
assert excinfo.value.status_code == 403
# ----- Regression tests for allow_team_guardrail_config -----
@pytest.mark.asyncio
async def test_team_create_guardrail_forbidden_when_allow_team_guardrail_config_false(
mocker, mock_prisma_client, mock_team_user_auth
):
"""Team user gets 403 when team has allow_team_guardrail_config=False."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_team = _make_mock_team(mocker, allow_team_guardrail_config=False)
mocker.patch(
"litellm.proxy.auth.auth_checks.get_team_object",
AsyncMock(return_value=mock_team),
)
request = CreateGuardrailRequest(
guardrail={
"guardrail_name": "Forbidden Guardrail",
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
"guardrail_info": {},
}
)
with pytest.raises(HTTPException) as excinfo:
await create_guardrail(request=request, user_api_key_dict=mock_team_user_auth)
assert excinfo.value.status_code == 403
assert "Guardrail configuration is not enabled" in str(excinfo.value.detail)
@pytest.mark.asyncio
async def test_proxy_admin_can_create_guardrail_without_team_guardrail_config_permission(
mocker, mock_prisma_client, mock_proxy_admin_auth
):
"""Proxy admin can create guardrails without team allow_team_guardrail_config (check is skipped)."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_registry = mocker.Mock()
mock_registry.get_guardrail_by_name_from_db = AsyncMock(return_value=None)
mock_registry.add_guardrail_to_db = AsyncMock(
return_value={
"guardrail_id": "admin-g",
"guardrail_name": "Admin Guardrail",
"team_id": "team-123",
"created_at": "2024-01-01T00:00:00.000Z",
"updated_at": "2024-01-01T00:00:00.000Z",
}
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_registry,
)
mock_in_memory = mocker.Mock()
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory,
)
# Do NOT mock get_team_object: proxy admin skips _check_team_can_configure_guardrails
request = CreateGuardrailRequest(
guardrail={
"guardrail_name": "Admin Guardrail",
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
"guardrail_info": {},
}
)
response = await create_guardrail(
request=request, user_api_key_dict=mock_proxy_admin_auth
)
mock_registry.add_guardrail_to_db.assert_called_once()
assert response["guardrail_id"] == "admin-g"
assert response["team_id"] == "team-123"
ui/litellm-dashboard/src/components/team/team_info.tsx
+16 -71
View File
@@ -96,7 +96,6 @@ export interface TeamData {
model_aliases: Record<string, string>;
} | null;
created_at: string;
allow_team_guardrail_config?: boolean;
guardrails?: string[];
policies?: string[];
object_permission?: {
@@ -468,7 +467,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
},
policies: values.policies || [],
organization_id: values.organization_id,
...(values.allow_team_guardrail_config !== undefined ? { allow_team_guardrail_config: values.allow_team_guardrail_config } : {}),
};
updateData.max_budget = mapEmptyStringToNull(updateData.max_budget);
@@ -660,15 +658,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<Card>
<Text className="font-semibold text-gray-900 mb-3">Guardrails</Text>
{info.allow_team_guardrail_config === true ? (
<div className="mb-3">
<Badge color="green">Team can configure guardrails</Badge>
</div>
) : (
<div className="mb-3">
<Badge color="gray">Only proxy admin can configure guardrails for this team</Badge>
</div>
)}
{info.guardrails && info.guardrails.length > 0 ? (
<div className="flex flex-wrap gap-2">
{info.guardrails.map((guardrail: string, index: number) => (
@@ -772,7 +761,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
team_member_budget_duration: info.team_member_budget_table?.budget_duration,
guardrails: info.metadata?.guardrails || [],
policies: info.policies || [],
allow_team_guardrail_config: info.allow_team_guardrail_config ?? false,
disable_global_guardrails: info.metadata?.disable_global_guardrails || false,
metadata: info.metadata
? JSON.stringify(
@@ -888,72 +876,29 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<NumericalInput step={1} style={{ width: "100%" }} />
</Form.Item>
<Form.Item
noStyle
shouldUpdate={(prev, cur) =>
prev.allow_team_guardrail_config !== cur.allow_team_guardrail_config
}
>
{() => {
const allowTeamGuardrailConfig =
form.getFieldValue("allow_team_guardrail_config");
const guardrailsDisabled =
!is_proxy_admin && !allowTeamGuardrailConfig;
return (
<Form.Item
label={
<span>
Guardrails{" "}
<Tooltip title="Setup your first guardrail">
<a
href="https://docs.litellm.ai/docs/proxy/guardrails/quick_start"
target="_blank"
rel="noopener noreferrer"
onClick={(e) => e.stopPropagation()}
>
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</a>
</Tooltip>
</span>
}
name="guardrails"
help="Select existing guardrails or enter new ones"
>
<Select
mode="tags"
placeholder="Select or enter guardrails"
options={guardrailsList.map((name) => ({
value: name,
label: name,
}))}
disabled={guardrailsDisabled}
/>
</Form.Item>
);
}}
</Form.Item>
<Form.Item
label={
<span>
Allow team to configure guardrails{" "}
<Tooltip title="When enabled, team admins can create, update, delete, and assign guardrails for this team. When disabled, only proxy admin can manage guardrails for this team. Only proxy admin can change this setting.">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
Guardrails{" "}
<Tooltip title="Setup your first guardrail">
<a
href="https://docs.litellm.ai/docs/proxy/guardrails/quick_start"
target="_blank"
rel="noopener noreferrer"
onClick={(e) => e.stopPropagation()}
>
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</a>
</Tooltip>
</span>
}
name="allow_team_guardrail_config"
valuePropName="checked"
help={
is_proxy_admin
? "Let team admins manage guardrails for this team"
: "Only proxy admin can change this setting"
}
name="guardrails"
help="Select existing guardrails or enter new ones"
>
<Switch
checkedChildren="Yes"
unCheckedChildren="No"
disabled={!is_proxy_admin}
<Select
mode="tags"
placeholder="Select or enter guardrails"
options={guardrailsList.map((name) => ({ value: name, label: name }))}
/>
</Form.Item>