From 76a399ba69239899a1ec2cff0a9e235f75944889 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 3 Feb 2026 08:44:55 +0530 Subject: [PATCH 01/15] feat(guardrails): implement team-based isolation guardrails mgmnt (#19889) * feat(guardrails): implement team-based isolation guardrails mgmnt * fix lint errors * add allow_team_guardrail_config for admin permissions --- .../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, 1293 insertions(+), 618 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql create 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 new file mode 100644 index 00000000000..706a2f25e95 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql @@ -0,0 +1,8 @@ +-- 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 045d2fd5f14..e84eabde6be 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -358,7 +358,6 @@ 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,6 +629,9 @@ 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 ## @@ -1477,6 +1479,9 @@ 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 @@ -1574,6 +1579,9 @@ 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 4df773dec2b..3bd56177d71 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -34,6 +34,58 @@ 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. @@ -245,6 +297,11 @@ 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( @@ -271,7 +328,7 @@ async def authenticate_user( # noqa: PLR0915 "config": {}, "spend": 0, "user_id": user_id, - "team_id": "litellm-dashboard", + "team_id": user_team_id, }, ) else: @@ -340,4 +397,3 @@ 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 3ce819439cb..b80b02fdc6f 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 UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, 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,6 +32,8 @@ from litellm.types.guardrails import ( PresidioPresidioConfigModelUserInterface, SupportedGuardrailIntegrations, ToolPermissionGuardrailConfigModel, + CreateGuardrailRequest, + UpdateGuardrailRequest, ) #### GUARDRAILS ENDPOINTS #### @@ -40,6 +42,37 @@ 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: @@ -64,9 +97,19 @@ def _get_guardrails_list_response( dependencies=[Depends(user_api_key_auth)], response_model=ListGuardrailsResponse, ) -async def list_guardrails(): +@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), +): """ - List the guardrails that are available on the proxy server + List the guardrails that are available in the database using GuardrailRegistry + + Now supports both DB-persisted and In-Memory (Config) guardrails. 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) @@ -75,59 +118,6 @@ 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 { @@ -139,12 +129,14 @@ async def list_guardrails_v2(): "guardrail": "bedrock", "mode": "pre_call", "guardrailIdentifier": "ff6ujrregl1q", - "guardrailVersion": "DRAFT", + "guardrailVersion": "1.0", "default_on": true }, "guardrail_info": { - "description": "Bedrock content moderation guardrail" - } + "description": "Updated Bedrock content moderation guardrail" + }, + "created_at": "2023-11-09T12:34:56.789Z", + "updated_at": "2023-11-09T13:45:12.345Z" } ] } @@ -153,12 +145,29 @@ async def list_guardrails_v2(): 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: - raise HTTPException(status_code=500, detail="Prisma client not initialized") + # 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) 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 + prisma_client=prisma_client, team_id=filter_team_id ) guardrail_configs: List[GuardrailInfoResponse] = [] @@ -173,6 +182,7 @@ async def list_guardrails_v2(): 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")) @@ -180,6 +190,15 @@ async def list_guardrails_v2(): # get guardrails initialized on litellm config.yaml in_memory_guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails() for guardrail in in_memory_guardrails: + # 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( @@ -189,6 +208,7 @@ async def list_guardrails_v2(): 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")) @@ -199,16 +219,15 @@ async def list_guardrails_v2(): 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): +async def create_guardrail( + request: CreateGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Create a new guardrail @@ -257,14 +276,21 @@ async def create_guardrail(request: CreateGuardrailRequest): ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache 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 + guardrail=request.guardrail, prisma_client=prisma_client, team_id=team_id ) guardrail_name = result.get("guardrail_name", "Unknown") @@ -283,21 +309,23 @@ async def create_guardrail(request: CreateGuardrailRequest): ) 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): +async def update_guardrail( + guardrail_id: str, + request: UpdateGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Update an existing guardrail @@ -346,7 +374,7 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -362,6 +390,25 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): 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, @@ -394,7 +441,9 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def delete_guardrail(guardrail_id: str): +async def delete_guardrail( + guardrail_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth) +): """ Delete a guardrail @@ -414,7 +463,7 @@ async def delete_guardrail(guardrail_id: str): ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -430,6 +479,25 @@ async def delete_guardrail(guardrail_id: str): 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 ) @@ -460,7 +528,11 @@ async def delete_guardrail(guardrail_id: str): tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): +async def patch_guardrail( + guardrail_id: str, + request: PatchGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Partially update an existing guardrail @@ -507,7 +579,7 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): ``` """ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -523,6 +595,25 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): 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 @@ -594,7 +685,10 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): tags=["Guardrails"], dependencies=[Depends(user_api_key_auth)], ) -async def get_guardrail_info(guardrail_id: str): +async def get_guardrail_info( + guardrail_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Get detailed information about a specific guardrail by ID @@ -653,6 +747,19 @@ async def get_guardrail_info(guardrail_id: str): 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" ) @@ -675,6 +782,7 @@ async def get_guardrail_info(guardrail_id: str): 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 @@ -1209,9 +1317,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 = { @@ -1250,10 +1358,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 c3da6892209..487a9d481c9 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -233,7 +233,10 @@ class GuardrailRegistry: ########### DB management helpers for guardrails ########### ############################################################ async def add_guardrail_to_db( - self, guardrail: Guardrail, prisma_client: PrismaClient + self, + guardrail: Guardrail, + prisma_client: PrismaClient, + team_id: Optional[str] = None, ): """ Add a guardrail to the database @@ -248,6 +251,8 @@ 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", {})) @@ -257,6 +262,7 @@ 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), } @@ -265,18 +271,32 @@ 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 + self, + guardrail_id: str, + prisma_client: PrismaClient, + team_id: Optional[str] = None, ): """ 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} @@ -287,12 +307,26 @@ 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 + self, + guardrail_id: str, + guardrail: Guardrail, + prisma_client: PrismaClient, + team_id: Optional[str] = None, ): """ 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", {}) @@ -302,6 +336,7 @@ 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", {})) @@ -324,27 +359,67 @@ 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: - guardrails_from_db = ( - await prisma_client.db.litellm_guardrailstable.find_many( - order={"created_at": "desc"}, - ) + # Normal ORM Fetch + all_guardrails = 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: - guardrails.append(Guardrail(**(dict(guardrail)))) # type: ignore + 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 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 + self, + guardrail_id: str, + prisma_client: PrismaClient, + team_id: Optional[str] = None, ) -> Optional[Guardrail]: """ Get a guardrail by its ID from the database @@ -357,6 +432,11 @@ 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)}") @@ -473,6 +553,7 @@ class InMemoryGuardrailHandler: guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail["guardrail_name"], litellm_params=litellm_params, + team_id=guardrail.get("team_id"), ) # store references to the guardrail in memory diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 63db2d72fe4..43c32121949 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -753,12 +753,16 @@ 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 @@ -918,12 +922,16 @@ 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 ) @@ -1099,7 +1107,9 @@ 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, ) @@ -1107,7 +1117,9 @@ 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. @@ -1158,7 +1170,9 @@ 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 @@ -1204,7 +1218,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), @@ -1288,19 +1302,25 @@ 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( @@ -1367,6 +1387,22 @@ 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) @@ -1411,16 +1447,19 @@ 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: @@ -1429,7 +1468,9 @@ 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()), @@ -1771,113 +1812,6 @@ 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"], @@ -1953,27 +1887,16 @@ async def team_member_add( complete_team_data=complete_team_data, ) - # 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, - ) + ( + 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 @@ -2152,15 +2075,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, @@ -2539,10 +2462,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 @@ -2601,10 +2524,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: @@ -2643,7 +2566,6 @@ async def delete_team( return deleted_teams - def _transform_teams_to_deleted_records( teams: List[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, @@ -2666,7 +2588,13 @@ 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]) @@ -2685,9 +2613,7 @@ 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( @@ -2707,6 +2633,7 @@ async def _persist_deleted_team_records( prisma_client=prisma_client, ) + def validate_membership( user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable ): @@ -2830,11 +2757,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 @@ -3297,7 +3224,9 @@ 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 3b81da10923..cc33b9bed8a 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -128,6 +128,7 @@ 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]) @@ -159,8 +160,9 @@ 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") @@ -784,6 +786,7 @@ 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 ca22049720e..ed463491043 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -730,6 +730,7 @@ 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] @@ -761,6 +762,7 @@ class GuardrailInfoResponse(BaseModel): guardrail_name: str litellm_params: Optional[BaseLitellmParams] = None guardrail_info: Optional[Dict] = None + team_id: Optional[str] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = ( @@ -807,3 +809,11 @@ 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 3b81da10923..3886f4db199 100644 --- a/schema.prisma +++ b/schema.prisma @@ -129,6 +129,7 @@ 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]) @@ -160,7 +161,8 @@ 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") @@ -784,6 +786,7 @@ 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 88f56c24067..62439327278 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1,15 +1,14 @@ -import json import os import sys from datetime import datetime -from typing import Dict, List, Optional +from typing import List from unittest.mock import AsyncMock import pytest sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path + 0, os.path.abspath("../../../../..") +) # Adds the repo root directory to the system path from fastapi import HTTPException @@ -21,12 +20,12 @@ from litellm.proxy.guardrails.guardrail_endpoints import ( create_guardrail, delete_guardrail, get_guardrail_info, - list_guardrails_v2, + list_guardrails, 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 ( @@ -63,7 +62,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) @@ -71,7 +70,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"}, ) @@ -102,22 +101,35 @@ 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 + mocker, mock_prisma_client, mock_in_memory_handler, mock_admin_user_auth ): """Test listing guardrails from both DB and config""" # Mock the prisma client @@ -128,7 +140,7 @@ async def test_list_guardrails_v2_with_db_and_config( mock_in_memory_handler, ) - response = await list_guardrails_v2() + response = await list_guardrails(user_api_key_dict=mock_admin_user_auth) assert len(response.guardrails) == 2 @@ -150,11 +162,17 @@ 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): +async def test_get_guardrail_info_from_db( + mocker, mock_prisma_client, mock_admin_user_auth +): """Test getting guardrail info from DB""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - response = await get_guardrail_info("test-db-guardrail") + 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 + ) assert response.guardrail_id == "test-db-guardrail" assert response.guardrail_name == "Test DB Guardrail" @@ -164,7 +182,7 @@ async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): @pytest.mark.asyncio async def test_get_guardrail_info_from_config( - mocker, mock_prisma_client, mock_in_memory_handler + mocker, mock_prisma_client, mock_in_memory_handler, mock_admin_user_auth ): """Test getting guardrail info from config when not found in DB""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -178,7 +196,9 @@ async def test_get_guardrail_info_from_config( return_value=None ) - response = await get_guardrail_info("test-config-guardrail") + response = await get_guardrail_info( + "test-config-guardrail", user_api_key_dict=mock_admin_user_auth + ) assert response.guardrail_id == "test-config-guardrail" assert response.guardrail_name == "Test Config Guardrail" @@ -188,7 +208,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 + mocker, mock_prisma_client, mock_in_memory_handler, mock_admin_user_auth ): """Test getting guardrail info when not found in either DB or config""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -204,7 +224,9 @@ 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") + await get_guardrail_info( + "non-existent-guardrail", user_api_key_dict=mock_admin_user_auth + ) assert exc_info.value.status_code == 404 assert "not found" in str(exc_info.value.detail) @@ -222,7 +244,6 @@ 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) @@ -238,12 +259,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"] == False - assert fields["api_key"]["type"] == "string" # Should be string, not None + assert fields["api_key"]["required"] is False + assert fields["api_key"]["type"] == "string" # Check the structure of the nested optional_params field assert fields["optional_params"]["type"] == "nested" - assert fields["optional_params"]["required"] == True + assert fields["optional_params"]["required"] is True assert "fields" in fields["optional_params"] # Check nested fields within optional_params @@ -260,7 +281,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"] == False + assert nested_fields["severity_threshold"]["required"] is False assert ( nested_fields["severity_threshold"]["type"] == "number" ) # Should be number, not None @@ -269,16 +290,14 @@ 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" - ) # Literal type should be select + assert nested_fields["outputType"]["type"] == "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 BaseModel, Field + from pydantic import Field from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -299,7 +318,6 @@ 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 @@ -337,14 +355,13 @@ 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, patch + from unittest.mock import Mock from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, @@ -352,27 +369,23 @@ 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 @@ -389,45 +402,44 @@ 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 - prepared_request = guardrail_hook._prepare_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() @@ -442,34 +454,30 @@ 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 - - prepared_request = guardrail_hook._prepare_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 @@ -485,45 +493,51 @@ 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 @@ -531,336 +545,419 @@ 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 + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, + mock_admin_user_auth, ): """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) - + await create_guardrail( + MOCK_CREATE_REQUEST, user_api_key_dict=mock_admin_user_auth + ) + 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) - + result = await create_guardrail( + MOCK_CREATE_REQUEST, user_api_key_dict=mock_admin_user_auth + ) + 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 + prisma_client=mocker.ANY, + team_id=None, ) - + 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 + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, + mock_admin_user_auth, ): """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) - + await update_guardrail( + "test-guardrail-id", + MOCK_UPDATE_REQUEST, + user_api_key_dict=mock_admin_user_auth, + ) + 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) - + result = await update_guardrail( + "test-guardrail-id", + MOCK_UPDATE_REQUEST, + user_api_key_dict=mock_admin_user_auth, + ) + 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 + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, + mock_admin_user_auth, ): """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) - + await patch_guardrail( + "test-guardrail-id", + MOCK_PATCH_REQUEST, + user_api_key_dict=mock_admin_user_auth, + ) + 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) - + result = await patch_guardrail( + "test-guardrail-id", + MOCK_PATCH_REQUEST, + user_api_key_dict=mock_admin_user_auth, + ) + 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 + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, + mock_admin_user_auth, ): """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) + await delete_guardrail( + guardrail_id=expected_result, user_api_key_dict=mock_admin_user_auth + ) else: - result = await delete_guardrail(guardrail_id=expected_result) - + result = await delete_guardrail( + guardrail_id=expected_result, user_api_key_dict=mock_admin_user_auth + ) + 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() @@ -877,21 +974,22 @@ 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() @@ -909,30 +1007,34 @@ 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): +async def test_get_guardrail_info_endpoint_config_guardrail( + mocker, mock_admin_user_auth +): """ Test get_guardrail_info endpoint returns proper response when guardrail is found in config. """ @@ -945,21 +1047,28 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker): # 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") + result = await get_guardrail_info( + guardrail_id="test-config-guardrail", user_api_key_dict=mock_admin_user_auth + ) # Verify the response is of the correct type assert isinstance(result, GuardrailInfoResponse) @@ -967,8 +1076,9 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker): 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): +async def test_get_guardrail_info_endpoint_db_guardrail(mocker, mock_admin_user_auth): """ Test get_guardrail_info endpoint returns proper response when guardrail is found in DB. """ @@ -980,19 +1090,28 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker): # 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") + result = await get_guardrail_info( + guardrail_id="test-db-guardrail", user_api_key_dict=mock_admin_user_auth + ) # 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" \ No newline at end of file + assert result.guardrail_definition_location == "db" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py b/tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py new file mode 100644 index 00000000000..e09a9ae2f0f --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_team_access.py @@ -0,0 +1,295 @@ +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 193d056fdd4..8996fee6637 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.tsx @@ -96,6 +96,7 @@ export interface TeamData { model_aliases: Record; } | null; created_at: string; + allow_team_guardrail_config?: boolean; guardrails?: string[]; policies?: string[]; object_permission?: { @@ -466,6 +467,7 @@ 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); @@ -658,6 +660,15 @@ 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) => ( @@ -761,6 +772,7 @@ 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( @@ -876,29 +888,72 @@ 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" + > + ({ value: name, label: name }))} + From 04c348e7bb1a65078052fea4085ebeb586d5f5fb Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 3 Feb 2026 08:50:14 +0530 Subject: [PATCH 02/15] fixes failure metrics labels (#20152) Co-authored-by: Krish Dholakia --- litellm/integrations/prometheus.py | 169 ++++++++++++++++-- .../test_prometheus_logging_callbacks.py | 144 ++++++++------- 2 files changed, 230 insertions(+), 83 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2c897cb0692..00c38eac188 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1683,6 +1683,108 @@ class PrometheusLogger(CustomLogger): ) pass + def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any: + """Get value from dict or Pydantic model.""" + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + def _extract_deployment_failure_label_values( + self, request_kwargs: dict + ) -> Dict[str, Optional[str]]: + """ + Extract label values for deployment failure metrics from all available + sources in request_kwargs. Falls back to litellm_params metadata and + user_api_key_auth when standard_logging_payload has None values. + """ + standard_logging_payload = ( + request_kwargs.get("standard_logging_object", {}) or {} + ) + _litellm_params = request_kwargs.get("litellm_params", {}) or {} + _metadata_raw = self._safe_get(standard_logging_payload, "metadata") or {} + if isinstance(_metadata_raw, dict): + _metadata = _metadata_raw + else: + _metadata = { + "user_api_key_alias": getattr( + _metadata_raw, "user_api_key_alias", None + ), + "user_api_key_team_id": getattr( + _metadata_raw, "user_api_key_team_id", None + ), + "user_api_key_team_alias": getattr( + _metadata_raw, "user_api_key_team_alias", None + ), + "user_api_key_hash": getattr(_metadata_raw, "user_api_key_hash", None), + "requester_ip_address": getattr( + _metadata_raw, "requester_ip_address", None + ), + "user_agent": getattr(_metadata_raw, "user_agent", None), + } + _litellm_params_metadata = _litellm_params.get("metadata", {}) or {} + + # Extract user_api_key_auth if present (proxy injects this, skipped in merge) + user_api_key_auth = _litellm_params_metadata.get("user_api_key_auth") + + def _get_api_key_alias() -> Optional[str]: + val = _metadata.get("user_api_key_alias") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_alias") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "key_alias", None) + return None + + def _get_team_id() -> Optional[str]: + val = _metadata.get("user_api_key_team_id") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_team_id") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "team_id", None) + return None + + def _get_team_alias() -> Optional[str]: + val = _metadata.get("user_api_key_team_alias") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_team_alias") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "team_alias", None) + return None + + def _get_hashed_api_key() -> Optional[str]: + val = _metadata.get("user_api_key_hash") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_hash") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "api_key", None) or getattr( + user_api_key_auth, "api_key_hash", None + ) + return None + + return { + "api_key_alias": _get_api_key_alias(), + "team": _get_team_id(), + "team_alias": _get_team_alias(), + "hashed_api_key": _get_hashed_api_key(), + "client_ip": _metadata.get("requester_ip_address") + or _litellm_params_metadata.get("requester_ip_address"), + "user_agent": _metadata.get("user_agent") + or _litellm_params_metadata.get("user_agent"), + } + def set_llm_deployment_failure_metrics(self, request_kwargs: dict): """ Sets Failure metrics when an LLM API call fails @@ -1707,6 +1809,21 @@ class PrometheusLogger(CustomLogger): model_id = standard_logging_payload.get("model_id", None) exception = request_kwargs.get("exception", None) + # Fallback: model_id from litellm_metadata.model_info + if model_id is None: + _model_info = ( + (_litellm_params.get("litellm_metadata") or {}).get("model_info") + or (_litellm_params.get("metadata") or {}).get("model_info") + or {} + ) + model_id = _model_info.get("id") + + # Fallback: model_group from litellm_metadata + if model_group is None: + model_group = (_litellm_params.get("litellm_metadata") or {}).get( + "model_group" + ) or (_litellm_params.get("metadata") or {}).get("model_group") + llm_provider = _litellm_params.get("custom_llm_provider", None) if self._should_skip_metrics_for_invalid_key( @@ -1714,9 +1831,37 @@ class PrometheusLogger(CustomLogger): standard_logging_payload=standard_logging_payload, ): return - hashed_api_key = standard_logging_payload.get("metadata", {}).get( + + # Extract context labels from all available sources (fix for None labels) + fallback_values = self._extract_deployment_failure_label_values( + request_kwargs + ) + _metadata = standard_logging_payload.get("metadata", {}) or {} + hashed_api_key = fallback_values.get("hashed_api_key") or _metadata.get( "user_api_key_hash" ) + api_key_alias = fallback_values.get("api_key_alias") or _metadata.get( + "user_api_key_alias" + ) + team = fallback_values.get("team") or _metadata.get("user_api_key_team_id") + team_alias = fallback_values.get("team_alias") or _metadata.get( + "user_api_key_team_alias" + ) + client_ip = fallback_values.get("client_ip") or _metadata.get( + "requester_ip_address" + ) + user_agent = fallback_values.get("user_agent") or _metadata.get( + "user_agent" + ) + + # exception_status: prefer status_code, fallback to exception class for known types + exception_status = None + if exception is not None: + exception_status = str(getattr(exception, "status_code", None)) + if exception_status == "None" or not exception_status: + code = getattr(exception, "code", None) + if code is not None: + exception_status = str(code) # Create enum_values for the label factory (always create for use in different metrics) enum_values = UserAPIKeyLabelValues( @@ -1724,26 +1869,18 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, api_provider=llm_provider, - exception_status=( - str(getattr(exception, "status_code", None)) if exception else None - ), + exception_status=exception_status, exception_class=( self._get_exception_class_name(exception) if exception else None ), - requested_model=model_group, + requested_model=model_group or litellm_model_name, hashed_api_key=hashed_api_key, - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], + api_key_alias=api_key_alias, + team=team, + team_alias=team_alias, tags=standard_logging_payload.get("request_tags", []), - client_ip=standard_logging_payload["metadata"].get( - "requester_ip_address" - ), - user_agent=standard_logging_payload["metadata"].get("user_agent"), + client_ip=client_ip, + user_agent=user_agent, ) """ diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 0a57d046c72..c39454728a8 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -1,4 +1,3 @@ -import io import os import sys @@ -10,13 +9,10 @@ from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, call, patch import pytest -from prometheus_client import REGISTRY, CollectorRegistry +from prometheus_client import REGISTRY import litellm -from litellm import completion from litellm._logging import verbose_logger -from litellm._uuid import uuid -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -37,7 +33,6 @@ from litellm.proxy._types import UserAPIKeyAuth verbose_logger.setLevel(logging.DEBUG) litellm.set_verbose = True -import time @pytest.fixture @@ -293,7 +288,6 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): ) as mock_get_team, patch( "litellm.proxy.auth.auth_checks.get_key_object" ) as mock_get_key: - mock_get_team.return_value = MagicMock(budget_reset_at=future_reset_time_team) mock_get_key.return_value = MagicMock(budget_reset_at=future_reset_time_key) @@ -648,25 +642,16 @@ async def test_async_log_failure_event(prometheus_logger): ) # litellm_llm_api_failed_requests_metric incremented - """ - Expected metrics - end_user_id, - user_api_key, - user_api_key_alias, - model, - user_api_team, - user_api_team_alias, - user_id, - """ + # Labels: end_user, api_key_hash, api_key_alias, model, team, team_alias, user, model_id prometheus_logger.litellm_llm_api_failed_requests_metric.labels.assert_called_once_with( - None, + None, # end_user_id "test_hash", "test_alias", "gpt-3.5-turbo", "test_team", "test_team_alias", "test_user", - "model-123", + "model-123", # model_id from standard_logging_payload ) prometheus_logger.litellm_llm_api_failed_requests_metric.labels().inc.assert_called_once() @@ -678,38 +663,54 @@ async def test_async_log_failure_event(prometheus_logger): api_provider="openai", ) - # deployment failure responses incremented - prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once_with( - litellm_model_name="gpt-3.5-turbo", - model_id="model-123", - api_base="https://api.openai.com", - api_provider="openai", - exception_status="None", - exception_class="Exception", - requested_model="openai-gpt", # passed in standard logging payload - hashed_api_key="test_hash", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - client_ip="127.0.0.1", # from standard logging payload - user_agent=None, + # deployment failure responses incremented - verify key labels are populated + prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once() + actual_failure_labels = ( + prometheus_logger.litellm_deployment_failure_responses.labels.call_args.kwargs ) + expected_failure_labels = { + "litellm_model_name": "gpt-3.5-turbo", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "api_provider": "openai", + "exception_class": "Exception", + "requested_model": "openai-gpt", + "hashed_api_key": "test_hash", + "api_key_alias": "test_alias", + "team": "test_team", + "team_alias": "test_team_alias", + } + for key, expected_val in expected_failure_labels.items(): + assert key in actual_failure_labels, f"Missing label {key}" + assert ( + actual_failure_labels[key] == expected_val + ), f"Label {key}: expected {expected_val!r}, got {actual_failure_labels[key]!r}" + assert actual_failure_labels.get("exception_status") in ("None", None) + assert actual_failure_labels.get("client_ip") == "127.0.0.1" prometheus_logger.litellm_deployment_failure_responses.labels().inc.assert_called_once() - # deployment total requests incremented - prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once_with( - litellm_model_name="gpt-3.5-turbo", - model_id="model-123", - api_base="https://api.openai.com", - api_provider="openai", - requested_model="openai-gpt", # passed in standard logging payload - hashed_api_key="test_hash", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - client_ip="127.0.0.1", # from standard logging payload - user_agent=None, + # deployment total requests incremented - verify key labels are populated + prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once() + actual_total_labels = ( + prometheus_logger.litellm_deployment_total_requests.labels.call_args.kwargs ) + expected_total_labels = { + "litellm_model_name": "gpt-3.5-turbo", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "api_provider": "openai", + "requested_model": "openai-gpt", + "hashed_api_key": "test_hash", + "api_key_alias": "test_alias", + "team": "test_team", + "team_alias": "test_team_alias", + } + for key, expected_val in expected_total_labels.items(): + assert key in actual_total_labels, f"Missing label {key}" + assert ( + actual_total_labels[key] == expected_val + ), f"Label {key}: expected {expected_val!r}, got {actual_total_labels[key]!r}" + assert actual_total_labels.get("client_ip") == "127.0.0.1" prometheus_logger.litellm_deployment_total_requests.labels().inc.assert_called_once() @@ -1095,7 +1096,7 @@ def test_increment_deployment_cooled_down(prometheus_logger): import inspect method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down) - expected_label_count = len([p for p in method_sig.parameters.keys() if p != 'self']) + expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"]) mock_chain = MagicMock() @@ -1103,11 +1104,15 @@ def test_increment_deployment_cooled_down(prometheus_logger): """Validate label count matches metric definition""" total = len(label_values) + len(label_kwargs) if total != expected_label_count: - raise ValueError(f"Incorrect label count: expected {expected_label_count}, got {total}") + raise ValueError( + f"Incorrect label count: expected {expected_label_count}, got {total}" + ) return mock_chain prometheus_logger.litellm_deployment_cooled_down = MagicMock() - prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(side_effect=validating_labels) + prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock( + side_effect=validating_labels + ) prometheus_logger.increment_deployment_cooled_down( litellm_model_name="gpt-3.5-turbo", @@ -1179,8 +1184,12 @@ def test_get_custom_labels_from_top_level_metadata(monkeypatch): metadata = { "requester_ip_address": "10.48.203.20", # Top-level field "user_api_key_alias": "TestAlias", # Top-level field - "requester_metadata": {"nested_field": "nested_value"}, # Nested dict (excluded) - "user_api_key_auth_metadata": {"another_nested": "value"}, # Nested dict (excluded) + "requester_metadata": { + "nested_field": "nested_value" + }, # Nested dict (excluded) + "user_api_key_auth_metadata": { + "another_nested": "value" + }, # Nested dict (excluded) } result = get_custom_labels_from_metadata(metadata) assert result == { @@ -1217,7 +1226,9 @@ def test_get_custom_labels_from_top_level_and_nested_metadata(monkeypatch): } -async def test_async_log_success_event_with_top_level_metadata(prometheus_logger, monkeypatch): +async def test_async_log_success_event_with_top_level_metadata( + prometheus_logger, monkeypatch +): """ Test that async_log_success_event correctly extracts custom labels from top-level metadata fields like requester_ip_address, not just from nested dictionaries. @@ -1231,7 +1242,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger standard_logging_object = create_standard_logging_payload() standard_logging_object["metadata"]["requester_ip_address"] = "10.48.203.20" standard_logging_object["metadata"]["requester_metadata"] = {} # Empty nested dict - standard_logging_object["metadata"]["user_api_key_auth_metadata"] = {} # Empty nested dict + standard_logging_object["metadata"][ + "user_api_key_auth_metadata" + ] = {} # Empty nested dict kwargs = { "model": "gpt-3.5-turbo", @@ -1273,7 +1286,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger prometheus_logger.litellm_remaining_user_budget_metric = create_mock_metric() prometheus_logger.litellm_user_max_budget_metric = create_mock_metric() prometheus_logger.litellm_user_budget_remaining_hours_metric = create_mock_metric() - prometheus_logger.litellm_remaining_api_key_requests_for_model = create_mock_metric() + prometheus_logger.litellm_remaining_api_key_requests_for_model = ( + create_mock_metric() + ) prometheus_logger.litellm_remaining_api_key_tokens_for_model = create_mock_metric() prometheus_logger.litellm_llm_api_time_to_first_token_metric = create_mock_metric() prometheus_logger.litellm_llm_api_latency_metric = create_mock_metric() @@ -1302,7 +1317,7 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger # This confirms that the custom label extraction logic ran without errors assert prometheus_logger.litellm_requests_metric.labels.called assert prometheus_logger.litellm_spend_metric.labels.called - + # Verify that the labels() method was called with some arguments (either positional or keyword) # This ensures the custom label extraction happened and didn't cause a "Incorrect label names" error call_args = prometheus_logger.litellm_requests_metric.labels.call_args @@ -1494,7 +1509,6 @@ async def test_initialize_remaining_budget_metrics(prometheus_logger): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.team_endpoints.get_paginated_teams" ) as mock_get_teams: - # Create mock team data with proper datetime objects for budget_reset_at future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now mock_teams = [ @@ -1592,21 +1606,22 @@ async def test_initialize_remaining_budget_metrics_exception_handling( ) as mock_get_teams, patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: - # Make get_paginated_teams raise an exception mock_get_teams.side_effect = Exception("Database error") mock_list_keys.side_effect = Exception("Key listing error") - + # Mock prisma_client structure to raise an exception for user budget metrics # The code accesses prisma_client.db.litellm_usertable.find_many and count mock_usertable = MagicMock() - mock_usertable.find_many = MagicMock(side_effect=Exception("User database error")) + mock_usertable.find_many = MagicMock( + side_effect=Exception("User database error") + ) mock_usertable.count = MagicMock(side_effect=Exception("User count error")) - + # Mock litellm_teamtable to raise an exception for team count metrics mock_teamtable = MagicMock() mock_teamtable.count = MagicMock(side_effect=Exception("Team count error")) - + mock_db = MagicMock() mock_db.litellm_usertable = mock_usertable mock_db.litellm_teamtable = mock_teamtable @@ -1661,7 +1676,6 @@ async def test_initialize_api_key_budget_metrics(prometheus_logger): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: - # Create mock key data with proper datetime objects for budget_reset_at future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now key1 = UserAPIKeyAuth( @@ -1916,7 +1930,6 @@ def test_prometheus_label_factory_with_custom_tags(monkeypatch): Test that prometheus_label_factory correctly handles custom tags """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -1954,7 +1967,6 @@ def test_prometheus_label_factory_with_no_custom_tags(monkeypatch): Test that prometheus_label_factory works when no custom tags are configured """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -2179,9 +2191,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): All three metrics should be properly incremented when making a successful completion request. """ - from prometheus_client import CollectorRegistry, Counter - import litellm from litellm.types.integrations.prometheus import PrometheusMetricsConfig # Clear registry before test From ec279eb426d26e74c14b92101c61851537847579 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 3 Feb 2026 08:56:50 +0530 Subject: [PATCH 03/15] fix: proxy failure cases, now log ip and user agent, key hash, name (#20145) --- litellm/integrations/prometheus.py | 30 +--- litellm/proxy/auth/auth_checks.py | 8 +- litellm/proxy/auth/auth_exception_handler.py | 65 ++++++-- litellm/proxy/auth/auth_utils.py | 145 ++++++++++++------ litellm/proxy/auth/user_api_key_auth.py | 14 +- litellm/proxy/utils.py | 10 +- .../test_prometheus_invalid_key_filtering.py | 134 ++++++++++------ 7 files changed, 272 insertions(+), 134 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 00c38eac188..d751e12460b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1478,35 +1478,19 @@ class PrometheusLogger(CustomLogger): """ Determine if a request has an invalid API key based on status code and exception. - This method prevents invalid authentication attempts from being recorded in - Prometheus metrics. A 401 status code is the definitive indicator of authentication - failure. Additionally, we check exception messages for authentication error patterns - to catch cases where the exception hasn't been converted to a ProxyException yet. + Returns True only when we truly cannot record useful metrics (e.g. missing required + data). We no longer skip 401/invalid-key requests - all requests including + authentication failures and bad requests must be tracked for debugging, security + auditing, abuse detection, and capacity planning. Args: - status_code: HTTP status code (401 indicates authentication error) - exception: Exception object to check for auth-related error messages + status_code: HTTP status code + exception: Exception object (unused, kept for API compatibility) Returns: - True if the request has an invalid API key and metrics should be skipped, + True if metrics should be skipped (currently always False - track all requests), False otherwise """ - if status_code == 401: - return True - - # Handle cases where AssertionError is raised before conversion to ProxyException - if exception is not None: - exception_str = str(exception).lower() - auth_error_patterns = [ - "virtual key expected", - "expected to start with 'sk-'", - "authentication error", - "invalid api key", - "api key not valid", - ] - if any(pattern in exception_str for pattern in auth_error_patterns): - return True - return False def _should_skip_metrics_for_invalid_key( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 359bb944546..85e44ebc827 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -198,7 +198,7 @@ async def common_checks( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", type=ProxyErrorTypes.team_model_access_denied, param="model", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) ## 2.1 If user can call model (if personal key) @@ -2056,7 +2056,7 @@ def _can_object_call_model( object_type=object_type ), param="model", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) @@ -2157,7 +2157,7 @@ async def can_user_call_model( message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}", type=ProxyErrorTypes.key_model_access_denied, param="model", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) return _can_object_call_model( @@ -2739,7 +2739,7 @@ def _can_object_call_vector_stores( object_type ), param="vector_store", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) return True diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c6..90378bd8f3a 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -9,7 +9,10 @@ from fastapi import HTTPException, Request, status import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth -from litellm.proxy.auth.auth_utils import _get_request_ip_address +from litellm.proxy.auth.auth_utils import ( + _get_request_ip_address, + add_client_context_to_request_data, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -30,6 +33,7 @@ class UserAPIKeyAuthExceptionHandler: route: str, parent_otel_span: Optional[Span], api_key: str, + valid_token: Optional[UserAPIKeyAuth] = None, ) -> UserAPIKeyAuth: """ Handles Connection Errors when reading a Virtual Key from LiteLLM DB @@ -71,30 +75,63 @@ class UserAPIKeyAuthExceptionHandler: ) else: # raise the exception to the caller + use_x_forwarded_for = general_settings.get("use_x_forwarded_for", False) requester_ip = _get_request_ip_address( request=request, - use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False), + use_x_forwarded_for=use_x_forwarded_for, ) - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format( - str(e), - requester_ip, - ), - extra={"requester_ip": requester_ip}, + user_agent = request.headers.get("user-agent", "") if request else "" + + # Ensure request_data has client context for callbacks (Prometheus, etc.) + add_client_context_to_request_data( + request=request, + request_data=request_data, + use_x_forwarded_for=use_x_forwarded_for, ) - # Log this exception to OTEL, Datadog etc - user_api_key_dict = UserAPIKeyAuth( - parent_otel_span=parent_otel_span, - api_key=api_key, - request_route=route, + key_name = ( + valid_token.key_alias or getattr(valid_token, "key_name", None) + if valid_token + else "" ) + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}\nUser-Agent:{}\nKey Hash:{}\nKey Name:{}".format( + str(e), + requester_ip or "", + user_agent or "", + api_key or "", + key_name, + ), + extra={ + "requester_ip": requester_ip, + "user_agent": user_agent, + "key_hash": api_key, + "key_name": key_name, + }, + ) + + # Log this exception to OTEL, Datadog etc - use valid_token when available (e.g. model access denied) + if valid_token is not None: + user_api_key_dict = valid_token + else: + user_api_key_dict = UserAPIKeyAuth( + parent_otel_span=parent_otel_span, + api_key=api_key, + request_route=route, + key_alias="", + ) # Allow callbacks to transform the error response + error_type = ProxyErrorTypes.auth_error + if isinstance(e, ProxyException) and hasattr(e, "type"): + try: + error_type = ProxyErrorTypes(e.type) + except (ValueError, TypeError): + pass transformed_exception = await proxy_logging_obj.post_call_failure_hook( request_data=request_data, original_exception=e, user_api_key_dict=user_api_key_dict, - error_type=ProxyErrorTypes.auth_error, + error_type=error_type, route=route, ) # Use transformed exception if callback returned one, otherwise use original diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 2bd84a1d98d..4049dcad14f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -26,6 +26,28 @@ def _get_request_ip_address( return client_ip +def add_client_context_to_request_data( + request: Request, request_data: dict, use_x_forwarded_for: bool = False +) -> None: + """ + Add client_ip (requester_ip_address) and User-Agent to request_data metadata early. + Call this at the start of the request pipeline so failures have client context for logging. + """ + if "metadata" not in request_data: + request_data["metadata"] = {} + metadata = request_data["metadata"] + + requester_ip = _get_request_ip_address( + request=request, use_x_forwarded_for=use_x_forwarded_for + ) + metadata["requester_ip_address"] = requester_ip or "" + + user_agent = "" + if hasattr(request, "headers") and "user-agent" in request.headers: + user_agent = request.headers.get("user-agent", "") + metadata["user_agent"] = user_agent + + def _check_valid_ip( allowed_ips: Optional[List[str]], request: Request, @@ -314,17 +336,17 @@ def get_request_route(request: Request) -> str: def normalize_request_route(route: str) -> str: """ Normalize request routes by replacing dynamic path parameters with placeholders. - + This prevents high cardinality in Prometheus metrics by collapsing routes like: - /v1/responses/1234567890 -> /v1/responses/{response_id} - /v1/threads/thread_123 -> /v1/threads/{thread_id} - + Args: route: The request route path - + Returns: Normalized route with dynamic parameters replaced by placeholders - + Examples: >>> normalize_request_route("/v1/responses/abc123") '/v1/responses/{response_id}' @@ -337,58 +359,90 @@ def normalize_request_route(route: str) -> str: # Format: (regex_pattern, replacement_template) patterns = [ # Responses API - must come before generic patterns - (r'^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), - (r'^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), - (r'^(/(?:openai/)?v1/responses)/([^/]+)$', r'\1/{response_id}'), - (r'^(/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), - (r'^(/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), - (r'^(/responses)/([^/]+)$', r'\1/{response_id}'), - + (r"^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"), + (r"^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"), + (r"^(/(?:openai/)?v1/responses)/([^/]+)$", r"\1/{response_id}"), + (r"^(/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"), + (r"^(/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"), + (r"^(/responses)/([^/]+)$", r"\1/{response_id}"), # Threads API - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$', r'\1/{thread_id}\3/{run_id}\5/{step_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$', r'\1/{thread_id}\3/{run_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$', r'\1/{thread_id}\3'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$', r'\1/{thread_id}\3/{message_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$', r'\1/{thread_id}\3'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)$', r'\1/{thread_id}'), - + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$", + r"\1/{thread_id}\3/{run_id}\5/{step_id}", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$", + r"\1/{thread_id}\3/{run_id}", + ), + (r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$", r"\1/{thread_id}\3"), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$", + r"\1/{thread_id}\3/{message_id}", + ), + (r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$", r"\1/{thread_id}\3"), + (r"^(/(?:openai/)?v1/threads)/([^/]+)$", r"\1/{thread_id}"), # Vector Stores API - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$', r'\1/{vector_store_id}\3/{file_id}'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$', r'\1/{vector_store_id}\3'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$', r'\1/{vector_store_id}\3/{batch_id}'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$', r'\1/{vector_store_id}\3'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)$', r'\1/{vector_store_id}'), - + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$", + r"\1/{vector_store_id}\3/{file_id}", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$", + r"\1/{vector_store_id}\3", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$", + r"\1/{vector_store_id}\3/{batch_id}", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$", + r"\1/{vector_store_id}\3", + ), + (r"^(/(?:openai/)?v1/vector_stores)/([^/]+)$", r"\1/{vector_store_id}"), # Assistants API - (r'^(/(?:openai/)?v1/assistants)/([^/]+)$', r'\1/{assistant_id}'), - + (r"^(/(?:openai/)?v1/assistants)/([^/]+)$", r"\1/{assistant_id}"), # Files API - (r'^(/(?:openai/)?v1/files)/([^/]+)(/content)$', r'\1/{file_id}\3'), - (r'^(/(?:openai/)?v1/files)/([^/]+)$', r'\1/{file_id}'), - + (r"^(/(?:openai/)?v1/files)/([^/]+)(/content)$", r"\1/{file_id}\3"), + (r"^(/(?:openai/)?v1/files)/([^/]+)$", r"\1/{file_id}"), # Batches API - (r'^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$', r'\1/{batch_id}\3'), - (r'^(/(?:openai/)?v1/batches)/([^/]+)$', r'\1/{batch_id}'), - + (r"^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$", r"\1/{batch_id}\3"), + (r"^(/(?:openai/)?v1/batches)/([^/]+)$", r"\1/{batch_id}"), # Fine-tuning API - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$', r'\1/{fine_tuning_job_id}'), - + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$", + r"\1/{fine_tuning_job_id}\3", + ), + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$", + r"\1/{fine_tuning_job_id}\3", + ), + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$", + r"\1/{fine_tuning_job_id}\3", + ), + (r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$", r"\1/{fine_tuning_job_id}"), # Models API - (r'^(/(?:openai/)?v1/models)/([^/]+)$', r'\1/{model}'), + (r"^(/(?:openai/)?v1/models)/([^/]+)$", r"\1/{model}"), ] - + # Apply patterns in order for pattern, replacement in patterns: normalized = re.sub(pattern, replacement, route) if normalized != route: return normalized - + # Return original route if no pattern matched return route @@ -644,6 +698,7 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: return header_name return None + def _get_customer_id_from_standard_headers( request_headers: Optional[dict], ) -> Optional[str]: @@ -679,7 +734,9 @@ def get_end_user_id_from_request_body( from litellm.proxy.proxy_server import general_settings # Check 1: Standard customer ID headers (always checked, no configuration required) - customer_id = _get_customer_id_from_standard_headers(request_headers=request_headers) + customer_id = _get_customer_id_from_standard_headers( + request_headers=request_headers + ) if customer_id is not None: return customer_id diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a153c6e51cc..e562b134818 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -42,6 +42,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, + add_client_context_to_request_data, get_end_user_id_from_request_body, get_model_from_request, get_request_route, @@ -247,7 +248,9 @@ async def get_global_proxy_spend( proxy_logging_obj: ProxyLogging, ) -> Optional[float]: global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget + if ( + litellm.max_budget > 0 and prisma_client is not None + ): # user set proxy max budget # Use event-driven coordination to prevent cache stampede cache_key = "{}:spend".format(litellm_proxy_admin_name) global_proxy_spend = await _fetch_global_spend_with_event_coordination( @@ -1251,6 +1254,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, parent_otel_span=parent_otel_span, api_key=api_key, + valid_token=valid_token, ) @@ -1278,6 +1282,14 @@ async def user_api_key_auth( request_data = populate_request_with_path_params( request_data=request_data, request=request ) + # Capture client context early so failures have IP/User-Agent for logging and metrics + from litellm.proxy.proxy_server import general_settings + + add_client_context_to_request_data( + request=request, + request_data=request_data, + use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False), + ) route: str = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6bbf0df74de..9877825f671 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1633,8 +1633,16 @@ class ProxyLogging: if RouteChecks.is_llm_api_route(route) is not True: return False + proxy_only_error_types = { + ProxyErrorTypes.auth_error, + ProxyErrorTypes.key_model_access_denied, + ProxyErrorTypes.team_model_access_denied, + ProxyErrorTypes.user_model_access_denied, + ProxyErrorTypes.org_model_access_denied, + ProxyErrorTypes.token_not_found_in_db, + } return isinstance(original_exception, HTTPException) or ( - error_type == ProxyErrorTypes.auth_error + error_type in proxy_only_error_types ) async def _handle_logging_proxy_only_error( diff --git a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py index ff433480d5e..565b7f3ef7b 100644 --- a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py @@ -1,8 +1,8 @@ """ -Unit tests for Prometheus invalid API key request filtering. +Unit tests for Prometheus request tracking. -Tests functionality that prevents invalid API key requests (401 status codes) -from being recorded in Prometheus metrics. +Tests that all requests including 401/invalid-key failures are tracked in metrics +for debugging, security auditing, abuse detection, and capacity planning. """ import os @@ -29,12 +29,14 @@ def prometheus_logger(): class ExceptionWithCode: """Exception-like object with 'code' attribute (ProxyException pattern).""" + def __init__(self, code): self.code = code class ExceptionWithStatusCode: """Exception-like object with 'status_code' attribute.""" + def __init__(self, status_code): self.status_code = status_code @@ -42,17 +44,25 @@ class ExceptionWithStatusCode: class TestExtractStatusCode: """Test status code extraction from various sources.""" - @pytest.mark.parametrize("exception_class,code_value,expected", [ - (ExceptionWithCode, "401", 401), - (ExceptionWithStatusCode, 401, 401), - ]) - def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected): + @pytest.mark.parametrize( + "exception_class,code_value,expected", + [ + (ExceptionWithCode, "401", 401), + (ExceptionWithStatusCode, 401, 401), + ], + ) + def test_extract_from_exception( + self, prometheus_logger, exception_class, code_value, expected + ): exception = exception_class(code_value) assert prometheus_logger._extract_status_code(exception=exception) == expected def test_extract_from_kwargs(self, prometheus_logger): exception = ExceptionWithCode("401") - assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401 + assert ( + prometheus_logger._extract_status_code(kwargs={"exception": exception}) + == 401 + ) def test_extract_from_enum_values(self, prometheus_logger): enum_values = Mock(status_code="401") @@ -60,45 +70,62 @@ class TestExtractStatusCode: class TestInvalidAPIKeyDetection: - """Test invalid API key request detection logic.""" + """Test that we no longer skip metrics - all requests are tracked.""" - @pytest.mark.parametrize("status_code,expected", [ - (401, True), - (200, False), - (500, False), - (None, False), - ]) - def test_status_code_detection(self, prometheus_logger, status_code, expected): - assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected + @pytest.mark.parametrize("status_code", [401, 200, 500, None]) + def test_no_skip_for_any_status_code(self, prometheus_logger, status_code): + """All status codes are tracked - never skip metrics.""" + assert ( + prometheus_logger._is_invalid_api_key_request(status_code=status_code) + is False + ) - def test_auth_error_message_detection(self, prometheus_logger): - exception = AssertionError("LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'.") - assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True + def test_auth_error_message_not_skipped(self, prometheus_logger): + exception = AssertionError( + "LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'." + ) + assert ( + prometheus_logger._is_invalid_api_key_request( + status_code=None, exception=exception + ) + is False + ) - def test_non_auth_exception_not_detected(self, prometheus_logger): + def test_non_auth_exception_not_skipped(self, prometheus_logger): exception = ValueError("Some other error") - assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False + assert ( + prometheus_logger._is_invalid_api_key_request( + status_code=None, exception=exception + ) + is False + ) class TestSkipMetricsValidation: - """Test high-level validation method that orchestrates detection and extraction.""" + """Test that we never skip metrics - all requests are tracked.""" - def test_skip_for_401_exception(self, prometheus_logger): - """Test full flow: extraction -> detection -> skip decision.""" + def test_no_skip_for_401_exception(self, prometheus_logger): + """401 requests are now tracked for security auditing and abuse detection.""" exception = ExceptionWithCode("401") - assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True + assert ( + prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) + is False + ) - def test_skip_for_auth_error_message(self, prometheus_logger): - """Test full flow: exception message -> detection -> skip decision.""" + def test_no_skip_for_auth_error_message(self, prometheus_logger): + """Auth error messages are now tracked.""" exception = AssertionError("expected to start with 'sk-'") - assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True + assert ( + prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) + is False + ) def test_no_skip_for_valid_request(self, prometheus_logger): assert prometheus_logger._should_skip_metrics_for_invalid_key() is False class TestAsyncHooks: - """Test async hook methods skip metrics for invalid API keys.""" + """Test async hook methods record metrics for all requests including 401s.""" @pytest.fixture def mock_user_api_key(self): @@ -115,24 +142,33 @@ class TestAsyncHooks: return user_key @pytest.mark.asyncio - async def test_post_call_failure_hook_skips_401(self, prometheus_logger, mock_user_api_key): + async def test_post_call_failure_hook_records_401( + self, prometheus_logger, mock_user_api_key + ): + """401 failures are now recorded for security auditing and abuse detection.""" exception = ExceptionWithCode("401") exception.__class__.__name__ = "ProxyException" - with patch.object(prometheus_logger, 'litellm_proxy_failed_requests_metric') as mock_failed, \ - patch.object(prometheus_logger, 'litellm_proxy_total_requests_metric') as mock_total: + with patch.object( + prometheus_logger, "litellm_proxy_failed_requests_metric" + ) as mock_failed, patch.object( + prometheus_logger, "litellm_proxy_total_requests_metric" + ) as mock_total: + mock_failed.labels.return_value = Mock(inc=Mock()) + mock_total.labels.return_value = Mock(inc=Mock()) await prometheus_logger.async_post_call_failure_hook( - request_data={"model": "test-model"}, + request_data={"model": "test-model", "metadata": {}}, original_exception=exception, - user_api_key_dict=mock_user_api_key + user_api_key_dict=mock_user_api_key, ) - mock_failed.labels.assert_not_called() - mock_total.labels.assert_not_called() + mock_failed.labels.assert_called_once() + mock_total.labels.assert_called_once() @pytest.mark.asyncio - async def test_log_failure_event_skips_401(self, prometheus_logger): + async def test_log_failure_event_records_401(self, prometheus_logger): + """401 failures in log_failure_event are now recorded.""" exception = ExceptionWithCode("401") kwargs = { "model": "test-model", @@ -140,6 +176,9 @@ class TestAsyncHooks: "metadata": { "user_api_key_hash": "test-key", "user_api_key_user_id": "test-user", + "user_api_key_alias": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, }, "model_group": "test-model", }, @@ -147,15 +186,16 @@ class TestAsyncHooks: "litellm_params": {}, } - with patch.object(prometheus_logger, 'litellm_llm_api_failed_requests_metric') as mock_failed, \ - patch.object(prometheus_logger, 'set_llm_deployment_failure_metrics') as mock_deployment: + with patch.object( + prometheus_logger, "litellm_llm_api_failed_requests_metric" + ) as mock_failed, patch.object( + prometheus_logger, "set_llm_deployment_failure_metrics" + ) as mock_deployment: + mock_failed.labels.return_value = Mock(inc=Mock()) await prometheus_logger.async_log_failure_event( - kwargs=kwargs, - response_obj=None, - start_time=None, - end_time=None + kwargs=kwargs, response_obj=None, start_time=None, end_time=None ) - mock_failed.labels.assert_not_called() - mock_deployment.assert_not_called() + mock_failed.labels.assert_called_once() + mock_deployment.assert_called_once() From b33e1e80198247ec717cbe9fab1fac61754b4db8 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Feb 2026 03:04:08 -0300 Subject: [PATCH 04/15] feat(sdk): add proxy_auth for auto OAuth2/JWT token management (#20238) Adds litellm.proxy_auth to automatically obtain and refresh OAuth2/JWT tokens when connecting to LiteLLM Proxy or any OAuth2-protected endpoint. - Add ProxyAuthHandler for token lifecycle (obtain, cache, refresh) - Add AzureADCredential wrapper for azure-identity credentials - Add GenericOAuth2Credential for any OAuth2 provider (Okta, Auth0, etc) - Auto-inject Authorization headers in completion() and embedding() Closes #19834 --- litellm/__init__.py | 2 + litellm/main.py | 14 ++ litellm/proxy_auth/__init__.py | 30 ++++ litellm/proxy_auth/credentials.py | 240 ++++++++++++++++++++++++++++++ tests/litellm/test_proxy_auth.py | 204 +++++++++++++++++++++++++ 5 files changed, 490 insertions(+) create mode 100644 litellm/proxy_auth/__init__.py create mode 100644 litellm/proxy_auth/credentials.py create mode 100644 tests/litellm/test_proxy_auth.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 112d58d49d8..d9811b07da4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -261,6 +261,8 @@ extra_spend_tag_headers: Optional[List[str]] = None in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False enable_azure_ad_token_refresh: Optional[bool] = False +# Proxy Authentication - auto-obtain/refresh OAuth2/JWT tokens for LiteLLM Proxy +proxy_auth: Optional[Any] = None ### DEFAULT AZURE API VERSION ### AZURE_DEFAULT_API_VERSION = "2025-02-01-preview" # this is updated to the latest ### DEFAULT WATSONX API VERSION ### diff --git a/litellm/main.py b/litellm/main.py index 7d591f76882..e7991ae1f17 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1199,6 +1199,13 @@ def completion( # type: ignore # noqa: PLR0915 headers = {} if extra_headers is not None: headers.update(extra_headers) + # Inject proxy auth headers if configured + if litellm.proxy_auth is not None: + try: + proxy_headers = litellm.proxy_auth.get_auth_headers() + headers.update(proxy_headers) + except Exception as e: + verbose_logger.warning(f"Failed to get proxy auth headers: {e}") num_retries = kwargs.get( "num_retries", None ) ## alt. param for 'max_retries'. Use this to pass retries w/ instructor. @@ -4555,6 +4562,13 @@ def embedding( # noqa: PLR0915 headers = {} if extra_headers is not None: headers.update(extra_headers) + # Inject proxy auth headers if configured + if litellm.proxy_auth is not None: + try: + proxy_headers = litellm.proxy_auth.get_auth_headers() + headers.update(proxy_headers) + except Exception as e: + verbose_logger.warning(f"Failed to get proxy auth headers: {e}") ### CUSTOM MODEL COST ### input_cost_per_token = kwargs.get("input_cost_per_token", None) output_cost_per_token = kwargs.get("output_cost_per_token", None) diff --git a/litellm/proxy_auth/__init__.py b/litellm/proxy_auth/__init__.py new file mode 100644 index 00000000000..27624a94fb9 --- /dev/null +++ b/litellm/proxy_auth/__init__.py @@ -0,0 +1,30 @@ +""" +Proxy Authentication module for LiteLLM SDK. + +This module provides OAuth2/JWT token management for authenticating +with LiteLLM Proxy or any OAuth2-protected endpoint. + +Usage: + from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler + + litellm.proxy_auth = ProxyAuthHandler( + credential=AzureADCredential(), + scope="api://my-proxy/.default" + ) +""" + +from .credentials import ( + AccessToken, + TokenCredential, + AzureADCredential, + GenericOAuth2Credential, + ProxyAuthHandler, +) + +__all__ = [ + "AccessToken", + "TokenCredential", + "AzureADCredential", + "GenericOAuth2Credential", + "ProxyAuthHandler", +] diff --git a/litellm/proxy_auth/credentials.py b/litellm/proxy_auth/credentials.py new file mode 100644 index 00000000000..cddaf1278f9 --- /dev/null +++ b/litellm/proxy_auth/credentials.py @@ -0,0 +1,240 @@ +""" +Credential providers for proxy authentication. + +This module provides a provider-agnostic interface for obtaining OAuth2/JWT tokens. +It follows the same TokenCredential protocol used by Azure SDK. +""" + +import time +from dataclasses import dataclass +from typing import Any, Optional, Protocol, runtime_checkable + + +@dataclass +class AccessToken: + """ + Represents an OAuth2 access token with expiration. + + This matches the structure used by azure.core.credentials.AccessToken. + + Attributes: + token: The access token string (typically a JWT). + expires_on: Unix timestamp when the token expires. + """ + + token: str + expires_on: int + + +@runtime_checkable +class TokenCredential(Protocol): + """ + Protocol for credential providers. + + This matches the azure.core.credentials.TokenCredential interface, + allowing any Azure SDK credential to be used directly. + + Any class implementing get_token(scope) -> AccessToken can be used. + """ + + def get_token(self, scope: str) -> AccessToken: + """ + Get an access token for the specified scope. + + Args: + scope: The OAuth2 scope to request (e.g., "api://my-app/.default") + + Returns: + AccessToken with the token string and expiration timestamp. + """ + ... + + +class AzureADCredential: + """ + Wrapper for Azure Identity credentials. + + This wraps any azure-identity credential (DefaultAzureCredential, + ClientSecretCredential, ManagedIdentityCredential, etc.) and converts + the token to our AccessToken format. + + If no credential is provided, it will use DefaultAzureCredential + which tries multiple authentication methods automatically. + + Example: + # Use default credential chain (env vars, managed identity, CLI, etc.) + cred = AzureADCredential() + + # Or provide a specific credential + from azure.identity import ClientSecretCredential + azure_cred = ClientSecretCredential(tenant_id, client_id, client_secret) + cred = AzureADCredential(credential=azure_cred) + """ + + def __init__(self, credential: Optional[Any] = None): + """ + Initialize with an optional Azure credential. + + Args: + credential: An azure-identity credential object. If None, + DefaultAzureCredential will be used on first token request. + """ + self._credential = credential + self._initialized = credential is not None + + def get_token(self, scope: str) -> AccessToken: + """ + Get an access token from Azure AD. + + Args: + scope: The OAuth2 scope (e.g., "api://my-app/.default") + + Returns: + AccessToken with the JWT and expiration. + + Raises: + ImportError: If azure-identity is not installed. + """ + if not self._initialized: + try: + from azure.identity import DefaultAzureCredential + + self._credential = DefaultAzureCredential() + self._initialized = True + except ImportError: + raise ImportError( + "azure-identity is required for AzureADCredential. " + "Install it with: pip install azure-identity" + ) + + result = self._credential.get_token(scope) + return AccessToken(token=result.token, expires_on=result.expires_on) + + +class GenericOAuth2Credential: + """ + Generic OAuth2 client credentials flow. + + This works with any OAuth2 provider (Okta, Auth0, Keycloak, etc.) + that supports the client_credentials grant type. + + Example: + cred = GenericOAuth2Credential( + client_id="my-client-id", + client_secret="my-client-secret", + token_url="https://my-idp.com/oauth2/token" + ) + """ + + def __init__(self, client_id: str, client_secret: str, token_url: str): + """ + Initialize OAuth2 client credentials. + + Args: + client_id: OAuth2 client ID + client_secret: OAuth2 client secret + token_url: Token endpoint URL (e.g., "https://idp.com/oauth2/token") + """ + self.client_id = client_id + self.client_secret = client_secret + self.token_url = token_url + self._cached_token: Optional[AccessToken] = None + + def get_token(self, scope: str) -> AccessToken: + """ + Get an access token using OAuth2 client credentials flow. + + Tokens are cached and reused until they expire (with 60s buffer). + + Args: + scope: The OAuth2 scope to request + + Returns: + AccessToken with the token and expiration. + """ + # Return cached token if still valid (with 60s buffer) + if self._cached_token and self._cached_token.expires_on > time.time() + 60: + return self._cached_token + + import httpx + + response = httpx.post( + self.token_url, + data={ + "grant_type": "client_credentials", + "client_id": self.client_id, + "client_secret": self.client_secret, + "scope": scope, + }, + ) + response.raise_for_status() + data = response.json() + + self._cached_token = AccessToken( + token=data["access_token"], + expires_on=int(time.time()) + data.get("expires_in", 3600), + ) + return self._cached_token + + +class ProxyAuthHandler: + """ + Manages OAuth2/JWT token lifecycle for proxy authentication. + + This handler: + - Obtains tokens from the configured credential provider + - Caches tokens to avoid unnecessary requests + - Automatically refreshes tokens before they expire (60s buffer) + - Generates Authorization headers for HTTP requests + + Set this as litellm.proxy_auth to automatically inject auth headers + into all requests to your LiteLLM Proxy. + + Example: + import litellm + from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler + + litellm.proxy_auth = ProxyAuthHandler( + credential=AzureADCredential(), + scope="api://my-litellm-proxy/.default" + ) + litellm.api_base = "https://my-proxy.example.com" + + # Auth headers are now automatically injected + response = litellm.completion(model="gpt-4", messages=[...]) + """ + + def __init__(self, credential: TokenCredential, scope: str): + """ + Initialize the proxy auth handler. + + Args: + credential: A TokenCredential implementation (AzureADCredential, + GenericOAuth2Credential, or any custom implementation) + scope: The OAuth2 scope to request tokens for + """ + self.credential = credential + self.scope = scope + self._cached_token: Optional[AccessToken] = None + + def get_token(self) -> AccessToken: + """ + Get a valid access token, refreshing if necessary. + + Returns: + AccessToken that is valid for at least 60 more seconds. + """ + # Refresh if no token or token expires within 60 seconds + if not self._cached_token or self._cached_token.expires_on <= time.time() + 60: + self._cached_token = self.credential.get_token(self.scope) + return self._cached_token + + def get_auth_headers(self) -> dict: + """ + Get HTTP headers for authentication. + + Returns: + Dict with Authorization header containing Bearer token. + """ + token = self.get_token() + return {"Authorization": f"Bearer {token.token}"} diff --git a/tests/litellm/test_proxy_auth.py b/tests/litellm/test_proxy_auth.py new file mode 100644 index 00000000000..1d73e143e10 --- /dev/null +++ b/tests/litellm/test_proxy_auth.py @@ -0,0 +1,204 @@ +""" +Unit tests for litellm.proxy_auth module. + +Tests the OAuth2/JWT token management for LiteLLM Proxy authentication. +""" + +import time +from unittest.mock import Mock, patch + +import pytest + +from litellm.proxy_auth import ( + AccessToken, + AzureADCredential, + GenericOAuth2Credential, + ProxyAuthHandler, +) + + +class TestAccessToken: + """Tests for AccessToken dataclass.""" + + def test_access_token_creation(self): + """Test AccessToken can be created with required fields.""" + token = AccessToken(token="test-token", expires_on=1234567890) + assert token.token == "test-token" + assert token.expires_on == 1234567890 + + def test_access_token_equality(self): + """Test AccessToken equality comparison.""" + token1 = AccessToken(token="test", expires_on=123) + token2 = AccessToken(token="test", expires_on=123) + assert token1 == token2 + + +class MockCredential: + """Mock credential for testing.""" + + def __init__(self, expires_in_seconds: int = 3600): + self.call_count = 0 + self.expires_in = expires_in_seconds + + def get_token(self, scope: str) -> AccessToken: + self.call_count += 1 + return AccessToken( + token=f"mock-token-{self.call_count}", + expires_on=int(time.time()) + self.expires_in, + ) + + +class TestProxyAuthHandler: + """Tests for ProxyAuthHandler.""" + + def test_get_auth_headers_returns_bearer_token(self): + """Test that get_auth_headers returns correct Authorization header.""" + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + headers = handler.get_auth_headers() + + assert "Authorization" in headers + assert headers["Authorization"].startswith("Bearer ") + assert "mock-token-1" in headers["Authorization"] + + def test_token_caching(self): + """Test that tokens are cached and not re-requested.""" + cred = MockCredential(expires_in_seconds=3600) # Long expiry + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + # Multiple calls should only request token once + handler.get_auth_headers() + handler.get_auth_headers() + handler.get_auth_headers() + + assert cred.call_count == 1 + + def test_token_refresh_when_about_to_expire(self): + """Test that tokens are refreshed when about to expire (within 60s buffer).""" + cred = MockCredential(expires_in_seconds=30) # Expires in 30s (< 60s buffer) + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + # First call gets token + handler.get_auth_headers() + # Second call should refresh because token expires within 60s buffer + handler.get_auth_headers() + + assert cred.call_count == 2 + + def test_get_token_method(self): + """Test the get_token method returns AccessToken.""" + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + token = handler.get_token() + + assert isinstance(token, AccessToken) + assert token.token == "mock-token-1" + + +class TestAzureADCredential: + """Tests for AzureADCredential.""" + + def test_lazy_initialization(self): + """Test that azure-identity is not imported until get_token is called.""" + # This should not raise ImportError even if azure-identity is not installed + cred = AzureADCredential(credential=None) + # _initialized should be False until get_token is called + assert cred._initialized is False + + def test_wraps_azure_credential(self): + """Test that AzureADCredential wraps an azure-identity credential.""" + # Mock Azure credential + mock_azure_cred = Mock() + mock_azure_cred.get_token.return_value = Mock( + token="azure-token", expires_on=9999999999 + ) + + cred = AzureADCredential(credential=mock_azure_cred) + token = cred.get_token("https://graph.microsoft.com/.default") + + assert token.token == "azure-token" + assert token.expires_on == 9999999999 + mock_azure_cred.get_token.assert_called_once_with( + "https://graph.microsoft.com/.default" + ) + + +class TestGenericOAuth2Credential: + """Tests for GenericOAuth2Credential.""" + + def test_token_request(self): + """Test that GenericOAuth2Credential makes correct OAuth2 request.""" + with patch("httpx.post") as mock_post: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "oauth2-token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + cred = GenericOAuth2Credential( + client_id="test-client", + client_secret="test-secret", + token_url="https://example.com/oauth2/token", + ) + token = cred.get_token("test-scope") + + assert token.token == "oauth2-token" + mock_post.assert_called_once() + call_kwargs = mock_post.call_args + assert call_kwargs[1]["data"]["grant_type"] == "client_credentials" + assert call_kwargs[1]["data"]["client_id"] == "test-client" + assert call_kwargs[1]["data"]["client_secret"] == "test-secret" + assert call_kwargs[1]["data"]["scope"] == "test-scope" + + def test_token_caching(self): + """Test that GenericOAuth2Credential caches tokens.""" + with patch("httpx.post") as mock_post: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "oauth2-token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + cred = GenericOAuth2Credential( + client_id="test-client", + client_secret="test-secret", + token_url="https://example.com/oauth2/token", + ) + + # Multiple calls should only make one HTTP request + cred.get_token("test-scope") + cred.get_token("test-scope") + cred.get_token("test-scope") + + assert mock_post.call_count == 1 + + +class TestLiteLLMIntegration: + """Tests for integration with litellm module.""" + + def test_proxy_auth_variable_exists(self): + """Test that litellm.proxy_auth variable exists.""" + import litellm + + # Should be None by default + assert hasattr(litellm, "proxy_auth") + + def test_proxy_auth_can_be_set(self): + """Test that litellm.proxy_auth can be set to a ProxyAuthHandler.""" + import litellm + + original_value = litellm.proxy_auth + try: + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test") + litellm.proxy_auth = handler + + assert litellm.proxy_auth is handler + finally: + litellm.proxy_auth = original_value From a904c3f40d7d539d56e0c6eea4e6c906673d7e28 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Feb 2026 03:05:44 -0300 Subject: [PATCH 05/15] fix(github_copilot): preserve system prompts and auto-inject headers (#20113) - Remove system-to-assistant message conversion (API now supports system prompts) - Auto-inject required Copilot headers in chat completions (same as /responses) - Deprecate disable_copilot_system_to_assistant flag - Update docs to remove manual extra_headers requirement Fixes #19873 --- .../docs/providers/github_copilot.md | 35 ++++++------------- docs/my-website/docs/proxy/config_settings.md | 4 +-- .../github_copilot/chat/transformation.py | 21 +++++------ litellm/main.py | 14 ++++++++ 4 files changed, 38 insertions(+), 36 deletions(-) diff --git a/docs/my-website/docs/providers/github_copilot.md b/docs/my-website/docs/providers/github_copilot.md index 306c9f949ec..e9fd3444f5f 100644 --- a/docs/my-website/docs/providers/github_copilot.md +++ b/docs/my-website/docs/providers/github_copilot.md @@ -35,11 +35,10 @@ from litellm import completion response = completion( model="github_copilot/gpt-4", - messages=[{"role": "user", "content": "Write a Python function to calculate fibonacci numbers"}], - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + messages=[ + {"role": "system", "content": "You are a helpful coding assistant"}, + {"role": "user", "content": "Write a Python function to calculate fibonacci numbers"} + ] ) print(response) ``` @@ -50,11 +49,7 @@ from litellm import completion stream = completion( model="github_copilot/gpt-4", messages=[{"role": "user", "content": "Explain async/await in Python"}], - stream=True, - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + stream=True ) for chunk in stream: @@ -134,11 +129,7 @@ client = OpenAI( # Non-streaming response response = client.chat.completions.create( model="github_copilot/gpt-4", - messages=[{"role": "user", "content": "How do I optimize this SQL query?"}], - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + messages=[{"role": "user", "content": "How do I optimize this SQL query?"}] ) print(response.choices[0].message.content) @@ -156,11 +147,7 @@ response = litellm.completion( model="litellm_proxy/github_copilot/gpt-4", messages=[{"role": "user", "content": "Review this code for bugs"}], api_base="http://localhost:4000", - api_key="your-proxy-api-key", - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + api_key="your-proxy-api-key" ) print(response.choices[0].message.content) @@ -174,8 +161,6 @@ print(response.choices[0].message.content) curl http://localhost:4000/v1/chat/completions \ -H "Content-Type: application/json" \ -H "Authorization: Bearer your-proxy-api-key" \ - -H "editor-version: vscode/1.85.1" \ - -H "Copilot-Integration-Id: vscode-chat" \ -d '{ "model": "github_copilot/gpt-4", "messages": [{"role": "user", "content": "Explain this error message"}] @@ -211,9 +196,11 @@ export GITHUB_COPILOT_API_KEY_FILE="api-key.json" ### Headers -GitHub Copilot supports various editor-specific headers: +LiteLLM automatically injects the required GitHub Copilot headers (simulating VSCode). You don't need to specify them manually. -```python showLineNumbers title="Common Headers" +If you want to override the defaults (e.g., to simulate a different editor), you can use `extra_headers`: + +```python showLineNumbers title="Custom Headers (Optional)" extra_headers = { "editor-version": "vscode/1.85.1", # Editor version "editor-plugin-version": "copilot/1.155.0", # Plugin version diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 264c7d765b3..80dfe11742a 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -94,7 +94,7 @@ litellm_settings: # /chat/completions, /completions, /embeddings, /audio/transcriptions mode: default_off # if default_off, you need to opt in to caching on a per call basis ttl: 600 # ttl for caching - disable_copilot_system_to_assistant: False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. + disable_copilot_system_to_assistant: False # DEPRECATED - GitHub Copilot API supports system prompts. callback_settings: otel: @@ -197,7 +197,7 @@ router_settings: | disable_add_transform_inline_image_block | boolean | For Fireworks AI models - if true, turns off the auto-add of `#transform=inline` to the url of the image_url, if the model is not a vision model. | | disable_hf_tokenizer_download | boolean | If true, it defaults to using the openai tokenizer for all models (including huggingface models). | | enable_json_schema_validation | boolean | If true, enables json schema validation for all requests. | -| disable_copilot_system_to_assistant | boolean | If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. Useful for tools (like Claude Code) that send system messages, which Copilot does not support. | +| disable_copilot_system_to_assistant | boolean | **DEPRECATED** - GitHub Copilot API supports system prompts. | ### general_settings - Reference diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 50f18cedf9b..f001bb65f8b 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -5,7 +5,7 @@ from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues from ..authenticator import Authenticator -from ..common_utils import GetAPIKeyError, GITHUB_COPILOT_API_BASE +from ..common_utils import GetAPIKeyError, GITHUB_COPILOT_API_BASE, get_copilot_default_headers class GithubCopilotConfig(OpenAIConfig): @@ -43,15 +43,8 @@ class GithubCopilotConfig(OpenAIConfig): messages, model: str, ): - import litellm - - disable_copilot_system_to_assistant = ( - litellm.disable_copilot_system_to_assistant - ) - if not disable_copilot_system_to_assistant: - for message in messages: - if "role" in message and message["role"] == "system": - cast(Any, message)["role"] = "assistant" + # GitHub Copilot API now supports system prompts for all models (Claude, GPT, etc.) + # No conversion needed - just return messages as-is return messages def validate_environment( @@ -69,6 +62,14 @@ class GithubCopilotConfig(OpenAIConfig): headers, model, messages, optional_params, litellm_params, api_key, api_base ) + # Add Copilot-specific headers (editor-version, user-agent, etc.) + try: + copilot_api_key = self.authenticator.get_api_key() + copilot_headers = get_copilot_default_headers(copilot_api_key) + validated_headers = {**copilot_headers, **validated_headers} + except GetAPIKeyError: + pass # Will be handled later in the request flow + # Add X-Initiator header based on message roles initiator = self._determine_initiator(messages) validated_headers["X-Initiator"] = initiator diff --git a/litellm/main.py b/litellm/main.py index e7991ae1f17..440b7e52e70 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2462,6 +2462,20 @@ def completion( # type: ignore # noqa: PLR0915 headers = headers or litellm.headers + # Add GitHub Copilot headers (same as /responses endpoint does) + if custom_llm_provider == "github_copilot": + from litellm.llms.github_copilot.common_utils import ( + get_copilot_default_headers, + ) + from litellm.llms.github_copilot.authenticator import Authenticator + + copilot_auth = Authenticator() + copilot_api_key = copilot_auth.get_api_key() + copilot_headers = get_copilot_default_headers(copilot_api_key) + if extra_headers: + copilot_headers.update(extra_headers) + extra_headers = copilot_headers + if extra_headers is not None: optional_params["extra_headers"] = extra_headers From 17c0a88a60550c86771dea75386c6b6b7657a4ce Mon Sep 17 00:00:00 2001 From: krauckbot Date: Tue, 3 Feb 2026 07:07:17 +0100 Subject: [PATCH 06/15] fix: add missing capability flags to vercel_ai_gateway models (#20276) 67 vercel_ai_gateway models were missing capability flags (supports_vision, supports_function_calling, supports_tool_choice, supports_response_schema). These capabilities were inferred from the corresponding direct provider entries for the same models (e.g., vercel_ai_gateway/anthropic/claude-3.5-sonnet now has the same capabilities as anthropic/claude-3.5-sonnet). Models fixed include: - Claude 3/3.5/3.7 (Anthropic) - GPT-4/5 variants (OpenAI) - Gemini 2.0/2.5 (Google) - Grok 3/4 (xAI) - Mistral/Mixtral variants - Qwen models - DeepSeek models - And more This ensures consistent capability reporting across providers for the same underlying models. Co-authored-by: krauckbot --- model_prices_and_context_window.json | 333 +++++++++++++++++++++------ 1 file changed, 261 insertions(+), 72 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 485bee4f191..801aec4ceb1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27829,7 +27829,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/alibaba/qwen3-coder": { "input_cost_per_token": 4e-07, @@ -27838,7 +27840,9 @@ "max_output_tokens": 66536, "max_tokens": 66536, "mode": "chat", - "output_cost_per_token": 1.6e-06 + "output_cost_per_token": 1.6e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/amazon/nova-lite": { "input_cost_per_token": 6e-08, @@ -27847,7 +27851,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.4e-07 + "output_cost_per_token": 2.4e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/nova-micro": { "input_cost_per_token": 3.5e-08, @@ -27856,7 +27863,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.4e-07 + "output_cost_per_token": 1.4e-07, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/nova-pro": { "input_cost_per_token": 8e-07, @@ -27865,7 +27874,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3.2e-06 + "output_cost_per_token": 3.2e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/titan-embed-text-v2": { "input_cost_per_token": 2e-08, @@ -27885,7 +27897,11 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.25e-06 + "output_cost_per_token": 1.25e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3-opus": { "cache_creation_input_token_cost": 1.875e-05, @@ -27896,7 +27912,11 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 7.5e-05 + "output_cost_per_token": 7.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.5-haiku": { "cache_creation_input_token_cost": 1e-06, @@ -27907,7 +27927,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 4e-06 + "output_cost_per_token": 4e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.5-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -27918,7 +27942,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.7-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -27929,7 +27957,11 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-4-opus": { "cache_creation_input_token_cost": 1.875e-05, @@ -27940,7 +27972,11 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 7.5e-05 + "output_cost_per_token": 7.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-4-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -27951,7 +27987,9 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-a": { "input_cost_per_token": 2.5e-06, @@ -27960,7 +27998,9 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-r": { "input_cost_per_token": 1.5e-07, @@ -27969,7 +28009,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-r-plus": { "input_cost_per_token": 2.5e-06, @@ -27978,7 +28020,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, @@ -27996,7 +28040,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.19e-06 + "output_cost_per_token": 2.19e-06, + "supports_tool_choice": true }, "vercel_ai_gateway/deepseek/deepseek-r1-distill-llama-70b": { "input_cost_per_token": 7.5e-07, @@ -28005,7 +28050,10 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.9e-07 + "output_cost_per_token": 9.9e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/deepseek/deepseek-v3": { "input_cost_per_token": 9e-07, @@ -28014,7 +28062,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 9e-07 + "output_cost_per_token": 9e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/google/gemini-2.0-flash": { "deprecation_date": "2026-03-31", @@ -28024,7 +28073,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.0-flash-lite": { "deprecation_date": "2026-03-31", @@ -28034,7 +28087,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.5-flash": { "input_cost_per_token": 3e-07, @@ -28043,7 +28100,11 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.5e-06 + "output_cost_per_token": 2.5e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { "input_cost_per_token": 2.5e-06, @@ -28052,7 +28113,11 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -28070,7 +28135,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2e-07 + "output_cost_per_token": 2e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/google/text-embedding-005": { "input_cost_per_token": 2.5e-08, @@ -28106,7 +28174,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.9e-07 + "output_cost_per_token": 7.9e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3-8b": { "input_cost_per_token": 5e-08, @@ -28115,7 +28184,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8e-08 + "output_cost_per_token": 8e-08, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.1-70b": { "input_cost_per_token": 7.2e-07, @@ -28124,7 +28194,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.1-8b": { "input_cost_per_token": 5e-08, @@ -28133,7 +28204,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-08 + "output_cost_per_token": 8e-08, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/meta/llama-3.2-11b": { "input_cost_per_token": 1.6e-07, @@ -28142,7 +28215,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.6e-07 + "output_cost_per_token": 1.6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.2-1b": { "input_cost_per_token": 1e-07, @@ -28160,7 +28236,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.5e-07 + "output_cost_per_token": 1.5e-07, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/meta/llama-3.2-90b": { "input_cost_per_token": 7.2e-07, @@ -28169,7 +28247,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.3-70b": { "input_cost_per_token": 7.2e-07, @@ -28178,7 +28259,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-4-maverick": { "input_cost_per_token": 2e-07, @@ -28187,7 +28270,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-4-scout": { "input_cost_per_token": 1e-07, @@ -28196,7 +28280,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/codestral": { "input_cost_per_token": 3e-07, @@ -28205,7 +28292,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 9e-07 + "output_cost_per_token": 9e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/codestral-embed": { "input_cost_per_token": 1.5e-07, @@ -28223,7 +28312,10 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.8e-07 + "output_cost_per_token": 2.8e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/magistral-medium": { "input_cost_per_token": 2e-06, @@ -28232,7 +28324,10 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 5e-06 + "output_cost_per_token": 5e-06, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/magistral-small": { "input_cost_per_token": 5e-07, @@ -28241,7 +28336,8 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-06 + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true }, "vercel_ai_gateway/mistral/ministral-3b": { "input_cost_per_token": 4e-08, @@ -28250,7 +28346,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 4e-08 + "output_cost_per_token": 4e-08, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/ministral-8b": { "input_cost_per_token": 1e-07, @@ -28259,7 +28357,10 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1e-07 + "output_cost_per_token": 1e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/mistral-embed": { "input_cost_per_token": 1e-07, @@ -28277,7 +28378,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 6e-06 + "output_cost_per_token": 6e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/mistral-saba-24b": { "input_cost_per_token": 7.9e-07, @@ -28295,7 +28398,10 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/mixtral-8x22b-instruct": { "input_cost_per_token": 1.2e-06, @@ -28304,7 +28410,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.2e-06 + "output_cost_per_token": 1.2e-06, + "supports_function_calling": true }, "vercel_ai_gateway/mistral/pixtral-12b": { "input_cost_per_token": 1.5e-07, @@ -28313,7 +28420,11 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1.5e-07 + "output_cost_per_token": 1.5e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/pixtral-large": { "input_cost_per_token": 2e-06, @@ -28322,7 +28433,11 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 6e-06 + "output_cost_per_token": 6e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/moonshotai/kimi-k2": { "input_cost_per_token": 5.5e-07, @@ -28331,7 +28446,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 2.2e-06 + "output_cost_per_token": 2.2e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/morph/morph-v3-fast": { "input_cost_per_token": 8e-07, @@ -28358,7 +28475,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.5e-06 + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/openai/gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, @@ -28376,7 +28495,10 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 3e-05 + "output_cost_per_token": 3e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/openai/gpt-4.1": { "cache_creation_input_token_cost": 0.0, @@ -28387,7 +28509,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 8e-06 + "output_cost_per_token": 8e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4.1-mini": { "cache_creation_input_token_cost": 0.0, @@ -28398,7 +28524,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.6e-06 + "output_cost_per_token": 1.6e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4.1-nano": { "cache_creation_input_token_cost": 0.0, @@ -28409,7 +28539,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 4e-07 + "output_cost_per_token": 4e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4o": { "cache_creation_input_token_cost": 0.0, @@ -28420,7 +28554,11 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4o-mini": { "cache_creation_input_token_cost": 0.0, @@ -28431,7 +28569,11 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o1": { "cache_creation_input_token_cost": 0.0, @@ -28442,7 +28584,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 6e-05 + "output_cost_per_token": 6e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o3": { "cache_creation_input_token_cost": 0.0, @@ -28453,7 +28599,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 8e-06 + "output_cost_per_token": 8e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o3-mini": { "cache_creation_input_token_cost": 0.0, @@ -28464,7 +28614,10 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 4.4e-06 + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o4-mini": { "cache_creation_input_token_cost": 0.0, @@ -28475,7 +28628,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 4.4e-06 + "output_cost_per_token": 4.4e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/text-embedding-3-large": { "input_cost_per_token": 1.3e-07, @@ -28547,7 +28704,10 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/vercel/v0-1.5-md": { "input_cost_per_token": 3e-06, @@ -28556,7 +28716,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-2": { "input_cost_per_token": 2e-06, @@ -28565,7 +28728,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-2-vision": { "input_cost_per_token": 2e-06, @@ -28574,7 +28739,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3": { "input_cost_per_token": 3e-06, @@ -28583,7 +28751,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3-fast": { "input_cost_per_token": 5e-06, @@ -28592,7 +28762,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05 + "output_cost_per_token": 2.5e-05, + "supports_function_calling": true }, "vercel_ai_gateway/xai/grok-3-mini": { "input_cost_per_token": 3e-07, @@ -28601,7 +28772,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07 + "output_cost_per_token": 5e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3-mini-fast": { "input_cost_per_token": 6e-07, @@ -28610,7 +28783,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06 + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-4": { "input_cost_per_token": 3e-06, @@ -28619,7 +28794,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.5": { "input_cost_per_token": 6e-07, @@ -28628,7 +28805,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.2e-06 + "output_cost_per_token": 2.2e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.5-air": { "input_cost_per_token": 2e-07, @@ -28637,7 +28816,9 @@ "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", - "output_cost_per_token": 1.1e-06 + "output_cost_per_token": 1.1e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.6": { "litellm_provider": "vercel_ai_gateway", @@ -29799,7 +29980,9 @@ "mode": "chat", "output_cost_per_token": 1e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29812,7 +29995,9 @@ "mode": "chat", "output_cost_per_token": 4e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29825,7 +30010,9 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29838,7 +30025,9 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -34787,4 +34976,4 @@ "output_cost_per_token": 0, "supports_reasoning": true } -} \ No newline at end of file +} From cc76f95555b5e5576df903b05f8c7b4a1abc74ab Mon Sep 17 00:00:00 2001 From: Alexander Grattan <51346343+agrattan0820@users.noreply.github.com> Date: Tue, 3 Feb 2026 17:41:13 -0500 Subject: [PATCH 07/15] fix: check for model_response_choices before guardrail input (#19784) * fix: check for model_response_choices before guardrail input * test: add tests for responses api translation * fix: protect other guardrail translations * refactor: remove type ignores * anthropic request body got mutated fix * add warning when extra_body is provided but user is non premium * fix: resolve mypy union-attr errors in anthropic guardrail handler Cast choices[0] to Choices type before accessing .message attribute to satisfy mypy's union type checking for Choices | StreamingChoices. Co-Authored-By: Claude Opus 4.5 * add logger when model response has no choices for streaming /response and /messages * update pyproject.toml as requested * Revert "update pyproject.toml as requested" This reverts commit 541a2b075a91b1b2d9efaf0407572f35bf5d4324. * update pyproject.toml as requested * Revert "update pyproject.toml as requested" This reverts commit 716ea0caa1fee5e5f028d3f86479fedea2fac68b. --------- Co-authored-by: Xiaohan Fu Co-authored-by: Claude Opus 4.5 --- litellm/integrations/custom_guardrail.py | 9 +- .../chat/guardrail_translation/handler.py | 72 ++++-- .../chat/guardrail_translation/handler.py | 21 +- .../guardrail_translation/handler.py | 51 ++-- .../test_anthropic_guardrail_handler.py | 236 ++++++++++++++++++ .../test_openai_guardrail_handler.py | 148 +++++++++++ ...test_openai_responses_guardrail_handler.py | 178 +++++++++++++ 7 files changed, 655 insertions(+), 60 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a5bb530fc56..1652ec2aa0c 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -475,11 +475,18 @@ class CustomGuardrail(CustomLogger): guardrail_config: DynamicGuardrailParams = DynamicGuardrailParams( **guardrail[self.guardrail_name] ) + extra_body = guardrail_config.get("extra_body", {}) if self._validate_premium_user() is not True: + if isinstance(extra_body, dict) and extra_body: + verbose_logger.warning( + "Guardrail %s: ignoring dynamic extra_body keys %s because premium_user is False", + self.guardrail_name, + list(extra_body.keys()), + ) return {} # Return the extra_body if it exists, otherwise empty dict - return guardrail_config.get("extra_body", {}) + return extra_body return {} diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 8e1016bd5bd..a14e7d118e8 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -34,6 +34,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ( ChatCompletionMessageToolCall, + Choices, GenericGuardrailAPIInputs, ModelResponse, ) @@ -76,7 +77,8 @@ class AnthropicMessagesHandler(BaseTranslation): chat_completion_compatible_request, tool_name_mapping = ( LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - anthropic_message_request=cast(AnthropicMessagesRequest, data) + # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). + anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) ) ) @@ -84,9 +86,9 @@ class AnthropicMessagesHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] - tools_to_check: List[ChatCompletionToolParam] = ( - chat_completion_compatible_request.get("tools", []) - ) + tools_to_check: List[ + ChatCompletionToolParam + ] = chat_completion_compatible_request.get("tools", []) task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each text # content_index is None for string content, int for list content @@ -282,7 +284,10 @@ class AnthropicMessagesHandler(BaseTranslation): if hasattr(content_block, "model_dump"): block_dict = content_block.model_dump() else: - block_dict = {"type": block_type, "text": getattr(content_block, "text", None)} + block_dict = { + "type": block_type, + "text": getattr(content_block, "text", None), + } else: continue @@ -358,30 +363,40 @@ class AnthropicMessagesHandler(BaseTranslation): """ has_ended = self._check_streaming_has_ended(responses_so_far) if has_ended: - # build the model response from the responses_so_far - model_response = cast( - ModelResponse, - AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=responses_so_far, - litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), - model="", - ), + built_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=responses_so_far, + litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), + model="", ) - tool_calls_list = cast(Optional[List[ChatCompletionMessageToolCall]], model_response.choices[0].message.tool_calls) # type: ignore - string_so_far = model_response.choices[0].message.content # type: ignore - guardrail_inputs = GenericGuardrailAPIInputs() - if string_so_far: - guardrail_inputs["texts"] = [string_so_far] - if tool_calls_list: - guardrail_inputs["tool_calls"] = tool_calls_list - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid - inputs=guardrail_inputs, - request_data={}, - input_type="response", - logging_obj=litellm_logging_obj, - ) + # Check if model_response is valid and has choices before accessing + if ( + built_response is not None + and hasattr(built_response, "choices") + and built_response.choices + ): + model_response = cast(ModelResponse, built_response) + first_choice = cast(Choices, model_response.choices[0]) + tool_calls_list = cast( + Optional[List[ChatCompletionMessageToolCall]], + first_choice.message.tool_calls, + ) + string_so_far = first_choice.message.content + guardrail_inputs = GenericGuardrailAPIInputs() + if string_so_far: + guardrail_inputs["texts"] = [string_so_far] + if tool_calls_list: + guardrail_inputs["tool_calls"] = tool_calls_list + + _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid + inputs=guardrail_inputs, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + else: + verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far string_so_far = self.get_streaming_string_so_far(responses_so_far) @@ -648,7 +663,10 @@ class AnthropicMessagesHandler(BaseTranslation): if isinstance(content_block, dict): if content_block.get("type") == "text": cast(Dict[str, Any], content_block)["text"] = guardrail_response - elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": + elif ( + hasattr(content_block, "type") + and getattr(content_block, "type", None) == "text" + ): # Update Pydantic object's text attribute if hasattr(content_block, "text"): content_block.text = guardrail_response diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index fb00aa28f45..c406f502b45 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -21,7 +21,13 @@ from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.main import stream_chunk_builder from litellm.types.llms.openai import ChatCompletionToolParam -from litellm.types.utils import Choices, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + Choices, + GenericGuardrailAPIInputs, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -80,9 +86,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore if messages: - inputs["structured_messages"] = ( - messages # pass the openai /chat/completions messages to the guardrail, as-is - ) + inputs[ + "structured_messages" + ] = messages # pass the openai /chat/completions messages to the guardrail, as-is # Pass tools (function definitions) to the guardrail tools = data.get("tools") if tools: @@ -362,14 +368,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # check if the stream has ended has_stream_ended = False for chunk in responses_so_far: - if chunk.choices[0].finish_reason is not None: + if chunk.choices and chunk.choices[0].finish_reason is not None: has_stream_ended = True break if has_stream_ended: # convert to model response model_response = cast( - ModelResponse, stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj) + ModelResponse, + stream_chunk_builder( + chunks=responses_so_far, logging_obj=litellm_logging_obj + ), ) # run process_output_response await self.process_output_response( diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index d943662f9e4..ad3d4c932d4 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -319,9 +319,7 @@ class OpenAIResponsesHandler(BaseTranslation): return response if not response_output: - verbose_proxy_logger.debug( - "OpenAI Responses API: Empty output in response" - ) + verbose_proxy_logger.debug("OpenAI Responses API: Empty output in response") return response # Step 1: Extract all text content and tool calls from response output @@ -427,27 +425,30 @@ class OpenAIResponsesHandler(BaseTranslation): handle_raw_dict_callback=None, ) - tool_calls = model_response_choices[0].message.tool_calls - text = model_response_choices[0].message.content - guardrail_inputs = GenericGuardrailAPIInputs() - if text: - guardrail_inputs["texts"] = [text] - if tool_calls: - guardrail_inputs["tool_calls"] = cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - # Include model information from the response if available - response_model = final_chunk.get("response", {}).get("model") - if response_model: - guardrail_inputs["model"] = response_model - if tool_calls or text: - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=guardrail_inputs, - request_data={}, - input_type="response", - logging_obj=litellm_logging_obj, - ) - return responses_so_far + if model_response_choices: + tool_calls = model_response_choices[0].message.tool_calls + text = model_response_choices[0].message.content + guardrail_inputs = GenericGuardrailAPIInputs() + if text: + guardrail_inputs["texts"] = [text] + if tool_calls: + guardrail_inputs["tool_calls"] = cast( + List[ChatCompletionToolCallChunk], tool_calls + ) + # Include model information from the response if available + response_model = final_chunk.get("response", {}).get("model") + if response_model: + guardrail_inputs["model"] = response_model + if tool_calls or text: + _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=guardrail_inputs, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + else: + verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") # model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) # tool_calls = model_response_stream.choices[0].tool_calls # convert openai response to model response @@ -513,11 +514,9 @@ class OpenAIResponsesHandler(BaseTranslation): # Check if it's an OutputText with text if isinstance(content_item, OutputText): if content_item.text: - return True elif isinstance(content_item, dict): if content_item.get("text"): - return True return False diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py new file mode 100644 index 00000000000..82517b7af9e --- /dev/null +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -0,0 +1,236 @@ +""" +Unit tests for Anthropic Messages Guardrail Translation Handler + +Tests the handler's ability to process streaming output for Anthropic Messages API +with guardrail transformations, specifically testing edge cases with empty choices. +""" + +import os +import sys +from typing import Any, List, Literal, Optional +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../../..") +) # Adds the parent directory to the system path + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, +) +from litellm.types.utils import GenericGuardrailAPIInputs + + +class MockPassThroughGuardrail(CustomGuardrail): + """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + """Simply return inputs unchanged""" + return inputs + + +class MockDynamicGuardrail(CustomGuardrail): + """Mock guardrail that records dynamic params from request metadata.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.dynamic_params: Optional[dict] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.dynamic_params = self.get_guardrail_dynamic_request_body_params( + request_data + ) + return inputs + + +class TestAnthropicMessagesHandlerStreamingOutputProcessing: + """Test streaming output processing functionality""" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_model_response(self): + """Test that streaming response with None model_response doesn't raise error + + This test verifies the fix for the bug where accessing model_response.choices[0] + would raise an error when _build_complete_streaming_response returns None. + """ + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Mock _check_streaming_has_ended to return True (stream ended) + # and _build_complete_streaming_response to return None + with patch.object( + handler, "_check_streaming_has_ended", return_value=True + ), patch( + "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response", + return_value=None, + ): + responses_so_far = [b"data: some chunk"] + + # This should not raise an error + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses unchanged + assert result == responses_so_far + + +class TestAnthropicMessagesHandlerInputProcessing: + """Test input processing preserves litellm_metadata for dynamic guardrails.""" + + @pytest.mark.asyncio + async def test_process_input_messages_preserves_litellm_metadata_guardrails(self): + handler = AnthropicMessagesHandler() + guardrail = MockDynamicGuardrail(guardrail_name="cygnal-monitor") + + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "hello"}], + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": {"policy_id": "policy-123"} + } + } + ] + }, + } + + with patch("litellm.proxy.proxy_server.premium_user", True): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data.get("litellm_metadata", {}).get("guardrails") + assert guardrail.dynamic_params == {"policy_id": "policy-123"} + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_choices(self): + """Test that streaming response with empty choices doesn't raise IndexError + + This test verifies the fix for the bug where accessing model_response.choices[0] + would raise IndexError when the response has an empty choices list. + """ + from litellm.types.utils import ModelResponse + + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create a mock response with empty choices + mock_response = ModelResponse( + id="msg_123", + created=1234567890, + model="claude-3", + object="chat.completion", + choices=[], # Empty choices + ) + + # Mock _check_streaming_has_ended to return True (stream ended) + # and _build_complete_streaming_response to return the mock response + with patch.object( + handler, "_check_streaming_has_ended", return_value=True + ), patch( + "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response", + return_value=mock_response, + ): + responses_so_far = [b"data: some chunk"] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_with_valid_choices(self): + """Test that streaming response with valid choices still works correctly""" + from litellm.types.utils import Choices, Message, ModelResponse + + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create a mock response with valid choices + mock_response = ModelResponse( + id="msg_123", + created=1234567890, + model="claude-3", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="Hello world", + role="assistant", + ), + ) + ], + ) + + # Mock _check_streaming_has_ended to return True (stream ended) + # and _build_complete_streaming_response to return the mock response + with patch.object( + handler, "_check_streaming_has_ended", return_value=True + ), patch( + "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response", + return_value=mock_response, + ): + responses_so_far = [b"data: some chunk"] + + # This should process successfully + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_stream_not_ended(self): + """Test that streaming response falls back to text processing when stream hasn't ended""" + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Mock _check_streaming_has_ended to return False (stream not ended) + with patch.object( + handler, "_check_streaming_has_ended", return_value=False + ), patch.object( + handler, "get_streaming_string_so_far", return_value="partial text" + ): + responses_so_far = [b"data: some chunk"] + + # This should process successfully using text-based guardrail + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses + assert result == responses_so_far + + +if __name__ == "__main__": + # Run the tests + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 6c0195d2831..1f5f53d0f0c 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -733,6 +733,154 @@ class TestOpenAIChatCompletionsHandlerToolCallsOutput: assert response.choices[0].finish_reason == "tool_calls" +class MockPassThroughGuardrail(CustomGuardrail): + """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + """Simply return inputs unchanged""" + return inputs + + +class TestOpenAIChatCompletionsHandlerStreamingOutput: + """Test streaming output processing functionality""" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_choices(self): + """Test that streaming response with empty choices doesn't raise IndexError + + This test verifies the fix for the bug where accessing chunk.choices[0] + would raise IndexError when a streaming chunk has an empty choices list. + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create a streaming chunk with empty choices + chunk_with_empty_choices = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[], # Empty choices - this was causing the IndexError + ) + + responses_so_far = [chunk_with_empty_choices] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_with_valid_choices(self): + """Test that streaming response with valid choices still works correctly""" + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create streaming chunks with valid choices + chunk1 = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello"), + finish_reason=None, + ) + ], + ) + + chunk2 = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=" world"), + finish_reason="stop", + ) + ], + ) + + responses_so_far = [chunk1, chunk2] + + # This should process successfully + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_mixed_empty_and_valid_choices_no_finish(self): + """Test streaming response with mix of empty and valid choices chunks (stream not finished) + + This tests the has_stream_ended check when iterating through chunks with mixed choices. + The stream hasn't finished yet (no finish_reason), so it won't trigger stream_chunk_builder. + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Mix of chunks - some with empty choices, some with valid choices + # Stream hasn't finished (no finish_reason) + chunk_empty = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[], + ) + + chunk_valid = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello"), + finish_reason=None, # Stream not finished + ) + ], + ) + + responses_so_far = [chunk_empty, chunk_valid] + + # This should not raise IndexError when checking has_stream_ended + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses + assert result == responses_so_far + + if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a2849ab91a2..ccece8018ff 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -817,3 +817,181 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert task_mappings[0] == (0, 0) assert task_mappings[1] == (0, 1) assert task_mappings[2] == (0, 2) + + +class MockPassThroughGuardrail(CustomGuardrail): + """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + """Simply return inputs unchanged""" + return inputs + + +class TestOpenAIResponsesHandlerStreamingOutputProcessing: + """Test streaming output processing functionality""" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_output(self): + """Test that streaming response with empty output doesn't raise IndexError + + This test verifies the fix for the bug where accessing model_response_choices[0] + would raise IndexError when the response.completed event has an empty output array. + """ + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with empty output + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [], # Empty output - this was causing the IndexError + "status": "completed", + }, + } + ] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_missing_output_key(self): + """Test that streaming response with missing output key doesn't raise IndexError + + This test verifies the handler gracefully handles when the response dict + doesn't contain an 'output' key at all. + """ + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with missing output key + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_123", + "status": "completed", + # No 'output' key - get() will return [] + }, + } + ] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_unrecognized_output_type(self): + """Test that streaming response with unrecognized output types doesn't raise IndexError + + This test verifies the handler gracefully handles when output items are of + unrecognized types that _convert_response_output_to_choices skips over. + """ + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with unrecognized output type + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "unknown_type", # Unrecognized type + "id": "item_123", + "data": "some data", + } + ], + "status": "completed", + }, + } + ] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_with_valid_output(self): + """Test that streaming response with valid output still works correctly""" + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with valid message output + responses_so_far = [ + { + "type": "response.created", + "response": {"id": "resp_123"}, + }, + { + "type": "response.output_item.added", + "item": {"type": "message", "id": "msg_123"}, + }, + { + "type": "response.content_part.added", + "part": {"type": "output_text", "text": ""}, + }, + { + "type": "response.output_text.delta", + "delta": "Hello", + }, + { + "type": "response.output_text.delta", + "delta": " world", + }, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Hello world"}, + ], + } + ], + "status": "completed", + }, + }, + ] + + # This should process successfully + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses + assert result == responses_so_far From 2b25d03046da9dbd42f3d2b9f85ed02400910c1c Mon Sep 17 00:00:00 2001 From: Xiaohan Fu Date: Tue, 3 Feb 2026 17:41:31 -0500 Subject: [PATCH 08/15] Fix fail-open for grayswan and pass metadata to cygnal api endpoint (#19837) * fix fail-open for grayswan; pass metadata to cygnal api endpoint; update docs * pass litellm_metadata to cygnal in payload * switch error msg to const, and clean exception handling. * update pyproject.toml as requested * Revert "update pyproject.toml as requested" This reverts commit 4eece154d056ba33689a5584c86c8fc352bb7cdd. --- .../docs/proxy/guardrails/grayswan.md | 197 +++++++++++------- .../guardrail_hooks/grayswan/grayswan.py | 42 +++- .../guardrail_hooks/test_grayswan.py | 124 +++++++++++ 3 files changed, 278 insertions(+), 85 deletions(-) diff --git a/docs/my-website/docs/proxy/guardrails/grayswan.md b/docs/my-website/docs/proxy/guardrails/grayswan.md index d6efaf15504..6c0ccbc293d 100644 --- a/docs/my-website/docs/proxy/guardrails/grayswan.md +++ b/docs/my-website/docs/proxy/guardrails/grayswan.md @@ -13,20 +13,26 @@ Cygnal returns a `violation` score between `0` and `1` (higher means more likely ### 1. Obtain Credentials -1. Create a Gray Swan account and generate a Cygnal API key. +1. Log in to our Gray Swan platform and generate a Cygnal API key. + + For existing customers, you should already have access to our [platform](https://platform.grayswan.ai). + + For new users, please register at this [page](https://hubs.ly/Q03-sX1J0) and we are more than happy to give you an onboarding! + + 2. Configure environment variables for the LiteLLM proxy host: -```bash -export GRAYSWAN_API_KEY="your-grayswan-key" -export GRAYSWAN_API_BASE="https://api.grayswan.ai" -``` + ```bash + export GRAYSWAN_API_KEY="your-grayswan-key" + export GRAYSWAN_API_BASE="https://api.grayswan.ai" + ``` ### 2. Configure `config.yaml` -Add a guardrail entry that references the Gray Swan integration. Below is a balanced example that monitors both input and output but only blocks once the violation score reaches the configured threshold. +Add a guardrail entry that references the Gray Swan integration. Below is our recommmended settings. ```yaml -model_list: +model_list: # this part is a standard litellm configuration for reference - model_name: openai/gpt-4.1-mini litellm_params: model: openai/gpt-4.1-mini @@ -40,13 +46,14 @@ guardrails: api_key: os.environ/GRAYSWAN_API_KEY api_base: os.environ/GRAYSWAN_API_BASE # optional optional_params: - on_flagged_action: monitor # or "block" + on_flagged_action: passthrough # or "block" or "monitor" violation_threshold: 0.5 # score >= threshold is flagged reasoning_mode: hybrid # off | hybrid | thinking - categories: - safety: "Detect jailbreaks and policy violations" - policy_id: "your-cygnal-policy-id" + policy_id: "your-cygnal-policy-id" # Optional: Your Cygnal policy ID. Defaults to a content safety policy if empty. + streaming_end_of_stream_only: true # For streaming API, only send the assembled message to Cygnal (post_call only). Defaults to false. default_on: true + guardrail_timeout: 30 # Defaults to 30 seconds. Change accordingly. + fail_open: true # Defaults to true; set to false to propagate guardrail errors. general_settings: master_key: "your-litellm-master-key" @@ -65,13 +72,13 @@ litellm --config config.yaml --port 4000 ## Choosing Guardrail Modes -Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements. +Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements. | Mode | When it Runs | Protects | Typical Use Case | |--------------|-------------------|-----------------------|------------------| | `pre_call` | Before LLM call | User input only | Block prompt injection before it reaches the model | | `during_call`| Parallel to call | User input only | Low-latency monitoring without blocking | -| `post_call` | After response | Full conversation | Scan output for policy violations, leaked secrets, or IPI | +| `post_call` | After response | Model Outputs | Scan output for policy violations, leaked secrets, or IPI | When using `during_call` with `on_flagged_action: block` or `on_flagged_action: passthrough`: @@ -81,87 +88,110 @@ When using `during_call` with `on_flagged_action: block` or `on_flagged_action: - The guardrail exception prevents the response from reaching the user, but **does not cancel the running LLM task** - This means you pay full LLM costs while returning an error/passthrough message to the user -**Recommendation:** For cost-sensitive applications, use `pre_call` and `post_call` instead of `during_call` for blocking or passthrough modes. Reserve `during_call` for `monitor` mode where you want low-latency logging without impacting the user experience. +**Recommendation:** Use `pre_call` and `post_call` instead of `during_call` for `passthrough` (or `block`) `on_flagged_action` (see our recommended configuration above). Reserve `during_call` for `monitor` mode ONLY when you want low-latency logging without impacting the user experience. - - +--- -```yaml -guardrails: - - guardrail_name: "cygnal-monitor-only" - litellm_params: - guardrail: grayswan - mode: "during_call" - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: monitor - violation_threshold: 0.6 - default_on: true +## Work with Claude Code + +Follow the official litellm [guide](https://docs.litellm.ai/docs/tutorials/claude_responses_api) on setting up Claude Code with litellm, with the guardrail part mentioned above added to your litellm configuration. Cygnal natively supports coding agent policies defense. Define your own policy or use the provided coding policies on the platform. The example config we show above is also the recommended setup for Claude Code (with the `policy_id` replaced with an appropriate one). + +--- + +## Per-request overrides via `extra_body` + +You can override parts of the Gray Swan guardrail configuration on a per-request basis by passing `litellm_metadata.guardrails[*].grayswan.extra_body`. + +`extra_body` is merged into the Cygnal request body and takes precedence over specific fields from `config.yaml`, which are `policy_id`, `violation_threshold`, and `reasoning_mode`. + +If you include a `metadata` field inside `extra_body`, it is forwarded to the Cygnal API as-is under the request body's `metadata` field. + +Example: + +```bash +curl -X POST "http://0.0.0.0:4000/v1/messages?beta=true" \ + -H "Authorization: Bearer token" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "openrouter/anthropic/claude-sonnet-4.5", + "messages": [{"role": "user", "content": "hello"}], + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": { + "policy_id": "specific policy id you want to use", + "metadata": { + "user": "health-check" + } + } + } + } + ] + } + }' ``` -Best for visibility without blocking. Alerts are logged via LiteLLM’s standard logging callbacks. +OpenAI client: - - +```python +from openai import OpenAI -```yaml -guardrails: - - guardrail_name: "cygnal-block-input" - litellm_params: - guardrail: grayswan - mode: "pre_call" - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: block - violation_threshold: 0.4 - categories: - pii: "Detect sensitive data" - default_on: true +client = OpenAI(api_key="anything", base_url="http://0.0.0.0:4000") + +resp = client.responses.create( + model="openrouter/anthropic/claude-sonnet-4.5", + input="hello", + extra_body={ + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": { + "policy_id": "69038214e5cdb6befc5e991e", + "metadata": {"trace_id": "trace-123"}, + } + } + } + ] + } + }, +) ``` -Stops malicious or sensitive prompts before any tokens are generated. +Anthropic client: - - +```python +from anthropic import Anthropic -```yaml -guardrails: - - guardrail_name: "cygnal-full-coverage" - litellm_params: - guardrail: grayswan - mode: [pre_call, post_call] - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: block - violation_threshold: 0.5 - reasoning_mode: thinking - policy_id: "policy-id-from-grayswan" - default_on: true +client = Anthropic(api_key="anything", base_url="http://0.0.0.0:4000") + +resp = client.messages.create( + model="openrouter/anthropic/claude-sonnet-4.5", + max_tokens=256, + messages=[{"role": "user", "content": "hello"}], + extra_body={ + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": { + "policy_id": "69038214e5cdb6befc5e991e", + "metadata": {"trace_id": "trace-123"}, + } + } + } + ] + } + }, +) ``` -Provides the strongest enforcement by inspecting both prompts and responses. +Notes: - - - -```yaml -guardrails: - - guardrail_name: "cygnal-passthrough" - litellm_params: - guardrail: grayswan - mode: [pre_call, post_call] - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: passthrough - violation_threshold: 0.5 - default_on: true -``` - -Allows requests to proceed without raising a 400 error when content is flagged. Instead of blocking, the model response content is replaced with a detailed violation message including violation score, violated rules, and detection flags (mutation, IPI). **Supported Response Formats:** OpenAI chat/text completions, Anthropic Messages API. Other response types (embeddings, images, etc.) will log a warning and return unchanged. - - - +- The guardrail name (for example, `cygnal-monitor`) must match the `guardrail_name` in `config.yaml`. +- Per-request guardrail overrides may require a premium license, depending on your proxy settings. --- @@ -170,9 +200,14 @@ Allows requests to proceed without raising a 400 error when content is flagged. | Parameter | Type | Description | |---------------------------------------|-----------------|-------------| | `api_key` | string | Gray Swan Cygnal API key. Reads from `GRAYSWAN_API_KEY` if omitted. | +| `api_base` | string | Override for the Gray Swan API base URL. Defaults to `https://api.grayswan.ai` or `GRAYSWAN_API_BASE`. | | `mode` | string or list | Guardrail stages (`pre_call`, `during_call`, `post_call`). | | `optional_params.on_flagged_action` | string | `monitor` (log only), `block` (raise `HTTPException`), or `passthrough` (replace response content with violation message, no 400 error). | -| `.optional_params.violation_threshold`| number (0-1) | Scores at or above this value are considered violations. | +| `optional_params.violation_threshold` | number (0-1) | Scores at or above this value are considered violations. | | `optional_params.reasoning_mode` | string | `off`, `hybrid`, or `thinking`. Enables Cygnal's reasoning capabilities. | | `optional_params.categories` | object | Map of custom category names to descriptions. | | `optional_params.policy_id` | string | Gray Swan policy identifier. | +| `guardrail_timeout` | number | Timeout in seconds for the Cygnal request. Defaults to 30. | +| `fail_open` | boolean | If true, errors contacting Cygnal are logged and the request proceeds; if false, errors propagate. Defaults to treu. | +| `streaming_end_of_stream_only` | boolean | For streaming `post_call`, only send the final assembled response to Cygnal. Defaults to false. | +| `default_on` | boolean | Run the guardrail on every request by default. | diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 2a852cbda08..90f689ed23c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -9,8 +9,10 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, + ModifyResponseException ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -21,6 +23,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +GRAYSWAN_BLOCK_ERROR_MSG = "Blocked by Gray Swan Guardrail" + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -205,9 +209,13 @@ class GraySwanGuardrail(CustomGuardrail): # Get dynamic params from request metadata dynamic_body = self.get_guardrail_dynamic_request_body_params(request_data) or {} + if dynamic_body: + verbose_proxy_logger.debug( + "Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body) + ) # Prepare and send payload - payload = self._prepare_payload(messages, dynamic_body) + payload = self._prepare_payload(messages, dynamic_body, request_data) if payload is None: return inputs @@ -223,6 +231,8 @@ class GraySwanGuardrail(CustomGuardrail): ) return result except Exception as exc: + if self._is_grayswan_exception(exc): + raise end_time = time.time() status_code = getattr(exc, "status_code", None) or getattr( exc, "exception_status_code", None @@ -240,8 +250,20 @@ class GraySwanGuardrail(CustomGuardrail): exc, ) return inputs + if isinstance(exc, GraySwanGuardrailAPIError): + raise exc raise GraySwanGuardrailAPIError(str(exc), status_code=status_code) from exc + def _is_grayswan_exception(self, exc: Exception) -> bool: + # Guardrail decision (passthrough) should always propagate, + # regardless of fail_open. + if isinstance(exc, ModifyResponseException): + return True + detail = getattr(exc, "detail", None) + if isinstance(detail, dict): + return detail.get("error") == GRAYSWAN_BLOCK_ERROR_MSG + return False + # ------------------------------------------------------------------ # Legacy Test Interface (for backward compatibility) # ------------------------------------------------------------------ @@ -324,7 +346,7 @@ class GraySwanGuardrail(CustomGuardrail): raise HTTPException( status_code=400, detail={ - "error": "Blocked by Gray Swan Guardrail", + "error": GRAYSWAN_BLOCK_ERROR_MSG, "violation_location": violation_location, "violation": violation_score, "violated_rules": violated_rules, @@ -445,7 +467,7 @@ class GraySwanGuardrail(CustomGuardrail): raise HTTPException( status_code=400, detail={ - "error": "Blocked by Gray Swan Guardrail", + "error": GRAYSWAN_BLOCK_ERROR_MSG, "violation_location": violation_location, "violation": violation_score, "violated_rules": violated_rules, @@ -494,7 +516,7 @@ class GraySwanGuardrail(CustomGuardrail): } def _prepare_payload( - self, messages: List[Dict[str, str]], dynamic_body: dict + self, messages: List[Dict[str, str]], dynamic_body: dict, request_data: dict ) -> Optional[Dict[str, Any]]: payload: Dict[str, Any] = {"messages": messages} @@ -510,6 +532,18 @@ class GraySwanGuardrail(CustomGuardrail): if reasoning_mode: payload["reasoning_mode"] = reasoning_mode + # Pass through arbitrary metadata when provided via dynamic extra_body. + if "metadata" in dynamic_body: + payload["metadata"] = dynamic_body["metadata"] + + litellm_metadata = request_data.get("litellm_metadata") + if isinstance(litellm_metadata, dict) and litellm_metadata: + cleaned_litellm_metadata = dict(litellm_metadata) + # cleaned_litellm_metadata.pop("user_api_key_auth", None) + sanitized = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + if isinstance(sanitized, dict) and sanitized: + payload["litellm_metadata"] = sanitized + return payload def _format_violation_message( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py index 6dc658827bc..2b00a099328 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -55,6 +55,17 @@ def test_prepare_payload_falls_back_to_guardrail_defaults( assert payload["reasoning_mode"] == "hybrid" +def test_prepare_payload_includes_dynamic_metadata( + grayswan_guardrail: GraySwanGuardrail, +) -> None: + messages = [{"role": "user", "content": "hello"}] + dynamic_body = {"metadata": {"trace_id": "trace-123", "tags": ["a", "b"]}} + + payload = grayswan_guardrail._prepare_payload(messages, dynamic_body) + + assert payload["metadata"] == dynamic_body["metadata"] + + def test_process_response_does_not_block_under_threshold( grayswan_guardrail: GraySwanGuardrail, ) -> None: @@ -160,6 +171,119 @@ async def test_run_guardrail_raises_api_error( await grayswan_guardrail.run_grayswan_guardrail(payload) +@pytest.mark.asyncio +async def test_apply_guardrail_passthrough_not_swallowed_by_fail_open( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-passthrough", + api_key="test-key", + on_flagged_action="passthrough", + violation_threshold=0.2, + fail_open=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.92, "violated_rule_descriptions": []} + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + + with pytest.raises(ModifyResponseException): + await guardrail.apply_guardrail( + inputs={"texts": ["bad"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_block_not_swallowed_by_fail_open( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-block", + api_key="test-key", + on_flagged_action="block", + violation_threshold=0.2, + fail_open=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.92, "violated_rule_descriptions": []} + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs={"texts": ["bad"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_grayswan_http_exception_fail_open_true( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-error", + api_key="test-key", + on_flagged_action="monitor", + violation_threshold=0.2, + fail_open=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.0, "violated_rule_descriptions": []} + + def _fake_process(**_kwargs): + raise HTTPException(status_code=500, detail={"error": "upstream failed"}) + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + monkeypatch.setattr(guardrail, "_process_response_internal", _fake_process) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["ok"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + assert result["texts"] == ["ok"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_grayswan_http_exception_fail_open_false( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-error", + api_key="test-key", + on_flagged_action="monitor", + violation_threshold=0.2, + fail_open=False, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.0, "violated_rule_descriptions": []} + + def _fake_process(**_kwargs): + raise HTTPException(status_code=500, detail={"error": "upstream failed"}) + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + monkeypatch.setattr(guardrail, "_process_response_internal", _fake_process) + + with pytest.raises(GraySwanGuardrailAPIError): + await guardrail.apply_guardrail( + inputs={"texts": ["ok"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + def test_process_response_passthrough_raises_exception_in_pre_call() -> None: """Test that passthrough mode raises ModifyResponseException in pre_call hook.""" guardrail = GraySwanGuardrail( From 0cb6b5876802884ca51038cc23c452dd342dee6a Mon Sep 17 00:00:00 2001 From: naaa760 Date: Wed, 4 Feb 2026 08:56:50 +0530 Subject: [PATCH 09/15] fix(proxy): forward extra_headers in chat --- .../transformation.py | 3 +++ ...responses_transformation_transformation.py | 22 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 57bd05124aa..753a94295b3 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -329,6 +329,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: request_data[key] = value + if headers: + request_data["extra_headers"] = headers + return request_data @staticmethod diff --git a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py index adbaf219079..57352eafaf1 100644 --- a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py @@ -134,3 +134,25 @@ def test_transform_request_with_response_format(): assert result["text"]["format"]["type"] == "json_schema" assert result["text"]["format"]["name"] == "person_schema" assert "schema" in result["text"]["format"] + + +def test_transform_request_includes_extra_headers(): + """Test that transform_request forwards headers as extra_headers for upstream call.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [{"role": "user", "content": "Hello"}] + optional_params = {} + litellm_params = {} + + class MockLoggingObj: + pass + + headers = {"cf-aig-authorization": "secret-token"} + result = handler.transform_request( + model="gpt-5-pro", + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + litellm_logging_obj=MockLoggingObj(), + ) + assert result.get("extra_headers") == headers From 7056d9984ea5f7146fb346d77e4e21ee222c2be1 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 3 Feb 2026 19:57:24 -0800 Subject: [PATCH 10/15] Custom Code Guardrails UI Playground (#20377) * feat(guardrails/): allow custom code execution for guardrails first step in allowing teams to submit custom code for guardrails * feat: custom_code_guardrail.md support passing custom code for guardrails * feat: initial commit adding ui for custom code guardrails allows users to write guardrails based on custom code * feat: expose new test custom code guardrail endpoint allows ui testing playground to sanity check if guardrail is working as expected * fix: fix linting errors * fix: fix max recursion check * fix: fix linting error --- .../simple_guardrail_tutorial.md | 7 +- .../proxy/guardrails/custom_code_guardrail.md | 278 ++++++++ docs/my-website/sidebars.js | 1 + .../out/{404.html => 404/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{budgets.html => budgets/index.html} | 0 .../{caching.html => caching/index.html} | 0 .../index.html} | 0 .../{old-usage.html => old-usage/index.html} | 0 .../{prompts.html => prompts/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{login.html => login/index.html} | 0 .../out/{logs.html => logs/index.html} | 0 .../{callback.html => callback/index.html} | 0 .../{model-hub.html => model-hub/index.html} | 0 .../{model_hub.html => model_hub/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{policies.html => policies/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{ui-theme.html => ui-theme/index.html} | 0 .../out/{teams.html => teams/index.html} | 0 .../{test-key.html => test-key/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{usage.html => usage/index.html} | 0 .../out/{users.html => users/index.html} | 0 .../index.html} | 0 litellm/proxy/_new_secret_config.yaml | 11 + .../proxy/guardrails/guardrail_endpoints.py | 269 ++++++++ .../guardrail_hooks/custom_code/__init__.py | 65 ++ .../custom_code/custom_code_guardrail.py | 372 +++++++++++ .../guardrail_hooks/custom_code/primitives.py | 602 ++++++++++++++++++ litellm/types/guardrails.py | 23 +- .../code_coverage_tests/recursive_detector.py | 1 + .../src/components/guardrails.tsx | 49 +- .../custom_code/CustomCodeEditor.tsx | 188 ++++++ .../custom_code/CustomCodeModal.tsx | 540 ++++++++++++++++ .../custom_code/CustomCodePlayground.tsx | 588 +++++++++++++++++ .../custom_code/custom_code_constants.ts | 194 ++++++ .../guardrails/custom_code/index.ts | 1 + .../src/components/networking.tsx | 83 +++ 49 files changed, 3254 insertions(+), 18 deletions(-) create mode 100644 docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%) rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%) rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%) rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%) rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%) rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%) rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%) rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%) rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%) rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%) rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%) rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%) rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%) rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%) rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeEditor.tsx create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodePlayground.tsx create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/custom_code_constants.ts create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/index.ts diff --git a/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md b/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md index 9c654cd1560..884a7397bde 100644 --- a/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md +++ b/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md @@ -101,12 +101,11 @@ model_list: - model_name: gpt-4 litellm_params: model: gpt-4 - api_key: os.environ/OPENAI_API_KEY + api_key: os.environ/OPENAI_API_KEY -litellm_settings: - guardrails: +guardrails: - guardrail_name: my_guardrail - litellm_params: + litellm_params: guardrail: my_guardrail mode: during_call api_key: os.environ/MY_GUARDRAIL_API_KEY diff --git a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md new file mode 100644 index 00000000000..cb246144497 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md @@ -0,0 +1,278 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Custom Code Guardrail + +Write custom guardrail logic using Python-like code that runs in a sandboxed environment. + +## Quick Start + +### 1. Define the guardrail in config + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: gpt-4 + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: block-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected") + return allow() +``` + +### 2. Start proxy + +```bash +litellm --config config.yaml +``` + +### 3. Test + +```bash +curl -X POST http://localhost:4000/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789"}], + "guardrails": ["block-ssn"] + }' +``` + +## Configuration + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `guardrail` | string | ✅ | Must be `custom_code` | +| `mode` | string | ✅ | When to run: `pre_call`, `post_call`, `during_call` | +| `custom_code` | string | ✅ | Python-like code with `apply_guardrail` function | +| `default_on` | bool | ❌ | Run on all requests (default: `false`) | + +## Writing Custom Code + +### Function Signature + +Your code must define an `apply_guardrail` function: + +```python +def apply_guardrail(inputs, request_data, input_type): + # inputs: see table below + # request_data: {"model": "...", "user_id": "...", "team_id": "...", "metadata": {...}} + # input_type: "request" or "response" + + return allow() # or block() or modify() +``` + +### `inputs` Parameter + +| Field | Type | Description | +|-------|------|-------------| +| `texts` | `List[str]` | Extracted text from the request/response | +| `images` | `List[str]` | Extracted images (for image guardrails) | +| `tools` | `List[dict]` | Tools sent to the LLM | +| `tool_calls` | `List[dict]` | Tool calls returned from the LLM | +| `structured_messages` | `List[dict]` | Full messages with role info (system/user/assistant) | +| `model` | `str` | The model being used | + +### `request_data` Parameter + +| Field | Type | Description | +|-------|------|-------------| +| `model` | `str` | Model name | +| `user_id` | `str` | User ID from API key | +| `team_id` | `str` | Team ID from API key | +| `end_user_id` | `str` | End user ID | +| `metadata` | `dict` | Request metadata | + +### Return Values + +| Function | Description | +|----------|-------------| +| `allow()` | Let request/response through | +| `block(reason)` | Reject with message | +| `modify(texts=[], images=[], tool_calls=[])` | Transform content | + +## Built-in Primitives + +### Regex + +| Function | Description | +|----------|-------------| +| `regex_match(text, pattern)` | Returns `True` if pattern found | +| `regex_replace(text, pattern, replacement)` | Replace all matches | +| `regex_find_all(text, pattern)` | Return list of matches | + +### JSON + +| Function | Description | +|----------|-------------| +| `json_parse(text)` | Parse JSON string, returns `None` on error | +| `json_stringify(obj)` | Convert to JSON string | +| `json_schema_valid(obj, schema)` | Validate against JSON schema | + +### URL + +| Function | Description | +|----------|-------------| +| `extract_urls(text)` | Extract all URLs from text | +| `is_valid_url(url)` | Check if URL is valid | +| `all_urls_valid(text)` | Check all URLs in text are valid | + +### Code Detection + +| Function | Description | +|----------|-------------| +| `detect_code(text)` | Returns `True` if code detected | +| `detect_code_languages(text)` | Returns list of detected languages | +| `contains_code_language(text, ["sql", "python"])` | Check for specific languages | + +### Text Utilities + +| Function | Description | +|----------|-------------| +| `contains(text, substring)` | Check if substring exists | +| `contains_any(text, [substr1, substr2])` | Check if any substring exists | +| `word_count(text)` | Count words | +| `char_count(text)` | Count characters | +| `lower(text)` / `upper(text)` / `trim(text)` | String transforms | + +## Examples + +### Block PII (SSN) + +```python +def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected") + return allow() +``` + +### Redact Email Addresses + +```python +def apply_guardrail(inputs, request_data, input_type): + pattern = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}" + modified = [] + for text in inputs["texts"]: + modified.append(regex_replace(text, pattern, "[EMAIL REDACTED]")) + return modify(texts=modified) +``` + +### Block SQL Injection + +```python +def apply_guardrail(inputs, request_data, input_type): + if input_type != "request": + return allow() + for text in inputs["texts"]: + if contains_code_language(text, ["sql"]): + return block("SQL code not allowed") + return allow() +``` + +### Validate JSON Response + +```python +def apply_guardrail(inputs, request_data, input_type): + if input_type != "response": + return allow() + + schema = { + "type": "object", + "required": ["name", "value"] + } + + for text in inputs["texts"]: + obj = json_parse(text) + if obj is None: + return block("Invalid JSON response") + if not json_schema_valid(obj, schema): + return block("Response missing required fields") + return allow() +``` + +### Check URLs in Response + +```python +def apply_guardrail(inputs, request_data, input_type): + if input_type != "response": + return allow() + for text in inputs["texts"]: + if not all_urls_valid(text): + return block("Response contains invalid URLs") + return allow() +``` + +### Combine Multiple Checks + +```python +def apply_guardrail(inputs, request_data, input_type): + modified = [] + + for text in inputs["texts"]: + # Redact SSN + text = regex_replace(text, r"\d{3}-\d{2}-\d{4}", "[SSN]") + # Redact credit cards + text = regex_replace(text, r"\d{16}", "[CARD]") + modified.append(text) + + # Block SQL in requests + if input_type == "request": + for text in inputs["texts"]: + if contains_code_language(text, ["sql"]): + return block("SQL injection blocked") + + return modify(texts=modified) +``` + +## Sandbox Restrictions + +Custom code runs in a restricted environment: + +- ❌ No `import` statements +- ❌ No file I/O +- ❌ No network access +- ❌ No `exec()` or `eval()` +- ✅ Only LiteLLM-provided primitives available + +## Per-Request Usage + +Enable guardrail per request: + +```bash +curl -X POST http://localhost:4000/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "guardrails": ["block-ssn"] + }' +``` + +## Default On + +Run guardrail on all requests: + +```yaml +litellm_settings: + guardrails: + - guardrail_name: block-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + default_on: true + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + ... +``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 98c8ee6eaa1..20e66fb532d 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -79,6 +79,7 @@ const sidebars = { "proxy/guardrails/panw_prisma_airs", "proxy/guardrails/secret_detection", "proxy/guardrails/custom_guardrail", + "proxy/guardrails/custom_code_guardrail", "proxy/guardrails/prompt_injection", "proxy/guardrails/tool_permission", "proxy/guardrails/zscaler_ai_guard", diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 13eeae14485..6f527e268b2 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -14,3 +14,14 @@ model_list: litellm_params: model: openai/gpt-4.1-mini +guardrails: + - guardrail_name: redact-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected in message") + return allow() \ No newline at end of file diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index b80b02fdc6f..aec5dc1edef 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1344,6 +1344,275 @@ async def get_provider_specific_params(): return provider_params +class TestCustomCodeGuardrailRequest(BaseModel): + """Request model for testing custom code guardrails.""" + + custom_code: str + """The Python-like code containing the apply_guardrail function.""" + + test_input: Dict[str, Any] + """The test input to pass to the guardrail. Should contain 'texts', optionally 'images', 'tools', etc.""" + + input_type: str = "request" + """Whether this is a 'request' or 'response' input type.""" + + request_data: Optional[Dict[str, Any]] = None + """Optional mock request_data (model, user_id, team_id, metadata, etc.).""" + + +class TestCustomCodeGuardrailResponse(BaseModel): + """Response model for testing custom code guardrails.""" + + success: bool + """Whether the test executed successfully (no errors).""" + + result: Optional[Dict[str, Any]] = None + """The guardrail result: action (allow/block/modify), reason, modified_texts, etc.""" + + error: Optional[str] = None + """Error message if execution failed.""" + + error_type: Optional[str] = None + """Type of error: 'compilation' or 'execution'.""" + + +@router.post( + "/guardrails/test_custom_code", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], + response_model=TestCustomCodeGuardrailResponse, +) +async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest): + """ + Test custom code guardrail logic without creating a guardrail. + + This endpoint allows admins to experiment with custom code guardrails by: + 1. Compiling the provided code in a sandbox + 2. Executing the apply_guardrail function with test input + 3. Returning the result (allow/block/modify) + + 👉 [Custom Code Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/custom_code_guardrail) + + Example Request: + ```bash + curl -X POST "http://localhost:4000/guardrails/test_custom_code" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "custom_code": "def apply_guardrail(inputs, request_data, input_type):\\n for text in inputs[\\"texts\\"]:\\n if regex_match(text, r\\"\\\\d{3}-\\\\d{2}-\\\\d{4}\\"):\\n return block(\\"SSN detected\\")\\n return allow()", + "test_input": { + "texts": ["My SSN is 123-45-6789"] + }, + "input_type": "request" + }' + ``` + + Example Success Response (blocked): + ```json + { + "success": true, + "result": { + "action": "block", + "reason": "SSN detected" + }, + "error": null, + "error_type": null + } + ``` + + Example Success Response (allowed): + ```json + { + "success": true, + "result": { + "action": "allow" + }, + "error": null, + "error_type": null + } + ``` + + Example Success Response (modified): + ```json + { + "success": true, + "result": { + "action": "modify", + "texts": ["My SSN is [REDACTED]"] + }, + "error": null, + "error_type": null + } + ``` + + Example Error Response (compilation error): + ```json + { + "success": false, + "result": null, + "error": "Syntax error in custom code: invalid syntax (, line 1)", + "error_type": "compilation" + } + ``` + """ + import concurrent.futures + import re + + from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( + get_custom_code_primitives, + ) + + # Security validation patterns + FORBIDDEN_PATTERNS = [ + # Import statements + (r"\bimport\s+", "import statements are not allowed"), + (r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"), + (r"__import__\s*\(", "__import__() is not allowed"), + # Dangerous builtins + (r"\bexec\s*\(", "exec() is not allowed"), + (r"\beval\s*\(", "eval() is not allowed"), + (r"\bcompile\s*\(", "compile() is not allowed"), + (r"\bopen\s*\(", "open() is not allowed"), + (r"\bgetattr\s*\(", "getattr() is not allowed"), + (r"\bsetattr\s*\(", "setattr() is not allowed"), + (r"\bdelattr\s*\(", "delattr() is not allowed"), + (r"\bglobals\s*\(", "globals() is not allowed"), + (r"\blocals\s*\(", "locals() is not allowed"), + (r"\bvars\s*\(", "vars() is not allowed"), + (r"\bdir\s*\(", "dir() is not allowed"), + (r"\bbreakpoint\s*\(", "breakpoint() is not allowed"), + (r"\binput\s*\(", "input() is not allowed"), + # Dangerous dunder access + (r"__builtins__", "__builtins__ access is not allowed"), + (r"__globals__", "__globals__ access is not allowed"), + (r"__code__", "__code__ access is not allowed"), + (r"__subclasses__", "__subclasses__ access is not allowed"), + (r"__bases__", "__bases__ access is not allowed"), + (r"__mro__", "__mro__ access is not allowed"), + (r"__class__", "__class__ access is not allowed"), + (r"__dict__", "__dict__ access is not allowed"), + (r"__getattribute__", "__getattribute__ access is not allowed"), + (r"__reduce__", "__reduce__ access is not allowed"), + (r"__reduce_ex__", "__reduce_ex__ access is not allowed"), + # OS/system access + (r"\bos\.", "os module access is not allowed"), + (r"\bsys\.", "sys module access is not allowed"), + (r"\bsubprocess\.", "subprocess module access is not allowed"), + ] + + EXECUTION_TIMEOUT_SECONDS = 5 + + try: + # Step 0: Security validation - check for forbidden patterns + code = request.custom_code + for pattern, error_msg in FORBIDDEN_PATTERNS: + if re.search(pattern, code): + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Security violation: {error_msg}", + error_type="compilation", + ) + + # Step 1: Compile the custom code with restricted environment + exec_globals = get_custom_code_primitives().copy() + + # Remove access to builtins to prevent escape + exec_globals["__builtins__"] = {} + + try: + exec(compile(request.custom_code, "", "exec"), exec_globals) + except SyntaxError as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Syntax error in custom code: {e}", + error_type="compilation", + ) + except Exception as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Failed to compile custom code: {e}", + error_type="compilation", + ) + + # Step 2: Verify apply_guardrail function exists + if "apply_guardrail" not in exec_globals: + return TestCustomCodeGuardrailResponse( + success=False, + error="Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)", + error_type="compilation", + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + return TestCustomCodeGuardrailResponse( + success=False, + error="'apply_guardrail' must be a callable function", + error_type="compilation", + ) + + # Step 3: Prepare test inputs + test_inputs = request.test_input + if "texts" not in test_inputs: + test_inputs["texts"] = [] + + # Prepare mock request_data + mock_request_data = request.request_data or {} + safe_request_data = { + "model": mock_request_data.get("model", "test-model"), + "user_id": mock_request_data.get("user_id"), + "team_id": mock_request_data.get("team_id"), + "end_user_id": mock_request_data.get("end_user_id"), + "metadata": mock_request_data.get("metadata", {}), + } + + # Step 4: Execute the function with timeout protection + + def execute_guardrail(): + return apply_fn(test_inputs, safe_request_data, request.input_type) + + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(execute_guardrail) + try: + result = future.result(timeout=EXECUTION_TIMEOUT_SECONDS) + except concurrent.futures.TimeoutError: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution timeout: code took longer than {EXECUTION_TIMEOUT_SECONDS} seconds", + error_type="execution", + ) + except Exception as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution error: {e}", + error_type="execution", + ) + + # Step 5: Validate and return result + if not isinstance(result, dict): + return TestCustomCodeGuardrailResponse( + success=True, + result={ + "action": "allow", + "warning": f"Expected dict result, got {type(result).__name__}. Treating as allow.", + }, + ) + + return TestCustomCodeGuardrailResponse( + success=True, + result=result, + ) + + except Exception as e: + verbose_proxy_logger.exception(f"Error testing custom code guardrail: {e}") + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Unexpected error: {e}", + error_type="execution", + ) + + @router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse) @router.post("/apply_guardrail", response_model=ApplyGuardrailResponse) async def apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py new file mode 100644 index 00000000000..747b188feea --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -0,0 +1,65 @@ +"""Custom code guardrail integration for LiteLLM. + +This module allows users to write custom guardrail logic using Python-like code +that runs in a sandboxed environment with access to LiteLLM-provided primitives. +""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .custom_code_guardrail import CustomCodeGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> CustomCodeGuardrail: + """ + Initialize a custom code guardrail. + + Args: + litellm_params: Configuration parameters including the custom code + guardrail: The guardrail configuration dict + + Returns: + CustomCodeGuardrail instance + """ + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Custom code guardrail requires a guardrail_name") + + # Get the custom code from litellm_params + custom_code = getattr(litellm_params, "custom_code", None) + if not custom_code: + raise ValueError( + "Custom code guardrail requires 'custom_code' in litellm_params" + ) + + custom_code_guardrail = CustomCodeGuardrail( + guardrail_name=guardrail_name, + custom_code=custom_code, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(custom_code_guardrail) + return custom_code_guardrail + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.CUSTOM_CODE.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.CUSTOM_CODE.value: CustomCodeGuardrail, +} + +__all__ = [ + "CustomCodeGuardrail", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py new file mode 100644 index 00000000000..a0ca324411c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -0,0 +1,372 @@ +""" +Custom code guardrail for LiteLLM. + +This module provides a guardrail that executes user-defined Python-like code +to implement custom guardrail logic. The code runs in a sandboxed environment +with access to LiteLLM-provided primitives for common guardrail operations. + +Example custom code: + + def apply_guardrail(inputs, request_data, input_type): + '''Block messages containing SSNs''' + for text in inputs["texts"]: + if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): + return block("Social Security Number detected") + return allow() +""" + +import threading +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +from litellm.types.utils import GenericGuardrailAPIInputs + +from .primitives import get_custom_code_primitives + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +class CustomCodeGuardrailError(Exception): + """Raised when custom code guardrail execution fails.""" + + def __init__(self, message: str, details: Optional[Dict[str, Any]] = None) -> None: + super().__init__(message) + self.details = details or {} + + +class CustomCodeCompilationError(CustomCodeGuardrailError): + """Raised when custom code fails to compile.""" + + +class CustomCodeExecutionError(CustomCodeGuardrailError): + """Raised when custom code fails during execution.""" + + +class CustomCodeGuardrailConfigModel(GuardrailConfigModel): + """Configuration parameters for the custom code guardrail.""" + + custom_code: str + """The Python-like code containing the apply_guardrail function.""" + + +class CustomCodeGuardrail(CustomGuardrail): + """ + Guardrail that executes user-defined Python-like code. + + The code runs in a sandboxed environment that provides: + - Access to LiteLLM primitives (regex_match, json_parse, etc.) + - No file I/O or network access + - No imports allowed + + Users write an `apply_guardrail(inputs, request_data, input_type)` function + that returns one of: + - allow() - let the request/response through + - block(reason) - reject with a message + - modify(texts=...) - transform the content + + Example: + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"password"): + return block("Sensitive content detected") + return allow() + """ + + def __init__( + self, + custom_code: str, + guardrail_name: Optional[str] = "custom_code", + **kwargs: Any, + ) -> None: + """ + Initialize the custom code guardrail. + + Args: + custom_code: The source code containing apply_guardrail function + guardrail_name: Name of this guardrail instance + **kwargs: Additional arguments passed to CustomGuardrail + """ + self.custom_code = custom_code + self._compiled_function: Optional[Any] = None + self._compile_lock = threading.Lock() + self._compile_error: Optional[str] = None + + supported_event_hooks = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ] + + super().__init__( + guardrail_name=guardrail_name, + supported_event_hooks=supported_event_hooks, + **kwargs, + ) + + # Compile the code on initialization + self._compile_custom_code() + + @staticmethod + def get_config_model() -> Optional[Type[GuardrailConfigModel]]: + """Returns the config model for the UI.""" + return CustomCodeGuardrailConfigModel + + def _compile_custom_code(self) -> None: + """ + Compile the custom code and extract the apply_guardrail function. + + The code runs in a sandboxed environment with only the allowed primitives. + """ + with self._compile_lock: + if self._compiled_function is not None: + return + + try: + # Create a restricted execution environment + # Only include our safe primitives + exec_globals = get_custom_code_primitives().copy() + + # Execute the user code in the restricted environment + exec(compile(self.custom_code, "", "exec"), exec_globals) + + # Extract the apply_guardrail function + if "apply_guardrail" not in exec_globals: + raise CustomCodeCompilationError( + "Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)" + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + raise CustomCodeCompilationError( + "'apply_guardrail' must be a callable function" + ) + + self._compiled_function = apply_fn + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}' compiled successfully" + ) + + except SyntaxError as e: + self._compile_error = f"Syntax error in custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e + except CustomCodeCompilationError: + raise + except Exception as e: + self._compile_error = f"Failed to compile custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Apply the custom code guardrail to the inputs. + + This method calls the user-defined apply_guardrail function and + processes its result to determine the appropriate action. + + Args: + inputs: Dictionary containing texts, images, tool_calls + request_data: The original request data with metadata + input_type: "request" for pre-call, "response" for post-call + logging_obj: Optional logging object + + Returns: + GenericGuardrailAPIInputs - possibly modified + + Raises: + HTTPException: If content is blocked + CustomCodeExecutionError: If execution fails + """ + if self._compiled_function is None: + if self._compile_error: + raise CustomCodeExecutionError( + f"Custom code guardrail not compiled: {self._compile_error}" + ) + raise CustomCodeExecutionError("Custom code guardrail not compiled") + + try: + # Prepare inputs dict for the function + + # Prepare request_data with safe subset of information + safe_request_data = self._prepare_safe_request_data(request_data) + + # Execute the custom function + result = self._compiled_function(inputs, safe_request_data, input_type) + + # Process the result + return self._process_result( + result=result, + inputs=inputs, + request_data=request_data, + input_type=input_type, + ) + + except HTTPException: + # Re-raise HTTP exceptions (from block action) + raise + except Exception as e: + verbose_proxy_logger.error( + f"Custom code guardrail '{self.guardrail_name}' execution error: {e}" + ) + raise CustomCodeExecutionError( + f"Custom code guardrail execution failed: {e}", + details={ + "guardrail_name": self.guardrail_name, + "input_type": input_type, + }, + ) from e + + def _prepare_safe_request_data(self, request_data: dict) -> Dict[str, Any]: + """ + Prepare a safe subset of request_data for code execution. + + This filters out sensitive information and provides only what's + needed for guardrail logic. + + Args: + request_data: The full request data + + Returns: + Safe subset of request data + """ + return { + "model": request_data.get("model"), + "user_id": request_data.get("user_api_key_user_id"), + "team_id": request_data.get("user_api_key_team_id"), + "end_user_id": request_data.get("user_api_key_end_user_id"), + "metadata": request_data.get("metadata", {}), + } + + def _process_result( + self, + result: Any, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + ) -> GenericGuardrailAPIInputs: + """ + Process the result from the custom code function. + + Args: + result: The return value from apply_guardrail + inputs: The original inputs + request_data: The request data + input_type: "request" or "response" + + Returns: + GenericGuardrailAPIInputs - possibly modified + + Raises: + HTTPException: If action is "block" + """ + if not isinstance(result, dict): + verbose_proxy_logger.warning( + f"Custom code guardrail '{self.guardrail_name}': " + f"Expected dict result, got {type(result).__name__}. Treating as allow." + ) + return inputs + + action = result.get("action", "allow") + + if action == "allow": + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}': Allowing {input_type}" + ) + return inputs + + elif action == "block": + reason = result.get("reason", "Blocked by custom code guardrail") + detection_info = result.get("detection_info", {}) + + verbose_proxy_logger.info( + f"Custom code guardrail '{self.guardrail_name}': Blocking {input_type} - {reason}" + ) + + is_output = input_type == "response" + + # For pre-call, raise passthrough exception to return synthetic response + if not is_output: + self.raise_passthrough_exception( + violation_message=reason, + request_data=request_data, + detection_info=detection_info, + ) + + # For post-call, raise HTTP exception + raise HTTPException( + status_code=400, + detail={ + "error": reason, + "guardrail": self.guardrail_name, + "detection_info": detection_info, + }, + ) + + elif action == "modify": + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}': Modifying {input_type}" + ) + + # Apply modifications + modified_inputs = dict(inputs) + + if "texts" in result and result["texts"] is not None: + modified_inputs["texts"] = result["texts"] + + if "images" in result and result["images"] is not None: + modified_inputs["images"] = result["images"] + + if "tool_calls" in result and result["tool_calls"] is not None: + modified_inputs["tool_calls"] = result["tool_calls"] + + return cast(GenericGuardrailAPIInputs, modified_inputs) + + else: + verbose_proxy_logger.warning( + f"Custom code guardrail '{self.guardrail_name}': " + f"Unknown action '{action}'. Treating as allow." + ) + return inputs + + def update_custom_code(self, new_code: str) -> None: + """ + Update the custom code and recompile. + + This method allows hot-reloading of guardrail logic without + restarting the server. + + Args: + new_code: The new source code + + Raises: + CustomCodeCompilationError: If the new code fails to compile + """ + with self._compile_lock: + # Reset state + old_function = self._compiled_function + old_code = self.custom_code + self._compiled_function = None + self._compile_error = None + + try: + self.custom_code = new_code + self._compile_custom_code() + verbose_proxy_logger.info( + f"Custom code guardrail '{self.guardrail_name}': Code updated successfully" + ) + except CustomCodeCompilationError: + # Rollback on failure + self.custom_code = old_code + self._compiled_function = old_function + raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py new file mode 100644 index 00000000000..695e59977c8 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -0,0 +1,602 @@ +""" +Built-in primitives provided to custom code guardrails. + +These functions are injected into the custom code execution environment +and provide safe, sandboxed functionality for common guardrail operations. +""" + +import json +import re +from typing import Any, Dict, List, Optional, Tuple, Type, Union +from urllib.parse import urlparse + +from litellm._logging import verbose_proxy_logger + +# ============================================================================= +# Result Types - Used by Starlark code to return guardrail decisions +# ============================================================================= + + +def allow() -> Dict[str, Any]: + """ + Allow the request/response to proceed unchanged. + + Returns: + Dict indicating the request should be allowed + """ + return {"action": "allow"} + + +def block( + reason: str, detection_info: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + """ + Block the request/response with a reason. + + Args: + reason: Human-readable reason for blocking + detection_info: Optional additional detection metadata + + Returns: + Dict indicating the request should be blocked + """ + result: Dict[str, Any] = {"action": "block", "reason": reason} + if detection_info: + result["detection_info"] = detection_info + return result + + +def modify( + texts: Optional[List[str]] = None, + images: Optional[List[Any]] = None, + tool_calls: Optional[List[Any]] = None, +) -> Dict[str, Any]: + """ + Modify the request/response content. + + Args: + texts: Modified text content (if None, keeps original) + images: Modified image content (if None, keeps original) + tool_calls: Modified tool calls (if None, keeps original) + + Returns: + Dict indicating the content should be modified + """ + result: Dict[str, Any] = {"action": "modify"} + if texts is not None: + result["texts"] = texts + if images is not None: + result["images"] = images + if tool_calls is not None: + result["tool_calls"] = tool_calls + return result + + +# ============================================================================= +# Regex Primitives +# ============================================================================= + + +def regex_match(text: str, pattern: str, flags: int = 0) -> bool: + """ + Check if a regex pattern matches anywhere in the text. + + Args: + text: The text to search in + pattern: The regex pattern to match + flags: Optional regex flags (default: 0) + + Returns: + True if pattern matches, False otherwise + """ + try: + return bool(re.search(pattern, text, flags)) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_match error: {e}") + return False + + +def regex_match_all(text: str, pattern: str, flags: int = 0) -> bool: + """ + Check if a regex pattern matches the entire text. + + Args: + text: The text to match + pattern: The regex pattern + flags: Optional regex flags + + Returns: + True if pattern matches entire text, False otherwise + """ + try: + return bool(re.fullmatch(pattern, text, flags)) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_match_all error: {e}") + return False + + +def regex_replace(text: str, pattern: str, replacement: str, flags: int = 0) -> str: + """ + Replace all occurrences of a pattern in text. + + Args: + text: The text to modify + pattern: The regex pattern to find + replacement: The replacement string + flags: Optional regex flags + + Returns: + The text with replacements applied + """ + try: + return re.sub(pattern, replacement, text, flags=flags) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_replace error: {e}") + return text + + +def regex_find_all(text: str, pattern: str, flags: int = 0) -> List[str]: + """ + Find all occurrences of a pattern in text. + + Args: + text: The text to search + pattern: The regex pattern to find + flags: Optional regex flags + + Returns: + List of all matches + """ + try: + return re.findall(pattern, text, flags) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_find_all error: {e}") + return [] + + +# ============================================================================= +# JSON Primitives +# ============================================================================= + + +def json_parse(text: str) -> Optional[Any]: + """ + Parse a JSON string into a Python object. + + Args: + text: The JSON string to parse + + Returns: + Parsed Python object, or None if parsing fails + """ + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError) as e: + verbose_proxy_logger.debug(f"Starlark json_parse error: {e}") + return None + + +def json_stringify(obj: Any) -> str: + """ + Convert a Python object to a JSON string. + + Args: + obj: The object to serialize + + Returns: + JSON string representation + """ + try: + return json.dumps(obj) + except (TypeError, ValueError) as e: + verbose_proxy_logger.warning(f"Starlark json_stringify error: {e}") + return "" + + +def json_schema_valid(obj: Any, schema: Dict[str, Any]) -> bool: + """ + Validate an object against a JSON schema. + + Args: + obj: The object to validate + schema: The JSON schema to validate against + + Returns: + True if valid, False otherwise + """ + try: + # Try to import jsonschema, fall back to basic validation if not available + try: + import jsonschema + + jsonschema.validate(instance=obj, schema=schema) + return True + except ImportError: + # Basic validation without jsonschema library + return _basic_json_schema_validate(obj, schema) + except Exception as validation_error: + # Catch jsonschema.ValidationError and other validation errors + if "ValidationError" in type(validation_error).__name__: + return False + raise + except Exception as e: + verbose_proxy_logger.warning(f"Custom code json_schema_valid error: {e}") + return False + + +def _basic_json_schema_validate( + obj: Any, schema: Dict[str, Any], max_depth: int = 50 +) -> bool: + """ + Basic JSON schema validation without external library. + Handles: type, required, properties + + Uses an iterative approach with a stack to avoid recursion limits. + max_depth limits nesting to prevent infinite loops from circular schemas. + """ + type_map: Dict[str, Union[Type, Tuple[Type, ...]]] = { + "object": dict, + "array": list, + "string": str, + "number": (int, float), + "integer": int, + "boolean": bool, + "null": type(None), + } + + # Stack of (obj, schema, depth) tuples to process + stack: List[Tuple[Any, Dict[str, Any], int]] = [(obj, schema, 0)] + + while stack: + current_obj, current_schema, depth = stack.pop() + + # Circuit breaker: stop if we've gone too deep + if depth > max_depth: + return False + + # Check type + schema_type = current_schema.get("type") + if schema_type: + expected_type = type_map.get(schema_type) + if expected_type is not None and not isinstance(current_obj, expected_type): + return False + + # Check required fields and properties for dicts + if isinstance(current_obj, dict): + required = current_schema.get("required", []) + for field in required: + if field not in current_obj: + return False + + # Queue property validations + properties = current_schema.get("properties", {}) + for prop_name, prop_schema in properties.items(): + if prop_name in current_obj: + stack.append((current_obj[prop_name], prop_schema, depth + 1)) + + return True + + +# ============================================================================= +# URL Primitives +# ============================================================================= + + +# Common URL pattern for extraction +_URL_PATTERN = re.compile( + r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[^\s]*", re.IGNORECASE +) + + +def extract_urls(text: str) -> List[str]: + """ + Extract all URLs from text. + + Args: + text: The text to search for URLs + + Returns: + List of URLs found in the text + """ + return _URL_PATTERN.findall(text) + + +def is_valid_url(url: str) -> bool: + """ + Check if a URL is syntactically valid. + + Args: + url: The URL to validate + + Returns: + True if the URL is valid, False otherwise + """ + try: + result = urlparse(url) + return all([result.scheme, result.netloc]) + except Exception: + return False + + +def all_urls_valid(text: str) -> bool: + """ + Check if all URLs in text are valid. + + Args: + text: The text containing URLs + + Returns: + True if all URLs are valid (or no URLs), False otherwise + """ + urls = extract_urls(text) + return all(is_valid_url(url) for url in urls) + + +def get_url_domain(url: str) -> Optional[str]: + """ + Extract the domain from a URL. + + Args: + url: The URL to parse + + Returns: + The domain, or None if invalid + """ + try: + result = urlparse(url) + return result.netloc if result.netloc else None + except Exception: + return None + + +# ============================================================================= +# Code Detection Primitives +# ============================================================================= + + +# Common code patterns for detection +_CODE_PATTERNS = { + "sql": [ + r"\b(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE)\b.*\b(FROM|INTO|TABLE|SET|WHERE)\b", + r"\b(SELECT)\s+[\w\*,\s]+\s+FROM\s+\w+", + r"\b(INSERT\s+INTO|UPDATE\s+\w+\s+SET|DELETE\s+FROM)\b", + ], + "python": [ + r"^\s*(def|class|import|from|if|for|while|try|except|with)\s+", + r"^\s*@\w+", # decorators + r"\b(print|len|range|str|int|float|list|dict|set)\s*\(", + ], + "javascript": [ + r"\b(function|const|let|var|class|import|export)\s+", + r"=>", # arrow functions + r"\b(console\.(log|error|warn))\s*\(", + ], + "typescript": [ + r":\s*(string|number|boolean|any|void|never)\b", + r"\b(interface|type|enum)\s+\w+", + r"<[A-Z]\w*>", # generics + ], + "java": [ + r"\b(public|private|protected)\s+(static\s+)?(class|void|int|String)\b", + r"\bSystem\.(out|err)\.print", + ], + "go": [ + r"\bfunc\s+\w+\s*\(", + r"\b(package|import)\s+", + r":=", # short variable declaration + ], + "rust": [ + r"\b(fn|let|mut|impl|struct|enum|pub|mod)\s+", + r"->", # return type + r"\b(println!|format!)\s*\(", + ], + "shell": [ + r"^#!.*\b(bash|sh|zsh)\b", + r"\b(echo|grep|sed|awk|cat|ls|cd|mkdir|rm)\s+", + r"\$\{?\w+\}?", # variable expansion + ], + "html": [ + r"<\s*(html|head|body|div|span|p|a|img|script|style)\b[^>]*>", + r"", + ], + "css": [ + r"\{[^}]*:\s*[^}]+;[^}]*\}", + r"@(media|keyframes|import|font-face)\b", + ], +} + + +def detect_code(text: str) -> bool: + """ + Check if text contains code of any language. + + Args: + text: The text to check + + Returns: + True if code is detected, False otherwise + """ + return len(detect_code_languages(text)) > 0 + + +def detect_code_languages(text: str) -> List[str]: + """ + Detect which programming languages are present in text. + + Args: + text: The text to analyze + + Returns: + List of detected language names + """ + detected = [] + for lang, patterns in _CODE_PATTERNS.items(): + for pattern in patterns: + try: + if re.search(pattern, text, re.IGNORECASE | re.MULTILINE): + detected.append(lang) + break # Only add each language once + except re.error: + continue + return detected + + +def contains_code_language(text: str, languages: List[str]) -> bool: + """ + Check if text contains code from specific languages. + + Args: + text: The text to check + languages: List of language names to check for + + Returns: + True if any of the specified languages are detected + """ + detected = detect_code_languages(text) + return any(lang.lower() in [d.lower() for d in detected] for lang in languages) + + +# ============================================================================= +# Text Utility Primitives +# ============================================================================= + + +def contains(text: str, substring: str) -> bool: + """ + Check if text contains a substring. + + Args: + text: The text to search in + substring: The substring to find + + Returns: + True if substring is found, False otherwise + """ + return substring in text + + +def contains_any(text: str, substrings: List[str]) -> bool: + """ + Check if text contains any of the given substrings. + + Args: + text: The text to search in + substrings: List of substrings to find + + Returns: + True if any substring is found, False otherwise + """ + return any(s in text for s in substrings) + + +def contains_all(text: str, substrings: List[str]) -> bool: + """ + Check if text contains all of the given substrings. + + Args: + text: The text to search in + substrings: List of substrings to find + + Returns: + True if all substrings are found, False otherwise + """ + return all(s in text for s in substrings) + + +def word_count(text: str) -> int: + """ + Count the number of words in text. + + Args: + text: The text to count words in + + Returns: + Number of words + """ + return len(text.split()) + + +def char_count(text: str) -> int: + """ + Count the number of characters in text. + + Args: + text: The text to count characters in + + Returns: + Number of characters + """ + return len(text) + + +def lower(text: str) -> str: + """Convert text to lowercase.""" + return text.lower() + + +def upper(text: str) -> str: + """Convert text to uppercase.""" + return text.upper() + + +def trim(text: str) -> str: + """Remove leading and trailing whitespace.""" + return text.strip() + + +# ============================================================================= +# Primitives Registry +# ============================================================================= + + +def get_custom_code_primitives() -> Dict[str, Any]: + """ + Get all primitives to inject into the custom code environment. + + Returns: + Dict of function name to function + """ + return { + # Result types + "allow": allow, + "block": block, + "modify": modify, + # Regex + "regex_match": regex_match, + "regex_match_all": regex_match_all, + "regex_replace": regex_replace, + "regex_find_all": regex_find_all, + # JSON + "json_parse": json_parse, + "json_stringify": json_stringify, + "json_schema_valid": json_schema_valid, + # URL + "extract_urls": extract_urls, + "is_valid_url": is_valid_url, + "all_urls_valid": all_urls_valid, + "get_url_domain": get_url_domain, + # Code detection + "detect_code": detect_code, + "detect_code_languages": detect_code_languages, + "contains_code_language": contains_code_language, + # Text utilities + "contains": contains, + "contains_any": contains_any, + "contains_all": contains_all, + "word_count": word_count, + "char_count": char_count, + "lower": lower, + "upper": upper, + "trim": trim, + # Python builtins (safe subset) + "len": len, + "str": str, + "int": int, + "float": float, + "bool": bool, + "list": list, + "dict": dict, + "True": True, + "False": False, + "None": None, + } diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ed463491043..3f7de10dde6 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -14,14 +14,14 @@ from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( + ContentFilterCategoryConfig, ) from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( - ContentFilterCategoryConfig, +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, ) """ @@ -68,6 +68,7 @@ class SupportedGuardrailIntegrations(Enum): PROMPT_SECURITY = "prompt_security" GENERIC_GUARDRAIL_API = "generic_guardrail_api" QUALIFIRE = "qualifire" + CUSTOM_CODE = "custom_code" class Role(Enum): @@ -296,13 +297,7 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = Field( default=None, description="Configuration for PII entity types and actions" ) - presidio_filter_scope: Literal["input", "output", "both"] = Field( - default="both", - description=( - "Where to apply Presidio checks: 'input' runs on user → model traffic, " - "'output' runs on model → user traffic, and 'both' applies to both." - ), - ) + presidio_score_thresholds: Optional[Dict[Union[PiiEntityType, str], float]] = Field( default=None, description=( @@ -656,6 +651,12 @@ class BaseLitellmParams( description="Additional provider-specific parameters for generic guardrail APIs", ) + # Custom code guardrail params + custom_code: Optional[str] = Field( + default=None, + description="Python-like code containing the apply_guardrail function for custom guardrail logic", + ) + model_config = ConfigDict(extra="allow", protected_namespaces=()) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index ed7595bb023..71e7798b09e 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -40,6 +40,7 @@ IGNORE_FUNCTIONS = [ "filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. "__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion. "_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation. + "_basic_json_schema_validate", # max depth set. "extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing. ] diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 26b9acb6a26..d0031b872f7 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -1,5 +1,7 @@ import React, { useState, useEffect } from "react"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; +import { Dropdown } from "antd"; +import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons"; import { getGuardrailsList, deleteGuardrailCall } from "./networking"; import AddGuardrailForm from "./guardrails/add_guardrail_form"; import GuardrailTable from "./guardrails/guardrail_table"; @@ -10,6 +12,7 @@ import NotificationsManager from "./molecules/notifications_manager"; import { Guardrail, GuardrailDefinitionLocation } from "./guardrails/types"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { getGuardrailLogoAndName } from "./guardrails/guardrail_info_helpers"; +import { CustomCodeModal } from "./guardrails/custom_code"; interface GuardrailsPanelProps { accessToken: string | null; @@ -37,6 +40,7 @@ interface GuardrailsResponse { const GuardrailsPanel: React.FC = ({ accessToken, userRole }) => { const [guardrailsList, setGuardrailsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); + const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); @@ -74,10 +78,21 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole setIsAddModalVisible(true); }; + const handleAddCustomCodeGuardrail = () => { + if (selectedGuardrailId) { + setSelectedGuardrailId(null); + } + setIsCustomCodeModalVisible(true); + }; + const handleCloseModal = () => { setIsAddModalVisible(false); }; + const handleCloseCustomCodeModal = () => { + setIsCustomCodeModalVisible(false); + }; + const handleSuccess = () => { fetchGuardrails(); }; @@ -128,9 +143,30 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole
- + , + label: "Add Provider Guardrail", + onClick: handleAddGuardrail, + }, + { + key: "custom_code", + icon: , + label: "Create Custom Code Guardrail", + onClick: handleAddCustomCodeGuardrail, + }, + ], + }} + trigger={["click"]} + disabled={!accessToken} + > + +
{selectedGuardrailId ? ( @@ -159,6 +195,13 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole onSuccess={handleSuccess} /> + + void; + height?: string; + placeholder?: string; + disabled?: boolean; +} + +const CustomCodeEditor: React.FC = ({ + value, + onChange, + height = "350px", + placeholder = `def apply_guardrail(inputs, request_data, input_type): + # inputs: contains texts, images, tools, tool_calls, structured_messages, model + # request_data: contains model, user_id, team_id, end_user_id, metadata + # input_type: "request" or "response" + + for text in inputs["texts"]: + # Example: Block if SSN pattern is detected + if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): + return block("SSN detected in message") + + return allow()`, + disabled = false, +}) => { + const textareaRef = useRef(null); + const [activeTab, setActiveTab] = useState("edit"); + const [cursorPosition, setCursorPosition] = useState({ line: 1, column: 1 }); + + // Calculate cursor position + const updateCursorPosition = () => { + if (textareaRef.current) { + const textarea = textareaRef.current; + const textBeforeCursor = value.substring(0, textarea.selectionStart); + const lines = textBeforeCursor.split("\n"); + const line = lines.length; + const column = lines[lines.length - 1].length + 1; + setCursorPosition({ line, column }); + } + }; + + // Handle tab key for indentation + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Tab") { + e.preventDefault(); + const textarea = e.currentTarget; + const start = textarea.selectionStart; + const end = textarea.selectionEnd; + + // Insert 4 spaces at cursor position + const newValue = value.substring(0, start) + " " + value.substring(end); + onChange(newValue); + + // Move cursor after the inserted spaces + setTimeout(() => { + textarea.selectionStart = textarea.selectionEnd = start + 4; + }, 0); + } + }; + + const lineCount = value.split("\n").length; + + const tabItems = [ + { + key: "edit", + label: ( + + + Edit + + ), + children: ( +
+ {/* Line numbers */} +
+ {Array.from({ length: Math.max(lineCount, 15) }, (_, i) => ( +
+ {i + 1} +
+ ))} +
+ + {/* Code editor */} +