From ca34b4ee13f2ec62e44f6ae7d3df5a98ffd0e5c0 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 3 Feb 2026 20:55:46 -0800 Subject: [PATCH] =?UTF-8?q?Revert=20"feat(guardrails):=20implement=20team-?= =?UTF-8?q?based=20isolation=20guardrails=20mgmnt=20(#1=E2=80=A6"=20(#2039?= =?UTF-8?q?3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 76a399ba69239899a1ec2cff0a9e235f75944889. --- .../migration.sql | 8 - litellm/proxy/_types.py | 10 +- litellm/proxy/auth/login_utils.py | 60 +- .../proxy/guardrails/guardrail_endpoints.py | 280 ++----- .../proxy/guardrails/guardrail_registry.py | 99 +-- .../management_endpoints/team_endpoints.py | 279 ++++--- litellm/proxy/schema.prisma | 5 +- litellm/types/guardrails.py | 10 - schema.prisma | 5 +- .../guardrails/test_guardrail_endpoints.py | 773 ++++++++---------- .../guardrails/test_guardrail_team_access.py | 295 ------- .../src/components/team/team_info.tsx | 87 +- 12 files changed, 618 insertions(+), 1293 deletions(-) delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql delete mode 100644 tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql deleted file mode 100644 index 706a2f25e9..0000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql +++ /dev/null @@ -1,8 +0,0 @@ --- AlterTable -ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN "team_id" TEXT; - --- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false; - --- AlterTable -ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false; \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 784e56a03f..131a6caab0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -359,6 +359,7 @@ class LiteLLMRoutes(enum.Enum): "/v1/vector_stores/{vector_store_id}/files/{file_id}/content", "/vector_store/list", "/v1/vector_store/list", + # search "/search", "/v1/search", @@ -630,9 +631,6 @@ class LiteLLMRoutes(enum.Enum): "/model/{model_id}/update", "/prompt/list", "/prompt/info", - "/guardrails", - "/guardrails/{guardrail_id}", - "/v2/guardrails/list", ] # routes that manage their own allowed/disallowed logic ## Org Admin Routes ## @@ -1484,9 +1482,6 @@ class TeamBase(LiteLLMPydanticObjectBase): members: list = [] members_with_roles: List[Member] = [] team_member_permissions: Optional[List[str]] = None - allow_team_guardrail_config: Optional[ - bool - ] = None # if True, team admin can configure guardrails for this team metadata: Optional[dict] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None @@ -1584,9 +1579,6 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): model_tpm_limit: Optional[Dict[str, int]] = None allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None router_settings: Optional[dict] = None - allow_team_guardrail_config: Optional[ - bool - ] = None # if True, team admin can configure guardrails for this team class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 3bd56177d7..4df773dec2 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -34,58 +34,6 @@ from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.ui_sso import ReturnedUITokenObject -async def expire_previous_ui_session_tokens( - user_id: str, prisma_client: Optional[PrismaClient] -) -> None: - """ - Expire (block) all other valid UI session tokens for a user. - - This prevents accumulation of multiple valid UI session tokens that - are supposed to be short-lived test keys. Only affects keys with - team_id = "litellm-dashboard" and that haven't expired yet. - - Args: - user_id: The user ID whose previous UI session tokens should be expired - prisma_client: Database client for performing the update - """ - if prisma_client is None: - return - - try: - from datetime import datetime, timezone - - current_time = datetime.now(timezone.utc) - - # Find all unblocked AND non-expired UI session tokens for this user - ui_session_tokens = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": user_id, - "team_id": "litellm-dashboard", - "OR": [ - {"blocked": None}, # Tokens that have never been blocked (null) - {"blocked": False}, # Tokens explicitly set to not blocked - ], - "expires": {"gt": current_time}, # Only get tokens that haven't expired - } - ) - - if not ui_session_tokens: - return - - # Block all the found tokens - tokens_to_block = [token.token for token in ui_session_tokens if token.token] - - if tokens_to_block: - await prisma_client.db.litellm_verificationtoken.update_many( - where={"token": {"in": tokens_to_block}}, data={"blocked": True} - ) - - except Exception: - # Silently fail - don't block login if cleanup fails - # This is a best-effort operation - pass - - def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]: """ Get UI username and password from environment variables or master key. @@ -297,11 +245,6 @@ async def authenticate_user( # noqa: PLR0915 ) user_email = getattr(_user_row, "user_email", "unknown") _password = getattr(_user_row, "password", "unknown") - user_team_id = getattr(_user_row, "team_id", "litellm-dashboard") - - # if user_team_id is None, set it to "litellm-dashboard" - if user_team_id is None: - user_team_id = "litellm-dashboard" if _password is None: raise ProxyException( @@ -328,7 +271,7 @@ async def authenticate_user( # noqa: PLR0915 "config": {}, "spend": 0, "user_id": user_id, - "team_id": user_team_id, + "team_id": "litellm-dashboard", }, ) else: @@ -397,3 +340,4 @@ def create_ui_token_object( disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) + diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index aec5dc1ede..a825ce22b2 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -11,7 +11,7 @@ from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.types.guardrails import ( @@ -32,8 +32,6 @@ from litellm.types.guardrails import ( PresidioPresidioConfigModelUserInterface, SupportedGuardrailIntegrations, ToolPermissionGuardrailConfigModel, - CreateGuardrailRequest, - UpdateGuardrailRequest, ) #### GUARDRAILS ENDPOINTS #### @@ -42,37 +40,6 @@ router = APIRouter() GUARDRAIL_REGISTRY = GuardrailRegistry() -async def _check_team_can_configure_guardrails( - user_api_key_dict: UserAPIKeyAuth, - team_id: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, -) -> None: - """ - If the user is not proxy admin and team_id is set, verify the team has - allow_team_guardrail_config enabled. Raise HTTPException 403 otherwise. - """ - if team_id is None or team_id == "litellm-dashboard": - return - user_role = getattr(user_api_key_dict, "user_role", None) - if user_role == LitellmUserRoles.PROXY_ADMIN: - return - from litellm.proxy.auth.auth_checks import get_team_object - - team_table = await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - allow_config = getattr(team_table, "allow_team_guardrail_config", False) - if not allow_config: - raise HTTPException( - status_code=403, - detail="Guardrail configuration is not enabled for this team. Contact your administrator to enable it.", - ) - - def _get_guardrails_list_response( guardrails_config: List[Dict], ) -> ListGuardrailsResponse: @@ -97,19 +64,9 @@ def _get_guardrails_list_response( dependencies=[Depends(user_api_key_auth)], response_model=ListGuardrailsResponse, ) -@router.get( - "/v2/guardrails/list", - tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], - response_model=ListGuardrailsResponse, -) -async def list_guardrails( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def list_guardrails(): """ - List the guardrails that are available in the database using GuardrailRegistry - - Now supports both DB-persisted and In-Memory (Config) guardrails. + List the guardrails that are available on the proxy server 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) @@ -118,6 +75,59 @@ async def list_guardrails( curl -X GET "http://localhost:4000/guardrails/list" -H "Authorization: Bearer " ``` + Example Response: + ```json + { + "guardrails": [ + { + "guardrail_name": "bedrock-pre-guard", + "guardrail_info": { + "params": [ + { + "name": "toxicity_score", + "type": "float", + "description": "Score between 0-1 indicating content toxicity level" + }, + { + "name": "pii_detection", + "type": "boolean" + } + ] + } + } + ] + } + ``` + """ + from litellm.proxy.proxy_server import proxy_config + + config = proxy_config.config + + _guardrails_config = cast(Optional[list[dict]], config.get("guardrails")) + + if _guardrails_config is None: + return _get_guardrails_list_response([]) + + return _get_guardrails_list_response(_guardrails_config) + + +@router.get( + "/v2/guardrails/list", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], + response_model=ListGuardrailsResponse, +) +async def list_guardrails_v2(): + """ + List the guardrails that are available in the database using GuardrailRegistry + + 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) + + Example Request: + ```bash + curl -X GET "http://localhost:4000/v2/guardrails/list" -H "Authorization: Bearer " + ``` + Example Response: ```json { @@ -129,14 +139,12 @@ async def list_guardrails( "guardrail": "bedrock", "mode": "pre_call", "guardrailIdentifier": "ff6ujrregl1q", - "guardrailVersion": "1.0", + "guardrailVersion": "DRAFT", "default_on": true }, "guardrail_info": { - "description": "Updated Bedrock content moderation guardrail" - }, - "created_at": "2023-11-09T12:34:56.789Z", - "updated_at": "2023-11-09T13:45:12.345Z" + "description": "Bedrock content moderation guardrail" + } } ] } @@ -145,29 +153,12 @@ async def list_guardrails( from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client - # Check if DB is connected if prisma_client is None: - # Fallback to config only if DB is missing (though prisma_client usually exists) - from litellm.proxy.proxy_server import proxy_config - - config = proxy_config.config - _guardrails_config = cast(Optional[list[dict]], config.get("guardrails")) - if _guardrails_config is None: - return _get_guardrails_list_response([]) - return _get_guardrails_list_response(_guardrails_config) + raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - team_id = getattr(user_api_key_dict, "team_id", None) - user_role = getattr(user_api_key_dict, "user_role", None) - - filter_team_id = None - if user_role != LitellmUserRoles.PROXY_ADMIN or ( - team_id is not None and team_id != "litellm-dashboard" - ): - filter_team_id = team_id - guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db( - prisma_client=prisma_client, team_id=filter_team_id + prisma_client=prisma_client ) guardrail_configs: List[GuardrailInfoResponse] = [] @@ -182,7 +173,6 @@ async def list_guardrails( created_at=guardrail.get("created_at"), updated_at=guardrail.get("updated_at"), guardrail_definition_location="db", - team_id=guardrail.get("team_id"), ) ) seen_guardrail_ids.add(guardrail.get("guardrail_id")) @@ -190,15 +180,6 @@ async def list_guardrails( # get guardrails initialized on litellm config.yaml in_memory_guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails() for guardrail in in_memory_guardrails: - # Check access for in-memory guardrails too - if filter_team_id: - g_team = guardrail.get("team_id") - - # If guardrail has a team_id, it must match. - # If guardrail has NO team_id, it is likely a global config guardrail, so we usually allow it. - if g_team and g_team != filter_team_id: - continue - # only add guardrails that are not in DB guardrail list already if guardrail.get("guardrail_id") not in seen_guardrail_ids: guardrail_configs.append( @@ -208,7 +189,6 @@ async def list_guardrails( litellm_params=dict(guardrail.get("litellm_params") or {}), guardrail_info=dict(guardrail.get("guardrail_info") or {}), guardrail_definition_location="config", - team_id=guardrail.get("team_id"), ) ) seen_guardrail_ids.add(guardrail.get("guardrail_id")) @@ -219,15 +199,16 @@ async def list_guardrails( raise HTTPException(status_code=500, detail=str(e)) +class CreateGuardrailRequest(BaseModel): + guardrail: Guardrail + + @router.post( "/guardrails", tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def create_guardrail( - request: CreateGuardrailRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def create_guardrail(request: CreateGuardrailRequest): """ Create a new guardrail @@ -276,21 +257,14 @@ async def create_guardrail( ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - team_id = getattr(user_api_key_dict, "team_id", None) - await _check_team_can_configure_guardrails( - user_api_key_dict=user_api_key_dict, - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) result = await GUARDRAIL_REGISTRY.add_guardrail_to_db( - guardrail=request.guardrail, prisma_client=prisma_client, team_id=team_id + guardrail=request.guardrail, prisma_client=prisma_client ) guardrail_name = result.get("guardrail_name", "Unknown") @@ -309,23 +283,21 @@ async def create_guardrail( ) return result - except HTTPException: - raise except Exception as e: verbose_proxy_logger.exception(f"Error adding guardrail to db: {e}") raise HTTPException(status_code=500, detail=str(e)) +class UpdateGuardrailRequest(BaseModel): + guardrail: Guardrail + + @router.put( "/guardrails/{guardrail_id}", tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def update_guardrail( - guardrail_id: str, - request: UpdateGuardrailRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): """ Update an existing guardrail @@ -374,7 +346,7 @@ async def update_guardrail( ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -390,25 +362,6 @@ async def update_guardrail( status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" ) - if existing_guardrail.get("team_id"): - team_id = getattr(user_api_key_dict, "team_id", None) - if getattr( - user_api_key_dict, "user_role", None - ) != LitellmUserRoles.PROXY_ADMIN or ( - team_id is not None and team_id != "litellm-dashboard" - ): - if existing_guardrail.get("team_id") != team_id: - raise HTTPException( - status_code=403, - detail="Not authorized to access this guardrail", - ) - await _check_team_can_configure_guardrails( - user_api_key_dict=user_api_key_dict, - team_id=existing_guardrail.get("team_id"), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - result = await GUARDRAIL_REGISTRY.update_guardrail_in_db( guardrail_id=guardrail_id, guardrail=request.guardrail, @@ -441,9 +394,7 @@ async def update_guardrail( tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def delete_guardrail( - guardrail_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth) -): +async def delete_guardrail(guardrail_id: str): """ Delete a guardrail @@ -463,7 +414,7 @@ async def delete_guardrail( ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -479,25 +430,6 @@ async def delete_guardrail( status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" ) - if existing_guardrail.get("team_id"): - team_id = getattr(user_api_key_dict, "team_id", None) - if getattr( - user_api_key_dict, "user_role", None - ) != LitellmUserRoles.PROXY_ADMIN or ( - team_id is not None and team_id != "litellm-dashboard" - ): - if existing_guardrail.get("team_id") != team_id: - raise HTTPException( - status_code=403, - detail="Not authorized to access this guardrail", - ) - await _check_team_can_configure_guardrails( - user_api_key_dict=user_api_key_dict, - team_id=existing_guardrail.get("team_id"), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db( guardrail_id=guardrail_id, prisma_client=prisma_client ) @@ -528,11 +460,7 @@ async def delete_guardrail( tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def patch_guardrail( - guardrail_id: str, - request: PatchGuardrailRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): """ Partially update an existing guardrail @@ -579,7 +507,7 @@ async def patch_guardrail( ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -595,25 +523,6 @@ async def patch_guardrail( status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" ) - if existing_guardrail.get("team_id"): - team_id = getattr(user_api_key_dict, "team_id", None) - if getattr( - user_api_key_dict, "user_role", None - ) != LitellmUserRoles.PROXY_ADMIN or ( - team_id is not None and team_id != "litellm-dashboard" - ): - if existing_guardrail.get("team_id") != team_id: - raise HTTPException( - status_code=403, - detail="Not authorized to access this guardrail", - ) - await _check_team_can_configure_guardrails( - user_api_key_dict=user_api_key_dict, - team_id=existing_guardrail.get("team_id"), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - # Create updated guardrail object guardrail_name = ( request.guardrail_name @@ -685,10 +594,7 @@ async def patch_guardrail( tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def get_guardrail_info( - guardrail_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def get_guardrail_info(guardrail_id: str): """ Get detailed information about a specific guardrail by ID @@ -747,19 +653,6 @@ async def get_guardrail_info( status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" ) - if result.get("team_id"): - team_id = getattr(user_api_key_dict, "team_id", None) - if getattr( - user_api_key_dict, "user_role", None - ) != LitellmUserRoles.PROXY_ADMIN or ( - team_id is not None and team_id != "litellm-dashboard" - ): - if result.get("team_id") != team_id: - raise HTTPException( - status_code=403, - detail="Not authorized to access this guardrail", - ) - litellm_params: Optional[Union[LitellmParams, dict]] = result.get( "litellm_params" ) @@ -782,7 +675,6 @@ async def get_guardrail_info( created_at=result.get("created_at"), updated_at=result.get("updated_at"), guardrail_definition_location=guardrail_definition_location, - team_id=result.get("team_id"), ) except HTTPException as e: raise e @@ -1317,9 +1209,9 @@ async def get_provider_specific_params(): lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel) - tool_permission_fields[ - "ui_friendly_name" - ] = ToolPermissionGuardrailConfigModel.ui_friendly_name() + tool_permission_fields["ui_friendly_name"] = ( + ToolPermissionGuardrailConfigModel.ui_friendly_name() + ) # Return the provider-specific parameters provider_params = { @@ -1627,10 +1519,10 @@ async def apply_guardrail( from litellm.proxy.utils import handle_exception_on_proxy try: - active_guardrail: Optional[ - CustomGuardrail - ] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( - guardrail_name=request.guardrail_name + active_guardrail: Optional[CustomGuardrail] = ( + GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + guardrail_name=request.guardrail_name + ) ) if active_guardrail is None: raise HTTPException( diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 487a9d481c..c3da689220 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -233,10 +233,7 @@ class GuardrailRegistry: ########### DB management helpers for guardrails ########### ############################################################ async def add_guardrail_to_db( - self, - guardrail: Guardrail, - prisma_client: PrismaClient, - team_id: Optional[str] = None, + self, guardrail: Guardrail, prisma_client: PrismaClient ): """ Add a guardrail to the database @@ -251,8 +248,6 @@ class GuardrailRegistry: litellm_params_dict = ( dict(litellm_params_obj) if litellm_params_obj else {} ) - - # Use safe_dumps to store as string, as Prisma client seems to prefer this for now litellm_params: str = safe_dumps(litellm_params_dict) guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) @@ -262,7 +257,6 @@ class GuardrailRegistry: "guardrail_name": guardrail_name, "litellm_params": litellm_params, "guardrail_info": guardrail_info, - "team_id": team_id, "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), } @@ -271,32 +265,18 @@ class GuardrailRegistry: # Add guardrail_id to the returned guardrail object guardrail_dict = dict(guardrail) guardrail_dict["guardrail_id"] = created_guardrail.guardrail_id - guardrail_dict["team_id"] = team_id return guardrail_dict except Exception as e: raise Exception(f"Error adding guardrail to DB: {str(e)}") async def delete_guardrail_from_db( - self, - guardrail_id: str, - prisma_client: PrismaClient, - team_id: Optional[str] = None, + self, guardrail_id: str, prisma_client: PrismaClient ): """ Delete a guardrail from the database """ try: - # Check ownership if team_id is provided - if team_id: - existing_guardrail = ( - await prisma_client.db.litellm_guardrailstable.find_unique( - where={"guardrail_id": guardrail_id} - ) - ) - if not existing_guardrail or existing_guardrail.team_id != team_id: - raise Exception("Guardrail not found or access denied") - # Delete from DB await prisma_client.db.litellm_guardrailstable.delete( where={"guardrail_id": guardrail_id} @@ -307,26 +287,12 @@ class GuardrailRegistry: raise Exception(f"Error deleting guardrail from DB: {str(e)}") async def update_guardrail_in_db( - self, - guardrail_id: str, - guardrail: Guardrail, - prisma_client: PrismaClient, - team_id: Optional[str] = None, + self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient ): """ Update a guardrail in the database """ try: - # Check ownership if team_id is provided - if team_id: - existing_guardrail = ( - await prisma_client.db.litellm_guardrailstable.find_unique( - where={"guardrail_id": guardrail_id} - ) - ) - if not existing_guardrail or existing_guardrail.team_id != team_id: - raise Exception("Guardrail not found or access denied") - guardrail_name = guardrail.get("guardrail_name") # Properly serialize LitellmParams Pydantic model to dict litellm_params_obj: Any = guardrail.get("litellm_params", {}) @@ -336,7 +302,6 @@ class GuardrailRegistry: litellm_params_dict = ( dict(litellm_params_obj) if litellm_params_obj else {} ) - litellm_params: str = safe_dumps(litellm_params_dict) guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) @@ -359,67 +324,27 @@ class GuardrailRegistry: @staticmethod async def get_all_guardrails_from_db( prisma_client: PrismaClient, - team_id: Optional[str] = None, ) -> List[Guardrail]: """ Get all guardrails from the database """ try: - # Normal ORM Fetch - all_guardrails = await prisma_client.db.litellm_guardrailstable.find_many( - order={"created_at": "desc"}, + guardrails_from_db = ( + await prisma_client.db.litellm_guardrailstable.find_many( + order={"created_at": "desc"}, + ) ) - # Filter in Python - if team_id: - guardrails_from_db = [g for g in all_guardrails if g.team_id == team_id] - else: - guardrails_from_db = all_guardrails - guardrails: List[Guardrail] = [] for guardrail in guardrails_from_db: - try: - # Deep copy to avoid mutating cache/original - g_dict = ( - guardrail.dict() - if hasattr(guardrail, "dict") - else dict(guardrail) - ) - - # Handle litellm_params - params = g_dict.get("litellm_params") - if isinstance(params, str): - import json - - try: - g_dict["litellm_params"] = json.loads(params) - except Exception: - g_dict["litellm_params"] = {} - - # Handle guardrail_info - info = g_dict.get("guardrail_info") - if isinstance(info, str): - import json - - try: - g_dict["guardrail_info"] = json.loads(info) - except Exception: - g_dict["guardrail_info"] = {} - - # Construct - guardrails.append(Guardrail(**g_dict)) # type: ignore - except Exception: - continue + guardrails.append(Guardrail(**(dict(guardrail)))) # type: ignore return guardrails except Exception as e: raise Exception(f"Error getting guardrails from DB: {str(e)}") async def get_guardrail_by_id_from_db( - self, - guardrail_id: str, - prisma_client: PrismaClient, - team_id: Optional[str] = None, + self, guardrail_id: str, prisma_client: PrismaClient ) -> Optional[Guardrail]: """ Get a guardrail by its ID from the database @@ -432,11 +357,6 @@ class GuardrailRegistry: if not guardrail: return None - if team_id and guardrail.team_id != team_id: - # Return None if not found or not owned by team - # Alternatively could raise exception, but returning None resembles "not found" - return None - return Guardrail(**(dict(guardrail))) # type: ignore except Exception as e: raise Exception(f"Error getting guardrail from DB: {str(e)}") @@ -553,7 +473,6 @@ class InMemoryGuardrailHandler: guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail["guardrail_name"], litellm_params=litellm_params, - team_id=guardrail.get("team_id"), ) # store references to the guardrail in memory diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 43c3212194..63db2d72fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -753,16 +753,12 @@ async def new_team( # noqa: PLR0915 if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"max_budget cannot be negative. Received: {data.max_budget}" - }, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} ) if data.team_member_budget is not None and data.team_member_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}" - }, + detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} ) # Check if license is over limit @@ -922,16 +918,12 @@ async def new_team( # noqa: PLR0915 complete_team_data.members_with_roles = [] complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) - + # Serialize router_settings to JSON (matching key creation pattern) router_settings_value = getattr(data, "router_settings", None) - router_settings_json = ( - safe_dumps(router_settings_value) - if router_settings_value is not None - else safe_dumps({}) - ) + router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({}) complete_team_data_dict["router_settings"] = router_settings_json - + complete_team_data_dict = prisma_client.jsonify_team_object( db_data=complete_team_data_dict ) @@ -1107,9 +1099,7 @@ async def fetch_and_validate_organization( validate_team_org_change( team=LiteLLM_TeamTable(**existing_team_row.model_dump()), - organization=LiteLLM_OrganizationTableWithMembers( - **organization_row.model_dump() - ), + organization=LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()), llm_router=llm_router, ) @@ -1117,9 +1107,7 @@ async def fetch_and_validate_organization( def validate_team_org_change( - team: LiteLLM_TeamTable, - organization: LiteLLM_OrganizationTableWithMembers, - llm_router: Router, + team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router ) -> bool: """ Validate that a team can be moved to an organization. @@ -1170,9 +1158,7 @@ def validate_team_org_change( # Check if the team's user_id is a member of the org team_members = [m.user_id for m in team.members_with_roles] - org_members = ( - [m.user_id for m in organization.members] if organization.members else [] - ) + org_members = [m.user_id for m in organization.members] if organization.members else [] not_in_org = [ m for m in team_members @@ -1218,7 +1204,7 @@ def validate_team_org_change( "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @management_endpoint_wrapper -async def update_team( # noqa: PLR0915 +async def update_team( # noqa: PLR0915 data: UpdateTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -1302,25 +1288,19 @@ async def update_team( # noqa: PLR0915 ) if data.team_id is None: - raise HTTPException( - status_code=400, detail={"error": "No team id passed in"} - ) + raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) verbose_proxy_logger.debug("/team/update - %s", data) # Validate budget values are not negative if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"max_budget cannot be negative. Received: {data.max_budget}" - }, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} ) if data.team_member_budget is not None and data.team_member_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}" - }, + detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} ) existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( @@ -1387,22 +1367,6 @@ async def update_team( # noqa: PLR0915 updated_kv = data.json(exclude_unset=True) - # Only proxy admin can change allow_team_guardrail_config - if ( - "allow_team_guardrail_config" in updated_kv - and user_api_key_dict.user_role - not in ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN.value, - ) - ): - raise HTTPException( - status_code=403, - detail={ - "error": "Only proxy admin can change 'Allow team to configure guardrails'. Contact your administrator." - }, - ) - # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) @@ -1447,19 +1411,16 @@ async def update_team( # noqa: PLR0915 updated_kv["model_id"] = _model_id # Serialize router_settings to JSON if present (matching key update pattern) - if ( - "router_settings" in updated_kv - and updated_kv["router_settings"] is not None - ): + if "router_settings" in updated_kv and updated_kv["router_settings"] is not None: updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore + ) ) if team_row is None or team_row.team_id is None: @@ -1468,9 +1429,7 @@ async def update_team( # noqa: PLR0915 detail={"error": "Team doesn't exist. Got={}".format(team_row)}, ) - verbose_proxy_logger.info( - "Successfully updated team - %s, info", team_row.team_id - ) + verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) await _cache_team_object( team_id=team_row.team_id, team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), @@ -1812,6 +1771,113 @@ async def _add_team_members_to_team( return updated_team, updated_users, updated_team_memberships +async def _validate_and_populate_member_user_info( + member: Member, + prisma_client: PrismaClient, +) -> Member: + """ + Validate and populate user_email/user_id for a member. + + Logic: + 1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth) + 2. If only user_email is provided, populate user_id from DB + 3. If only user_id is provided, populate user_email from DB (if user exists) + 4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later) + 5. If user_email and user_id mismatch, throw error + + Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist). + """ + if member.user_email is None and member.user_id is None: + raise HTTPException( + status_code=400, + detail={"error": "Either user_id or user_email must be provided"}, + ) + + # Case 1: Both user_email and user_id provided - verify they match + if member.user_email is not None and member.user_id is not None: + # Use user_email as source of truth + # Check for multiple users with same email first + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email is None or ( + isinstance(users_by_email, list) and len(users_by_email) == 0 + ): + # User doesn't exist yet - this is fine, will be created later + return member + + if isinstance(users_by_email, list) and len(users_by_email) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." + }, + ) + + # Get the single user + user_by_email = users_by_email[0] + + # Verify the user_id matches + if user_by_email.user_id != member.user_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user." + }, + ) + + # Both match, return as is + return member + + # Case 2: Only user_email provided - populate user_id from DB + if member.user_email is not None and member.user_id is None: + user_by_email = await prisma_client.db.litellm_usertable.find_first( + where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} + ) + + if user_by_email is None: + # User doesn't exist yet - this is fine, will be created later + return member + + # Check for multiple users with same email + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." + }, + ) + + # Populate user_id + member.user_id = user_by_email.user_id + return member + + # Case 3: Only user_id provided - populate user_email from DB if user exists + if member.user_id is not None and member.user_email is None: + user_by_id = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": member.user_id} + ) + + if user_by_id is None: + # User doesn't exist yet - allow it to pass with user_email as None + # Will be upserted later with just user_id and null email + return member + + # Populate user_email + member.user_email = user_by_id.user_email + return member + + return member + @router.post( "/team/member_add", tags=["team management"], @@ -1887,16 +1953,27 @@ async def team_member_add( complete_team_data=complete_team_data, ) - ( - updated_team, - updated_users, - updated_team_memberships, - ) = await _add_team_members_to_team( - data=data, - complete_team_data=complete_team_data, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + # Validate and populate user_email/user_id for members before processing + if isinstance(data.member, Member): + await _validate_and_populate_member_user_info( + member=data.member, + prisma_client=prisma_client, + ) + elif isinstance(data.member, List): + for m in data.member: + await _validate_and_populate_member_user_info( + member=m, + prisma_client=prisma_client, + ) + + updated_team, updated_users, updated_team_memberships = ( + await _add_team_members_to_team( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) ) # Check if updated_team is None @@ -2075,15 +2152,15 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) ) - + if keys_to_delete: await _persist_deleted_verification_tokens( keys=keys_to_delete, @@ -2462,10 +2539,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row_base: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) if team_row_base is None: raise Exception @@ -2524,10 +2601,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} + ) ) if keys_to_delete: @@ -2566,6 +2643,7 @@ async def delete_team( return deleted_teams + def _transform_teams_to_deleted_records( teams: List[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, @@ -2588,13 +2666,7 @@ def _transform_teams_to_deleted_records( ) record = deleted_record.model_dump() - for json_field in [ - "members_with_roles", - "metadata", - "model_spend", - "model_max_budget", - "router_settings", - ]: + for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]: if json_field in record and record[json_field] is not None: record[json_field] = json.dumps(record[json_field]) @@ -2613,7 +2685,9 @@ async def _save_deleted_team_records( """Save deleted team records to the database.""" if not records: return - await prisma_client.db.litellm_deletedteamtable.create_many(data=records) + await prisma_client.db.litellm_deletedteamtable.create_many( + data=records + ) async def _persist_deleted_team_records( @@ -2633,7 +2707,6 @@ async def _persist_deleted_team_records( prisma_client=prisma_client, ) - def validate_membership( user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable ): @@ -2757,11 +2830,11 @@ async def team_info( ) try: - team_info: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, + team_info: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, + ) ) if team_info is None: raise Exception @@ -3224,9 +3297,7 @@ async def list_team_v2( order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count( - where=where_conditions - ) + total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d5393bc32b..b118400b62 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -128,7 +128,6 @@ model LiteLLM_TeamTable { router_settings Json? @default("{}") team_member_permissions String[] @default([]) policies String[] @default([]) - allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) @@ -160,9 +159,8 @@ model LiteLLM_DeletedTeamTable { router_settings Json? @default("{}") team_member_permissions String[] @default([]) policies String[] @default([]) - allow_team_guardrail_config Boolean @default(false) model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases - + // Original timestamps from team creation/updates created_at DateTime? @map("created_at") updated_at DateTime? @map("updated_at") @@ -776,7 +774,6 @@ model LiteLLM_GuardrailsTable { guardrail_name String @unique litellm_params Json guardrail_info Json? - team_id String? created_at DateTime @default(now()) updated_at DateTime @updatedAt } diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 3f7de10dde..74ccb34ca6 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -731,7 +731,6 @@ class Guardrail(TypedDict, total=False): guardrail_name: Required[str] litellm_params: Required[LitellmParams] guardrail_info: Optional[Dict] - team_id: Optional[str] created_at: Optional[datetime] updated_at: Optional[datetime] @@ -763,7 +762,6 @@ class GuardrailInfoResponse(BaseModel): guardrail_name: str litellm_params: Optional[BaseLitellmParams] = None guardrail_info: Optional[Dict] = None - team_id: Optional[str] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = ( @@ -810,11 +808,3 @@ class PatchGuardrailRequest(BaseModel): guardrail_name: Optional[str] = None litellm_params: Optional[BaseLitellmParams] = None guardrail_info: Optional[Dict[str, Any]] = None - - -class CreateGuardrailRequest(BaseModel): - guardrail: Guardrail - - -class UpdateGuardrailRequest(BaseModel): - guardrail: Guardrail diff --git a/schema.prisma b/schema.prisma index dc49036cb1..b118400b62 100644 --- a/schema.prisma +++ b/schema.prisma @@ -129,7 +129,6 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) policies String[] @default([]) model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases - allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) @@ -161,8 +160,7 @@ model LiteLLM_DeletedTeamTable { team_member_permissions String[] @default([]) policies String[] @default([]) model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases - allow_team_guardrail_config Boolean @default(false) - + // Original timestamps from team creation/updates created_at DateTime? @map("created_at") updated_at DateTime? @map("updated_at") @@ -776,7 +774,6 @@ model LiteLLM_GuardrailsTable { guardrail_name String @unique litellm_params Json guardrail_info Json? - team_id String? created_at DateTime @default(now()) updated_at DateTime @updatedAt } diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 6243932727..88f56c2406 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1,14 +1,15 @@ +import json import os import sys from datetime import datetime -from typing import List +from typing import Dict, List, Optional from unittest.mock import AsyncMock import pytest sys.path.insert( - 0, os.path.abspath("../../../../..") -) # Adds the repo root directory to the system path + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path from fastapi import HTTPException @@ -20,12 +21,12 @@ from litellm.proxy.guardrails.guardrail_endpoints import ( create_guardrail, delete_guardrail, get_guardrail_info, - list_guardrails, + list_guardrails_v2, patch_guardrail, update_guardrail, ) -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_registry import ( + IN_MEMORY_GUARDRAIL_HANDLER, InMemoryGuardrailHandler, ) from litellm.types.guardrails import ( @@ -62,7 +63,7 @@ MOCK_CONFIG_GUARDRAIL = { MOCK_GUARDRAIL = Guardrail( guardrail_name=MOCK_CONFIG_GUARDRAIL["guardrail_name"], litellm_params=LitellmParams(**MOCK_CONFIG_GUARDRAIL["litellm_params"]), - guardrail_info=MOCK_CONFIG_GUARDRAIL["guardrail_info"], + guardrail_info=MOCK_CONFIG_GUARDRAIL["guardrail_info"] ) MOCK_CREATE_REQUEST = CreateGuardrailRequest(guardrail=MOCK_GUARDRAIL) @@ -70,7 +71,7 @@ MOCK_UPDATE_REQUEST = UpdateGuardrailRequest(guardrail=MOCK_GUARDRAIL) MOCK_PATCH_REQUEST = PatchGuardrailRequest( guardrail_name="Updated Test Guardrail", litellm_params={"guardrail": "updated.guardrail", "mode": "post_call"}, - guardrail_info={"description": "Updated test guardrail"}, + guardrail_info={"description": "Updated test guardrail"} ) @@ -101,35 +102,22 @@ def mock_in_memory_handler(mocker): mock_handler.delete_in_memory_guardrail = mocker.Mock() return mock_handler - @pytest.fixture def mock_guardrail_registry(mocker): """Mock GuardrailRegistry for testing""" mock_registry = mocker.Mock() - mock_registry.add_guardrail_to_db = AsyncMock( - return_value={**MOCK_DB_GUARDRAIL, "guardrail_id": "new-test-guardrail-id"} - ) + mock_registry.add_guardrail_to_db = AsyncMock(return_value={ + **MOCK_DB_GUARDRAIL, + "guardrail_id": "new-test-guardrail-id" + }) mock_registry.delete_guardrail_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) - mock_registry.get_guardrail_by_id_from_db = AsyncMock( - return_value=MOCK_DB_GUARDRAIL - ) + mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) mock_registry.update_guardrail_in_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) return mock_registry - -@pytest.fixture -def mock_admin_user_auth(): - return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) - - -@pytest.fixture -def mock_team_user_auth(): - return UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, team_id="team-123") - - @pytest.mark.asyncio async def test_list_guardrails_v2_with_db_and_config( - mocker, mock_prisma_client, mock_in_memory_handler, mock_admin_user_auth + mocker, mock_prisma_client, mock_in_memory_handler ): """Test listing guardrails from both DB and config""" # Mock the prisma client @@ -140,7 +128,7 @@ async def test_list_guardrails_v2_with_db_and_config( mock_in_memory_handler, ) - response = await list_guardrails(user_api_key_dict=mock_admin_user_auth) + response = await list_guardrails_v2() assert len(response.guardrails) == 2 @@ -162,17 +150,11 @@ async def test_list_guardrails_v2_with_db_and_config( @pytest.mark.asyncio -async def test_get_guardrail_info_from_db( - mocker, mock_prisma_client, mock_admin_user_auth -): +async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): """Test getting guardrail info from DB""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - response = await get_guardrail_info( - "test-db-guardrail", user_api_key_dict=mock_admin_user_auth - ) + response = await get_guardrail_info("test-db-guardrail") assert response.guardrail_id == "test-db-guardrail" assert response.guardrail_name == "Test DB Guardrail" @@ -182,7 +164,7 @@ async def test_get_guardrail_info_from_db( @pytest.mark.asyncio async def test_get_guardrail_info_from_config( - mocker, mock_prisma_client, mock_in_memory_handler, mock_admin_user_auth + mocker, mock_prisma_client, mock_in_memory_handler ): """Test getting guardrail info from config when not found in DB""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -196,9 +178,7 @@ async def test_get_guardrail_info_from_config( return_value=None ) - response = await get_guardrail_info( - "test-config-guardrail", user_api_key_dict=mock_admin_user_auth - ) + response = await get_guardrail_info("test-config-guardrail") assert response.guardrail_id == "test-config-guardrail" assert response.guardrail_name == "Test Config Guardrail" @@ -208,7 +188,7 @@ async def test_get_guardrail_info_from_config( @pytest.mark.asyncio async def test_get_guardrail_info_not_found( - mocker, mock_prisma_client, mock_in_memory_handler, mock_admin_user_auth + mocker, mock_prisma_client, mock_in_memory_handler ): """Test getting guardrail info when not found in either DB or config""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -224,9 +204,7 @@ async def test_get_guardrail_info_not_found( mock_in_memory_handler.get_guardrail_by_id.return_value = None with pytest.raises(HTTPException) as exc_info: - await get_guardrail_info( - "non-existent-guardrail", user_api_key_dict=mock_admin_user_auth - ) + await get_guardrail_info("non-existent-guardrail") assert exc_info.value.status_code == 404 assert "not found" in str(exc_info.value.detail) @@ -244,6 +222,7 @@ def test_get_provider_specific_params(): pytest.skip("Azure config model not available") fields = _get_fields_from_model(config_model) + print("FIELDS", fields) # Test that we get the expected nested structure assert isinstance(fields, dict) @@ -259,12 +238,12 @@ def test_get_provider_specific_params(): fields["api_key"]["description"] == "API key for the Azure Content Safety Prompt Shield guardrail" ) - assert fields["api_key"]["required"] is False - assert fields["api_key"]["type"] == "string" + assert fields["api_key"]["required"] == False + assert fields["api_key"]["type"] == "string" # Should be string, not None # Check the structure of the nested optional_params field assert fields["optional_params"]["type"] == "nested" - assert fields["optional_params"]["required"] is True + assert fields["optional_params"]["required"] == True assert "fields" in fields["optional_params"] # Check nested fields within optional_params @@ -281,7 +260,7 @@ def test_get_provider_specific_params(): nested_fields["severity_threshold"]["description"] == "Severity threshold for the Azure Content Safety Text Moderation guardrail across all categories" ) - assert nested_fields["severity_threshold"]["required"] is False + assert nested_fields["severity_threshold"]["required"] == False assert ( nested_fields["severity_threshold"]["type"] == "number" ) # Should be number, not None @@ -290,14 +269,16 @@ def test_get_provider_specific_params(): assert nested_fields["categories"]["type"] == "multiselect" assert nested_fields["blocklistNames"]["type"] == "array" assert nested_fields["haltOnBlocklistHit"]["type"] == "boolean" - assert nested_fields["outputType"]["type"] == "select" + assert ( + nested_fields["outputType"]["type"] == "select" + ) # Literal type should be select def test_optional_params_not_returned_when_not_overridden(): """Test that optional_params is not returned when the config model doesn't override it""" from typing import Optional - from pydantic import Field + from pydantic import BaseModel, Field from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -318,6 +299,7 @@ def test_optional_params_not_returned_when_not_overridden(): # Get fields from the model fields = _get_fields_from_model(TestGuardrailConfig) + print("FIELDS", fields) assert "optional_params" not in fields @@ -355,13 +337,14 @@ def test_optional_params_returned_when_properly_overridden(): # Get fields from the model fields = _get_fields_from_model(TestGuardrailConfigWithOptionalParams) + print("FIELDS", fields) assert "optional_params" in fields @pytest.mark.asyncio async def test_bedrock_guardrail_prepare_request_with_api_key(): """Test _prepare_request method uses Bearer token when api_key is provided in data""" - from unittest.mock import Mock + from unittest.mock import Mock, patch from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, @@ -369,23 +352,27 @@ async def test_bedrock_guardrail_prepare_request_with_api_key(): # Setup guardrail hook guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", + guardrailVersion="1" ) mock_credentials = Mock() - test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - + test_data = { + "source": "INPUT", + "content": [{"text": {"text": "test content"}}] + } + prepared_request = guardrail_hook._prepare_request( credentials=mock_credentials, data=test_data, optional_params={}, aws_region_name="us-east-1", - api_key="test-bearer-token-123", + api_key="test-bearer-token-123" ) - + # Verify Bearer token is used in Authorization header assert "Authorization" in prepared_request.headers assert prepared_request.headers["Authorization"] == "Bearer test-bearer-token-123" - + # Verify URL is correct expected_url = "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail/test-guardrail-id/version/1/apply" assert prepared_request.url == expected_url @@ -402,44 +389,45 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): # Setup guardrail hook guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", + guardrailVersion="1" ) - + # Mock credentials mock_credentials = Mock() - + # Test data without api_key - test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - - with patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, patch( - "botocore.awsrequest.AWSRequest" - ) as mock_aws_request: + test_data = { + "source": "INPUT", + "content": [{"text": {"text": "test content"}}] + } + + with patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str") as mock_get_secret, \ + patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, \ + patch("botocore.awsrequest.AWSRequest") as mock_aws_request: + # Mock no AWS_BEARER_TOKEN_BEDROCK mock_get_secret.return_value = None - + # Mock SigV4Auth mock_sigv4_instance = Mock() mock_sigv4_auth.return_value = mock_sigv4_instance - + # Mock AWSRequest mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance - + # Call _prepare_request - guardrail_hook._prepare_request( + prepared_request = guardrail_hook._prepare_request( credentials=mock_credentials, data=test_data, optional_params={}, - aws_region_name="us-east-1", + aws_region_name="us-east-1" ) - + # Verify SigV4 auth was used - mock_sigv4_auth.assert_called_once_with( - mock_credentials, "bedrock", "us-east-1" - ) + mock_sigv4_auth.assert_called_once_with(mock_credentials, "bedrock", "us-east-1") mock_sigv4_instance.add_auth.assert_called_once() @@ -454,30 +442,34 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): # Setup guardrail hook guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", + guardrailVersion="1" ) - + # Mock credentials mock_credentials = Mock() - + # Test data without api_key - test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} - - with patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, patch("botocore.awsrequest.AWSRequest") as mock_aws_request: + test_data = { + "source": "INPUT", + "content": [{"text": {"text": "test content"}}] + } + + with patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str") as mock_get_secret, \ + patch("botocore.awsrequest.AWSRequest") as mock_aws_request: + mock_get_secret.return_value = "env-bearer-token-456" mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance - - guardrail_hook._prepare_request( + + prepared_request = guardrail_hook._prepare_request( credentials=mock_credentials, data=test_data, optional_params={}, - aws_region_name="us-east-1", + aws_region_name="us-east-1" ) - + # Verify Bearer token from environment is used mock_aws_request.assert_called_once() call_args = mock_aws_request.call_args @@ -493,51 +485,45 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, ) - + guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", + guardrailVersion="1" ) - + guardrail_hook.async_handler = Mock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = {"action": "NONE", "outputs": []} - - test_request_data = {"api_key": "test-api-key-789"} - - with patch.object( - guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response) - ), patch.object( - guardrail_hook, "_load_credentials" - ) as mock_load_creds, patch.object( - guardrail_hook, "convert_to_bedrock_format" - ) as mock_convert, patch.object( - guardrail_hook, "get_guardrail_dynamic_request_body_params" - ) as mock_get_params, patch.object( - guardrail_hook, "add_standard_logging_guardrail_information_to_request_data" - ), patch( - "botocore.awsrequest.AWSRequest" - ) as mock_aws_request: + + test_request_data = { + "api_key": "test-api-key-789" + } + + with patch.object(guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)), \ + patch.object(guardrail_hook, "_load_credentials") as mock_load_creds, \ + patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert, \ + patch.object(guardrail_hook, "get_guardrail_dynamic_request_body_params") as mock_get_params, \ + patch.object(guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"), \ + patch("botocore.awsrequest.AWSRequest") as mock_aws_request: + mock_load_creds.return_value = (Mock(), "us-east-1") mock_convert.return_value = {"source": "INPUT", "content": []} mock_get_params.return_value = {} - + mock_request_instance = Mock() mock_request_instance.url = "test-url" mock_request_instance.body = b"test-body" - mock_request_instance.headers = { - "Content-Type": "application/json", - "Authorization": "Bearer test-api-key-789", - } + mock_request_instance.headers = {"Content-Type": "application/json", "Authorization": "Bearer test-api-key-789"} mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance - + await guardrail_hook.make_bedrock_api_request( source="INPUT", messages=[{"role": "user", "content": "test"}], - request_data=test_request_data, + request_data=test_request_data ) - + # Verify _prepare_request was invoked and used the api_key mock_aws_request.assert_called_once() call_args = mock_aws_request.call_args @@ -545,419 +531,336 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): assert headers["Authorization"] == "Bearer test-api-key-789" -@pytest.mark.parametrize( - "scenario,expected_result,expected_exception", - [ - ("success_with_sync", "new-test-guardrail-id", None), - ("success_sync_fails", "new-test-guardrail-id", None), - ("database_failure", None, HTTPException), - ("no_prisma_client", None, HTTPException), - ], - ids=[ - "success_with_immediate_sync", - "success_but_sync_fails", - "database_error", - "missing_prisma_client", - ], -) +@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ + ( + "success_with_sync", + "new-test-guardrail-id", + None + ), + ( + "success_sync_fails", + "new-test-guardrail-id", + None + ), + ( + "database_failure", + None, + HTTPException + ), + ( + "no_prisma_client", + None, + HTTPException + ), +], ids=[ + "success_with_immediate_sync", + "success_but_sync_fails", + "database_error", + "missing_prisma_client" +]) @pytest.mark.asyncio async def test_create_guardrail_endpoint( - scenario, - expected_result, - expected_exception, - mocker, - mock_guardrail_registry, - mock_in_memory_handler, - mock_admin_user_auth, + scenario, expected_result, expected_exception, + mocker, mock_guardrail_registry, mock_in_memory_handler ): """Test create_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.initialize_guardrail.side_effect = Exception( - "Sync failed" - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) - + mock_in_memory_handler.initialize_guardrail.side_effect = Exception("Sync failed") + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception( - "Database error" - ) - - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - + mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception("Database error") + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + elif scenario == "no_prisma_client": mocker.patch("litellm.proxy.proxy_server.prisma_client", None) - + # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await create_guardrail( - MOCK_CREATE_REQUEST, user_api_key_dict=mock_admin_user_auth - ) - + await create_guardrail(MOCK_CREATE_REQUEST) + if scenario == "database_failure": assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) - + else: - result = await create_guardrail( - MOCK_CREATE_REQUEST, user_api_key_dict=mock_admin_user_auth - ) - + result = await create_guardrail(MOCK_CREATE_REQUEST) + assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" - + mock_guardrail_registry.add_guardrail_to_db.assert_called_once_with( guardrail=MOCK_CREATE_REQUEST.guardrail, - prisma_client=mocker.ANY, - team_id=None, + prisma_client=mocker.ANY ) - + mock_in_memory_handler.initialize_guardrail.assert_called_once() - + if scenario == "success_sync_fails": assert mock_logger is not None mock_logger.warning.assert_called_once() - assert "Failed to initialize guardrail" in str( - mock_logger.warning.call_args - ) + assert "Failed to initialize guardrail" in str(mock_logger.warning.call_args) - -@pytest.mark.parametrize( - "scenario,expected_result,expected_exception", - [ - ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), - ("database_failure", None, HTTPException), - ("no_prisma_client", None, HTTPException), - ], - ids=[ - "success_with_immediate_sync", - "success_but_sync_fails", - "database_error", - "missing_prisma_client", - ], -) +@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ + ( + "success_with_sync", + "test-db-guardrail", + None + ), + ( + "success_sync_fails", + "test-db-guardrail", + None + ), + ( + "database_failure", + None, + HTTPException + ), + ( + "no_prisma_client", + None, + HTTPException + ), +], ids=[ + "success_with_immediate_sync", + "success_but_sync_fails", + "database_error", + "missing_prisma_client" +]) @pytest.mark.asyncio async def test_update_guardrail_endpoint( - scenario, - expected_result, - expected_exception, - mocker, - mock_guardrail_registry, - mock_in_memory_handler, - mock_admin_user_auth, + scenario, expected_result, expected_exception, + mocker, mock_guardrail_registry, mock_in_memory_handler ): """Test update_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception( - "Sync failed" - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) - + mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception("Sync failed") + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( - "Database error" - ) - - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - + mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error") + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + elif scenario == "no_prisma_client": mocker.patch("litellm.proxy.proxy_server.prisma_client", None) - + # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await update_guardrail( - "test-guardrail-id", - MOCK_UPDATE_REQUEST, - user_api_key_dict=mock_admin_user_auth, - ) - + await update_guardrail("test-guardrail-id", MOCK_UPDATE_REQUEST) + if scenario == "database_failure": assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) - + else: - result = await update_guardrail( - "test-guardrail-id", - MOCK_UPDATE_REQUEST, - user_api_key_dict=mock_admin_user_auth, - ) - + result = await update_guardrail("test-guardrail-id", MOCK_UPDATE_REQUEST) + assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" - + mock_guardrail_registry.update_guardrail_in_db.assert_called_once_with( guardrail_id="test-guardrail-id", guardrail=MOCK_UPDATE_REQUEST.guardrail, - prisma_client=mocker.ANY, + prisma_client=mocker.ANY ) - + mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with( - guardrail_id="test-guardrail-id", guardrail=mocker.ANY + guardrail_id="test-guardrail-id", + guardrail=mocker.ANY ) - + if scenario == "success_sync_fails": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) - -@pytest.mark.parametrize( - "scenario,expected_result,expected_exception", - [ - ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), - ("database_failure", None, HTTPException), - ("no_prisma_client", None, HTTPException), - ], - ids=[ - "success_with_immediate_sync", - "success_but_sync_fails", - "database_error", - "missing_prisma_client", - ], -) +@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ + ( + "success_with_sync", + "test-db-guardrail", + None + ), + ( + "success_sync_fails", + "test-db-guardrail", + None + ), + ( + "database_failure", + None, + HTTPException + ), + ( + "no_prisma_client", + None, + HTTPException + ), +], ids=[ + "success_with_immediate_sync", + "success_but_sync_fails", + "database_error", + "missing_prisma_client" +]) @pytest.mark.asyncio async def test_patch_guardrail_endpoint( - scenario, - expected_result, - expected_exception, - mocker, - mock_guardrail_registry, - mock_in_memory_handler, - mock_admin_user_auth, + scenario, expected_result, expected_exception, + mocker, mock_guardrail_registry, mock_in_memory_handler ): """Test patch_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( - side_effect=Exception("Sync failed") - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) - + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(side_effect=Exception("Sync failed")) + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( - "Database error" - ) - - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - + mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error") + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + elif scenario == "no_prisma_client": mocker.patch("litellm.proxy.proxy_server.prisma_client", None) - + # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await patch_guardrail( - "test-guardrail-id", - MOCK_PATCH_REQUEST, - user_api_key_dict=mock_admin_user_auth, - ) - + await patch_guardrail("test-guardrail-id", MOCK_PATCH_REQUEST) + if scenario == "database_failure": assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) - + else: - result = await patch_guardrail( - "test-guardrail-id", - MOCK_PATCH_REQUEST, - user_api_key_dict=mock_admin_user_auth, - ) - + result = await patch_guardrail("test-guardrail-id", MOCK_PATCH_REQUEST) + assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" - + mock_guardrail_registry.update_guardrail_in_db.assert_called_once() - + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( guardrail=mocker.ANY ) - + if scenario == "success_sync_fails": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) - -@pytest.mark.parametrize( - "scenario,expected_result,expected_exception", - [ - ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), - ], - ids=["success_with_immediate_sync", "success_but_sync_fails"], -) +@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ + ( + "success_with_sync", + "test-db-guardrail", + None + ), + ( + "success_sync_fails", + "test-db-guardrail", + None + ), +], ids=[ + "success_with_immediate_sync", + "success_but_sync_fails" +]) @pytest.mark.asyncio async def test_delete_guardrail_endpoint( - scenario, - expected_result, - expected_exception, - mocker, - mock_guardrail_registry, - mock_in_memory_handler, - mock_admin_user_auth, + scenario, expected_result, expected_exception, + mocker, mock_guardrail_registry, mock_in_memory_handler ): """Test delete_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario mock_prisma_client = mocker.Mock() mock_logger = None - + if scenario == "success_with_sync": mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + elif scenario == "success_sync_fails": - mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception( - "Sync failed" - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) + mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception("Sync failed") + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_guardrail_registry, - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + if expected_exception: with pytest.raises(expected_exception): - await delete_guardrail( - guardrail_id=expected_result, user_api_key_dict=mock_admin_user_auth - ) + await delete_guardrail(guardrail_id=expected_result) else: - result = await delete_guardrail( - guardrail_id=expected_result, user_api_key_dict=mock_admin_user_auth - ) - + result = await delete_guardrail(guardrail_id=expected_result) + assert result == MOCK_DB_GUARDRAIL - + mock_guardrail_registry.get_guardrail_by_id_from_db.assert_called_once_with( - guardrail_id=expected_result, prisma_client=mock_prisma_client + guardrail_id=expected_result, + prisma_client=mock_prisma_client ) mock_guardrail_registry.delete_guardrail_from_db.assert_called_once_with( - guardrail_id=expected_result, prisma_client=mock_prisma_client + guardrail_id=expected_result, + prisma_client=mock_prisma_client ) - + mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with( guardrail_id=expected_result ) - + if scenario == "success_sync_fails": assert mock_logger is not None mock_logger.warning.assert_called_once() @@ -974,22 +877,21 @@ async def test_apply_guardrail_not_found(mocker): # Mock the GUARDRAIL_REGISTRY to return None (guardrail not found) mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = None - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) + # Create request request = ApplyGuardrailRequest( - guardrail_name="non-existent-guardrail", text="Test input text" + guardrail_name="non-existent-guardrail", + text="Test input text" ) - + # Mock user auth mock_user_auth = UserAPIKeyAuth() - + # Call endpoint and expect ProxyException with pytest.raises(ProxyException) as exc_info: await apply_guardrail(request=request, user_api_key_dict=mock_user_auth) - + # Verify error details assert str(exc_info.value.code) == "404" assert "not found" in str(exc_info.value.message).lower() @@ -1007,34 +909,30 @@ async def test_apply_guardrail_execution_error(mocker): mock_guardrail.apply_guardrail = AsyncMock( side_effect=Exception("Bedrock guardrail failed: Violated guardrail policy") ) - + # Mock the GUARDRAIL_REGISTRY mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) - + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) + # Create request request = ApplyGuardrailRequest( - guardrail_name="test-guardrail", text="Test input text with forbidden content" + guardrail_name="test-guardrail", + text="Test input text with forbidden content" ) - + # Mock user auth mock_user_auth = UserAPIKeyAuth() - + # Call endpoint and expect ProxyException with pytest.raises(ProxyException) as exc_info: await apply_guardrail(request=request, user_api_key_dict=mock_user_auth) - + # Verify error is properly handled assert "Bedrock guardrail failed" in str(exc_info.value.message) - @pytest.mark.asyncio -async def test_get_guardrail_info_endpoint_config_guardrail( - mocker, mock_admin_user_auth -): +async def test_get_guardrail_info_endpoint_config_guardrail(mocker): """ Test get_guardrail_info endpoint returns proper response when guardrail is found in config. """ @@ -1047,28 +945,21 @@ async def test_get_guardrail_info_endpoint_config_guardrail( # Mock the GUARDRAIL_REGISTRY to return None from DB (so it checks config) mock_registry = mocker.Mock() mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=None) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) # Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) # Mock _get_masked_values to return values as-is mocker.patch( "litellm.litellm_core_utils.litellm_logging._get_masked_values", - side_effect=lambda x, **kwargs: x, + side_effect=lambda x, **kwargs: x ) # Call endpoint and expect GuardrailInfoResponse - result = await get_guardrail_info( - guardrail_id="test-config-guardrail", user_api_key_dict=mock_admin_user_auth - ) + result = await get_guardrail_info(guardrail_id="test-config-guardrail") # Verify the response is of the correct type assert isinstance(result, GuardrailInfoResponse) @@ -1076,9 +967,8 @@ async def test_get_guardrail_info_endpoint_config_guardrail( assert result.guardrail_name == "Test Config Guardrail" assert result.guardrail_definition_location == "config" - @pytest.mark.asyncio -async def test_get_guardrail_info_endpoint_db_guardrail(mocker, mock_admin_user_auth): +async def test_get_guardrail_info_endpoint_db_guardrail(mocker): """ Test get_guardrail_info endpoint returns proper response when guardrail is found in DB. """ @@ -1090,28 +980,19 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker, mock_admin_user_ # Mock the GUARDRAIL_REGISTRY to return a guardrail from DB mock_registry = mocker.Mock() - mock_registry.get_guardrail_by_id_from_db = AsyncMock( - return_value=MOCK_DB_GUARDRAIL - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) # Mock IN_MEMORY_GUARDRAIL_HANDLER to return None mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.get_guardrail_by_id.return_value = None - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) # Call endpoint and expect GuardrailInfoResponse - result = await get_guardrail_info( - guardrail_id="test-db-guardrail", user_api_key_dict=mock_admin_user_auth - ) + result = await get_guardrail_info(guardrail_id="test-db-guardrail") # Verify the response is of the correct type assert isinstance(result, GuardrailInfoResponse) assert result.guardrail_id == "test-db-guardrail" assert result.guardrail_name == "Test DB Guardrail" - assert result.guardrail_definition_location == "db" + assert result.guardrail_definition_location == "db" \ No newline at end of file diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py b/tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py deleted file mode 100644 index e09a9ae2f0..0000000000 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py +++ /dev/null @@ -1,295 +0,0 @@ -import pytest -import sys -import os -from unittest.mock import AsyncMock -from fastapi import HTTPException - -# Add repo root to path -sys.path.insert(0, os.path.abspath("../../../../..")) - -from litellm.proxy.guardrails.guardrail_endpoints import ( - create_guardrail, - update_guardrail, - list_guardrails, -) -from litellm.types.guardrails import ( - CreateGuardrailRequest, - UpdateGuardrailRequest, -) - -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - - -# Fixtures -@pytest.fixture -def mock_team_user_auth(): - return UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, team_id="team-123") - - -@pytest.fixture -def mock_other_team_user_auth(): - return UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, team_id="team-456") - - -@pytest.fixture -def mock_proxy_admin_auth(): - return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, team_id="team-123") - - -@pytest.fixture -def mock_prisma_client(mocker): - return mocker.Mock() - - -@pytest.fixture -def mock_in_memory_handler(mocker): - mock = mocker.Mock() - mock.get_guardrail_by_id.return_value = None - mock.list_in_memory_guardrails.return_value = [] - return mock - - -def _make_mock_team(mocker, allow_team_guardrail_config: bool): - """Create a mock team object for get_team_object. Used by guardrail permission checks.""" - team = mocker.Mock() - team.allow_team_guardrail_config = allow_team_guardrail_config - return team - - -@pytest.mark.asyncio -async def test_team_list_guardrails_v2_isolation( - mocker, - mock_prisma_client, - mock_in_memory_handler, - mock_team_user_auth, - mock_other_team_user_auth, -): - """Test that teams only see their own guardrails in list endpoint""" - # Setup mocks - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory_handler, - ) - - async def mock_get_all_guardrails(prisma_client=None, team_id=None): - all_guardrails = [ - { - "guardrail_id": "g1", - "guardrail_name": "G1", - "team_id": "team-123", - "guardrail_config": {}, - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - }, - { - "guardrail_id": "g2", - "guardrail_name": "G2", - "team_id": "team-456", - "guardrail_config": {}, - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - }, - { - "guardrail_id": "g3", - "guardrail_name": "G3", - "team_id": None, - "guardrail_config": {}, - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - }, - ] - if team_id: - return [g for g in all_guardrails if g["team_id"] == team_id] - return all_guardrails - - mock_registry = mocker.Mock() - mock_registry.get_all_guardrails_from_db = AsyncMock( - side_effect=mock_get_all_guardrails - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) - - # Test for Team 123 - response_123 = await list_guardrails(user_api_key_dict=mock_team_user_auth) - assert len(response_123.guardrails) == 1 - assert response_123.guardrails[0].guardrail_id == "g1" - - # Test for Team 456 - response_456 = await list_guardrails(user_api_key_dict=mock_other_team_user_auth) - assert len(response_456.guardrails) == 1 - assert response_456.guardrails[0].guardrail_id == "g2" - - -@pytest.mark.asyncio -async def test_team_create_guardrail_sets_team_id( - mocker, mock_prisma_client, mock_team_user_auth -): - """Test that creating a guardrail as a team user sets the team_id when team has allow_team_guardrail_config.""" - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Team user must have allow_team_guardrail_config=True to create guardrails - mock_team = _make_mock_team(mocker, allow_team_guardrail_config=True) - mocker.patch( - "litellm.proxy.auth.auth_checks.get_team_object", - AsyncMock(return_value=mock_team), - ) - - mock_registry = mocker.Mock() - mock_registry.get_guardrail_by_name_from_db = AsyncMock(return_value=None) - - async def mock_add_guardrail(guardrail, team_id=None, **kwargs): - return { - "guardrail_id": "new-g", - "guardrail_name": guardrail["guardrail_name"], - "team_id": team_id, - "created_at": "2024-01-01T00:00:00.000Z", - "updated_at": "2024-01-01T00:00:00.000Z", - } - - mock_registry.add_guardrail_to_db = AsyncMock(side_effect=mock_add_guardrail) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) - - # Mock initialize_guardrail which is called after success - # NOTE: The endpoint calls initialize_guardrail on IN_MEMORY_GUARDRAIL_HANDLER imported from guardrail_registry - mock_in_memory = mocker.Mock() - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory, - ) - - request = CreateGuardrailRequest( - guardrail={ - "guardrail_name": "New Team Guardrail", - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - "guardrail_info": {}, - } - ) - - response = await create_guardrail( - request=request, user_api_key_dict=mock_team_user_auth - ) - - mock_registry.add_guardrail_to_db.assert_called_once() - call_kwargs = mock_registry.add_guardrail_to_db.call_args[1] - assert call_kwargs["team_id"] == "team-123" - assert response["team_id"] == "team-123" - - -@pytest.mark.asyncio -async def test_team_update_other_team_guardrail_fails( - mocker, mock_prisma_client, mock_team_user_auth -): - """Test that a team user cannot update another team's guardrail""" - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - mock_registry = mocker.Mock() - - async def mock_get_by_id(guardrail_id, **kwargs): - g = { - "guardrail_id": "g2", - "guardrail_name": "G2", - "team_id": "team-456", - "guardrail_config": {}, - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - } - if g["guardrail_id"] == guardrail_id: - return g - return None - - mock_registry.get_guardrail_by_id_from_db = AsyncMock(side_effect=mock_get_by_id) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) - - request = UpdateGuardrailRequest( - guardrail={ - "guardrail_name": "Updated Name", - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - } - ) - - # Team 123 tries to update G2 (owned by 456) - with pytest.raises(HTTPException) as excinfo: - await update_guardrail( - guardrail_id="g2", request=request, user_api_key_dict=mock_team_user_auth - ) - - assert excinfo.value.status_code == 403 - - -# ----- Regression tests for allow_team_guardrail_config ----- - - -@pytest.mark.asyncio -async def test_team_create_guardrail_forbidden_when_allow_team_guardrail_config_false( - mocker, mock_prisma_client, mock_team_user_auth -): - """Team user gets 403 when team has allow_team_guardrail_config=False.""" - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - mock_team = _make_mock_team(mocker, allow_team_guardrail_config=False) - mocker.patch( - "litellm.proxy.auth.auth_checks.get_team_object", - AsyncMock(return_value=mock_team), - ) - - request = CreateGuardrailRequest( - guardrail={ - "guardrail_name": "Forbidden Guardrail", - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - "guardrail_info": {}, - } - ) - - with pytest.raises(HTTPException) as excinfo: - await create_guardrail(request=request, user_api_key_dict=mock_team_user_auth) - - assert excinfo.value.status_code == 403 - assert "Guardrail configuration is not enabled" in str(excinfo.value.detail) - - -@pytest.mark.asyncio -async def test_proxy_admin_can_create_guardrail_without_team_guardrail_config_permission( - mocker, mock_prisma_client, mock_proxy_admin_auth -): - """Proxy admin can create guardrails without team allow_team_guardrail_config (check is skipped).""" - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - mock_registry = mocker.Mock() - mock_registry.get_guardrail_by_name_from_db = AsyncMock(return_value=None) - mock_registry.add_guardrail_to_db = AsyncMock( - return_value={ - "guardrail_id": "admin-g", - "guardrail_name": "Admin Guardrail", - "team_id": "team-123", - "created_at": "2024-01-01T00:00:00.000Z", - "updated_at": "2024-01-01T00:00:00.000Z", - } - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", - mock_registry, - ) - mock_in_memory = mocker.Mock() - mocker.patch( - "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", - mock_in_memory, - ) - - # Do NOT mock get_team_object: proxy admin skips _check_team_can_configure_guardrails - request = CreateGuardrailRequest( - guardrail={ - "guardrail_name": "Admin Guardrail", - "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, - "guardrail_info": {}, - } - ) - - response = await create_guardrail( - request=request, user_api_key_dict=mock_proxy_admin_auth - ) - - mock_registry.add_guardrail_to_db.assert_called_once() - assert response["guardrail_id"] == "admin-g" - assert response["team_id"] == "team-123" diff --git a/ui/litellm-dashboard/src/components/team/team_info.tsx b/ui/litellm-dashboard/src/components/team/team_info.tsx index e5bbe4dc48..34ff903c86 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.tsx @@ -96,7 +96,6 @@ export interface TeamData { model_aliases: Record; } | null; created_at: string; - allow_team_guardrail_config?: boolean; guardrails?: string[]; policies?: string[]; object_permission?: { @@ -468,7 +467,6 @@ const TeamInfoView: React.FC = ({ }, policies: values.policies || [], organization_id: values.organization_id, - ...(values.allow_team_guardrail_config !== undefined ? { allow_team_guardrail_config: values.allow_team_guardrail_config } : {}), }; updateData.max_budget = mapEmptyStringToNull(updateData.max_budget); @@ -660,15 +658,6 @@ const TeamInfoView: React.FC = ({ Guardrails - {info.allow_team_guardrail_config === true ? ( -
- Team can configure guardrails -
- ) : ( -
- Only proxy admin can configure guardrails for this team -
- )} {info.guardrails && info.guardrails.length > 0 ? (
{info.guardrails.map((guardrail: string, index: number) => ( @@ -772,7 +761,6 @@ const TeamInfoView: React.FC = ({ team_member_budget_duration: info.team_member_budget_table?.budget_duration, guardrails: info.metadata?.guardrails || [], policies: info.policies || [], - allow_team_guardrail_config: info.allow_team_guardrail_config ?? false, disable_global_guardrails: info.metadata?.disable_global_guardrails || false, metadata: info.metadata ? JSON.stringify( @@ -888,72 +876,29 @@ const TeamInfoView: React.FC = ({ - - prev.allow_team_guardrail_config !== cur.allow_team_guardrail_config - } - > - {() => { - const allowTeamGuardrailConfig = - form.getFieldValue("allow_team_guardrail_config"); - const guardrailsDisabled = - !is_proxy_admin && !allowTeamGuardrailConfig; - return ( - - Guardrails{" "} - - e.stopPropagation()} - > - - - - - } - name="guardrails" - help="Select existing guardrails or enter new ones" - > -