Spaces:
Configuration error
Configuration error
File size: 3,331 Bytes
447ebeb |
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 |
from typing import Dict, List, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
# v2 implementation
from litellm.types.guardrails import Guardrail, GuardrailItem, GuardrailItemSpec
all_guardrails: List[GuardrailItem] = []
def initialize_guardrails(
guardrails_config: List[Dict[str, GuardrailItemSpec]],
premium_user: bool,
config_file_path: str,
litellm_settings: dict,
) -> Dict[str, GuardrailItem]:
try:
verbose_proxy_logger.debug(f"validating guardrails passed {guardrails_config}")
global all_guardrails
for item in guardrails_config:
"""
one item looks like this:
{'prompt_injection': {'callbacks': ['lakera_prompt_injection', 'prompt_injection_api_2'], 'default_on': True, 'enabled_roles': ['user']}}
"""
for k, v in item.items():
guardrail_item = GuardrailItem(**v, guardrail_name=k)
all_guardrails.append(guardrail_item)
litellm.guardrail_name_config_map[k] = guardrail_item
# set appropriate callbacks if they are default on
default_on_callbacks = set()
callback_specific_params = {}
for guardrail in all_guardrails:
verbose_proxy_logger.debug(guardrail.guardrail_name)
verbose_proxy_logger.debug(guardrail.default_on)
callback_specific_params.update(guardrail.callback_args)
if guardrail.default_on is True:
# add these to litellm callbacks if they don't exist
for callback in guardrail.callbacks:
if callback not in litellm.callbacks:
default_on_callbacks.add(callback)
if guardrail.logging_only is True:
if callback == "presidio":
callback_specific_params["presidio"] = {"logging_only": True} # type: ignore
default_on_callbacks_list = list(default_on_callbacks)
if len(default_on_callbacks_list) > 0:
initialize_callbacks_on_proxy(
value=default_on_callbacks_list,
premium_user=premium_user,
config_file_path=config_file_path,
litellm_settings=litellm_settings,
callback_specific_params=callback_specific_params,
)
return litellm.guardrail_name_config_map
except Exception as e:
verbose_proxy_logger.exception(
"error initializing guardrails {}".format(str(e))
)
raise e
"""
Map guardrail_name: <pre_call>, <post_call>, during_call
"""
def init_guardrails_v2(
all_guardrails: List[Dict],
config_file_path: Optional[str] = None,
):
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
guardrail_list: List[Guardrail] = []
for guardrail in all_guardrails:
initialized_guardrail = IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(
guardrail=guardrail,
config_file_path=config_file_path,
)
if initialized_guardrail:
guardrail_list.append(initialized_guardrail)
verbose_proxy_logger.info(f"\nGuardrail List:{guardrail_list}\n")
|