scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

feat: working team BYOG

Browse source
This commit is contained in:
Krrish Dholakia 2026-02-26 19:52:26 -08:00
parent d5ef6c7f93
commit 94c8705d0c
22 changed files with 767 additions and 275 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
docs/my-website/docs/proxy/guardrails/team_guardrails.md Normal file
View file

@ -0,0 +1,64 @@
# Team-based guardrails
:::info
This is an Enterprise feature.
[Enterprise Pricing](https://www.litellm.ai/#pricing)
[Contact us here to get a free trial](https://calendly.com/d/cx9p-5yf-2nm/litellm-introductions)
:::
Team admins can create guardrails scoped to their team. Those guardrails are only available to that team (and to proxy admins). This mirrors [team model onboarding](/proxy/team_model_add).
## Create a team guardrail
Use the same `/guardrails` POST endpoint with a team API key and include `team_id` in the body (top-level or inside `guardrail` / `guardrail_info`):
```bash
curl -X POST "http://localhost:4000/guardrails" \
-H "Authorization: Bearer <team_api_key>" \
-H "Content-Type: application/json" \
-d '{
"guardrail": {
"guardrail_name": "my-team-content-filter",
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
"default_on": true
},
"guardrail_info": { "description": "Team content filter" }
},
"team_id": "e59e2671-a064-436a-a0fa-16ae96e5a0a1"
}'
```
- **Proxy admin**: Can create global guardrails (omit `team_id`) or team-scoped guardrails (set `team_id`).
- **Team admin**: Must set `team_id` to a team they administer; the guardrail is then only available for that team.
## List guardrails by team
`GET /v2/guardrails/list` supports optional query params:
- **`team_id`** – When used with `view=current_team`, returns global guardrails plus guardrails for that team.
- **`view`** – `all` (default): all guardrails; `current_team`: global + guardrails for the given `team_id`.
```bash
# All guardrails (admin)
curl -X GET "http://localhost:4000/v2/guardrails/list" -H "Authorization: Bearer <admin_key>"
# Global + team guardrails for a team
curl -X GET "http://localhost:4000/v2/guardrails/list?team_id=e59e2671-a064-436a-a0fa-16ae96e5a0a1&view=current_team" \
-H "Authorization: Bearer <team_api_key>"
```
Each guardrail in the response includes `team_id` (null for global) so the UI can show scope.
## Request-time behavior
For a request with a team API key, guardrails are resolved by name with team precedence:
1. If a guardrail with that name exists for the request’s team, it is used.
2. Otherwise the global guardrail with that name (if any) is used.
So team-scoped guardrails override global ones for that team; they are only available to that team.

6
litellm-proxy-extras/litellm_proxy_extras/migrations/20260226195018_add_guardrail_team_unique/migration.sql Normal file
View file

@ -0,0 +1,6 @@
-- DropIndex
DROP INDEX "LiteLLM_GuardrailsTable_guardrail_name_key";
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_GuardrailsTable_guardrail_name_team_id_key" ON "LiteLLM_GuardrailsTable"("guardrail_name", "team_id");

4
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -863,12 +863,14 @@ model LiteLLM_ManagedVectorStoresTable {
// Guardrails table for storing guardrail configurations
model LiteLLM_GuardrailsTable {
guardrail_id String @id @default(uuid())
guardrail_name String @unique
guardrail_name String
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([guardrail_name, team_id])
}
// Daily guardrail metrics for usage dashboard (one row per guardrail per day)

86
litellm/proxy/_types.py
View file

@ -1,59 +1,40 @@
import enum
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
Optional, Union)
import httpx
from pydantic import (
BaseModel,
ConfigDict,
Field,
Json,
field_validator,
model_validator,
)
from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator,
model_validator)
from typing_extensions import Required, TypedDict
from litellm._uuid import uuid
from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIFileObject,
ResponsesAPIResponse,
)
from litellm.types.mcp import (
MCPAuthType,
MCPCredentials,
MCPTransport,
MCPTransportType,
)
from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject,
ResponsesAPIResponse)
from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport,
MCPTransportType)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo
from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CallTypes,
CostBreakdown,
EmbeddingResponse,
GenericBudgetConfigType,
ImageResponse,
LiteLLMBatch,
LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase,
ModelResponse,
ProviderField,
StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse,
)
from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse,
GenericBudgetConfigType, ImageResponse,
LiteLLMBatch, LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase, ModelResponse,
ProviderField, StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse)
from litellm.types.videos.main import VideoObject
from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
from .types_utils.utils import (get_instance_fn,
validate_custom_validate_return_type)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -607,6 +588,13 @@ class LiteLLMRoutes(enum.Enum):
"/health/services",
] + info_routes
guardrails_ui_routes = [
"/guardrails/ui/add_guardrail_settings",
"/guardrails/ui/category_yaml/{category_name}",
"/guardrails/ui/major_airlines",
"/guardrails/ui/provider_specific_params",
]
internal_user_routes = (
[
"/global/spend/tags",
@ -621,6 +609,7 @@ class LiteLLMRoutes(enum.Enum):
]
+ spend_tracking_routes
+ key_management_routes
+ guardrails_ui_routes
)
internal_user_view_only_routes = (
@ -644,6 +633,10 @@ class LiteLLMRoutes(enum.Enum):
# Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
"/invitation/new",
"/invitation/delete",
# Guardrail CRUD - proxy admin or team admin enforced in guardrail_endpoints
"/guardrails",
"/guardrails/*",
"/v2/guardrails/list",
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
@ -2373,7 +2366,8 @@ class UserAPIKeyAuth(
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
from litellm.constants import \
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@ -2405,7 +2399,8 @@ class UserAPIKeyAuth(
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
from litellm.constants import \
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
@ -2796,7 +2791,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
@model_validator(mode="after")
def mask_api_keys(self):
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.litellm_core_utils.sensitive_data_masker import \
SensitiveDataMasker
masker = SensitiveDataMasker(sensitive_patterns={"key"})

329
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -12,37 +12,33 @@ 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 (CommonProxyErrors, LiteLLM_TeamTable,
LitellmUserRoles, UserAPIKeyAuth)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
CustomCodeValidationError,
validate_custom_code,
)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
get_custom_code_primitives,
)
CustomCodeValidationError, validate_custom_code)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import \
get_custom_code_primitives
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
from litellm.types.guardrails import (
PII_ENTITY_CATEGORIES_MAP,
ApplyGuardrailRequest,
ApplyGuardrailResponse,
BaseLitellmParams,
BedrockGuardrailConfigModel,
Guardrail,
GuardrailEventHooks,
GuardrailInfoResponse,
GuardrailUIAddGuardrailSettings,
LakeraV2GuardrailConfigModel,
ListGuardrailsResponse,
LitellmParams,
PatchGuardrailRequest,
PiiAction,
PiiEntityType,
PresidioPresidioConfigModelUserInterface,
SupportedGuardrailIntegrations,
ToolPermissionGuardrailConfigModel,
)
from litellm.proxy.guardrails.usage_endpoints import \
router as guardrails_usage_router
from litellm.proxy.management_endpoints.model_management_endpoints import \
ModelManagementAuthChecks
from litellm.types.guardrails import (PII_ENTITY_CATEGORIES_MAP,
ApplyGuardrailRequest,
ApplyGuardrailResponse,
BaseLitellmParams,
BedrockGuardrailConfigModel, Guardrail,
GuardrailEventHooks,
GuardrailInfoResponse,
GuardrailUIAddGuardrailSettings,
LakeraV2GuardrailConfigModel,
ListGuardrailsResponse, LitellmParams,
PatchGuardrailRequest, PiiAction,
PiiEntityType,
PresidioPresidioConfigModelUserInterface,
SupportedGuardrailIntegrations,
ToolPermissionGuardrailConfigModel)
#### GUARDRAILS ENDPOINTS ####
@ -127,15 +123,24 @@ async def list_guardrails():
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
async def list_guardrails_v2():
async def list_guardrails_v2(
team_id: Optional[str] = None,
view: Optional[str] = "all",
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
List the guardrails that are available in the database using GuardrailRegistry
Query params:
- team_id: optional; when set with view "current_team", return global guardrails plus guardrails for this team
- view: "all" (default) return all guardrails; "current_team" return global + team_id guardrails when team_id is set
👉 [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>"
curl -X GET "http://localhost:4000/v2/guardrails/list?team_id=<team_id>&view=current_team" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
@ -154,22 +159,41 @@ async def list_guardrails_v2():
},
"guardrail_info": {
"description": "Bedrock content moderation guardrail"
}
},
"team_id": null
}
]
}
```
"""
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.management_endpoints.common_utils import \
get_team_ids_where_user_is_team_admin
from litellm.proxy.proxy_server import (prisma_client, proxy_logging_obj,
user_api_key_cache)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
view_val = view if view is not None else "all"
allowed_team_ids: Optional[List[str]] = None
if view_val == "all" and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
allowed_team_ids = await get_team_ids_where_user_is_team_admin(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
try:
guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db(
prisma_client=prisma_client
guardrails = await GUARDRAIL_REGISTRY.get_guardrails_from_db(
prisma_client=prisma_client,
team_id=team_id,
view=view_val,
allowed_team_ids=allowed_team_ids,
)
guardrail_configs: List[GuardrailInfoResponse] = []
@ -199,6 +223,7 @@ async def list_guardrails_v2():
guardrail_name=guardrail.get("guardrail_name"),
litellm_params=masked_litellm_params,
guardrail_info=guardrail.get("guardrail_info"),
team_id=guardrail.get("team_id"),
created_at=guardrail.get("created_at"),
updated_at=guardrail.get("updated_at"),
guardrail_definition_location="db",
@ -209,6 +234,11 @@ async def list_guardrails_v2():
# get guardrails initialized on litellm config.yaml
in_memory_guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails()
for guardrail in in_memory_guardrails:
# for non-proxy-admin view=all, only include in-memory guardrails for their admin teams
if allowed_team_ids is not None:
g_team_id = guardrail.get("team_id")
if g_team_id not in allowed_team_ids:
continue
# only add guardrails that are not in DB guardrail list already
if guardrail.get("guardrail_id") not in seen_guardrail_ids:
in_memory_litellm_params_raw = guardrail.get("litellm_params")
@ -233,6 +263,7 @@ async def list_guardrails_v2():
guardrail_name=guardrail.get("guardrail_name"),
litellm_params=masked_in_memory_litellm_params_typed,
guardrail_info=dict(guardrail.get("guardrail_info") or {}),
team_id=guardrail.get("team_id"),
guardrail_definition_location="config",
)
)
@ -246,6 +277,7 @@ async def list_guardrails_v2():
class CreateGuardrailRequest(BaseModel):
guardrail: Guardrail
team_id: Optional[str] = None # top-level for UI; also allowed inside guardrail
@router.post(
@ -303,21 +335,56 @@ async def create_guardrail(
}
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import premium_user, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
# Extract team_id: top-level (UI sends it here) or inside guardrail / guardrail_info
guardrail_dict = (
request.guardrail
if isinstance(request.guardrail, dict)
else request.guardrail.model_dump(exclude_none=True)
)
team_id: Optional[str] = (
getattr(request, "team_id", None)
or guardrail_dict.get("team_id")
or (guardrail_dict.get("guardrail_info") or {}).get("team_id")
)
if team_id is not None:
if premium_user is not True:
raise HTTPException(
status_code=403,
detail=CommonProxyErrors.not_premium_user.value,
)
team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_row is None:
raise HTTPException(
status_code=400,
detail={"error": f"Team id={team_id} does not exist in db"},
)
team_obj = LiteLLM_TeamTable(**team_row.model_dump())
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=team_id,
user_api_key_dict=user_api_key_dict,
team_obj=team_obj,
premium_user=premium_user,
)
guardrail_dict["team_id"] = team_id
else:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
try:
result = await GUARDRAIL_REGISTRY.add_guardrail_to_db(
guardrail=request.guardrail, prisma_client=prisma_client
guardrail=cast(Guardrail, guardrail_dict), prisma_client=prisma_client
)
guardrail_name = result.get("guardrail_name", "Unknown")
@ -401,14 +468,9 @@ async def update_guardrail(
}
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import premium_user, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -424,6 +486,37 @@ async def update_guardrail(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
existing_team_id = existing_guardrail.get("team_id")
if existing_team_id is not None:
if premium_user is not True:
raise HTTPException(
status_code=403,
detail=CommonProxyErrors.not_premium_user.value,
)
team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": existing_team_id}
)
if team_row is None:
raise HTTPException(
status_code=400,
detail={
"error": f"Team id={existing_team_id} does not exist in db"
},
)
team_obj = LiteLLM_TeamTable(**team_row.model_dump())
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=existing_team_id,
user_api_key_dict=user_api_key_dict,
team_obj=team_obj,
premium_user=premium_user,
)
else:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
result = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=request.guardrail,
@ -477,14 +570,9 @@ async def delete_guardrail(
}
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import premium_user, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -500,6 +588,37 @@ async def delete_guardrail(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
existing_team_id = existing_guardrail.get("team_id")
if existing_team_id is not None:
if premium_user is not True:
raise HTTPException(
status_code=403,
detail=CommonProxyErrors.not_premium_user.value,
)
team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": existing_team_id}
)
if team_row is None:
raise HTTPException(
status_code=400,
detail={
"error": f"Team id={existing_team_id} does not exist in db"
},
)
team_obj = LiteLLM_TeamTable(**team_row.model_dump())
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=existing_team_id,
user_api_key_dict=user_api_key_dict,
team_obj=team_obj,
premium_user=premium_user,
)
else:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db(
guardrail_id=guardrail_id, prisma_client=prisma_client
)
@ -579,14 +698,9 @@ async def patch_guardrail(
}
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import premium_user, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -602,6 +716,37 @@ async def patch_guardrail(
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
existing_team_id = existing_guardrail.get("team_id")
if existing_team_id is not None:
if premium_user is not True:
raise HTTPException(
status_code=403,
detail=CommonProxyErrors.not_premium_user.value,
)
team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": existing_team_id}
)
if team_row is None:
raise HTTPException(
status_code=400,
detail={
"error": f"Team id={existing_team_id} does not exist in db"
},
)
team_obj = LiteLLM_TeamTable(**team_row.model_dump())
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=existing_team_id,
user_api_key_dict=user_api_key_dict,
team_obj=team_obj,
premium_user=premium_user,
)
else:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
# Create updated guardrail object
guardrail_name = (
request.guardrail_name
@ -673,7 +818,10 @@ async def patch_guardrail(
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_guardrail_info(guardrail_id: str):
async def get_guardrail_info(
guardrail_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get detailed information about a specific guardrail by ID
@ -707,8 +855,9 @@ async def get_guardrail_info(guardrail_id: str):
"""
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import premium_user, prisma_client
from litellm.types.guardrails import GUARDRAIL_DEFINITION_LOCATION
if prisma_client is None:
@ -732,6 +881,37 @@ async def get_guardrail_info(guardrail_id: str):
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
result_team_id = result.get("team_id")
if result_team_id is not None:
if premium_user is not True:
raise HTTPException(
status_code=403,
detail=CommonProxyErrors.not_premium_user.value,
)
team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": result_team_id}
)
if team_row is None:
raise HTTPException(
status_code=400,
detail={
"error": f"Team id={result_team_id} does not exist in db"
},
)
team_obj = LiteLLM_TeamTable(**team_row.model_dump())
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=result_team_id,
user_api_key_dict=user_api_key_dict,
team_obj=team_obj,
premium_user=premium_user,
)
else:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to view this guardrail",
)
litellm_params: Optional[Union[LitellmParams, dict]] = result.get(
"litellm_params"
)
@ -782,10 +962,8 @@ async def get_guardrail_ui_settings():
- Content filter settings (patterns and categories)
"""
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import (
PATTERN_CATEGORIES,
get_available_content_categories,
get_pattern_metadata,
)
PATTERN_CATEGORIES, get_available_content_categories,
get_pattern_metadata)
# Convert the PII_ENTITY_CATEGORIES_MAP to the format expected by the UI
category_maps = []
@ -1369,7 +1547,8 @@ async def get_provider_specific_params():
}
### get the config model for the guardrail - go through the registry and get the config model for the guardrail
from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
from litellm.proxy.guardrails.guardrail_registry import \
guardrail_class_registry
for guardrail_name, guardrail_class in guardrail_class_registry.items():
guardrail_config_model = guardrail_class.get_config_model()

24
litellm/proxy/guardrails/guardrail_helpers.py
View file

@ -1,6 +1,6 @@
import os
import sys
from typing import Dict
from typing import Dict, Optional
import litellm
from litellm._logging import verbose_proxy_logger
@ -27,7 +27,9 @@ def can_modify_guardrails(team_obj: Optional[LiteLLM_TeamTable]) -> bool:
return True
async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> bool:
async def should_proceed_based_on_metadata(
data: dict, guardrail_name: str, team_id: Optional[str] = None
) -> bool:
"""
checks if this guardrail should be applied to this call
"""
@ -75,6 +77,24 @@ async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> b
return True
def resolve_guardrail_for_request(
guardrail_name: str, team_id: Optional[str] = None
) -> Optional[Guardrail]:
"""
Resolve guardrail config by name and optional team_id from in-memory handler.
Returns the Guardrail dict or None if not found.
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
result = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_name_and_team(
guardrail_name, team_id
)
if result is None:
return None
guardrail, _ = result
return guardrail
async def should_proceed_based_on_api_key(
user_api_key_dict: UserAPIKeyAuth, guardrail_name: str
) -> bool:

15
litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py
View file

@ -13,7 +13,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
import json
import sys
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Type
from typing import TYPE_CHECKING, Any, List, Optional, Type
from fastapi import HTTPException
@ -31,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypesLiteral
GUARDRAIL_NAME = "aporia"
@ -181,16 +182,7 @@ class AporiaGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@ -209,6 +201,7 @@ class AporiaGuardrail(CustomGuardrail):
await should_proceed_based_on_metadata(
data=data,
guardrail_name=GUARDRAIL_NAME,
team_id=getattr(user_api_key_dict, "team_id", None),
)
is False
):

76
litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
View file

@ -20,23 +20,17 @@ from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.integrations.custom_guardrail import (CustomGuardrail,
log_guardrail_information)
from litellm.llms.custom_httpx.http_handler import (get_async_httpx_client,
httpxSpecialProvider)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.proxy.guardrails.guardrail_helpers import (
resolve_guardrail_for_request, should_proceed_based_on_metadata)
from litellm.secret_managers.main import get_secret
from litellm.types.guardrails import (
GuardrailItem,
LakeraCategoryThresholds,
Role,
default_roles,
)
from litellm.types.guardrails import (LakeraCategoryThresholds, Role,
default_roles)
from litellm.types.utils import CallTypesLiteral
GUARDRAIL_NAME = "lakera_prompt_injection"
@ -125,24 +119,13 @@ class lakeraAI_Moderation(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
if (
await should_proceed_based_on_metadata(
data=data,
guardrail_name=GUARDRAIL_NAME,
team_id=getattr(user_api_key_dict, "team_id", None),
)
is False
):
@ -150,14 +133,15 @@ class lakeraAI_Moderation(CustomGuardrail):
text = ""
_json_data: str = ""
if "messages" in data and isinstance(data["messages"], list):
prompt_injection_obj: Optional[
GuardrailItem
] = litellm.guardrail_name_config_map.get("prompt_injection")
if prompt_injection_obj is not None:
enabled_roles = prompt_injection_obj.enabled_roles
team_id = getattr(user_api_key_dict, "team_id", None)
resolved = resolve_guardrail_for_request("prompt_injection", team_id=team_id)
litellm_params = resolved.get("litellm_params") if resolved else None
if litellm_params is not None and hasattr(litellm_params, "enabled_roles"):
enabled_roles = getattr(litellm_params, "enabled_roles", None)
elif isinstance(litellm_params, dict):
enabled_roles = litellm_params.get("enabled_roles")
else:
enabled_roles = None
if enabled_roles is None:
enabled_roles = default_roles
@ -306,18 +290,7 @@ class lakeraAI_Moderation(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: litellm.DualCache,
data: Dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, Dict]]:
from litellm.types.guardrails import GuardrailEventHooks
@ -344,16 +317,7 @@ class lakeraAI_Moderation(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
):
if self.event_hook is None:
if self.moderation_check == "pre_call":

182
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -3,7 +3,7 @@
import importlib
import os
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Type, cast
from typing import Any, Dict, List, Optional, Tuple, Type, cast
import litellm
from litellm import Router
@ -11,29 +11,21 @@ from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy.utils import PrismaClient
from litellm.proxy.guardrails.guardrail_hooks.grayswan import GraySwanGuardrail
from litellm.proxy.guardrails.guardrail_hooks.grayswan import \
initialize_guardrail as initialize_grayswan
from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.proxy.utils import PrismaClient
from litellm.secret_managers.main import get_secret
from litellm.types.guardrails import (
Guardrail,
GuardrailEventHooks,
LakeraCategoryThresholds,
LitellmParams,
SupportedGuardrailIntegrations,
)
from litellm.proxy.guardrails.guardrail_hooks.grayswan import (
GraySwanGuardrail,
initialize_guardrail as initialize_grayswan,
)
from litellm.types.guardrails import (Guardrail, GuardrailEventHooks,
LakeraCategoryThresholds, LitellmParams,
SupportedGuardrailIntegrations)
from .guardrail_initializers import (
initialize_bedrock,
initialize_hide_secrets,
initialize_lakera,
initialize_lakera_v2,
initialize_presidio,
initialize_tool_permission,
)
from .guardrail_initializers import (initialize_bedrock,
initialize_hide_secrets,
initialize_lakera, initialize_lakera_v2,
initialize_presidio,
initialize_tool_permission)
guardrail_initializer_registry = {
SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock,
@ -251,21 +243,28 @@ class GuardrailRegistry:
)
litellm_params: str = safe_dumps(litellm_params_dict)
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
team_id: Optional[str] = guardrail.get("team_id")
# Create guardrail in DB
data: Dict[str, Any] = {
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
"guardrail_info": guardrail_info,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
if team_id is not None:
data["team_id"] = team_id
created_guardrail = await prisma_client.db.litellm_guardrailstable.create(
data={
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
"guardrail_info": guardrail_info,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
data=data
)
# Add guardrail_id to the returned guardrail object
# Add guardrail_id and team_id to the returned guardrail object
guardrail_dict = dict(guardrail)
guardrail_dict["guardrail_id"] = created_guardrail.guardrail_id
guardrail_dict["team_id"] = getattr(
created_guardrail, "team_id", guardrail.get("team_id")
)
return guardrail_dict
except Exception as e:
@ -305,16 +304,19 @@ class GuardrailRegistry:
)
litellm_params: str = safe_dumps(litellm_params_dict)
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
team_id: Optional[str] = guardrail.get("team_id")
# Update in DB
# Update in DB (include team_id so it can be updated or cleared)
update_data: Dict[str, Any] = {
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
"guardrail_info": guardrail_info,
"team_id": team_id,
"updated_at": datetime.now(timezone.utc),
}
updated_guardrail = await prisma_client.db.litellm_guardrailstable.update(
where={"guardrail_id": guardrail_id},
data={
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
"guardrail_info": guardrail_info,
"updated_at": datetime.now(timezone.utc),
},
data=update_data,
)
# Convert to dict and return
@ -323,17 +325,50 @@ class GuardrailRegistry:
raise Exception(f"Error updating guardrail in DB: {str(e)}")
@staticmethod
async def get_all_guardrails_from_db(
async def get_guardrails_from_db(
prisma_client: PrismaClient,
team_id: Optional[str] = None,
view: str = "all",
allowed_team_ids: Optional[List[str]] = None,
) -> List[Guardrail]:
"""
Get all guardrails from the database
Get guardrails from the database.
- If view == 'all' and allowed_team_ids is None: return all guardrails, order by created_at desc.
- If view == 'all' and allowed_team_ids is not None: return only guardrails where
team_id is in allowed_team_ids (for non-proxy-admin list filtering). Empty list returns [].
- If view == 'current_team' and team_id is not None: return guardrails where
team_id is null (global) or team_id == team_id, order by created_at desc.
- If view == 'current_team' and team_id is None: return guardrails where
team_id is null, order by created_at desc.
"""
try:
if view == "all":
if allowed_team_ids is not None:
if len(allowed_team_ids) == 0:
return []
where_filter = {"team_id": {"in": allowed_team_ids}}
else:
where_filter = None
elif view == "current_team":
if team_id is not None:
where_filter = {
"OR": [
{"team_id": None},
{"team_id": team_id},
]
}
else:
where_filter = {"team_id": None}
else:
where_filter = None
kwargs: Dict[str, Any] = {"order": {"created_at": "desc"}}
if where_filter is not None:
kwargs["where"] = where_filter
guardrails_from_db = (
await prisma_client.db.litellm_guardrailstable.find_many(
order={"created_at": "desc"},
)
await prisma_client.db.litellm_guardrailstable.find_many(**kwargs)
)
guardrails: List[Guardrail] = []
@ -344,6 +379,17 @@ class GuardrailRegistry:
except Exception as e:
raise Exception(f"Error getting guardrails from DB: {str(e)}")
@staticmethod
async def get_all_guardrails_from_db(
prisma_client: PrismaClient,
) -> List[Guardrail]:
"""
Get all guardrails from the database (backward-compatible wrapper).
"""
return await GuardrailRegistry.get_guardrails_from_db(
prisma_client, team_id=None, view="all"
)
async def get_guardrail_by_id_from_db(
self, guardrail_id: str, prisma_client: PrismaClient
) -> Optional[Guardrail]:
@ -363,20 +409,34 @@ class GuardrailRegistry:
raise Exception(f"Error getting guardrail from DB: {str(e)}")
async def get_guardrail_by_name_from_db(
self, guardrail_name: str, prisma_client: PrismaClient
self,
guardrail_name: str,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
) -> Optional[Guardrail]:
"""
Get a guardrail by its name from the database
Get a guardrail by its name from the database.
When team_id is provided, prefer the row with matching team_id; else use
the row with team_id null (global).
"""
try:
guardrail = await prisma_client.db.litellm_guardrailstable.find_unique(
guardrails = await prisma_client.db.litellm_guardrailstable.find_many(
where={"guardrail_name": guardrail_name}
)
if not guardrail:
if not guardrails:
return None
return Guardrail(**(dict(guardrail))) # type: ignore
# Prefer team_id match, then team_id is null
if team_id is not None:
for g in guardrails:
if g.team_id == team_id:
return Guardrail(**(dict(g))) # type: ignore
for g in guardrails:
if g.team_id is None:
return Guardrail(**(dict(g))) # type: ignore
# Fallback: return first (e.g. another team's guardrail if only those exist)
return Guardrail(**(dict(guardrails[0]))) # type: ignore
except Exception as e:
raise Exception(f"Error getting guardrail from DB: {str(e)}")
@ -474,6 +534,7 @@ 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
@ -588,6 +649,35 @@ class InMemoryGuardrailHandler:
"""
return self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
def get_guardrail_by_name_and_team(
self,
guardrail_name: str,
team_id: Optional[str] = None,
) -> Optional[Tuple[Guardrail, Optional[CustomGuardrail]]]:
"""
Resolve guardrail by name and optional team_id.
Prefer: guardrail with guardrail_name and team_id == team_id.
Else: guardrail with guardrail_name and team_id is null (global).
Returns (Guardrail, callback) or None if not found.
"""
candidates_team: List[Tuple[Guardrail, Optional[CustomGuardrail]]] = []
candidates_global: List[Tuple[Guardrail, Optional[CustomGuardrail]]] = []
for g_id, guardrail in self.IN_MEMORY_GUARDRAILS.items():
if guardrail.get("guardrail_name") != guardrail_name:
continue
g_team_id = guardrail.get("team_id")
callback = self.guardrail_id_to_custom_guardrail.get(g_id)
pair = (guardrail, callback)
if g_team_id == team_id and team_id is not None:
candidates_team.append(pair)
elif g_team_id is None:
candidates_global.append(pair)
if team_id is not None and candidates_team:
return candidates_team[0]
if candidates_global:
return candidates_global[0]
return None
def _has_guardrail_params_changed(
self, guardrail_id: str, new_guardrail: Guardrail
) -> bool:

51
litellm/proxy/management_endpoints/common_utils.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
@ -57,6 +57,55 @@ def _team_member_has_permission(
return False
async def get_team_ids_where_user_is_team_admin(
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Optional["PrismaClient"] = None,
user_api_key_cache: Optional["DualCache"] = None,
proxy_logging_obj: Optional["ProxyLogging"] = None,
) -> List[str]:
"""
Return team_ids for which the user is a team admin.
Returns empty list if user is proxy admin (caller should not restrict in that case),
or if no DB / user_id, or if user is not admin of any team.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return []
if prisma_client is None or user_api_key_dict.user_id is None:
return []
from litellm.caching import DualCache as DualCacheImport
from litellm.proxy.auth.auth_checks import get_user_object
try:
user_obj = await get_user_object(
user_id=user_api_key_dict.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache or DualCacheImport(),
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
)
if user_obj is None or not user_obj.teams:
return []
teams = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": user_obj.teams}}
)
admin_team_ids: List[str] = []
for team in teams:
team_obj = LiteLLM_TeamTable(**team.model_dump())
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=team_obj
):
admin_team_ids.append(team_obj.team_id)
return admin_team_ids
except Exception as e:
verbose_proxy_logger.debug(
f"Error getting team admin list for user {user_api_key_dict.user_id}: {e}"
)
return []
async def _user_has_admin_privileges(
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Optional["PrismaClient"] = None,

39
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -10,19 +10,13 @@ from typing import Any, List, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
)
from litellm.integrations.custom_guardrail import (CustomGuardrail,
ModifyResponseException)
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import \
UnifiedLLMGuardrails
from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineExecutionResult,
PipelineStep,
PipelineStepResult,
)
PipelineExecutionResult, PipelineStep, PipelineStepResult)
try:
from fastapi.exceptions import HTTPException
@ -143,7 +137,10 @@ class PipelineExecutor:
- modified_data: dict if guardrail returned modified data, else None
- error_detail: error message string if fail/error, else None
"""
callback = PipelineExecutor._find_guardrail_callback(step.guardrail)
team_id = getattr(user_api_key_dict, "team_id", None)
callback = PipelineExecutor._find_guardrail_callback(
step.guardrail, team_id=team_id
)
if callback is None:
verbose_proxy_logger.warning(
f"Pipeline: guardrail '{step.guardrail}' not found in callbacks"
@ -196,8 +193,22 @@ class PipelineExecutor:
return ("error", None, str(e))
@staticmethod
def _find_guardrail_callback(guardrail_name: str) -> Optional[CustomGuardrail]:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""
def _find_guardrail_callback(
guardrail_name: str,
team_id: Optional[str] = None,
) -> Optional[CustomGuardrail]:
"""Look up an initialized guardrail callback by name. When team_id is set,
prefer team-scoped then global from in-memory handler; else search litellm.callbacks."""
if team_id is not None:
from litellm.proxy.guardrails.guardrail_registry import \
IN_MEMORY_GUARDRAIL_HANDLER
in_memory = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_name_and_team(
guardrail_name, team_id
)
if in_memory is not None:
_, callback = in_memory
return callback
for callback in litellm.callbacks:
if isinstance(callback, CustomGuardrail):
if callback.guardrail_name == guardrail_name:

4
litellm/proxy/schema.prisma
View file

@ -863,12 +863,14 @@ model LiteLLM_ManagedVectorStoresTable {
// Guardrails table for storing guardrail configurations
model LiteLLM_GuardrailsTable {
guardrail_id String @id @default(uuid())
guardrail_name String @unique
guardrail_name String
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([guardrail_name, team_id])
}
// Daily guardrail metrics for usage dashboard (one row per guardrail per day)

2
litellm/types/guardrails.py
View file

@ -780,6 +780,7 @@ class Guardrail(TypedDict, total=False):
litellm_params: Required[LitellmParams]
guardrail_info: Optional[Dict]
policy_template: Optional[str]
team_id: Optional[str]
created_at: Optional[datetime]
updated_at: Optional[datetime]
@ -812,6 +813,7 @@ 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 = (

4
schema.prisma
View file

@ -863,12 +863,14 @@ model LiteLLM_ManagedVectorStoresTable {
// Guardrails table for storing guardrail configurations
model LiteLLM_GuardrailsTable {
guardrail_id String @id @default(uuid())
guardrail_name String @unique
guardrail_name String
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([guardrail_name, team_id])
}
// Daily guardrail metrics for usage dashboard (one row per guardrail per day)

11
ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts
View file

@ -40,15 +40,18 @@ const mockPage2: PaginatedKeyAliasResponse = {
size: 2,
};
const createWrapper = () => {
function createWrapper() {
const queryClient = new QueryClient({
defaultOptions: {
queries: { retry: false },
},
});
return ({ children }: { children: ReactNode }) =>
React.createElement(QueryClientProvider, { client: queryClient }, children);
};
function Wrapper({ children }: { children: ReactNode }) {
return React.createElement(QueryClientProvider, { client: queryClient }, children);
}
Wrapper.displayName = "QueryClientWrapper";
return Wrapper;
}
describe("useInfiniteKeyAliases", () => {
beforeEach(() => {

31
ui/litellm-dashboard/src/components/guardrails.tsx
View file

@ -1,4 +1,4 @@
import React, { useState, useEffect } from "react";
import React, { useState, useEffect, useCallback } from "react";
import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
import { Dropdown } from "antd";
import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons";
@ -14,6 +14,9 @@ import DeleteResourceModal from "./common_components/DeleteResourceModal";
import { getGuardrailLogoAndName } from "./guardrails/guardrail_info_helpers";
import { CustomCodeModal } from "./guardrails/custom_code";
import GuardrailGarden from "./guardrails/guardrail_garden";
import { Team } from "./key_team_helpers/key_list";
import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
interface GuardrailsPanelProps {
accessToken: string | null;
@ -38,7 +41,13 @@ interface GuardrailsResponse {
guardrails: Guardrail[];
}
const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole }) => {
type GuardrailViewMode = "all" | "current_team";
const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken: accessTokenProp, userRole: userRoleProp }) => {
const { accessToken: accessTokenFromAuth, userRole: userRoleFromAuth } = useAuthorized();
const accessToken = accessTokenProp ?? accessTokenFromAuth;
const userRole = userRoleProp ?? userRoleFromAuth;
const { data: teams, isLoading: isLoadingTeams } = useTeams();
const [guardrailsList, setGuardrailsList] = useState<Guardrail[]>([]);
const [isAddModalVisible, setIsAddModalVisible] = useState(false);
const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false);
@ -48,17 +57,25 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [selectedGuardrailId, setSelectedGuardrailId] = useState<string | null>(null);
const [activeTab, setActiveTab] = useState<number>(0);
const [currentTeam, setCurrentTeam] = useState<Team | "personal" | null>("personal");
const [modelViewMode, setModelViewMode] = useState<GuardrailViewMode>("all");
const isAdmin = userRole ? isAdminRole(userRole) : false;
const fetchGuardrails = async () => {
const teamIdForQuery = currentTeam === "personal" || currentTeam === null ? undefined : currentTeam.team_id;
const fetchGuardrails = useCallback(async () => {
if (!accessToken) {
return;
}
setIsLoading(true);
try {
const response: GuardrailsResponse = await getGuardrailsList(accessToken);
const response: GuardrailsResponse = await getGuardrailsList(
accessToken,
teamIdForQuery,
modelViewMode,
);
console.log(`guardrails: ${JSON.stringify(response)}`);
setGuardrailsList(response.guardrails);
} catch (error) {
@ -66,11 +83,11 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
} finally {
setIsLoading(false);
}
};
}, [accessToken, teamIdForQuery, modelViewMode]);
useEffect(() => {
fetchGuardrails();
}, [accessToken]);
}, [fetchGuardrails]);
const handleAddGuardrail = () => {
if (selectedGuardrailId) {
@ -203,6 +220,8 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
onClose={handleCloseModal}
accessToken={accessToken}
onSuccess={handleSuccess}
teams={teams ?? null}
userRole={userRole}
/>
<CustomCodeModal

65
ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx
View file

@ -1,7 +1,11 @@
import { Form, Input, Modal, Select, Tag, Typography, Button } from "antd";
import { Alert, Form, Input, Modal, Select, Tag, Typography, Button } from "antd";
import React, { useEffect, useMemo, useState } from "react";
import NotificationsManager from "../molecules/notifications_manager";
import { createGuardrailCall, getGuardrailProviderSpecificParams, getGuardrailUISettings } from "../networking";
import { all_admin_roles, isUserTeamAdminForAnyTeam } from "@/utils/roles";
import TeamDropdown from "../common_components/team_dropdown";
import type { Team } from "../key_team_helpers/key_list";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import ContentFilterConfiguration from "./content_filter/ContentFilterConfiguration";
import {
getGuardrailProviders,
@ -44,6 +48,8 @@ interface AddGuardrailFormProps {
accessToken: string | null;
onSuccess: () => void;
preset?: GuardrailPreset;
teams?: Team[] | null;
userRole?: string;
}
interface GuardrailSettings {
@ -96,10 +102,22 @@ interface ProviderParamsResponse {
[provider: string]: { [key: string]: ProviderParam };
}
const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, accessToken, onSuccess, preset }) => {
const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({
visible,
onClose,
accessToken,
onSuccess,
preset,
teams = null,
userRole,
}) => {
const [form] = Form.useForm();
const [loading, setLoading] = useState(false);
const [selectedProvider, setSelectedProvider] = useState<string | null>(null);
const [teamAdminSelectedTeam, setTeamAdminSelectedTeam] = useState<string | null>(null);
const { userId } = useAuthorized();
const isAdmin = userRole ? all_admin_roles.includes(userRole) : false;
const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId ?? undefined);
const [guardrailSettings, setGuardrailSettings] = useState<GuardrailSettings | null>(null);
const [selectedEntities, setSelectedEntities] = useState<string[]>([]);
const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({});
@ -161,7 +179,7 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
populateGuardrailProviderMap(providerParamsResp);
} catch (error) {
console.error("Error fetching guardrail data:", error);
NotificationsManager.fromBackend("Failed to load guardrail configuration");
NotificationsManager.fromBackend("Failed to load guardrail configuration for the selected provider");
}
};
@ -353,6 +371,7 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
const resetForm = () => {
form.resetFields();
setTeamAdminSelectedTeam(null);
setSelectedProvider(null);
setSelectedEntities([]);
setSelectedActions({});
@ -583,8 +602,9 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
throw new Error("No access token available");
}
const teamId = values.team_id ?? null;
console.log("Sending guardrail data:", JSON.stringify(guardrailData));
await createGuardrailCall(accessToken, guardrailData);
await createGuardrailCall(accessToken, guardrailData, teamId);
NotificationsManager.success("Guardrail created successfully");
@ -605,6 +625,41 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
const renderBasicInfo = () => {
return (
<>
{isTeamAdmin && !isAdmin && (
<>
<Form.Item
name="team_id"
label="Select Team"
rules={[{ required: true, message: "Please select a team to continue" }]}
tooltip="Select the team for which you want to add this guardrail"
>
<TeamDropdown
teams={teams}
onChange={(value) => setTeamAdminSelectedTeam(value ?? null)}
/>
</Form.Item>
{!teamAdminSelectedTeam && (
<Alert
message="Team Selection Required"
description="As a team admin, you need to select your team first before adding guardrails."
type="info"
showIcon
className="mb-4"
/>
)}
</>
)}
{(isAdmin || !isTeamAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && (
<>
{isAdmin && (
<Form.Item
name="team_id"
label="Team (optional)"
tooltip="Leave empty for a global guardrail, or select a team to create a team-scoped guardrail"
>
<TeamDropdown teams={teams} />
</Form.Item>
)}
<Form.Item
name="guardrail_name"
label="Guardrail Name"
@ -757,6 +812,8 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
providerParams={providerParams}
/>
)}
</>
)}
</>
);
};

19
ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx
View file

@ -133,6 +133,25 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
);
},
},
{
header: "Scope",
accessorKey: "team_id",
cell: ({ row }) => {
const guardrail = row.original;
if (guardrail.team_id == null || guardrail.team_id === "") {
return (
<Badge color="gray" className="text-xs font-normal" size="xs">
Global
</Badge>
);
}
return (
<span className="text-xs">
Team: {guardrail.team_alias || guardrail.team_id}
</span>
);
},
},
{
header: "Created At",
accessorKey: "created_at",

2
ui/litellm-dashboard/src/components/guardrails/types.ts
View file

@ -32,6 +32,8 @@ export interface Guardrail {
created_at?: string;
updated_at?: string;
guardrail_definition_location: GuardrailDefinitionLocation;
team_id?: string | null;
team_alias?: string | null;
}
export enum GuardrailDefinitionLocation {

2
ui/litellm-dashboard/src/components/leftnav.tsx
View file

@ -103,7 +103,7 @@ const menuGroups: MenuGroup[] = [
page: "guardrails",
label: "Guardrails",
icon: <SafetyOutlined />,
roles: all_admin_roles,
roles: [...all_admin_roles, ...internalUserRoles],
},
{
key: "policies",

24
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -5458,9 +5458,17 @@ export const testMCPSemanticFilter = async (accessToken: string, model: string,
}
};
export const getGuardrailsList = async (accessToken: string) => {
export const getGuardrailsList = async (
accessToken: string,
teamId?: string | null,
view: "all" | "current_team" = "all",
) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/v2/guardrails/list` : `/v2/guardrails/list`;
let url = proxyBaseUrl ? `${proxyBaseUrl}/v2/guardrails/list` : `/v2/guardrails/list`;
const params = new URLSearchParams();
if (teamId != null && teamId !== "") params.append("team_id", teamId);
params.append("view", view);
if (params.toString()) url += `?${params.toString()}`;
const response = await fetch(url, {
method: "GET",
headers: {
@ -6694,9 +6702,15 @@ export const createAgentCall = async (accessToken: string, agentData: any) => {
}
};
export const createGuardrailCall = async (accessToken: string, guardrailData: any) => {
export const createGuardrailCall = async (
accessToken: string,
guardrailData: any,
teamId?: string | null,
) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails` : `/guardrails`;
const body: { guardrail: any; team_id?: string } = { guardrail: guardrailData };
if (teamId != null && teamId !== "") body.team_id = teamId;
const response = await fetch(url, {
method: "POST",
@ -6704,9 +6718,7 @@ export const createGuardrailCall = async (accessToken: string, guardrailData: an
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
guardrail: guardrailData,
}),
body: JSON.stringify(body),
});
if (!response.ok) {

2
ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
View file

@ -5,7 +5,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
import { InfoCircleOutlined } from "@ant-design/icons";
import { useQueryClient } from "@tanstack/react-query";
import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
import { Button as Button2, Form, Input, message, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
import debounce from "lodash/debounce";
import React, { useCallback, useEffect, useState } from "react";
import { rolesWithWriteAccess } from "../../utils/roles";