diff --git a/litellm/proxy/guardrails/guardrail_helpers.py b/litellm/proxy/guardrails/guardrail_helpers.py index 8a25abf3a9..682428cc9d 100644 --- a/litellm/proxy/guardrails/guardrail_helpers.py +++ b/litellm/proxy/guardrails/guardrail_helpers.py @@ -1,5 +1,6 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.guardrails.init_guardrails import guardrail_name_config_map +from litellm.proxy.proxy_server import UserAPIKeyAuth from litellm.types.guardrails import * @@ -47,3 +48,44 @@ async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> b return False return True + + +async def should_proceed_based_on_api_key( + user_api_key_dict: UserAPIKeyAuth, guardrail_name: str +) -> bool: + """ + checks if this guardrail should be applied to this call + """ + if user_api_key_dict.permissions is not None: + # { prompt_injection: true, rail_2: false } + verbose_proxy_logger.debug( + "Guardrails valid for API Key= %s - checking which to apply", + user_api_key_dict.permissions, + ) + + if not isinstance(user_api_key_dict.permissions, dict): + verbose_proxy_logger.error( + "API Key permissions must be a dict - %s running guardrail %s", + user_api_key_dict, + guardrail_name, + ) + return True + + for _guardrail_name, should_run in user_api_key_dict.permissions.items(): + if should_run is False: + verbose_proxy_logger.debug( + "Guardrail %s skipped because request set to False", + _guardrail_name, + ) + continue + + # lookup the guardrail in guardrail_name_config_map + guardrail_item: GuardrailItem = guardrail_name_config_map[_guardrail_name] + + guardrail_callbacks = guardrail_item.callbacks + if guardrail_name in guardrail_callbacks: + return True + + # Do not proceeed if - "metadata": { "guardrails": { "lakera_prompt_injection": false } } + return False + return True