scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83
litellm/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
yucheng-berri ecabfbd5af
fix(guardrail): hide-secrets playground redaction and guardrail telemetry (#39398) 
* Fix hide-secrets guardrail: playground redaction, UI dropdown entry, spend-log telemetry

The hide-secrets guardrail never implemented apply_guardrail, so the UI test
playground echoed secrets verbatim; it was missing from the Add Guardrail
dropdown; and it recorded no guardrail_information, so Spend Logs could not
distinguish a redacted request from a clean one.

- implement apply_guardrail (unified interface) with use_native_lifecycle_hooks
  so proxied traffic stays on async_pre_call_hook (per-key opt-out and
  data["prompt"] handling live only there)
- record standard_logging_guardrail_information (allow/mask + masked_entity_count)
  via _process_response/_process_error; opted-out keys and legacy nameless
  callback instances record nothing
- advertise hide-secrets in /guardrails/ui/add_guardrail_settings (pre_call only)
  and /guardrails/ui/provider_specific_params with a config model

Resolves LIT-3548

* Fix hide-secrets passthrough telemetry and JSON config input

* fix(guardrails): validate hide-secrets object config before submit

- apply_guardrail treats empty-string-only texts as no input, so no
  false allow is recorded
- the UI object field keeps raw text while editing and blocks submission
  until it parses to a JSON object, instead of posting a string to an
  object-only API
- supported_modes_by_provider keeps its dict[str, list[str]] value type

* fix(guardrails): record no hide-secrets telemetry when nothing was inspected

walk_user_text and the prompt redaction now report how many non-empty
strings they visited; when neither inspected anything (image-only
content, empty strings), the run records no guardrail entry instead of
an 'allow' row that counts a check which never saw any text.
2026-09-03 00:01:03 -07:00

2694 lines
100 KiB
Python
Raw Permalink Blame History

import json
from datetime import datetime
from typing import Dict, List, Optional
from unittest.mock import AsyncMock
import pytest
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_endpoints import (
CreateGuardrailRequest,
PatchGuardrailRequest,
RegisterGuardrailRequest,
UpdateGuardrailRequest,
apply_guardrail,
approve_guardrail_submission,
create_guardrail,
delete_guardrail,
get_guardrail_info,
get_guardrail_submission,
get_guardrail_ui_settings,
list_guardrail_submissions,
list_guardrails_v2,
patch_guardrail,
register_guardrail,
reject_guardrail_submission,
update_guardrail,
)
MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
from litellm.proxy.guardrails.guardrail_registry import (
IN_MEMORY_GUARDRAIL_HANDLER,
InMemoryGuardrailHandler,
)
from litellm.types.guardrails import (
ApplyGuardrailRequest,
BaseLitellmParams,
Guardrail,
GuardrailInfoResponse,
LitellmParams,
)
# Mock data for testing
MOCK_DB_GUARDRAIL = {
"guardrail_id": "test-db-guardrail",
"guardrail_name": "Test DB Guardrail",
"litellm_params": {
"guardrail": "test.guardrail",
"mode": "pre_call",
},
"guardrail_info": {"description": "Test guardrail from DB"},
"created_at": datetime.now(),
"updated_at": datetime.now(),
}
MOCK_CONFIG_GUARDRAIL = {
"guardrail_id": "test-config-guardrail",
"guardrail_name": "Test Config Guardrail",
"litellm_params": {
"guardrail": "custom_guardrail.myCustomGuardrail",
"mode": "during_call",
},
"guardrail_info": {"description": "Test guardrail from config"},
}
MOCK_GUARDRAIL = Guardrail(
guardrail_name=MOCK_CONFIG_GUARDRAIL["guardrail_name"],
litellm_params=LitellmParams(**MOCK_CONFIG_GUARDRAIL["litellm_params"]),
guardrail_info=MOCK_CONFIG_GUARDRAIL["guardrail_info"],
)
MOCK_CREATE_REQUEST = CreateGuardrailRequest(guardrail=MOCK_GUARDRAIL)
MOCK_UPDATE_REQUEST = UpdateGuardrailRequest(guardrail=MOCK_GUARDRAIL)
MOCK_PATCH_REQUEST = PatchGuardrailRequest(
guardrail_name="Updated Test Guardrail",
litellm_params={"guardrail": "updated.guardrail", "mode": "post_call"},
guardrail_info={"description": "Updated test guardrail"},
)
@pytest.fixture
def mock_prisma_client(mocker):
"""Mock Prisma client for testing"""
mock_client = mocker.Mock()
# Create async mocks for the database methods
mock_client.db = mocker.Mock()
mock_client.db.litellm_guardrailstable = mocker.Mock()
mock_client.db.litellm_guardrailstable.find_many = AsyncMock(
return_value=[MOCK_DB_GUARDRAIL]
)
mock_client.db.litellm_guardrailstable.find_unique = AsyncMock(
return_value=MOCK_DB_GUARDRAIL
)
return mock_client
@pytest.fixture
def mock_in_memory_handler(mocker):
"""Mock InMemoryGuardrailHandler for testing"""
mock_handler = mocker.Mock(spec=InMemoryGuardrailHandler)
mock_handler.list_in_memory_guardrails.return_value = [MOCK_CONFIG_GUARDRAIL]
mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
mock_handler.get_source.return_value = "config"
mock_handler.initialize_guardrail = mocker.Mock()
mock_handler.sync_guardrail_from_db = mocker.Mock()
mock_handler.delete_in_memory_guardrail = mocker.Mock()
mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[])
return mock_handler
@pytest.fixture
def mock_guardrail_registry(mocker):
"""Mock GuardrailRegistry for testing"""
mock_registry = mocker.Mock()
mock_registry.add_guardrail_to_db = AsyncMock(
return_value={**MOCK_DB_GUARDRAIL, "guardrail_id": "new-test-guardrail-id"}
)
mock_registry.delete_guardrail_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
mock_registry.get_guardrail_by_id_from_db = AsyncMock(
return_value=MOCK_DB_GUARDRAIL
)
mock_registry.update_guardrail_in_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
return mock_registry
@pytest.mark.asyncio
async def test_list_guardrails_v2_with_db_and_config(
mocker, mock_prisma_client, mock_in_memory_handler
):
"""Test listing guardrails from both DB and config"""
# Mock the prisma client
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Mock the in-memory handler
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
response = await list_guardrails_v2(user_api_key_dict=admin_auth)
assert len(response.guardrails) == 2
# Check DB guardrail
db_guardrail = next(
g for g in response.guardrails if g.guardrail_id == "test-db-guardrail"
)
assert db_guardrail.guardrail_name == "Test DB Guardrail"
assert db_guardrail.guardrail_definition_location == "db"
assert isinstance(db_guardrail.litellm_params, BaseLitellmParams)
# Check config guardrail
config_guardrail = next(
g for g in response.guardrails if g.guardrail_id == "test-config-guardrail"
)
assert config_guardrail.guardrail_name == "Test Config Guardrail"
assert config_guardrail.guardrail_definition_location == "config"
assert isinstance(config_guardrail.litellm_params, BaseLitellmParams)
@pytest.mark.asyncio
async def test_list_guardrails_v2_skips_stale_db_backed_in_memory_entries(mocker):
"""
A guardrail that's still in this pod's memory tagged source='db' but is no
longer in the DB result (deleted on another pod, awaiting reconcile) must
NOT surface in the list response — pre-fix it leaked as 'config'.
"""
stale_guardrail = {
"guardrail_id": "stale-db-id",
"guardrail_name": "Stale DB Guardrail",
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
"guardrail_info": {},
}
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 = [stale_guardrail]
mock_in_memory_handler.get_source.return_value = "db"
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,
)
admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
response = await list_guardrails_v2(user_api_key_dict=admin_auth)
assert response.guardrails == []
mock_in_memory_handler.get_source.assert_called_with("stale-db-id")
@pytest.mark.asyncio
async def test_get_guardrail_info_404s_stale_db_backed_entry(
mocker, mock_prisma_client, mock_in_memory_handler
):
"""
Stale DB-backed entry (in-memory but not in DB) must 404 instead of being
returned as if it were a config-loaded guardrail.
"""
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,
)
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
return_value=None
)
# In-memory still has it, but it's tagged as 'db' (stale, awaiting reconcile)
mock_in_memory_handler.get_source.return_value = "db"
with pytest.raises(HTTPException) as exc_info:
await get_guardrail_info("stale-db-id")
assert exc_info.value.status_code == 404
assert "not found" in str(exc_info.value.detail)
@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,
)
admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
response = await list_guardrails_v2(user_api_key_dict=admin_auth)
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,
)
admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
response = await list_guardrails_v2(user_api_key_dict=admin_auth)
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_list_guardrails_v2_admin_viewer_sees_guardrails_of_teams_they_are_not_in(
mocker,
):
"""
proxy_admin_viewer reads the same unscoped list as proxy_admin: a team-owned
guardrail must surface even though the viewer belongs to no teams.
"""
other_team_guardrail = {
"guardrail_id": "other-team-guardrail",
"guardrail_name": "Other Team Guardrail",
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
"guardrail_info": {"description": "owned by a team the viewer is not in"},
"team_id": "team-viewer-is-not-in",
"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=[other_team_guardrail]
)
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,
)
mock_get_user_team_ids = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=[]),
)
viewer_auth = UserAPIKeyAuth(
user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
)
response = await list_guardrails_v2(user_api_key_dict=viewer_auth)
assert [g.guardrail_id for g in response.guardrails] == ["other-team-guardrail"]
mock_get_user_team_ids.assert_not_called()
@pytest.mark.asyncio
async def test_list_guardrails_v2_masks_sensitive_data_for_admin_viewer(mocker):
"""
Read parity for proxy_admin_viewer must not also hand out unmasked secrets.
The guardrail is team-owned so it only reaches the viewer via the admin path.
"""
other_team_guardrail_with_secrets = {
"guardrail_id": "other-team-secret-guardrail",
"guardrail_name": "Other Team Guardrail with Secrets",
"litellm_params": {
"guardrail": "azure/text_moderations",
"mode": "pre_call",
"api_key": "sk-viewer-must-not-see-this",
},
"guardrail_info": {},
"team_id": "team-viewer-is-not-in",
"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=[other_team_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,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=[]),
)
viewer_auth = UserAPIKeyAuth(
user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
)
response = await list_guardrails_v2(user_api_key_dict=viewer_auth)
guardrail = next(
g
for g in response.guardrails
if g.guardrail_id == "other-team-secret-guardrail"
)
params = guardrail.litellm_params.model_dump()
assert params["api_key"] != "sk-viewer-must-not-see-this"
assert "****" in str(params["api_key"])
assert params["guardrail"] == "azure/text_moderations"
@pytest.mark.asyncio
async def test_get_guardrail_info_from_db(mocker, mock_prisma_client):
"""Test getting guardrail info from DB"""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
response = await get_guardrail_info("test-db-guardrail")
assert response.guardrail_id == "test-db-guardrail"
assert response.guardrail_name == "Test DB Guardrail"
assert isinstance(response.litellm_params, BaseLitellmParams)
assert response.guardrail_info == {"description": "Test guardrail from DB"}
@pytest.mark.asyncio
async def test_get_guardrail_info_from_config(
mocker, mock_prisma_client, mock_in_memory_handler
):
"""Test getting guardrail info from config when not found in DB"""
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,
)
# Mock DB to return None
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
return_value=None
)
response = await get_guardrail_info("test-config-guardrail")
assert response.guardrail_id == "test-config-guardrail"
assert response.guardrail_name == "Test Config Guardrail"
assert isinstance(response.litellm_params, BaseLitellmParams)
assert response.guardrail_info == {"description": "Test guardrail from config"}
@pytest.mark.asyncio
async def test_get_guardrail_info_not_found(
mocker, mock_prisma_client, mock_in_memory_handler
):
"""Test getting guardrail info when not found in either DB or config"""
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,
)
# Mock both DB and in-memory handler to return None
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
return_value=None
)
mock_in_memory_handler.get_guardrail_by_id.return_value = None
with pytest.raises(HTTPException) as exc_info:
await get_guardrail_info("non-existent-guardrail")
assert exc_info.value.status_code == 404
assert "not found" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_list_guardrails_v2_without_prisma_returns_config_guardrails(
mocker, mock_in_memory_handler
):
"""
A proxy without a DB must still list config-defined guardrails instead of
raising 500 'Prisma client not initialized'.
"""
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
response = await list_guardrails_v2(user_api_key_dict=MOCK_ADMIN_USER)
assert len(response.guardrails) == 1
config_guardrail = response.guardrails[0]
assert config_guardrail.guardrail_id == "test-config-guardrail"
assert config_guardrail.guardrail_name == "Test Config Guardrail"
assert config_guardrail.guardrail_definition_location == "config"
@pytest.mark.asyncio
async def test_list_guardrails_v2_without_prisma_non_admin_sees_unrestricted_config_guardrails(
mocker, mock_in_memory_handler
):
"""
A non-admin caller on a no-DB proxy must see config guardrails that carry
no team_id restriction; the team lookup must not blow up without a DB.
"""
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
non_admin_auth = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal-user-1"
)
response = await list_guardrails_v2(user_api_key_dict=non_admin_auth)
assert [g.guardrail_id for g in response.guardrails] == ["test-config-guardrail"]
@pytest.mark.asyncio
async def test_get_guardrail_info_without_prisma_returns_config_guardrail(
mocker, mock_in_memory_handler
):
"""
The info endpoint must serve config-defined guardrails from the in-memory
registry when no DB is attached instead of raising 500.
"""
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
response = await get_guardrail_info("test-config-guardrail")
assert response.guardrail_id == "test-config-guardrail"
assert response.guardrail_name == "Test Config Guardrail"
assert response.guardrail_definition_location == "config"
@pytest.mark.asyncio
async def test_get_guardrail_info_without_prisma_404s_unknown_id(
mocker, mock_in_memory_handler
):
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
mock_in_memory_handler.get_guardrail_by_id.return_value = None
with pytest.raises(HTTPException) as exc_info:
await get_guardrail_info("non-existent-guardrail")
assert exc_info.value.status_code == 404
def test_get_guardrails_list_response_includes_guardrail_id():
"""
The v1 list response is the UI's fallback when v2 fails; without ids every
row click requests /guardrails/undefined/info.
"""
from litellm.proxy.guardrails.guardrail_endpoints import (
_get_guardrails_list_response,
)
response = _get_guardrails_list_response(
[
{
"guardrail_id": "stable-config-id",
"guardrail_name": "tooling",
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
},
}
]
)
assert response.guardrails[0].guardrail_id == "stable-config-id"
def test_get_provider_specific_params():
"""Test getting provider-specific parameters"""
from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model
from litellm.proxy.guardrails.guardrail_hooks.azure import (
AzureContentSafetyTextModerationGuardrail,
)
config_model = AzureContentSafetyTextModerationGuardrail.get_config_model()
if config_model is None:
pytest.skip("Azure config model not available")
fields = _get_fields_from_model(config_model)
print("FIELDS", fields)
# Test that we get the expected nested structure
assert isinstance(fields, dict)
# Check that we have the expected top-level fields
assert "api_key" in fields
assert "api_base" in fields
assert "api_version" in fields
assert "optional_params" in fields
# Check the structure of a simple field
assert (
fields["api_key"]["description"]
== "API key for the Azure Content Safety Prompt Shield guardrail"
)
assert fields["api_key"]["required"] == False
assert fields["api_key"]["type"] == "string" # Should be string, not None
# Check the structure of the nested optional_params field
assert fields["optional_params"]["type"] == "nested"
assert fields["optional_params"]["required"] == True
assert "fields" in fields["optional_params"]
# Check nested fields within optional_params
nested_fields = fields["optional_params"]["fields"]
assert "severity_threshold" in nested_fields
assert "severity_threshold_by_category" in nested_fields
assert "categories" in nested_fields
assert "blocklistNames" in nested_fields
assert "haltOnBlocklistHit" in nested_fields
assert "outputType" in nested_fields
# Check structure of a nested field
assert (
nested_fields["severity_threshold"]["description"]
== "Severity threshold for the Azure Content Safety Text Moderation guardrail across all categories"
)
assert nested_fields["severity_threshold"]["required"] == False
assert (
nested_fields["severity_threshold"]["type"] == "number"
) # Should be number, not None
# Check other field types
assert nested_fields["categories"]["type"] == "multiselect"
assert nested_fields["blocklistNames"]["type"] == "array"
assert nested_fields["haltOnBlocklistHit"]["type"] == "boolean"
assert (
nested_fields["outputType"]["type"] == "select"
) # Literal type should be select
@pytest.mark.asyncio
async def test_provider_specific_params_includes_hide_secrets():
"""hide-secrets lives in the enterprise package so it is not in
guardrail_class_registry; the endpoint must still advertise it or the
Add Guardrail UI dropdown never offers it (LIT-3548)."""
from litellm.proxy.guardrails.guardrail_endpoints import (
get_provider_specific_params,
)
provider_params = await get_provider_specific_params()
assert "hide-secrets" in provider_params
# populateGuardrailProviders() in the dashboard only lists providers whose
# entry carries a ui_friendly_name.
assert provider_params["hide-secrets"]["ui_friendly_name"] == "Hide Secrets"
assert provider_params["hide-secrets"]["detect_secrets_config"]["required"] is False
@pytest.mark.asyncio
async def test_add_guardrail_settings_restricts_hide_secrets_to_pre_call():
"""hide-secrets only implements async_pre_call_hook, so offering the other
modes in the UI would create configs that boot clean and never run."""
from litellm.proxy.guardrails.guardrail_endpoints import (
get_guardrail_ui_settings,
)
settings = await get_guardrail_ui_settings()
assert settings.supported_modes_by_provider["hide-secrets"] == ["pre_call"]
def test_optional_params_not_returned_when_not_overridden():
"""Test that optional_params is not returned when the config model doesn't override it"""
from typing import Optional
from pydantic import BaseModel, Field
from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
class TestGuardrailConfig(GuardrailConfigModel):
api_key: Optional[str] = Field(
default=None,
description="Test API key",
)
api_base: Optional[str] = Field(
default=None,
description="Test API base",
)
@staticmethod
def ui_friendly_name() -> str:
return "Test Guardrail"
# Get fields from the model
fields = _get_fields_from_model(TestGuardrailConfig)
print("FIELDS", fields)
assert "optional_params" not in fields
def test_optional_params_returned_when_properly_overridden():
"""Test that optional_params IS returned when the config model properly overrides it"""
from typing import Optional
from pydantic import BaseModel, Field
from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
# Create specific optional params model
class SpecificOptionalParams(BaseModel):
threshold: Optional[float] = Field(
default=0.5, description="Detection threshold"
)
categories: Optional[List[str]] = Field(
default=None, description="Categories to check"
)
# Create a config model that DOES override optional_params with a specific type
class TestGuardrailConfigWithOptionalParams(
GuardrailConfigModel[SpecificOptionalParams]
):
api_key: Optional[str] = Field(
default=None,
description="Test API key",
)
@staticmethod
def ui_friendly_name() -> str:
return "Test Guardrail With Optional Params"
# Get fields from the model
fields = _get_fields_from_model(TestGuardrailConfigWithOptionalParams)
print("FIELDS", fields)
assert "optional_params" in fields
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_with_api_key():
"""Test _prepare_request method uses Bearer token when api_key is provided in data"""
from unittest.mock import Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
# Setup guardrail hook
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
)
mock_credentials = Mock()
test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
prepared_request = guardrail_hook._prepare_request(
credentials=mock_credentials,
data=test_data,
optional_params={},
aws_region_name="us-east-1",
api_key="test-bearer-token-123",
)
# Verify Bearer token is used in Authorization header
assert "Authorization" in prepared_request.headers
assert prepared_request.headers["Authorization"] == "Bearer test-bearer-token-123"
# Verify URL is correct
expected_url = "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail/test-guardrail-id/version/1/apply"
assert prepared_request.url == expected_url
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_without_api_key():
"""Test _prepare_request method falls back to SigV4 when no api_key is provided"""
from unittest.mock import Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
# Setup guardrail hook
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
)
# Mock credentials
mock_credentials = Mock()
# Test data without api_key
test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
with (
patch(
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str"
) as mock_get_secret,
patch("botocore.auth.SigV4Auth") as mock_sigv4_auth,
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
):
# Mock no AWS_BEARER_TOKEN_BEDROCK
mock_get_secret.return_value = None
# Mock SigV4Auth
mock_sigv4_instance = Mock()
mock_sigv4_auth.return_value = mock_sigv4_instance
# Mock AWSRequest
mock_request_instance = Mock()
mock_request_instance.prepare.return_value = Mock()
mock_aws_request.return_value = mock_request_instance
# Call _prepare_request
prepared_request = guardrail_hook._prepare_request(
credentials=mock_credentials,
data=test_data,
optional_params={},
aws_region_name="us-east-1",
)
# Verify SigV4 auth was used
mock_sigv4_auth.assert_called_once_with(
mock_credentials, "bedrock", "us-east-1"
)
mock_sigv4_instance.add_auth.assert_called_once()
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_with_bearer_token_env():
"""Test _prepare_request method uses Bearer token from environment when available"""
from unittest.mock import Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
# Setup guardrail hook
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
)
# Mock credentials
mock_credentials = Mock()
# Test data without api_key
test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
with (
patch(
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str"
) as mock_get_secret,
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
):
mock_get_secret.return_value = "env-bearer-token-456"
mock_request_instance = Mock()
mock_request_instance.prepare.return_value = Mock()
mock_aws_request.return_value = mock_request_instance
prepared_request = guardrail_hook._prepare_request(
credentials=mock_credentials,
data=test_data,
optional_params={},
aws_region_name="us-east-1",
)
# Verify Bearer token from environment is used
mock_aws_request.assert_called_once()
call_args = mock_aws_request.call_args
headers = call_args[1]["headers"]
assert headers["Authorization"] == "Bearer env-bearer-token-456"
@pytest.mark.asyncio
async def test_bedrock_guardrail_make_api_request_passes_api_key():
"""Test make_bedrock_api_request method correctly passes api_key from request_data"""
from unittest.mock import AsyncMock, Mock, patch
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
guardrail_hook = BedrockGuardrail(
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
)
guardrail_hook.async_handler = Mock()
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"action": "NONE", "outputs": []}
test_request_data = {"api_key": "test-api-key-789"}
with (
patch.object(
guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)
),
patch.object(guardrail_hook, "_load_credentials") as mock_load_creds,
patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert,
patch.object(
guardrail_hook, "get_guardrail_dynamic_request_body_params"
) as mock_get_params,
patch.object(
guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"
),
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
):
mock_load_creds.return_value = (Mock(), "us-east-1")
mock_convert.return_value = {"source": "INPUT", "content": [{"text": {"text": "test"}}]}
mock_get_params.return_value = {}
mock_request_instance = Mock()
mock_request_instance.url = "test-url"
mock_request_instance.body = b"test-body"
mock_request_instance.headers = {
"Content-Type": "application/json",
"Authorization": "Bearer test-api-key-789",
}
mock_request_instance.prepare.return_value = Mock(
headers=mock_request_instance.headers
)
mock_aws_request.return_value = mock_request_instance
await guardrail_hook.make_bedrock_api_request(
source="INPUT",
messages=[{"role": "user", "content": "test"}],
request_data=test_request_data,
)
# Verify _prepare_request was invoked and used the api_key
mock_aws_request.assert_called_once()
call_args = mock_aws_request.call_args
headers = call_args[1]["headers"]
assert headers["Authorization"] == "Bearer test-api-key-789"
@pytest.mark.parametrize(
"scenario,expected_result,expected_exception",
[
("success_with_sync", "new-test-guardrail-id", None),
("success_sync_fails", "new-test-guardrail-id", None),
("database_failure", None, HTTPException),
("no_prisma_client", None, HTTPException),
],
ids=[
"success_with_immediate_sync",
"success_but_sync_fails",
"database_error",
"missing_prisma_client",
],
)
@pytest.mark.asyncio
async def test_create_guardrail_endpoint(
scenario,
expected_result,
expected_exception,
mocker,
mock_guardrail_registry,
mock_in_memory_handler,
):
"""Test create_guardrail endpoint with different scenarios"""
# Configure mocks based on scenario
mock_logger = None
if scenario == "success_with_sync":
mock_prisma_client = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "success_sync_fails":
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.initialize_guardrail.side_effect = Exception(
"Sync failed"
)
mock_logger = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "database_failure":
mock_prisma_client = mocker.Mock()
mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception(
"Database error"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
elif scenario == "no_prisma_client":
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
# Run the test
if expected_exception:
with pytest.raises(expected_exception) as exc_info:
await create_guardrail(
MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
)
if scenario == "database_failure":
assert "Database error" in str(exc_info.value.detail)
elif scenario == "no_prisma_client":
assert "Prisma client not initialized" in str(exc_info.value.detail)
else:
result = await create_guardrail(
MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
)
assert result["guardrail_id"] == expected_result
assert result["guardrail_name"] == "Test DB Guardrail"
mock_guardrail_registry.add_guardrail_to_db.assert_called_once_with(
guardrail=MOCK_CREATE_REQUEST.guardrail, prisma_client=mocker.ANY
)
mock_in_memory_handler.initialize_guardrail.assert_called_once()
if scenario == "success_sync_fails":
assert mock_logger is not None
mock_logger.warning.assert_called_once()
assert "Failed to initialize guardrail" in str(
mock_logger.warning.call_args
)
@pytest.mark.parametrize(
"scenario,expected_result,expected_exception",
[
("success_with_sync", "test-db-guardrail", None),
("success_sync_fails_unexpected_error", "test-db-guardrail", None),
("sync_fails_invalid_config", None, HTTPException),
("database_failure", None, HTTPException),
("no_prisma_client", None, HTTPException),
],
ids=[
"success_with_immediate_sync",
"success_but_sync_fails_with_unexpected_error",
"sync_rejects_invalid_config",
"database_error",
"missing_prisma_client",
],
)
@pytest.mark.asyncio
async def test_update_guardrail_endpoint(
scenario,
expected_result,
expected_exception,
mocker,
mock_guardrail_registry,
mock_in_memory_handler,
):
"""Test update_guardrail endpoint with different scenarios"""
# Configure mocks based on scenario
mock_logger = None
if scenario == "success_with_sync":
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "success_sync_fails_unexpected_error":
# A non-ValueError/TypeError failure is not a config-rejection signal,
# so it keeps the pre-existing swallow-and-warn behavior rather than
# rolling back the DB write.
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=Exception("Sync failed")
)
mock_logger = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "sync_fails_invalid_config":
# Regression for the PUT half of the fix: a TypeError from the sync (the
# in-place update_in_memory_guardrail raised exactly this on every PUT)
# must roll back the DB write and surface a 422, not persist the
# rejected config with a 200.
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=TypeError("vars() argument must have __dict__ attribute")
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "database_failure":
mock_prisma_client = mocker.Mock()
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception(
"Database error"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
elif scenario == "no_prisma_client":
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
# Run the test
if expected_exception:
with pytest.raises(expected_exception) as exc_info:
await update_guardrail(
"test-guardrail-id",
MOCK_UPDATE_REQUEST,
user_api_key_dict=MOCK_ADMIN_USER,
)
if scenario == "database_failure":
assert "Database error" in str(exc_info.value.detail)
elif scenario == "no_prisma_client":
assert "Prisma client not initialized" in str(exc_info.value.detail)
elif scenario == "sync_fails_invalid_config":
assert exc_info.value.status_code == 422
assert "update rejected" in str(exc_info.value.detail)
# Rolled back: update_guardrail_in_db is called once for the
# rejected write and once more to restore the previous config.
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
assert (
mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
== MOCK_DB_GUARDRAIL
)
else:
result = await update_guardrail(
"test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
)
assert result["guardrail_id"] == expected_result
assert result["guardrail_name"] == "Test DB Guardrail"
mock_guardrail_registry.update_guardrail_in_db.assert_called_once_with(
guardrail_id="test-guardrail-id",
guardrail=MOCK_UPDATE_REQUEST.guardrail,
prisma_client=mocker.ANY,
)
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
guardrail=mocker.ANY
)
if scenario == "success_sync_fails_unexpected_error":
assert mock_logger is not None
mock_logger.warning.assert_called_once()
assert "Failed to update" in str(mock_logger.warning.call_args)
@pytest.mark.parametrize(
"scenario,expected_result,expected_exception",
[
("success_with_sync", "test-db-guardrail", None),
("success_sync_fails_unexpected_error", "test-db-guardrail", None),
("sync_fails_invalid_config", None, HTTPException),
("database_failure", None, HTTPException),
("no_prisma_client", None, HTTPException),
],
ids=[
"success_with_immediate_sync",
"success_but_sync_fails_with_unexpected_error",
"sync_rejects_invalid_config",
"database_error",
"missing_prisma_client",
],
)
@pytest.mark.asyncio
async def test_patch_guardrail_endpoint(
scenario,
expected_result,
expected_exception,
mocker,
mock_guardrail_registry,
mock_in_memory_handler,
):
"""Test patch_guardrail endpoint with different scenarios"""
# Configure mocks based on scenario
mock_logger = None
if scenario == "success_with_sync":
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "success_sync_fails_unexpected_error":
# A non-ValueError/TypeError failure (e.g. a transient bug) is not a
# config-rejection signal, so it keeps the pre-existing swallow-and-warn
# behavior rather than rolling back the DB write.
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=Exception("Sync failed")
)
mock_logger = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "sync_fails_invalid_config":
# Maintainer finding on BerriAI/litellm#34940: a ValueError from
# sync_guardrail_from_db (e.g. an invalid on_flagged combination) must
# roll back the DB write and surface a 422, not persist the rejected
# config with a 200.
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=ValueError("on_flagged='inject_system_message' requires payload=True and breakdown=True")
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "database_failure":
mock_prisma_client = mocker.Mock()
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception(
"Database error"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
elif scenario == "no_prisma_client":
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
# Run the test
if expected_exception:
with pytest.raises(expected_exception) as exc_info:
await patch_guardrail(
"test-guardrail-id",
MOCK_PATCH_REQUEST,
user_api_key_dict=MOCK_ADMIN_USER,
)
if scenario == "database_failure":
assert "Database error" in str(exc_info.value.detail)
elif scenario == "no_prisma_client":
assert "Prisma client not initialized" in str(exc_info.value.detail)
elif scenario == "sync_fails_invalid_config":
assert exc_info.value.status_code == 422
assert "update rejected" in str(exc_info.value.detail)
# Rolled back: update_guardrail_in_db is called once for the
# rejected write and once more to restore the previous config.
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
else:
result = await patch_guardrail(
"test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
)
assert result["guardrail_id"] == expected_result
assert result["guardrail_name"] == "Test DB Guardrail"
mock_guardrail_registry.update_guardrail_in_db.assert_called_once()
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
guardrail=mocker.ANY
)
if scenario == "success_sync_fails_unexpected_error":
assert mock_logger is not None
mock_logger.warning.assert_called_once()
assert "Failed to update" in str(mock_logger.warning.call_args)
@pytest.mark.parametrize(
"scenario,expected_result,expected_exception",
[
("success_with_sync", "test-db-guardrail", None),
("success_sync_fails", "test-db-guardrail", None),
],
ids=["success_with_immediate_sync", "success_but_sync_fails"],
)
@pytest.mark.asyncio
async def test_delete_guardrail_endpoint(
scenario,
expected_result,
expected_exception,
mocker,
mock_guardrail_registry,
mock_in_memory_handler,
):
"""Test delete_guardrail endpoint with different scenarios"""
# Configure mocks based on scenario
mock_prisma_client = mocker.Mock()
mock_logger = None
if scenario == "success_with_sync":
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "success_sync_fails":
mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception(
"Sync failed"
)
mock_logger = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
if expected_exception:
with pytest.raises(expected_exception):
await delete_guardrail(
guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER
)
else:
result = await delete_guardrail(
guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER
)
assert result == MOCK_DB_GUARDRAIL
mock_guardrail_registry.get_guardrail_by_id_from_db.assert_called_once_with(
guardrail_id=expected_result, prisma_client=mock_prisma_client
)
mock_guardrail_registry.delete_guardrail_from_db.assert_called_once_with(
guardrail_id=expected_result, prisma_client=mock_prisma_client
)
mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with(
guardrail_id=expected_result
)
if scenario == "success_sync_fails":
assert mock_logger is not None
mock_logger.warning.assert_called_once()
assert "Failed to remove guardrail" in str(mock_logger.warning.call_args)
@pytest.mark.asyncio
async def test_apply_guardrail_not_found(mocker):
"""
Test apply_guardrail endpoint returns proper error when guardrail is not found.
"""
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
# Mock the GUARDRAIL_REGISTRY to return None (guardrail not found)
mock_registry = mocker.Mock()
mock_registry.get_initialized_guardrail_callback.return_value = None
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
mock_proxy_logging = mocker.Mock()
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
mocker.patch("litellm.proxy.proxy_server.general_settings", {})
mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock())
mocker.patch("litellm.proxy.proxy_server.version", "test")
# Create request
request = ApplyGuardrailRequest(
guardrail_name="non-existent-guardrail", text="Test input text"
)
# Mock user auth
mock_user_auth = UserAPIKeyAuth()
# Call endpoint and expect ProxyException
with pytest.raises(ProxyException) as exc_info:
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=mock_user_auth,
)
# Verify error details
assert str(exc_info.value.code) == "404"
assert "not found" in str(exc_info.value.message).lower()
@pytest.mark.asyncio
async def test_apply_guardrail_execution_error(mocker):
"""
Test apply_guardrail endpoint handles exceptions from guardrail execution properly.
"""
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
# Mock guardrail that raises an exception
mock_guardrail = mocker.Mock()
mock_guardrail.apply_guardrail = AsyncMock(
side_effect=Exception("Bedrock guardrail failed: Violated guardrail policy")
)
# Mock the GUARDRAIL_REGISTRY
mock_registry = mocker.Mock()
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
mock_logging_obj = mocker.Mock()
mock_logging_obj.async_failure_handler = AsyncMock()
mock_logging_obj.model_call_details = {}
mock_processor = mocker.Mock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
)
mocker.patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
return_value=mock_processor,
)
mock_proxy_logging = mocker.Mock()
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
mocker.patch("litellm.proxy.proxy_server.general_settings", {})
mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock())
mocker.patch("litellm.proxy.proxy_server.version", "test")
mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor")
# Create request
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail", text="Test input text with forbidden content"
)
# Mock user auth
mock_user_auth = UserAPIKeyAuth()
# Call endpoint and expect ProxyException
with pytest.raises(ProxyException) as exc_info:
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=mock_user_auth,
)
# Verify error is properly handled
assert "Bedrock guardrail failed" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_apply_guardrail_invokes_logging_pipeline(mocker):
mock_guardrail = mocker.Mock()
mock_guardrail.apply_guardrail = AsyncMock(return_value={"texts": ["masked"]})
mock_registry = mocker.Mock()
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
mock_logging_obj = mocker.Mock()
mock_logging_obj.async_success_handler = AsyncMock()
mock_logging_obj.model_call_details = {}
mock_processor = mocker.Mock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
)
mocker.patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
return_value=mock_processor,
)
mock_proxy_logging = mocker.Mock()
mock_proxy_logging.post_call_success_hook = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
mocker.patch("litellm.proxy.proxy_server.general_settings", {})
mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock())
mocker.patch("litellm.proxy.proxy_server.version", "test")
mock_executor = mocker.Mock()
mocker.patch(
"litellm.litellm_core_utils.thread_pool_executor.executor", mock_executor
)
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail", text="hello@example.com"
)
response = await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
assert response.response_text == "masked"
mock_processor.common_processing_pre_call_logic.assert_awaited_once()
mock_proxy_logging.post_call_success_hook.assert_awaited_once()
mock_logging_obj.async_success_handler.assert_awaited_once()
assert mock_logging_obj.call_type == "pass_through_endpoint"
mock_executor.submit.assert_called_once()
assert mock_logging_obj.async_success_handler.await_args.kwargs["result"] == {
"response": {"response_text": "masked"}
}
def _patch_apply_guardrail_env(mocker, guardrail_result):
mock_guardrail = mocker.Mock()
mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result)
mock_registry = mocker.Mock()
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
mock_logging_obj = mocker.Mock()
mock_logging_obj.async_success_handler = AsyncMock()
mock_logging_obj.model_call_details = {}
mock_processor = mocker.Mock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
)
mocker.patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
return_value=mock_processor,
)
mock_proxy_logging = mocker.Mock()
mock_proxy_logging.post_call_success_hook = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
mocker.patch("litellm.proxy.proxy_server.general_settings", {})
mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock())
mocker.patch("litellm.proxy.proxy_server.version", "test")
mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor")
return mock_guardrail
@pytest.mark.asyncio
async def test_apply_guardrail_forwards_metadata_to_guardrail(mocker):
"""Client-supplied metadata must reach apply_guardrail via request_data so
parameterized custom guardrails can read per-request configuration."""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail",
text="What are tax loopholes?",
metadata={"forbidden_topics": ["tax"]},
)
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["What are tax loopholes?"]},
request_data={"metadata": {"forbidden_topics": ["tax"]}},
input_type="request",
)
@pytest.mark.asyncio
async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
"""metadata and messages must coexist in request_data; the dict merge must
not clobber messages when both fields are sent."""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
messages = [{"role": "user", "content": "What are tax loopholes?"}]
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail",
text="What are tax loopholes?",
messages=messages,
metadata={"forbidden_topics": ["tax"]},
)
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["What are tax loopholes?"]},
request_data={
"messages": messages,
"metadata": {"forbidden_topics": ["tax"]},
},
input_type="request",
)
@pytest.mark.asyncio
async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
"""Without metadata, request_data stays empty (backward-compatible)."""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["hello"]},
request_data={},
input_type="request",
)
@pytest.mark.asyncio
async def test_apply_guardrail_forwards_explicit_empty_messages_and_metadata(mocker):
"""Explicitly-sent empty messages/metadata must be forwarded, not dropped;
only omitted fields stay out of request_data."""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail",
text="hello",
messages=[],
metadata={},
)
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["hello"]},
request_data={"messages": [], "metadata": {}},
input_type="request",
)
@pytest.mark.asyncio
async def test_get_guardrail_info_endpoint_config_guardrail(mocker):
"""
Test get_guardrail_info endpoint returns proper response when guardrail is found in config.
"""
from litellm.proxy.guardrails.guardrail_endpoints import get_guardrail_info
# Mock prisma_client to not be None (patch at the source where it's imported from)
mock_prisma = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock the GUARDRAIL_REGISTRY to return None from DB (so it checks config)
mock_registry = mocker.Mock()
mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=None)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
# Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail
mock_in_memory_handler = mocker.Mock()
mock_in_memory_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
mock_in_memory_handler.get_source.return_value = "config"
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
# Mock _get_masked_values to return values as-is
mocker.patch(
"litellm.litellm_core_utils.litellm_logging._get_masked_values",
side_effect=lambda x, **kwargs: x,
)
# Call endpoint and expect GuardrailInfoResponse
result = await get_guardrail_info(guardrail_id="test-config-guardrail")
# Verify the response is of the correct type
assert isinstance(result, GuardrailInfoResponse)
assert result.guardrail_id == "test-config-guardrail"
assert result.guardrail_name == "Test Config Guardrail"
assert result.guardrail_definition_location == "config"
@pytest.mark.asyncio
async def test_get_guardrail_info_endpoint_db_guardrail(mocker):
"""
Test get_guardrail_info endpoint returns proper response when guardrail is found in DB.
"""
from litellm.proxy.guardrails.guardrail_endpoints import get_guardrail_info
# Mock prisma_client to not be None (patch at the source where it's imported from)
mock_prisma = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock the GUARDRAIL_REGISTRY to return a guardrail from DB
mock_registry = mocker.Mock()
mock_registry.get_guardrail_by_id_from_db = AsyncMock(
return_value=MOCK_DB_GUARDRAIL
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
)
# Mock IN_MEMORY_GUARDRAIL_HANDLER to return None
mock_in_memory_handler = mocker.Mock()
mock_in_memory_handler.get_guardrail_by_id.return_value = None
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
# Call endpoint and expect GuardrailInfoResponse
result = await get_guardrail_info(guardrail_id="test-db-guardrail")
# Verify the response is of the correct type
assert isinstance(result, GuardrailInfoResponse)
assert result.guardrail_id == "test-db-guardrail"
assert result.guardrail_name == "Test DB Guardrail"
assert result.guardrail_definition_location == "db"
class TestBuildFieldDict:
"""Test _build_field_dict handles both enum and string ui_type values."""
def test_build_field_dict_with_string_ui_type(self):
"""Test that _build_field_dict works when ui_type is a plain string (e.g. BlockCodeExecutionGuardrailConfigModel)."""
from unittest.mock import MagicMock
from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict
field = MagicMock()
field.json_schema_extra = {
"ui_type": "multiselect",
"options": ["python", "javascript"],
}
result = _build_field_dict(
field=field,
field_annotation=str,
description="Test field",
required=False,
)
assert result["type"] == "multiselect"
assert result["description"] == "Test field"
def test_build_field_dict_with_enum_ui_type(self):
"""Test that _build_field_dict works when ui_type is a GuardrailParamUITypes enum."""
from unittest.mock import MagicMock
from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict
from litellm.types.guardrails import GuardrailParamUITypes
field = MagicMock()
field.json_schema_extra = {"ui_type": GuardrailParamUITypes.BOOL}
result = _build_field_dict(
field=field,
field_annotation=bool,
description="Test bool field",
required=True,
)
assert result["type"] == "bool"
assert result["required"] is True
# --- Team guardrail registration (register / submissions) ---
MOCK_REGISTER_REQUEST = RegisterGuardrailRequest(
guardrail_name="team-prompt-guard",
litellm_params={
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"api_base": "https://guardrails.example.com/validate",
},
guardrail_info={"description": "Team prompt injection detector"},
)
@pytest.mark.asyncio
async def test_register_guardrail_success(mocker):
"""Register creates a row with status pending_review and returns guardrail_id."""
mock_prisma = mocker.Mock()
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
created_row = mocker.Mock(
guardrail_id="reg-123",
guardrail_name=MOCK_REGISTER_REQUEST.guardrail_name,
status="pending_review",
submitted_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.create = AsyncMock(return_value=created_row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_id="u1", user_email="alice@co.com", team_id="team-1")
result = await register_guardrail(MOCK_REGISTER_REQUEST, user)
assert result.guardrail_id == "reg-123"
assert result.guardrail_name == MOCK_REGISTER_REQUEST.guardrail_name
assert result.status == "pending_review"
mock_prisma.db.litellm_guardrailstable.create.assert_called_once()
call_data = mock_prisma.db.litellm_guardrailstable.create.call_args[1]["data"]
assert call_data["status"] == "pending_review"
assert call_data["guardrail_name"] == MOCK_REGISTER_REQUEST.guardrail_name
@pytest.mark.asyncio
async def test_register_guardrail_rejects_non_generic_api(mocker):
"""Register returns 400 when litellm_params.guardrail is not generic_guardrail_api."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
req = RegisterGuardrailRequest(
guardrail_name="other-guard",
litellm_params={
"guardrail": "bedrock",
"mode": "pre_call",
"api_base": "https://x.com",
},
)
user = UserAPIKeyAuth(user_id="u1", user_email="a@b.com", team_id="team-1")
with pytest.raises(HTTPException) as exc_info:
await register_guardrail(req, user)
assert exc_info.value.status_code == 400
assert "generic_guardrail_api" in exc_info.value.detail
@pytest.mark.asyncio
async def test_register_guardrail_requires_team_id(mocker):
"""Register returns 400 when API key has no associated team_id."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
user = UserAPIKeyAuth(user_id="u1", user_email="a@b.com", team_id=None)
with pytest.raises(HTTPException) as exc_info:
await register_guardrail(MOCK_REGISTER_REQUEST, user)
assert exc_info.value.status_code == 400
assert "team" in exc_info.value.detail.lower()
@pytest.mark.asyncio
async def test_register_guardrail_non_admin_cross_team_allowed(mocker):
"""Non-admin may register for a team in their user.teams list even if the key's team_id differs."""
mock_prisma = mocker.Mock()
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
created = mocker.Mock(
guardrail_id="g1",
guardrail_name=MOCK_REGISTER_REQUEST.guardrail_name,
status="pending_review",
submitted_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.create = AsyncMock(return_value=created)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=["team-alpha", "team-beta"]),
)
req = RegisterGuardrailRequest(
guardrail_name=MOCK_REGISTER_REQUEST.guardrail_name,
team_id="team-beta",
litellm_params=MOCK_REGISTER_REQUEST.litellm_params,
)
user = UserAPIKeyAuth(
user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha"
)
result = await register_guardrail(req, user)
assert result.guardrail_id == "g1"
@pytest.mark.asyncio
async def test_register_guardrail_non_admin_cross_team_forbidden(mocker):
"""Non-admin gets 403 when registering for a team they are not a member of."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=["team-alpha"]),
)
req = RegisterGuardrailRequest(
guardrail_name=MOCK_REGISTER_REQUEST.guardrail_name,
team_id="team-other",
litellm_params=MOCK_REGISTER_REQUEST.litellm_params,
)
user = UserAPIKeyAuth(
user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha"
)
with pytest.raises(HTTPException) as exc_info:
await register_guardrail(req, user)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_register_guardrail_duplicate_name(mocker):
"""Register returns 400 when guardrail_name already exists."""
mock_prisma = mocker.Mock()
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(
return_value={"guardrail_name": MOCK_REGISTER_REQUEST.guardrail_name}
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_id="u1", user_email="a@b.com", team_id="team-1")
with pytest.raises(HTTPException) as exc_info:
await register_guardrail(MOCK_REGISTER_REQUEST, user)
assert exc_info.value.status_code == 400
assert "already exists" in exc_info.value.detail
@pytest.mark.asyncio
async def test_list_guardrail_submissions_non_admin_scoped_to_own_teams(mocker):
"""Non-admin callers see only submissions for teams they belong to."""
mock_prisma = mocker.Mock()
own_team_row = mocker.Mock(
guardrail_id="mine",
guardrail_name="mine-guard",
status="pending_review",
team_id="team-mine",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
find_many = AsyncMock(return_value=[own_team_row])
mock_prisma.db.litellm_guardrailstable.find_many = find_many
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=["team-mine"]),
)
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER)
result = await list_guardrail_submissions(user_api_key_dict=user)
# DB query scoped to visible teams
where_clause = find_many.call_args.kwargs["where"]
assert where_clause["team_id"] == {"in": ["team-mine"]}
assert len(result.submissions) == 1
assert result.submissions[0].team_id == "team-mine"
# Summary counts reflect only visible teams
assert result.summary.total == 1
assert result.summary.pending_review == 1
@pytest.mark.asyncio
async def test_list_guardrail_submissions_non_admin_no_teams(mocker):
"""Non-admin caller with no team memberships gets an empty list (not 403)."""
mock_prisma = mocker.Mock()
find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_guardrailstable.find_many = find_many
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=[]),
)
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER)
result = await list_guardrail_submissions(user_api_key_dict=user)
assert result.submissions == []
assert result.summary.total == 0
assert find_many.call_count == 0 # no DB query when user has no teams
@pytest.mark.asyncio
async def test_list_guardrail_submissions_non_admin_team_filter_forbidden(mocker):
"""Non-admin caller filtering by a team they're not in gets 403."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=["team-mine"]),
)
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER)
with pytest.raises(HTTPException) as exc_info:
await list_guardrail_submissions(team_id="team-other", user_api_key_dict=user)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_list_guardrail_submissions_success(mocker):
"""List submissions returns list and summary for admin."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="sub-1",
guardrail_name="pending-guard",
status="pending_review",
team_id="t1",
litellm_params={
"guardrail": "generic_guardrail_api",
"api_base": "https://x.com",
},
guardrail_info={
"description": "A guard",
"submitted_by_user_id": "u1",
"submitted_by_email": "alice@co.com",
},
submitted_at=datetime.now(),
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[row])
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
result = await list_guardrail_submissions(user_api_key_dict=user)
assert len(result.submissions) == 1
assert result.submissions[0].guardrail_id == "sub-1"
assert result.submissions[0].status == "pending_review"
assert result.submissions[0].team_guardrail is True # team_id is set
assert result.summary.total >= 1
assert result.summary.pending_review >= 1
@pytest.mark.asyncio
async def test_list_guardrail_submissions_returns_only_team_guardrails(mocker):
"""List submissions only returns team guardrails (team_id not null)."""
mock_prisma = mocker.Mock()
find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_guardrailstable.find_many = find_many
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
await list_guardrail_submissions(user_api_key_dict=user)
calls = find_many.call_args_list
assert len(calls) >= 1
first_where = calls[0].kwargs.get("where", {})
assert first_where.get("team_id") == {"not": None}
@pytest.mark.asyncio
async def test_list_guardrail_submissions_team_id_filter(mocker):
"""List submissions with team_id filter returns only that team's guardrails."""
mock_prisma = mocker.Mock()
row_abc = mocker.Mock(
guardrail_id="team-1",
guardrail_name="team-guard",
status="active",
team_id="team-abc",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
row_other = mocker.Mock(
guardrail_id="team-2",
guardrail_name="other-guard",
status="active",
team_id="team-xyz",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
find_many = AsyncMock(return_value=[row_abc, row_other])
mock_prisma.db.litellm_guardrailstable.find_many = find_many
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
result = await list_guardrail_submissions(
user_api_key_dict=user, team_id="team-abc"
)
assert len(result.submissions) == 1
assert result.submissions[0].guardrail_id == "team-1"
assert result.submissions[0].team_guardrail is True
assert result.summary.total == 2 # summary counts all team guardrails
@pytest.mark.asyncio
async def test_get_guardrail_submission_not_found(mocker):
"""Get submission returns 404 when guardrail_id does not exist."""
mock_prisma = mocker.Mock()
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as exc_info:
await get_guardrail_submission("nonexistent-id", user)
assert exc_info.value.status_code == 404
@pytest.mark.asyncio
async def test_get_guardrail_submission_non_admin_own_team(mocker):
"""Non-admin caller can fetch a submission belonging to one of their teams."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="sub-1",
guardrail_name="team-guard",
status="pending_review",
team_id="team-mine",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=["team-mine"]),
)
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER)
result = await get_guardrail_submission("sub-1", user)
assert result.guardrail_id == "sub-1"
assert result.team_id == "team-mine"
@pytest.mark.asyncio
async def test_get_guardrail_submission_non_admin_other_team_forbidden(mocker):
"""Non-admin caller gets 403 when fetching a submission for a team they're not in."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="sub-1",
guardrail_name="team-guard",
status="pending_review",
team_id="team-other",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=["team-mine"]),
)
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER)
with pytest.raises(HTTPException) as exc_info:
await get_guardrail_submission("sub-1", user)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_get_guardrail_submission_admin_viewer_other_team_allowed(mocker):
"""proxy_admin_viewer reads any team's submission without the membership check."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="sub-1",
guardrail_name="team-guard",
status="pending_review",
team_id="team-other",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mock_get_user_team_ids = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
AsyncMock(return_value=[]),
)
user = UserAPIKeyAuth(
user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
)
result = await get_guardrail_submission("sub-1", user)
assert result.guardrail_id == "sub-1"
assert result.team_id == "team-other"
mock_get_user_team_ids.assert_not_called()
@pytest.mark.asyncio
async def test_approve_guardrail_submission_success(mocker):
"""Approve sets status to active and initializes guardrail in memory."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="approve-me",
guardrail_name="my-guard",
status="pending_review",
litellm_params={
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"api_base": "https://g.com",
},
guardrail_info={},
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mock_prisma.db.litellm_guardrailstable.update = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mock_handler = mocker.Mock()
mock_handler.initialize_guardrail = mocker.Mock()
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_handler,
)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
result = await approve_guardrail_submission("approve-me", user)
assert result["status"] == "active"
assert result["guardrail_id"] == "approve-me"
mock_prisma.db.litellm_guardrailstable.update.assert_called_once()
call_data = mock_prisma.db.litellm_guardrailstable.update.call_args[1]["data"]
assert call_data["status"] == "active"
@pytest.mark.asyncio
async def test_approve_guardrail_submission_not_pending(mocker):
"""Approve returns 400 when status is not pending_review."""
mock_prisma = mocker.Mock()
row = mocker.Mock(guardrail_id="x", guardrail_name="y", status="active")
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as exc_info:
await approve_guardrail_submission("x", user)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_reject_guardrail_submission_success(mocker):
"""Reject sets status to rejected."""
mock_prisma = mocker.Mock()
row = mocker.Mock(guardrail_id="rej-1", guardrail_name="r", status="pending_review")
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mock_prisma.db.litellm_guardrailstable.update = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
result = await reject_guardrail_submission("rej-1", user)
assert result["status"] == "rejected"
mock_prisma.db.litellm_guardrailstable.update.assert_called_once()
call_data = mock_prisma.db.litellm_guardrailstable.update.call_args[1]["data"]
assert call_data["status"] == "rejected"
@pytest.mark.asyncio
async def test_reject_guardrail_submission_not_pending(mocker):
"""Reject returns 400 when status is not pending_review (e.g. already active)."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="already-active", guardrail_name="g", status="active"
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as exc_info:
await reject_guardrail_submission("already-active", user)
assert exc_info.value.status_code == 400
assert "not pending review" in exc_info.value.detail.lower()
# --- Tests for review fixes ---
@pytest.mark.asyncio
@pytest.mark.parametrize(
"api_base,expected_detail",
[
("file:///etc/passwd", "http or https scheme"),
("ftp://internal.host/data", "http or https scheme"),
("javascript:alert(1)", "http or https scheme"),
("://missing-scheme", "http or https scheme"),
("https://", "valid hostname"),
],
ids=[
"file_scheme",
"ftp_scheme",
"javascript_scheme",
"no_scheme",
"no_hostname",
],
)
async def test_register_guardrail_rejects_bad_api_base(
mocker, api_base, expected_detail
):
"""Register returns 400 when api_base has invalid scheme or missing hostname."""
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
req = RegisterGuardrailRequest(
guardrail_name="bad-url-guard",
litellm_params={
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"api_base": api_base,
},
)
user = UserAPIKeyAuth(user_id="u1", user_email="a@b.com", team_id="team-1")
with pytest.raises(HTTPException) as exc_info:
await register_guardrail(req, user)
assert exc_info.value.status_code == 400
assert expected_detail in exc_info.value.detail
@pytest.mark.asyncio
async def test_register_guardrail_accepts_valid_https_url(mocker):
"""Register accepts valid https api_base URLs."""
mock_prisma = mocker.Mock()
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
created_row = mocker.Mock(
guardrail_id="valid-url-123",
guardrail_name="valid-guard",
status="pending_review",
submitted_at=datetime.now(),
)
mock_prisma.db.litellm_guardrailstable.create = AsyncMock(return_value=created_row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
req = RegisterGuardrailRequest(
guardrail_name="valid-guard",
litellm_params={
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"api_base": "https://guardrails.example.com/v1/check",
},
)
user = UserAPIKeyAuth(user_id="u1", user_email="a@b.com", team_id="team-1")
result = await register_guardrail(req, user)
assert result.guardrail_id == "valid-url-123"
assert result.status == "pending_review"
@pytest.mark.asyncio
async def test_approve_guardrail_init_failure_returns_warning(mocker):
"""Approve returns a warning field when in-memory initialization fails."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="warn-me",
guardrail_name="fragile-guard",
status="pending_review",
litellm_params={
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"api_base": "https://g.com",
},
guardrail_info={},
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mock_prisma.db.litellm_guardrailstable.update = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mock_handler = mocker.Mock()
mock_handler.initialize_guardrail = mocker.Mock(
side_effect=Exception("missing dependency")
)
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_handler,
)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
result = await approve_guardrail_submission("warn-me", user)
assert result["status"] == "active"
assert "warning" in result
assert "failed to initialize" in result["warning"].lower()
assert "missing dependency" in result["warning"]
@pytest.mark.asyncio
async def test_approve_guardrail_no_warning_on_success(mocker):
"""Approve does NOT include a warning field when init succeeds."""
mock_prisma = mocker.Mock()
row = mocker.Mock(
guardrail_id="ok-guard",
guardrail_name="good-guard",
status="pending_review",
litellm_params={
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"api_base": "https://g.com",
},
guardrail_info={},
)
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
mock_prisma.db.litellm_guardrailstable.update = AsyncMock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mock_handler = mocker.Mock()
mock_handler.initialize_guardrail = mocker.Mock() # no exception
mocker.patch(
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_handler,
)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
result = await approve_guardrail_submission("ok-guard", user)
assert result["status"] == "active"
assert "warning" not in result
@pytest.mark.asyncio
async def test_list_submissions_single_db_query(mocker):
"""List submissions makes exactly one find_many call (no redundant query)."""
mock_prisma = mocker.Mock()
find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_guardrailstable.find_many = find_many
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
await list_guardrail_submissions(user_api_key_dict=user)
assert find_many.call_count == 1
@pytest.mark.asyncio
async def test_list_submissions_summary_counts_unaffected_by_filters(mocker):
"""Summary counts reflect all team guardrails regardless of status filter."""
mock_prisma = mocker.Mock()
pending_row = mocker.Mock(
guardrail_id="p1",
guardrail_name="p",
status="pending_review",
team_id="t1",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
active_row = mocker.Mock(
guardrail_id="a1",
guardrail_name="a",
status="active",
team_id="t1",
litellm_params={},
guardrail_info={},
submitted_at=None,
reviewed_at=None,
created_at=datetime.now(),
updated_at=datetime.now(),
)
all_rows = [pending_row, active_row]
mock_prisma.db.litellm_guardrailstable.find_many = AsyncMock(return_value=all_rows)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
# Filter to only pending, but summary should still show both
result = await list_guardrail_submissions(
status="pending_review", user_api_key_dict=user
)
assert len(result.submissions) == 1 # filtered
assert result.summary.total == 2 # unfiltered
assert result.summary.pending_review == 1
assert result.summary.active == 1
@pytest.mark.asyncio
async def test_get_guardrail_ui_settings_returns_per_provider_supported_modes():
"""
Regression test for LIT-4226. The Admin UI used to render `pre_mcp_call` as a
selectable mode for every guardrail because the settings endpoint returned a
single global `supported_modes` list. The proxy then rejected the save because
Content Filter and Tool Permission do not accept `pre_mcp_call`. The endpoint
must now return per-provider modes so the UI can filter its dropdown.
"""
result = await get_guardrail_ui_settings()
modes_by_provider = result.supported_modes_by_provider
# Content Filter now supports pre_mcp_call (LIT-4226 feature half) but not
# during_mcp_call; Tool Permission still supports neither, and the settings
# endpoint must reflect both so the UI shows exactly the savable modes.
assert "pre_mcp_call" in modes_by_provider["litellm_content_filter"]
assert "during_mcp_call" not in modes_by_provider["litellm_content_filter"]
assert modes_by_provider["tool_permission"] == ["pre_call", "post_call"]
# MCP-capable guardrails must still advertise the MCP hooks so users who
# picked one of them can actually configure pre_mcp_call / during_mcp_call.
for provider in ("bedrock", "panw_prisma_airs", "cisco_ai_defense", "custom_code", "pillar"):
assert "pre_mcp_call" in modes_by_provider[provider], provider
assert "during_mcp_call" in modes_by_provider[provider], provider
# The union list stays exhaustive for legacy clients that ignore the
# per-provider map; it must cover every declared GuardrailEventHooks value.
from litellm.types.guardrails import GuardrailEventHooks
assert set(result.supported_modes) == {m.value for m in GuardrailEventHooks}
@pytest.mark.asyncio
async def test_ui_settings_map_matches_runtime_supported_event_hooks():
"""
Regression guard against the two-copy-of-the-list drift risk. The map the
UI reads must agree with what CustomGuardrail._validate_event_hook accepts
at save time, otherwise the bug in LIT-4226 comes back one classname at a
time as future guardrails drift.
"""
from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
result = await get_guardrail_ui_settings()
for provider, guardrail_class in guardrail_class_registry.items():
declared = guardrail_class.get_supported_event_hooks()
if declared is None:
assert (
provider not in result.supported_modes_by_provider
), f"{provider} returned None from classmethod but appears in map"
continue
assert provider in result.supported_modes_by_provider, provider
assert result.supported_modes_by_provider[provider] == [
hook.value for hook in declared
], provider
def test_content_filter_runtime_rejects_unsupported_mcp_hook():
"""
Locks the runtime side of the LIT-4226 contract: the ContentFilterGuardrail
validator must reject a hook missing from its supported list at
construction. pre_mcp_call is supported since the LIT-4226 feature half, so
during_mcp_call is the unsupported example now. If someone widens the
UI classmethod but forgets to widen the runtime supported_event_hooks (or
vice versa), the two-lists-must-agree test above catches the drift and this
test catches the specific bug the ticket reported.
"""
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
from litellm.types.guardrails import GuardrailEventHooks
with pytest.raises(ValueError, match="not in the supported event hooks"):
ContentFilterGuardrail(
guardrail_name="lit4226-runtime-check",
event_hook=GuardrailEventHooks.during_mcp_call,
)
def test_model_armor_runtime_supported_event_hooks_match_classmethod():
"""
Regression for the drift Round 2 caught: the ModelArmorGuardrail classmethod
declared its supported hooks for the UI, but __init__ did not seed the
runtime instance's `supported_event_hooks` from that classmethod, so the
runtime validator accepted any hook (including nonsense like logging_only)
while the UI hid them. Ensures the two sides agree at instantiation time.
"""
from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import (
ModelArmorGuardrail,
)
instance = ModelArmorGuardrail(
guardrail_name="lit4226-model-armor-drift",
template_id="t",
project_id="p",
)
assert instance.supported_event_hooks == ModelArmorGuardrail.get_supported_event_hooks()
def test_strict_guardrail_modes_flag_controls_raise_vs_warn(monkeypatch, caplog):
"""
Escape hatch for the boot-time behavior change. Deployments upgrading from
a build where a guardrail previously silently no-op'd on an unsupported
mode should be able to set LITELLM_STRICT_GUARDRAIL_MODES=false and boot
with a warning instead of a hard failure while they fix their config.
"""
import logging
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
from litellm.types.guardrails import GuardrailEventHooks
monkeypatch.delenv("LITELLM_STRICT_GUARDRAIL_MODES", raising=False)
with pytest.raises(ValueError, match="not in the supported event hooks"):
ContentFilterGuardrail(
guardrail_name="lit4226-strict-default",
event_hook=GuardrailEventHooks.during_mcp_call,
)
monkeypatch.setenv("LITELLM_STRICT_GUARDRAIL_MODES", "false")
with caplog.at_level(logging.WARNING):
instance = ContentFilterGuardrail(
guardrail_name="lit4226-strict-off",
event_hook=GuardrailEventHooks.during_mcp_call,
)
assert instance is not None
assert any("not in the supported event hooks" in rec.message for rec in caplog.records)
def test_field_type_inference_handles_pep604_unions():
from litellm.proxy.guardrails.guardrail_endpoints import (
_get_field_type_from_annotation,
_unwrap_optional_type,
)
assert _get_field_type_from_annotation(Optional[int]) == "number"
assert _get_field_type_from_annotation(int | None) == "number"
assert _get_field_type_from_annotation(list[str] | None) == "array"
assert _get_field_type_from_annotation(bool | None) == "boolean"
assert _unwrap_optional_type(str | None) is str
Reference in a new issue View git blame Copy permalink