scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

guardrails list sensitive values fix

Browse source
This commit is contained in:
yuneng-jiang 2026-02-09 17:11:54 -08:00
parent 35eb303098
commit 6b938f81d1
2 changed files with 129 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -150,6 +150,7 @@ async def list_guardrails_v2():
}
```
"""
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
@ -164,11 +165,24 @@ async def list_guardrails_v2():
guardrail_configs: List[GuardrailInfoResponse] = []
seen_guardrail_ids = set()
for guardrail in guardrails:
litellm_params: Optional[Union[LitellmParams, dict]] = guardrail.get(
"litellm_params"
)
litellm_params_dict = (
litellm_params.model_dump(exclude_none=True)
if isinstance(litellm_params, LitellmParams)
else litellm_params
) or {}
masked_litellm_params_dict = _get_masked_values(
litellm_params_dict,
unmasked_length=4,
number_of_asterisks=4,
)
guardrail_configs.append(
GuardrailInfoResponse(
guardrail_id=guardrail.get("guardrail_id"),
guardrail_name=guardrail.get("guardrail_name"),
litellm_params=guardrail.get("litellm_params"),
litellm_params=masked_litellm_params_dict,
guardrail_info=guardrail.get("guardrail_info"),
created_at=guardrail.get("created_at"),
updated_at=guardrail.get("updated_at"),
@ -182,11 +196,19 @@ async def list_guardrails_v2():
for guardrail in in_memory_guardrails:
# only add guardrails that are not in DB guardrail list already
if guardrail.get("guardrail_id") not in seen_guardrail_ids:
in_memory_litellm_params = dict(
guardrail.get("litellm_params") or {}
)
masked_in_memory_litellm_params = _get_masked_values(
in_memory_litellm_params,
unmasked_length=4,
number_of_asterisks=4,
)
guardrail_configs.append(
GuardrailInfoResponse(
guardrail_id=guardrail.get("guardrail_id"),
guardrail_name=guardrail.get("guardrail_name"),
litellm_params=dict(guardrail.get("litellm_params") or {}),
litellm_params=masked_in_memory_litellm_params,
guardrail_info=dict(guardrail.get("guardrail_info") or {}),
guardrail_definition_location="config",
)

105
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -149,6 +149,111 @@ async def test_list_guardrails_v2_with_db_and_config(
assert isinstance(config_guardrail.litellm_params, BaseLitellmParams)
@pytest.mark.asyncio
async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker):
"""Test that sensitive litellm_params are masked for DB guardrails in list response"""
db_guardrail_with_secrets = {
"guardrail_id": "secret-db-guardrail",
"guardrail_name": "DB Guardrail with Secrets",
"litellm_params": {
"guardrail": "azure/text_moderations",
"mode": "pre_call",
"api_key": "sk-1234567890abcdef",
"api_base": "https://api.secret.example.com",
},
"guardrail_info": {"description": "Test guardrail"},
"created_at": datetime.now(),
"updated_at": datetime.now(),
}
mock_prisma_client = mocker.Mock()
mock_prisma_client.db = mocker.Mock()
mock_prisma_client.db.litellm_guardrailstable = mocker.Mock()
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
return_value=[db_guardrail_with_secrets]
)
mock_in_memory_handler = mocker.Mock()
mock_in_memory_handler.list_in_memory_guardrails.return_value = []
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,
)
response = await list_guardrails_v2()
assert len(response.guardrails) == 1
guardrail = response.guardrails[0]
litellm_params = guardrail.litellm_params
if isinstance(litellm_params, dict):
params = litellm_params
else:
params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params)
# Sensitive keys (containing "key", "secret", "token", etc.) should be masked
assert params["api_key"] != "sk-1234567890abcdef"
assert "****" in str(params["api_key"])
# Non-sensitive keys should remain unchanged
assert params["guardrail"] == "azure/text_moderations"
assert params["mode"] == "pre_call"
assert params["api_base"] == "https://api.secret.example.com"
@pytest.mark.asyncio
async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mocker):
"""Test that sensitive litellm_params are masked for in-memory/config guardrails in list response"""
config_guardrail_with_secrets = {
"guardrail_id": "secret-config-guardrail",
"guardrail_name": "Config Guardrail with Secrets",
"litellm_params": {
"guardrail": "bedrock",
"mode": "during_call",
"api_key": "my-secret-bedrock-key",
"vertex_credentials": "{sensitive_creds}",
},
"guardrail_info": {"description": "Test guardrail from config"},
}
mock_prisma_client = mocker.Mock()
mock_prisma_client.db = mocker.Mock()
mock_prisma_client.db.litellm_guardrailstable = mocker.Mock()
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
return_value=[]
)
mock_in_memory_handler = mocker.Mock()
mock_in_memory_handler.list_in_memory_guardrails.return_value = [
config_guardrail_with_secrets
]
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,
)
response = await list_guardrails_v2()
assert len(response.guardrails) == 1
guardrail = response.guardrails[0]
litellm_params = guardrail.litellm_params
if isinstance(litellm_params, dict):
params = litellm_params
else:
params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params)
# Sensitive keys should be masked
assert params["api_key"] != "my-secret-bedrock-key"
assert "****" in str(params["api_key"])
assert params["vertex_credentials"] != "{sensitive_creds}"
assert "****" in str(params["vertex_credentials"])
# Non-sensitive keys should remain unchanged
assert params["guardrail"] == "bedrock"
assert params["mode"] == "during_call"
@pytest.mark.asyncio
async def test_get_guardrail_info_from_db(mocker, mock_prisma_client):
"""Test getting guardrail info from DB"""