mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
feat(guardrails): runtime enable/disable toggle for guardrails via API and UI
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI
parent
13df85cceb
commit
e81c8b7eb2
18 changed files with 718 additions and 17 deletions
|
|
@ -222,6 +222,7 @@ class CustomGuardrail(CustomLogger):
|
|||
self.run_in_parallel: bool = run_in_parallel
|
||||
self.scan_raw_request: bool = scan_raw_request
|
||||
self.only_scan_new_messages: bool = only_scan_new_messages
|
||||
self.enabled: bool = True
|
||||
|
||||
if supported_event_hooks:
|
||||
## validate event_hook is in supported_event_hooks
|
||||
|
|
@ -973,6 +974,8 @@ class CustomGuardrail(CustomLogger):
|
|||
"""
|
||||
Returns True if the guardrail should be run on the event_type
|
||||
"""
|
||||
if not self.enabled:
|
||||
return False
|
||||
if self._suppressed_by_auto_router_compression():
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -48,6 +48,8 @@ from litellm.types.guardrails import (
|
|||
PiiAction,
|
||||
PiiEntityType,
|
||||
PresidioPresidioConfigModelUserInterface,
|
||||
SetGuardrailEnabledRequest,
|
||||
SetGuardrailEnabledResponse,
|
||||
SupportedGuardrailIntegrations,
|
||||
ToolPermissionGuardrailConfigModel,
|
||||
)
|
||||
|
|
@ -269,6 +271,7 @@ async def list_guardrails_v2(
|
|||
created_at=guardrail.get("created_at"),
|
||||
updated_at=guardrail.get("updated_at"),
|
||||
guardrail_definition_location="db",
|
||||
enabled=IN_MEMORY_GUARDRAIL_HANDLER.is_enabled(guardrail.get("guardrail_id") or ""),
|
||||
)
|
||||
)
|
||||
seen_guardrail_ids.add(guardrail.get("guardrail_id"))
|
||||
|
|
@ -308,6 +311,7 @@ async def list_guardrails_v2(
|
|||
litellm_params=masked_in_memory_litellm_params_typed,
|
||||
guardrail_info=dict(guardrail.get("guardrail_info") or {}),
|
||||
guardrail_definition_location="config",
|
||||
enabled=IN_MEMORY_GUARDRAIL_HANDLER.is_enabled(gid),
|
||||
)
|
||||
)
|
||||
seen_guardrail_ids.add(gid)
|
||||
|
|
@ -1281,6 +1285,82 @@ async def patch_guardrail(
|
|||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/guardrails/{guardrail_id}/enabled",
|
||||
tags=["Guardrails"],
|
||||
response_model=SetGuardrailEnabledResponse,
|
||||
)
|
||||
async def set_guardrail_enabled(
|
||||
guardrail_id: str,
|
||||
request: SetGuardrailEnabledRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
Turn a guardrail off or back on at runtime without editing or deleting it
|
||||
|
||||
Works for guardrails defined in config.yaml as well as ones created in the DB.
|
||||
A disabled guardrail stays registered but skips every request until it is
|
||||
re-enabled. The state is persisted, so it survives restarts and applies to
|
||||
all proxy instances sharing the DB.
|
||||
|
||||
Example Request:
|
||||
```bash
|
||||
curl -X PATCH "http://localhost:4000/guardrails/123e4567-e89b-12d3-a456-426614174000/enabled" \\
|
||||
-H "Authorization: Bearer <your_api_key>" \\
|
||||
-H "Content-Type: application/json" \\
|
||||
-d '{"enabled": false}'
|
||||
```
|
||||
|
||||
Example Response:
|
||||
```json
|
||||
{
|
||||
"guardrail_id": "123e4567-e89b-12d3-a456-426614174000",
|
||||
"guardrail_name": "headroom-compression",
|
||||
"enabled": false
|
||||
}
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Admin access required to manage guardrails",
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Prisma client not initialized")
|
||||
|
||||
db_guardrail: Final = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db(
|
||||
guardrail_id=guardrail_id, prisma_client=prisma_client
|
||||
)
|
||||
config_guardrail: Final = (
|
||||
IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(guardrail_id=guardrail_id)
|
||||
if IN_MEMORY_GUARDRAIL_HANDLER.get_source(guardrail_id) == "config"
|
||||
else None
|
||||
)
|
||||
guardrail: Final = db_guardrail if db_guardrail is not None else config_guardrail
|
||||
if guardrail is None:
|
||||
raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found")
|
||||
|
||||
disabled_guardrail_ids: Final = await GUARDRAIL_REGISTRY.set_guardrail_enabled_in_db(
|
||||
guardrail_id=guardrail_id, enabled=request.enabled, prisma_client=prisma_client
|
||||
)
|
||||
IN_MEMORY_GUARDRAIL_HANDLER.set_disabled_guardrails(disabled_guardrail_ids)
|
||||
verbose_proxy_logger.info(
|
||||
"Guardrail '%s' (ID: %s) %s",
|
||||
guardrail.get("guardrail_name"),
|
||||
guardrail_id,
|
||||
"enabled" if request.enabled else "disabled",
|
||||
)
|
||||
return SetGuardrailEnabledResponse(
|
||||
guardrail_id=guardrail_id,
|
||||
guardrail_name=guardrail.get("guardrail_name") or "",
|
||||
enabled=request.enabled,
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/guardrails/{guardrail_id}",
|
||||
tags=["Guardrails"],
|
||||
|
|
@ -1369,6 +1449,7 @@ async def get_guardrail_info(guardrail_id: str):
|
|||
created_at=result.get("created_at"),
|
||||
updated_at=result.get("updated_at"),
|
||||
guardrail_definition_location=guardrail_definition_location,
|
||||
enabled=IN_MEMORY_GUARDRAIL_HANDLER.is_enabled(guardrail_id),
|
||||
)
|
||||
except HTTPException as e:
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tool_permission import (
|
|||
)
|
||||
from litellm.proxy.types_utils.utils import get_instance_fn
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import GuardrailsRepository
|
||||
from litellm.secret_managers.main import get_secret
|
||||
|
|
@ -90,6 +91,8 @@ guardrail_initializer_registry: Final = {
|
|||
|
||||
CONFIG_GUARDRAIL_ID_NAMESPACE: Final = uuid.UUID("625f63f4-935a-50e5-98b5-fbe77babc74a")
|
||||
|
||||
DISABLED_GUARDRAILS_CONFIG_PARAM: Final = "disabled_guardrails"
|
||||
|
||||
GuardrailCallbacks: TypeAlias = tuple[CustomGuardrail, ...]
|
||||
|
||||
guardrail_class_registry: Final[dict[str, type[CustomGuardrail]]] = {
|
||||
|
|
@ -414,6 +417,22 @@ class GuardrailRegistry:
|
|||
except Exception as e:
|
||||
raise Exception(f"Error getting guardrail from DB: {e}")
|
||||
|
||||
@staticmethod
|
||||
async def get_disabled_guardrail_ids_from_db(prisma_client: PrismaClient) -> frozenset[str]:
|
||||
param: Final = await ConfigRepository(prisma_client).get_param(DISABLED_GUARDRAILS_CONFIG_PARAM)
|
||||
if param is None or not isinstance(param.param_value, list):
|
||||
return frozenset()
|
||||
return frozenset(item for item in param.param_value if isinstance(item, str))
|
||||
|
||||
@staticmethod
|
||||
async def set_guardrail_enabled_in_db(
|
||||
guardrail_id: str, enabled: bool, prisma_client: PrismaClient
|
||||
) -> frozenset[str]:
|
||||
current: Final = await GuardrailRegistry.get_disabled_guardrail_ids_from_db(prisma_client)
|
||||
updated: Final = current - {guardrail_id} if enabled else current | {guardrail_id}
|
||||
await ConfigRepository(prisma_client).set_param(DISABLED_GUARDRAILS_CONFIG_PARAM, sorted(updated))
|
||||
return updated
|
||||
|
||||
|
||||
def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: LitellmParams) -> None:
|
||||
"""Override the parallel/raw-scan flags only when ``litellm_params`` explicitly
|
||||
|
|
@ -486,6 +505,8 @@ class InMemoryGuardrailHandler:
|
|||
and never deleted by reconciliation.
|
||||
"""
|
||||
|
||||
self._disabled_guardrail_ids: frozenset[str] = frozenset()
|
||||
|
||||
def _stable_guardrail_id(self, guardrail_name: str) -> str:
|
||||
seeds: Final = chain((guardrail_name,), (f"{guardrail_name}:{occurrence}" for occurrence in count(1)))
|
||||
candidate_ids: Final = (str(uuid.uuid5(CONFIG_GUARDRAIL_ID_NAMESPACE, seed.encode("utf-8"))) for seed in seeds)
|
||||
|
|
@ -545,6 +566,7 @@ class InMemoryGuardrailHandler:
|
|||
)
|
||||
for custom_guardrail_callback in created_callbacks:
|
||||
_configure_callback_scoping(custom_guardrail_callback, guardrail["guardrail_name"], litellm_params)
|
||||
custom_guardrail_callback.enabled = guardrail_id not in self._disabled_guardrail_ids
|
||||
|
||||
parsed_guardrail: Final = Guardrail(
|
||||
guardrail_id=guardrail.get("guardrail_id"),
|
||||
|
|
@ -702,6 +724,15 @@ class InMemoryGuardrailHandler:
|
|||
"""
|
||||
return self._sources.get(guardrail_id)
|
||||
|
||||
def is_enabled(self, guardrail_id: str) -> bool:
|
||||
return guardrail_id not in self._disabled_guardrail_ids
|
||||
|
||||
def set_disabled_guardrails(self, disabled_guardrail_ids: frozenset[str]) -> None:
|
||||
self._disabled_guardrail_ids = disabled_guardrail_ids
|
||||
for guardrail_id in self.IN_MEMORY_GUARDRAILS:
|
||||
for callback in self._tracked_callbacks(guardrail_id):
|
||||
callback.enabled = guardrail_id not in disabled_guardrail_ids
|
||||
|
||||
def list_config_guardrails(self) -> list[Guardrail]:
|
||||
"""
|
||||
List in-memory guardrails owned by config.yaml.
|
||||
|
|
|
|||
|
|
@ -7727,6 +7727,9 @@ class ProxyConfig:
|
|||
# Drop in-memory DB-backed entries whose row was deleted on another
|
||||
# pod. Config-loaded entries are never touched.
|
||||
IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids)
|
||||
IN_MEMORY_GUARDRAIL_HANDLER.set_disabled_guardrails(
|
||||
await GuardrailRegistry.get_disabled_guardrail_ids_from_db(prisma_client=prisma_client)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - %s", e)
|
||||
|
||||
|
|
|
|||
|
|
@ -1173,6 +1173,7 @@ class GuardrailInfoResponse(BaseModel):
|
|||
created_at: datetime | None = None
|
||||
updated_at: datetime | None = None
|
||||
guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.CONFIG
|
||||
enabled: bool = True
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -1218,3 +1219,13 @@ class PatchGuardrailRequest(BaseModel):
|
|||
guardrail_name: str | None = None
|
||||
litellm_params: BaseLitellmParams | None = None
|
||||
guardrail_info: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class SetGuardrailEnabledRequest(BaseModel):
|
||||
enabled: bool
|
||||
|
||||
|
||||
class SetGuardrailEnabledResponse(BaseModel):
|
||||
guardrail_id: str
|
||||
guardrail_name: str
|
||||
enabled: bool
|
||||
|
|
|
|||
|
|
@ -2966,3 +2966,37 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float):
|
|||
|
||||
assert guardrail.timeout.read == 60.0
|
||||
assert guardrail.timeout.connect == 5.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_headroom_guardrail_skips_compress_and_forwards_request_unchanged():
|
||||
"""
|
||||
A Headroom guardrail switched off at runtime (PATCH /guardrails/{id}/enabled)
|
||||
must not call /v1/compress at all, even though it is default_on and the
|
||||
request names it explicitly; the proxy pipeline forwards the request as is.
|
||||
"""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
guardrail = _make_guardrail(guardrail_name="headroom-compression-global")
|
||||
data = {"model": "gpt-4o", "messages": ORIGINAL_MESSAGES, "metadata": {"guardrails": ["headroom-compression-global"]}}
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=UserApiKeyCache())
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as post:
|
||||
assert guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
guardrail.enabled = False
|
||||
skipped = await proxy_logging._process_guardrail_callback(
|
||||
callback=guardrail,
|
||||
data=data,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
call_type="acompletion",
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
assert skipped is None
|
||||
post.assert_not_called()
|
||||
assert data["messages"] == ORIGINAL_MESSAGES
|
||||
assert _recorded_guardrail_entries(data) == []
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ def mock_prisma_client(mocker):
|
|||
def mock_in_memory_handler(mocker):
|
||||
"""Mock InMemoryGuardrailHandler for testing"""
|
||||
mock_handler = mocker.Mock(spec=InMemoryGuardrailHandler)
|
||||
mock_handler.is_enabled.return_value = True
|
||||
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"
|
||||
|
|
@ -179,6 +180,7 @@ async def test_list_guardrails_v2_skips_stale_db_backed_in_memory_entries(mocker
|
|||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = [stale_guardrail]
|
||||
mock_in_memory_handler.get_source.return_value = "db"
|
||||
|
||||
|
|
@ -246,6 +248,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker):
|
|||
)
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = []
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
|
@ -299,6 +302,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock
|
|||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = [
|
||||
config_guardrail_with_secrets
|
||||
]
|
||||
|
|
@ -360,6 +364,7 @@ async def test_list_guardrails_v2_admin_viewer_sees_guardrails_of_teams_they_are
|
|||
)
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = []
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
|
@ -409,6 +414,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_for_admin_viewer(mocker):
|
|||
)
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = []
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
|
@ -1762,6 +1768,7 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker):
|
|||
|
||||
# Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
|
||||
mock_in_memory_handler.get_source.return_value = "config"
|
||||
mocker.patch(
|
||||
|
|
@ -1807,6 +1814,7 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker):
|
|||
|
||||
# Mock IN_MEMORY_GUARDRAIL_HANDLER to return None
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.is_enabled.return_value = True
|
||||
mock_in_memory_handler.get_guardrail_by_id.return_value = None
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
|
||||
|
|
@ -2709,3 +2717,161 @@ def test_field_type_inference_handles_pep604_unions():
|
|||
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
|
||||
|
||||
|
||||
def _prisma_with_config_store(mocker, stored: dict, db_guardrails: list):
|
||||
async def find_unique(*, where):
|
||||
value = stored.get(where["param_name"])
|
||||
return None if value is None else mocker.Mock(param_name=where["param_name"], param_value=value)
|
||||
|
||||
async def upsert(*, where, data):
|
||||
stored[where["param_name"]] = data["update"]["param_value"]
|
||||
return mocker.Mock()
|
||||
|
||||
prisma_client = mocker.Mock()
|
||||
prisma_client.db.litellm_config.find_unique = find_unique
|
||||
prisma_client.db.litellm_config.upsert = upsert
|
||||
prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=db_guardrails)
|
||||
prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
|
||||
side_effect=lambda where: next(
|
||||
(g for g in db_guardrails if g["guardrail_id"] == where["guardrail_id"]), None
|
||||
)
|
||||
)
|
||||
return prisma_client
|
||||
|
||||
|
||||
def _real_handler_with_config_guardrail(mocker, name: str):
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.guardrails import guardrail_registry as registry_module
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
registry_module.guardrail_initializer_registry["toggle_endpoint_test"] = lambda litellm_params, guardrail: (
|
||||
CustomGuardrail(
|
||||
guardrail_name=guardrail["guardrail_name"], event_hook=GuardrailEventHooks.pre_call, default_on=True
|
||||
)
|
||||
)
|
||||
handler = InMemoryGuardrailHandler()
|
||||
result = handler.initialize_guardrail(
|
||||
guardrail={
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {"guardrail": "toggle_endpoint_test", "mode": "pre_call"},
|
||||
}
|
||||
)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler)
|
||||
return handler, result["guardrail_id"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_guardrail_enabled_requires_admin(mocker):
|
||||
from litellm.proxy.guardrails.guardrail_endpoints import set_guardrail_enabled
|
||||
from litellm.types.guardrails import SetGuardrailEnabledRequest
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await set_guardrail_enabled(
|
||||
guardrail_id="any",
|
||||
request=SetGuardrailEnabledRequest(enabled=False),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_guardrail_enabled_toggles_config_guardrail_and_surfaces_in_list_and_info(mocker):
|
||||
"""
|
||||
End to end through the endpoints: a config.yaml guardrail (no DB row) is
|
||||
disabled, the disabled id is persisted to LiteLLM_Config, the live callback
|
||||
stops running, /v2/guardrails/list and /guardrails/{id}/info report
|
||||
enabled=false, and re-enabling reverses all of it.
|
||||
"""
|
||||
from litellm.proxy.guardrails import guardrail_registry as registry_module
|
||||
from litellm.proxy.guardrails.guardrail_endpoints import set_guardrail_enabled
|
||||
from litellm.types.guardrails import GuardrailEventHooks, SetGuardrailEnabledRequest
|
||||
|
||||
stored: dict = {}
|
||||
prisma_client = _prisma_with_config_store(mocker, stored, db_guardrails=[])
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
try:
|
||||
handler, gid = _real_handler_with_config_guardrail(mocker, "headroom-compression")
|
||||
instance = handler.guardrail_id_to_custom_guardrail[gid]
|
||||
|
||||
disabled = await set_guardrail_enabled(
|
||||
guardrail_id=gid, request=SetGuardrailEnabledRequest(enabled=False), user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
|
||||
assert disabled.model_dump() == {"guardrail_id": gid, "guardrail_name": "headroom-compression", "enabled": False}
|
||||
assert stored["disabled_guardrails"] == json.dumps([gid])
|
||||
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
|
||||
listed = await list_guardrails_v2(user_api_key_dict=MOCK_ADMIN_USER)
|
||||
assert [(g.guardrail_id, g.enabled, g.guardrail_definition_location) for g in listed.guardrails] == [
|
||||
(gid, False, "config")
|
||||
]
|
||||
assert (await get_guardrail_info(gid)).enabled is False
|
||||
|
||||
enabled = await set_guardrail_enabled(
|
||||
guardrail_id=gid, request=SetGuardrailEnabledRequest(enabled=True), user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
|
||||
assert enabled.enabled is True
|
||||
assert stored["disabled_guardrails"] == "[]"
|
||||
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True
|
||||
assert (await list_guardrails_v2(user_api_key_dict=MOCK_ADMIN_USER)).guardrails[0].enabled is True
|
||||
assert (await get_guardrail_info(gid)).enabled is True
|
||||
finally:
|
||||
registry_module.guardrail_initializer_registry.pop("toggle_endpoint_test", None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_guardrail_enabled_disables_db_guardrail_without_touching_its_row(mocker):
|
||||
from litellm.proxy.guardrails.guardrail_endpoints import set_guardrail_enabled
|
||||
from litellm.types.guardrails import SetGuardrailEnabledRequest
|
||||
|
||||
stored: dict = {}
|
||||
prisma_client = _prisma_with_config_store(mocker, stored, db_guardrails=[MOCK_DB_GUARDRAIL])
|
||||
prisma_client.db.litellm_guardrailstable.update = AsyncMock()
|
||||
prisma_client.db.litellm_guardrailstable.delete = AsyncMock()
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
handler = InMemoryGuardrailHandler()
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler)
|
||||
|
||||
result = await set_guardrail_enabled(
|
||||
guardrail_id="test-db-guardrail",
|
||||
request=SetGuardrailEnabledRequest(enabled=False),
|
||||
user_api_key_dict=MOCK_ADMIN_USER,
|
||||
)
|
||||
|
||||
assert result.enabled is False
|
||||
assert result.guardrail_name == "Test DB Guardrail"
|
||||
assert stored["disabled_guardrails"] == '["test-db-guardrail"]'
|
||||
assert handler.is_enabled("test-db-guardrail") is False
|
||||
prisma_client.db.litellm_guardrailstable.update.assert_not_called()
|
||||
prisma_client.db.litellm_guardrailstable.delete.assert_not_called()
|
||||
listed = await list_guardrails_v2(user_api_key_dict=MOCK_ADMIN_USER)
|
||||
assert [(g.guardrail_id, g.enabled) for g in listed.guardrails] == [("test-db-guardrail", False)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_guardrail_enabled_404s_unknown_and_stale_db_backed_ids(mocker):
|
||||
from litellm.proxy.guardrails.guardrail_endpoints import set_guardrail_enabled
|
||||
from litellm.types.guardrails import SetGuardrailEnabledRequest
|
||||
|
||||
stored: dict = {}
|
||||
prisma_client = _prisma_with_config_store(mocker, stored, db_guardrails=[])
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
handler = mocker.Mock(spec=InMemoryGuardrailHandler)
|
||||
handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
|
||||
handler.get_source.return_value = "db"
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await set_guardrail_enabled(
|
||||
guardrail_id="stale-db-id",
|
||||
request=SetGuardrailEnabledRequest(enabled=False),
|
||||
user_api_key_dict=MOCK_ADMIN_USER,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert stored == {}
|
||||
handler.set_disabled_guardrails.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -1084,3 +1084,109 @@ def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance():
|
|||
finally:
|
||||
for cb_list, snapshot in zip(lists, snapshots):
|
||||
cb_list[:] = snapshot
|
||||
|
||||
|
||||
def test_set_disabled_guardrails_gates_should_run_and_re_enables():
|
||||
"""
|
||||
A disabled guardrail must stay registered and listed but never run: the
|
||||
runtime gate returns False even for a default_on guardrail, and flipping it
|
||||
back on takes effect immediately without re-initializing anything.
|
||||
"""
|
||||
from litellm.proxy.guardrails import guardrail_registry as registry_module
|
||||
|
||||
def _initializer(litellm_params, guardrail):
|
||||
return CustomGuardrail(
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
registry_module.guardrail_initializer_registry["toggle_test"] = _initializer
|
||||
lists = _all_callback_lists()
|
||||
snapshots = [list(cb_list) for cb_list in lists]
|
||||
try:
|
||||
handler = InMemoryGuardrailHandler()
|
||||
result = handler.initialize_guardrail(guardrail=_config_guardrail("headroom-compression", "toggle_test"))
|
||||
gid = result["guardrail_id"]
|
||||
instance = handler.guardrail_id_to_custom_guardrail[gid]
|
||||
|
||||
assert handler.is_enabled(gid) is True
|
||||
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
handler.set_disabled_guardrails(frozenset({gid}))
|
||||
assert handler.is_enabled(gid) is False
|
||||
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
|
||||
assert handler.get_guardrail_by_id(gid) is not None
|
||||
assert instance in handler._tracked_callbacks(gid)
|
||||
|
||||
handler.set_disabled_guardrails(frozenset())
|
||||
assert handler.is_enabled(gid) is True
|
||||
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True
|
||||
finally:
|
||||
registry_module.guardrail_initializer_registry.pop("toggle_test", None)
|
||||
for cb_list, snapshot in zip(lists, snapshots):
|
||||
cb_list[:] = snapshot
|
||||
|
||||
|
||||
def test_guardrail_initialized_after_disable_starts_disabled():
|
||||
"""
|
||||
Persisted disabled ids are loaded before DB guardrails are (re)initialized on
|
||||
a pod, so a callback created later for a disabled id must come up disabled
|
||||
rather than silently re-enabling itself on the next sync.
|
||||
"""
|
||||
registry_module = _register_noop_initializer("late_init_test")
|
||||
lists = _all_callback_lists()
|
||||
snapshots = [list(cb_list) for cb_list in lists]
|
||||
try:
|
||||
handler = InMemoryGuardrailHandler()
|
||||
handler.set_disabled_guardrails(frozenset({"db-guardrail-1"}))
|
||||
handler.initialize_guardrail(
|
||||
guardrail=_config_guardrail("late", "late_init_test", guardrail_id="db-guardrail-1")
|
||||
)
|
||||
|
||||
instance = handler.guardrail_id_to_custom_guardrail["db-guardrail-1"]
|
||||
assert instance.enabled is False
|
||||
assert instance.should_run_guardrail(
|
||||
data={"metadata": {"guardrails": ["late"]}}, event_type=GuardrailEventHooks.pre_call
|
||||
) is False
|
||||
finally:
|
||||
registry_module.guardrail_initializer_registry.pop("late_init_test", None)
|
||||
for cb_list, snapshot in zip(lists, snapshots):
|
||||
cb_list[:] = snapshot
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_guardrail_enabled_in_db_round_trips_through_config_table():
|
||||
"""
|
||||
The disabled set lives in LiteLLM_Config so config.yaml guardrails, which have
|
||||
no row in the guardrails table, can be switched off too. Disable adds the id,
|
||||
enable removes it, and the stored JSON reads back as the same set.
|
||||
"""
|
||||
stored: dict = {}
|
||||
|
||||
async def find_unique(*, where):
|
||||
value = stored.get(where["param_name"])
|
||||
if value is None:
|
||||
return None
|
||||
return MagicMock(param_name=where["param_name"], param_value=value)
|
||||
|
||||
async def upsert(*, where, data):
|
||||
stored[where["param_name"]] = data["update"]["param_value"]
|
||||
return MagicMock()
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_config.find_unique = find_unique
|
||||
prisma_client.db.litellm_config.upsert = upsert
|
||||
|
||||
assert await GuardrailRegistry.get_disabled_guardrail_ids_from_db(prisma_client) == frozenset()
|
||||
|
||||
assert await GuardrailRegistry.set_guardrail_enabled_in_db("b", enabled=False, prisma_client=prisma_client) == {"b"}
|
||||
assert await GuardrailRegistry.set_guardrail_enabled_in_db("a", enabled=False, prisma_client=prisma_client) == {
|
||||
"a",
|
||||
"b",
|
||||
}
|
||||
assert stored["disabled_guardrails"] == '["a", "b"]'
|
||||
assert await GuardrailRegistry.get_disabled_guardrail_ids_from_db(prisma_client) == frozenset({"a", "b"})
|
||||
|
||||
assert await GuardrailRegistry.set_guardrail_enabled_in_db("b", enabled=True, prisma_client=prisma_client) == {"a"}
|
||||
assert await GuardrailRegistry.get_disabled_guardrail_ids_from_db(prisma_client) == frozenset({"a"})
|
||||
|
|
|
|||
|
|
@ -0,0 +1,62 @@
|
|||
import React from "react";
|
||||
|
||||
import { setGuardrailEnabledCall } from "@/components/networking";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Card } from "@/components/ui/card";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
import { formatGuardrailMode } from "./guardrail_info_helpers";
|
||||
|
||||
interface GuardrailModeCardProps {
|
||||
guardrailId: string;
|
||||
mode: unknown;
|
||||
defaultOn: boolean;
|
||||
enabled: boolean;
|
||||
accessToken: string | null;
|
||||
isAdmin: boolean;
|
||||
onChanged: () => Promise<void>;
|
||||
}
|
||||
|
||||
const GuardrailModeCard: React.FC<GuardrailModeCardProps> = ({
|
||||
guardrailId,
|
||||
mode,
|
||||
defaultOn,
|
||||
enabled,
|
||||
accessToken,
|
||||
isAdmin,
|
||||
onChanged,
|
||||
}) => {
|
||||
const handleToggle = async (nextEnabled: boolean) => {
|
||||
if (!accessToken) return;
|
||||
try {
|
||||
await setGuardrailEnabledCall(accessToken, guardrailId, nextEnabled);
|
||||
toast.success(`Guardrail ${nextEnabled ? "enabled" : "disabled"}`);
|
||||
} catch (error) {
|
||||
console.error("Error updating guardrail enabled state:", error);
|
||||
toast.fromError(`Failed to ${nextEnabled ? "enable" : "disable"} guardrail`);
|
||||
}
|
||||
await onChanged();
|
||||
};
|
||||
|
||||
return (
|
||||
<Card className="block p-6">
|
||||
<p>Mode</p>
|
||||
<div className="mt-2">
|
||||
<h3 className="text-lg font-medium">{formatGuardrailMode(mode) || "-"}</h3>
|
||||
<div className="mt-1 flex items-center gap-2">
|
||||
<Badge variant={defaultOn ? "secondary" : "outline"}>{defaultOn ? "Default On" : "Default Off"}</Badge>
|
||||
<Badge variant={enabled ? "secondary" : "outline"}>{enabled ? "Enabled" : "Disabled"}</Badge>
|
||||
</div>
|
||||
</div>
|
||||
{isAdmin && (
|
||||
<div className="mt-4 flex items-center gap-3 text-sm">
|
||||
<Switch checked={enabled} aria-label="Enable guardrail" onCheckedChange={handleToggle} />
|
||||
<span>{enabled ? "Guardrail runs on requests" : "Guardrail is skipped for every request"}</span>
|
||||
</div>
|
||||
)}
|
||||
</Card>
|
||||
);
|
||||
};
|
||||
|
||||
export default GuardrailModeCard;
|
||||
|
|
@ -1,12 +1,13 @@
|
|||
import { type UrlUpdateEvent } from "nuqs/adapters/testing";
|
||||
import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import GuardrailsPanel from "./GuardrailsPanel";
|
||||
import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking";
|
||||
import { getGuardrailsList, deleteGuardrailCall, setGuardrailEnabledCall } from "@/components/networking";
|
||||
import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
getGuardrailsList: vi.fn(),
|
||||
deleteGuardrailCall: vi.fn(),
|
||||
setGuardrailEnabledCall: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("./add_guardrail_form", () => ({
|
||||
|
|
@ -16,7 +17,7 @@ vi.mock("./add_guardrail_form", () => ({
|
|||
|
||||
vi.mock("./guardrail_table", () => ({
|
||||
__esModule: true,
|
||||
default: ({ guardrailsList, onDeleteClick, onGuardrailClick }: any) => (
|
||||
default: ({ guardrailsList, onDeleteClick, onGuardrailClick, onToggleEnabled }: any) => (
|
||||
<div>
|
||||
<div>Mock Guardrail Table</div>
|
||||
{guardrailsList.length > 0 && (
|
||||
|
|
@ -30,6 +31,12 @@ vi.mock("./guardrail_table", () => ({
|
|||
<button data-testid="open-button" onClick={() => onGuardrailClick(guardrailsList[0].guardrail_id)}>
|
||||
Open
|
||||
</button>
|
||||
<button
|
||||
data-testid="toggle-button"
|
||||
onClick={() => onToggleEnabled(guardrailsList[0].guardrail_id, !guardrailsList[0].enabled)}
|
||||
>
|
||||
Toggle
|
||||
</button>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
|
|
@ -100,6 +107,7 @@ describe("GuardrailsPanel", () => {
|
|||
|
||||
const mockGetGuardrailsList = vi.mocked(getGuardrailsList);
|
||||
const mockDeleteGuardrailCall = vi.mocked(deleteGuardrailCall);
|
||||
const mockSetGuardrailEnabledCall = vi.mocked(setGuardrailEnabledCall);
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
|
@ -117,11 +125,27 @@ describe("GuardrailsPanel", () => {
|
|||
created_at: "2024-01-01T00:00:00Z",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
guardrail_definition_location: "database" as any,
|
||||
enabled: true,
|
||||
},
|
||||
],
|
||||
});
|
||||
});
|
||||
|
||||
it("should disable a guardrail through the toggle and refetch the list", async () => {
|
||||
mockSetGuardrailEnabledCall.mockResolvedValue({ guardrail_id: "test-guardrail-1", enabled: false });
|
||||
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
|
||||
fireEvent.click(screen.getByText("Guardrails"));
|
||||
|
||||
fireEvent.click(await screen.findByTestId("toggle-button"));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockSetGuardrailEnabledCall).toHaveBeenCalledWith("test-token", "test-guardrail-1", false);
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
it("should render the component", async () => {
|
||||
renderWithProviders(<GuardrailsPanel {...defaultProps} />);
|
||||
expect(screen.getByText("Guardrails")).toBeInTheDocument();
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import { parseAsString, useQueryState } from "nuqs";
|
|||
import React, { useState, useEffect } from "react";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { ChevronDown, Code, Plus } from "lucide-react";
|
||||
import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking";
|
||||
import { getGuardrailsList, deleteGuardrailCall, setGuardrailEnabledCall } from "@/components/networking";
|
||||
import { buttonVariants } from "@/components/ui/button";
|
||||
import {
|
||||
DropdownMenu,
|
||||
|
|
@ -126,6 +126,20 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
|
|||
setGuardrailToDelete(null);
|
||||
};
|
||||
|
||||
const handleToggleEnabled = async (guardrailId: string, enabled: boolean) => {
|
||||
if (!accessToken) return;
|
||||
const guardrail = guardrailsList.find((g) => g.guardrail_id === guardrailId);
|
||||
const label = guardrail?.guardrail_name || guardrailId;
|
||||
try {
|
||||
await setGuardrailEnabledCall(accessToken, guardrailId, enabled);
|
||||
toast.success(`Guardrail "${label}" ${enabled ? "enabled" : "disabled"}`);
|
||||
} catch (error) {
|
||||
console.error("Error updating guardrail enabled state:", error);
|
||||
toast.fromError(`Failed to ${enabled ? "enable" : "disable"} guardrail`);
|
||||
}
|
||||
await fetchGuardrails();
|
||||
};
|
||||
|
||||
const providerDisplayName =
|
||||
guardrailToDelete && guardrailToDelete.litellm_params
|
||||
? getGuardrailLogoAndName(guardrailToDelete.litellm_params.guardrail).displayName
|
||||
|
|
@ -193,6 +207,7 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
|
|||
isLoading={isLoading}
|
||||
onDeleteClick={handleDeleteClick}
|
||||
onGuardrailClick={(id) => void setSelectedGuardrailId(id)}
|
||||
onToggleEnabled={handleToggleEnabled}
|
||||
/>
|
||||
)}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import { DataTableSortHeader } from "@/components/shared/DataTable";
|
|||
import { DateCell, IdentityCell, StatusBadge } from "@/components/shared/table_cells";
|
||||
import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types";
|
||||
import { buttonVariants } from "@/components/ui/button";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import {
|
||||
DropdownMenu,
|
||||
DropdownMenuContent,
|
||||
|
|
@ -66,11 +67,13 @@ function GuardrailRowActions({ guardrail, onDeleteClick }: GuardrailRowActionsPr
|
|||
interface GuardrailTableColumnsDeps {
|
||||
onGuardrailClick: (guardrailId: string) => void;
|
||||
onDeleteClick: (guardrailId: string, guardrailName: string) => void;
|
||||
onToggleEnabled: (guardrailId: string, enabled: boolean) => void;
|
||||
}
|
||||
|
||||
export const getGuardrailTableColumns = ({
|
||||
onGuardrailClick,
|
||||
onDeleteClick,
|
||||
onToggleEnabled,
|
||||
}: GuardrailTableColumnsDeps): ColumnDef<Guardrail>[] => [
|
||||
{
|
||||
id: "guardrail_id",
|
||||
|
|
@ -139,6 +142,22 @@ export const getGuardrailTableColumns = ({
|
|||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "enabled",
|
||||
accessorKey: "enabled",
|
||||
meta: { title: "Enabled" },
|
||||
header: "Enabled",
|
||||
size: 100,
|
||||
enableSorting: false,
|
||||
cell: ({ row }) => (
|
||||
<Switch
|
||||
size="sm"
|
||||
checked={row.original.enabled}
|
||||
aria-label={`Enable ${row.original.guardrail_name || row.original.guardrail_id}`}
|
||||
onCheckedChange={(enabled) => onToggleEnabled(row.original.guardrail_id, enabled)}
|
||||
/>
|
||||
),
|
||||
},
|
||||
{
|
||||
id: "created_at",
|
||||
accessorKey: "created_at",
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ vi.mock("@/components/networking", () => ({
|
|||
getGuardrailUISettings: vi.fn(),
|
||||
getGuardrailProviderSpecificParams: vi.fn(),
|
||||
updateGuardrailCall: vi.fn(),
|
||||
setGuardrailEnabledCall: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock ContentFilterManager
|
||||
|
|
@ -41,6 +42,68 @@ describe("Guardrail Info", () => {
|
|||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("lets an admin re-enable a disabled config guardrail from the overview and refetches it", async () => {
|
||||
vi.mocked(networking.getGuardrailInfo)
|
||||
.mockResolvedValueOnce({
|
||||
guardrail_id: "cfg-1",
|
||||
guardrail_name: "Headroom",
|
||||
litellm_params: { guardrail: "headroom", mode: "pre_call", default_on: true },
|
||||
guardrail_definition_location: "config",
|
||||
enabled: false,
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
guardrail_id: "cfg-1",
|
||||
guardrail_name: "Headroom",
|
||||
litellm_params: { guardrail: "headroom", mode: "pre_call", default_on: true },
|
||||
guardrail_definition_location: "config",
|
||||
enabled: true,
|
||||
});
|
||||
vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({
|
||||
supported_entities: [],
|
||||
supported_actions: [],
|
||||
pii_entity_categories: [],
|
||||
supported_modes: ["pre_call"],
|
||||
});
|
||||
vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({});
|
||||
vi.mocked(networking.setGuardrailEnabledCall).mockResolvedValue({ guardrail_id: "cfg-1", enabled: true });
|
||||
|
||||
render(<GuardrailInfoView guardrailId="cfg-1" onClose={() => {}} accessToken="tok" isAdmin={true} />);
|
||||
|
||||
const toggle = await screen.findByRole("switch", { name: "Enable guardrail" });
|
||||
expect(toggle).not.toBeChecked();
|
||||
expect(screen.getByText("Disabled")).toBeInTheDocument();
|
||||
|
||||
fireEvent.click(toggle);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.setGuardrailEnabledCall).toHaveBeenCalledWith("tok", "cfg-1", true);
|
||||
});
|
||||
expect(await screen.findByText("Enabled", { selector: "[data-slot='badge']" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("switch", { name: "Enable guardrail" })).toBeChecked();
|
||||
});
|
||||
|
||||
it("hides the enable switch from non-admins but still shows the state", async () => {
|
||||
vi.mocked(networking.getGuardrailInfo).mockResolvedValue({
|
||||
guardrail_id: "123",
|
||||
guardrail_name: "Test Guardrail",
|
||||
litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true },
|
||||
guardrail_definition_location: "db",
|
||||
enabled: false,
|
||||
});
|
||||
vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({
|
||||
supported_entities: [],
|
||||
supported_actions: [],
|
||||
pii_entity_categories: [],
|
||||
supported_modes: ["pre_call"],
|
||||
});
|
||||
vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({});
|
||||
|
||||
render(<GuardrailInfoView guardrailId="123" onClose={() => {}} accessToken="tok" isAdmin={false} />);
|
||||
|
||||
expect(await screen.findByText("Disabled")).toBeInTheDocument();
|
||||
expect(screen.queryByRole("switch", { name: "Enable guardrail" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should render the guardrail info after loading", async () => {
|
||||
// Mock the network responses
|
||||
vi.mocked(networking.getGuardrailInfo).mockResolvedValue({
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ import {
|
|||
} from "./GuardrailFormField";
|
||||
import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager";
|
||||
import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal";
|
||||
import GuardrailModeCard from "./GuardrailModeCard";
|
||||
import {
|
||||
formatGuardrailMode,
|
||||
getGuardrailLogoAndName,
|
||||
|
|
@ -556,17 +557,15 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
|
|||
</div>
|
||||
</Card>
|
||||
|
||||
<Card className="block p-6">
|
||||
<p>Mode</p>
|
||||
<div className="mt-2">
|
||||
<h3 className="text-lg font-medium">
|
||||
{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}
|
||||
</h3>
|
||||
<Badge variant={guardrailData.litellm_params?.default_on ? "secondary" : "outline"}>
|
||||
{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"}
|
||||
</Badge>
|
||||
</div>
|
||||
</Card>
|
||||
<GuardrailModeCard
|
||||
guardrailId={guardrailId}
|
||||
mode={guardrailData.litellm_params?.mode}
|
||||
defaultOn={Boolean(guardrailData.litellm_params?.default_on)}
|
||||
enabled={guardrailData.enabled !== false}
|
||||
accessToken={accessToken}
|
||||
isAdmin={isAdmin}
|
||||
onChanged={fetchGuardrailInfo}
|
||||
/>
|
||||
|
||||
<Card className="block p-6">
|
||||
<p>Created At</p>
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ const baseProps = {
|
|||
isLoading: false,
|
||||
onDeleteClick: vi.fn(),
|
||||
onGuardrailClick: vi.fn(),
|
||||
onToggleEnabled: vi.fn(),
|
||||
};
|
||||
|
||||
const makeGuardrail = (overrides: Partial<Guardrail> = {}): Guardrail => ({
|
||||
|
|
@ -19,13 +20,23 @@ const makeGuardrail = (overrides: Partial<Guardrail> = {}): Guardrail => ({
|
|||
created_at: "2021-01-01",
|
||||
updated_at: "2021-01-02",
|
||||
guardrail_definition_location: GuardrailDefinitionLocation.DB,
|
||||
enabled: true,
|
||||
...overrides,
|
||||
});
|
||||
|
||||
describe("GuardrailTable", () => {
|
||||
it("renders every column header", () => {
|
||||
render(<GuardrailTable guardrailsList={[]} {...baseProps} />);
|
||||
for (const header of ["Guardrail ID", "Name", "Provider", "Mode", "Default On", "Created At", "Updated At"]) {
|
||||
for (const header of [
|
||||
"Guardrail ID",
|
||||
"Name",
|
||||
"Provider",
|
||||
"Mode",
|
||||
"Default On",
|
||||
"Enabled",
|
||||
"Created At",
|
||||
"Updated At",
|
||||
]) {
|
||||
expect(screen.getByText(header)).toBeInTheDocument();
|
||||
}
|
||||
});
|
||||
|
|
@ -86,4 +97,53 @@ describe("GuardrailTable", () => {
|
|||
expect(deleteItem).toHaveAttribute("data-disabled");
|
||||
expect(onDeleteClick).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reflects the enabled state in the row switch", () => {
|
||||
render(
|
||||
<GuardrailTable
|
||||
guardrailsList={[
|
||||
makeGuardrail({ guardrail_id: "on", guardrail_name: "On", enabled: true }),
|
||||
makeGuardrail({ guardrail_id: "off", guardrail_name: "Off", enabled: false }),
|
||||
]}
|
||||
{...baseProps}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByRole("switch", { name: "Enable On" })).toBeChecked();
|
||||
expect(screen.getByRole("switch", { name: "Enable Off" })).not.toBeChecked();
|
||||
});
|
||||
|
||||
it("toggles a DB guardrail off from the row switch", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onToggleEnabled = vi.fn();
|
||||
render(
|
||||
<GuardrailTable
|
||||
guardrailsList={[makeGuardrail({ guardrail_id: "gr-2", guardrail_name: "Headroom" })]}
|
||||
{...baseProps}
|
||||
onToggleEnabled={onToggleEnabled}
|
||||
/>,
|
||||
);
|
||||
|
||||
await user.click(screen.getByRole("switch", { name: "Enable Headroom" }));
|
||||
|
||||
expect(onToggleEnabled).toHaveBeenCalledWith("gr-2", false);
|
||||
});
|
||||
|
||||
it("lets config guardrails be toggled even though they cannot be deleted", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onToggleEnabled = vi.fn();
|
||||
const guardrail = makeGuardrail({
|
||||
guardrail_id: "cfg-2",
|
||||
guardrail_name: "Config Headroom",
|
||||
guardrail_definition_location: GuardrailDefinitionLocation.CONFIG,
|
||||
enabled: false,
|
||||
});
|
||||
render(<GuardrailTable guardrailsList={[guardrail]} {...baseProps} onToggleEnabled={onToggleEnabled} />);
|
||||
|
||||
const toggle = screen.getByRole("switch", { name: "Enable Config Headroom" });
|
||||
expect(toggle).toBeEnabled();
|
||||
await user.click(toggle);
|
||||
|
||||
expect(onToggleEnabled).toHaveBeenCalledWith("cfg-2", true);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ interface GuardrailTableProps {
|
|||
isLoading: boolean;
|
||||
onDeleteClick: (guardrailId: string, guardrailName: string) => void;
|
||||
onGuardrailClick: (id: string) => void;
|
||||
onToggleEnabled: (guardrailId: string, enabled: boolean) => void;
|
||||
}
|
||||
|
||||
const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }];
|
||||
|
|
@ -35,12 +36,13 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
|
|||
isLoading,
|
||||
onDeleteClick,
|
||||
onGuardrailClick,
|
||||
onToggleEnabled,
|
||||
}) => {
|
||||
const [sorting, setSorting] = useState<SortingState>(DEFAULT_SORTING);
|
||||
|
||||
const columns = useMemo(
|
||||
() => getGuardrailTableColumns({ onGuardrailClick, onDeleteClick }),
|
||||
[onGuardrailClick, onDeleteClick],
|
||||
() => getGuardrailTableColumns({ onGuardrailClick, onDeleteClick, onToggleEnabled }),
|
||||
[onGuardrailClick, onDeleteClick, onToggleEnabled],
|
||||
);
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ export interface Guardrail {
|
|||
created_at?: string;
|
||||
updated_at?: string;
|
||||
guardrail_definition_location: GuardrailDefinitionLocation;
|
||||
enabled: boolean;
|
||||
}
|
||||
|
||||
export enum GuardrailDefinitionLocation {
|
||||
|
|
|
|||
|
|
@ -6210,6 +6210,27 @@ export const getAgentInfo = async (accessToken: string, agentId: string) => {
|
|||
}
|
||||
};
|
||||
|
||||
export const setGuardrailEnabledCall = async (accessToken: string, guardrailId: string, enabled: boolean) => {
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails/${guardrailId}/enabled` : `/guardrails/${guardrailId}/enabled`;
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "PATCH",
|
||||
headers: {
|
||||
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({ enabled }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.text();
|
||||
handleError(errorData);
|
||||
throw new Error(errorData);
|
||||
}
|
||||
|
||||
return response.json();
|
||||
};
|
||||
|
||||
export const getGuardrailInfo = async (accessToken: string, guardrailId: string) => {
|
||||
try {
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails/${guardrailId}/info` : `/guardrails/${guardrailId}/info`;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue