File size: 4,278 Bytes
e3278e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import os
import sys
from typing import Dict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.proxy_server import LiteLLM_TeamTable, UserAPIKeyAuth
from litellm.types.guardrails import *
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
def can_modify_guardrails(team_obj: Optional[LiteLLM_TeamTable]) -> bool:
if team_obj is None:
return True
team_metadata = team_obj.metadata or {}
if team_metadata.get("guardrails", None) is not None and isinstance(
team_metadata.get("guardrails"), Dict
):
if team_metadata.get("guardrails", {}).get("modify_guardrails", None) is False:
return False
return True
async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> bool:
"""
checks if this guardrail should be applied to this call
"""
if "metadata" in data and isinstance(data["metadata"], dict):
if "guardrails" in data["metadata"]:
# expect users to pass
# guardrails: { prompt_injection: true, rail_2: false }
request_guardrails = data["metadata"]["guardrails"]
verbose_proxy_logger.debug(
"Guardrails %s passed in request - checking which to apply",
request_guardrails,
)
requested_callback_names = []
# v1 implementation of this
if isinstance(request_guardrails, dict):
# get guardrail configs from `init_guardrails.py`
# for all requested guardrails -> get their associated callbacks
for _guardrail_name, should_run in request_guardrails.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 = litellm.guardrail_name_config_map[
_guardrail_name
]
guardrail_callbacks = guardrail_item.callbacks
requested_callback_names.extend(guardrail_callbacks)
verbose_proxy_logger.debug(
"requested_callback_names %s", requested_callback_names
)
if guardrail_name in requested_callback_names:
return True
# Do no proceeed if - "metadata": { "guardrails": { "lakera_prompt_injection": false } }
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 = litellm.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
|