From 94c8705d0cfb2a78132b809ca0ae18c5295657b3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 26 Feb 2026 19:52:26 -0800 Subject: [PATCH] feat: working team BYOG --- .../docs/proxy/guardrails/team_guardrails.md | 64 ++++ .../migration.sql | 6 + .../litellm_proxy_extras/schema.prisma | 4 +- litellm/proxy/_types.py | 86 +++-- .../proxy/guardrails/guardrail_endpoints.py | 329 ++++++++++++++---- litellm/proxy/guardrails/guardrail_helpers.py | 24 +- .../guardrail_hooks/aporia_ai/aporia_ai.py | 15 +- .../guardrails/guardrail_hooks/lakera_ai.py | 76 ++-- .../proxy/guardrails/guardrail_registry.py | 182 +++++++--- .../management_endpoints/common_utils.py | 51 ++- .../proxy/policy_engine/pipeline_executor.py | 39 ++- litellm/proxy/schema.prisma | 4 +- litellm/types/guardrails.py | 2 + schema.prisma | 4 +- .../hooks/keys/useKeyAliases.test.ts | 11 +- .../src/components/guardrails.tsx | 31 +- .../guardrails/add_guardrail_form.tsx | 65 +++- .../components/guardrails/guardrail_table.tsx | 19 + .../src/components/guardrails/types.ts | 2 + .../src/components/leftnav.tsx | 2 +- .../src/components/networking.tsx | 24 +- .../organisms/create_key_button.tsx | 2 +- 22 files changed, 767 insertions(+), 275 deletions(-) create mode 100644 docs/my-website/docs/proxy/guardrails/team_guardrails.md create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260226195018_add_guardrail_team_unique/migration.sql diff --git a/docs/my-website/docs/proxy/guardrails/team_guardrails.md b/docs/my-website/docs/proxy/guardrails/team_guardrails.md new file mode 100644 index 00000000000..21a416930a0 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/team_guardrails.md @@ -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 " \ + -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 " + +# 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 " +``` + +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. diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226195018_add_guardrail_team_unique/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226195018_add_guardrail_team_unique/migration.sql new file mode 100644 index 00000000000..74c3055285b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226195018_add_guardrail_team_unique/migration.sql @@ -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"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 440c9c1d829..b5c7811ded2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 53513f7f522..0c4399ae864 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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"}) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 5215fca0293..cf4bd5582c2 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -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 " + curl -X GET "http://localhost:4000/v2/guardrails/list?team_id=&view=current_team" -H "Authorization: Bearer " ``` 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() diff --git a/litellm/proxy/guardrails/guardrail_helpers.py b/litellm/proxy/guardrails/guardrail_helpers.py index e9703114603..aeb307b19b2 100644 --- a/litellm/proxy/guardrails/guardrail_helpers.py +++ b/litellm/proxy/guardrails/guardrail_helpers.py @@ -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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py index 35f776e5bd1..0506f8444f0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py @@ -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 ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index 28f0d830f12..c0fdf471713 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -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": diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index c0903a35b6d..d8602432e75 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -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: diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 999c53e6b82..0e2ffe3e065 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c1e2d76e7cd..9e33491e6fa 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -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: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 440c9c1d829..b5c7811ded2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 0e71f20700e..31c52a36eb8 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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 = ( diff --git a/schema.prisma b/schema.prisma index 440c9c1d829..b5c7811ded2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts index b382b1f2ad3..7b37ecbb919 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts @@ -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(() => { diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index a8de7dd2f4f..832d7a23283 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -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 = ({ accessToken, userRole }) => { +type GuardrailViewMode = "all" | "current_team"; + +const GuardrailsPanel: React.FC = ({ 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([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false); @@ -48,17 +57,25 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [selectedGuardrailId, setSelectedGuardrailId] = useState(null); const [activeTab, setActiveTab] = useState(0); + const [currentTeam, setCurrentTeam] = useState("personal"); + const [modelViewMode, setModelViewMode] = useState("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 = ({ accessToken, userRole } finally { setIsLoading(false); } - }; + }, [accessToken, teamIdForQuery, modelViewMode]); useEffect(() => { fetchGuardrails(); - }, [accessToken]); + }, [fetchGuardrails]); const handleAddGuardrail = () => { if (selectedGuardrailId) { @@ -203,6 +220,8 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole onClose={handleCloseModal} accessToken={accessToken} onSuccess={handleSuccess} + teams={teams ?? null} + userRole={userRole} /> 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 = ({ visible, onClose, accessToken, onSuccess, preset }) => { +const AddGuardrailForm: React.FC = ({ + visible, + onClose, + accessToken, + onSuccess, + preset, + teams = null, + userRole, +}) => { const [form] = Form.useForm(); const [loading, setLoading] = useState(false); const [selectedProvider, setSelectedProvider] = useState(null); + const [teamAdminSelectedTeam, setTeamAdminSelectedTeam] = useState(null); + const { userId } = useAuthorized(); + const isAdmin = userRole ? all_admin_roles.includes(userRole) : false; + const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId ?? undefined); const [guardrailSettings, setGuardrailSettings] = useState(null); const [selectedEntities, setSelectedEntities] = useState([]); const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({}); @@ -161,7 +179,7 @@ const AddGuardrailForm: React.FC = ({ 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 = ({ visible, onClose, a const resetForm = () => { form.resetFields(); + setTeamAdminSelectedTeam(null); setSelectedProvider(null); setSelectedEntities([]); setSelectedActions({}); @@ -583,8 +602,9 @@ const AddGuardrailForm: React.FC = ({ 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 = ({ visible, onClose, a const renderBasicInfo = () => { return ( <> + {isTeamAdmin && !isAdmin && ( + <> + + setTeamAdminSelectedTeam(value ?? null)} + /> + + {!teamAdminSelectedTeam && ( + + )} + + )} + {(isAdmin || !isTeamAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && ( + <> + {isAdmin && ( + + + + )} = ({ visible, onClose, a providerParams={providerParams} /> )} + + )} ); }; diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx index 352f3148e7b..5f7e1bbe973 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx @@ -133,6 +133,25 @@ const GuardrailTable: React.FC = ({ ); }, }, + { + header: "Scope", + accessorKey: "team_id", + cell: ({ row }) => { + const guardrail = row.original; + if (guardrail.team_id == null || guardrail.team_id === "") { + return ( + + Global + + ); + } + return ( + + Team: {guardrail.team_alias || guardrail.team_id} + + ); + }, + }, { header: "Created At", accessorKey: "created_at", diff --git a/ui/litellm-dashboard/src/components/guardrails/types.ts b/ui/litellm-dashboard/src/components/guardrails/types.ts index e8ed27d9e45..5321dfbffed 100644 --- a/ui/litellm-dashboard/src/components/guardrails/types.ts +++ b/ui/litellm-dashboard/src/components/guardrails/types.ts @@ -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 { diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 2cbeb22ec81..8d59f3cda4e 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -103,7 +103,7 @@ const menuGroups: MenuGroup[] = [ page: "guardrails", label: "Guardrails", icon: , - roles: all_admin_roles, + roles: [...all_admin_roles, ...internalUserRoles], }, { key: "policies", diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 0df6d813d7c..d2d3716d276 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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) { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 961a5d4d460..7566edbcae8 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -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";