test guardrails with API Key

This commit is contained in:
Ishaan Jaff 2024-08-20 08:38:14 -07:00
parent 4ec6d8ff50
commit aceab2669f
2 changed files with 106 additions and 7 deletions

View file

@ -318,13 +318,33 @@ async def add_litellm_data_to_request(
# Guardrails
move_guardrails_to_metadata(
data=data, _metadata_variable_name=_metadata_variable_name
data=data,
_metadata_variable_name=_metadata_variable_name,
user_api_key_dict=user_api_key_dict,
)
return data
def move_guardrails_to_metadata(data: dict, _metadata_variable_name: str):
def move_guardrails_to_metadata(
data: dict,
_metadata_variable_name: str,
user_api_key_dict: UserAPIKeyAuth,
):
"""
Heper to add guardrails from request to metadata
- If guardrails set on API Key metadata then sets guardrails on request metadata
- If guardrails not set on API key, then checks request metadata
"""
if user_api_key_dict.metadata:
if "guardrails" in user_api_key_dict.metadata:
data[_metadata_variable_name]["guardrails"] = user_api_key_dict.metadata[
"guardrails"
]
return
if "guardrails" in data:
data[_metadata_variable_name]["guardrails"] = data["guardrails"]
del data["guardrails"]