mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(guardrails): make config guardrails viewable in UI on no-DB deployments and stable across restarts
Resolves the '#35256' UI 'Guardrail not found' failures for config.yaml-defined guardrails. Cause A: /v2/guardrails/list and /guardrails/{id}/info raised 500 when prisma_client was None, so config guardrails were unviewable on no-DB proxies; the v1 list fallback never carried guardrail_id. Both list surfaces now merge the in-memory registry without a DB, info consults the in-memory registry when prisma is absent, and the v1 list stamps each config guardrail with its in-memory id. Cause B: config guardrails without an explicit id got a fresh uuid4 per process, so the id changed on every restart and differed between replicas, 404ing live guardrails. Config ids are now derived deterministically from the guardrail name (uuid5). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
7c56317edf
commit
a4feb921fa
4 changed files with 177 additions and 10 deletions
|
|
@ -62,6 +62,12 @@ def _get_guardrails_list_response(
|
|||
Helper function to get the guardrails list response
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
|
||||
name_to_id = {
|
||||
g.get("guardrail_name"): g.get("guardrail_id")
|
||||
for g in IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails()
|
||||
}
|
||||
|
||||
guardrail_configs: List[GuardrailInfoResponse] = []
|
||||
for guardrail in guardrails_config:
|
||||
|
|
@ -73,6 +79,7 @@ def _get_guardrails_list_response(
|
|||
)
|
||||
guardrail_configs.append(
|
||||
GuardrailInfoResponse(
|
||||
guardrail_id=name_to_id.get(guardrail.get("guardrail_name")),
|
||||
guardrail_name=guardrail.get("guardrail_name"),
|
||||
litellm_params=masked_params,
|
||||
guardrail_info=guardrail.get("guardrail_info"),
|
||||
|
|
@ -178,13 +185,14 @@ async def list_guardrails_v2(
|
|||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Prisma client not initialized")
|
||||
|
||||
is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
try:
|
||||
guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db(prisma_client=prisma_client)
|
||||
guardrails = (
|
||||
await GUARDRAIL_REGISTRY.get_all_guardrails_from_db(prisma_client=prisma_client)
|
||||
if prisma_client is not None
|
||||
else []
|
||||
)
|
||||
|
||||
excluded_guardrail_ids: set = set()
|
||||
if not is_admin:
|
||||
|
|
@ -1228,13 +1236,14 @@ async def get_guardrail_info(guardrail_id: str):
|
|||
from litellm.proxy.proxy_server import prisma_client
|
||||
from litellm.types.guardrails import GUARDRAIL_DEFINITION_LOCATION
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Prisma client not initialized")
|
||||
|
||||
try:
|
||||
guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.DB
|
||||
result = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db(
|
||||
guardrail_id=guardrail_id, prisma_client=prisma_client
|
||||
result = (
|
||||
await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db(
|
||||
guardrail_id=guardrail_id, prisma_client=prisma_client
|
||||
)
|
||||
if prisma_client is not None
|
||||
else None
|
||||
)
|
||||
if result is None:
|
||||
in_memory = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(guardrail_id=guardrail_id)
|
||||
|
|
|
|||
|
|
@ -214,6 +214,21 @@ def get_guardrail_class_from_hooks():
|
|||
return discovered_classes
|
||||
|
||||
|
||||
CONFIG_GUARDRAIL_ID_NAMESPACE = uuid.UUID("a7c9f1e2-3b4d-5e6f-8a9b-0c1d2e3f4a5b")
|
||||
|
||||
|
||||
def get_stable_config_guardrail_id(guardrail_name: str) -> str:
|
||||
"""
|
||||
Derive a deterministic guardrail_id for a config.yaml-defined guardrail.
|
||||
|
||||
Config guardrails have no persisted id, so deriving it from the guardrail
|
||||
name keeps it stable across restarts and identical across replicas. That way
|
||||
the UI list and info lookups resolve to the same id no matter which pod booted
|
||||
when.
|
||||
"""
|
||||
return str(uuid.uuid5(CONFIG_GUARDRAIL_ID_NAMESPACE, guardrail_name.encode("utf-8")))
|
||||
|
||||
|
||||
guardrail_class_registry.update(get_guardrail_class_from_hooks())
|
||||
|
||||
|
||||
|
|
@ -419,7 +434,13 @@ class InMemoryGuardrailHandler:
|
|||
|
||||
Returns a Guardrail object if the guardrail is initialized successfully
|
||||
"""
|
||||
guardrail_id = guardrail.get("guardrail_id") or str(uuid.uuid4())
|
||||
provided_id = guardrail.get("guardrail_id")
|
||||
if provided_id:
|
||||
guardrail_id = provided_id
|
||||
elif source == "config":
|
||||
guardrail_id = get_stable_config_guardrail_id(guardrail["guardrail_name"])
|
||||
else:
|
||||
guardrail_id = str(uuid.uuid4())
|
||||
guardrail["guardrail_id"] = guardrail_id
|
||||
if guardrail_id in self.IN_MEMORY_GUARDRAILS:
|
||||
verbose_proxy_logger.debug("guardrail_id already exists in IN_MEMORY_GUARDRAILS")
|
||||
|
|
|
|||
|
|
@ -339,6 +339,81 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock
|
|||
assert params["mode"] == "during_call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_guardrails_v2_without_db_returns_config_guardrails(
|
||||
mocker, mock_in_memory_handler
|
||||
):
|
||||
"""
|
||||
#35256: on a no-DB deployment (prisma_client is None) the v2 list must still
|
||||
return config-defined guardrails from the in-memory registry instead of 500ing.
|
||||
"""
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
|
||||
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
|
||||
assert response.guardrails[0].guardrail_id == "test-config-guardrail"
|
||||
assert response.guardrails[0].guardrail_definition_location == "config"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_without_db_returns_config_guardrail(
|
||||
mocker, mock_in_memory_handler
|
||||
):
|
||||
"""
|
||||
#35256: /guardrails/{id}/info must consult the in-memory registry even when
|
||||
prisma_client is None instead of raising 500, so config guardrails are viewable
|
||||
on no-DB deployments.
|
||||
"""
|
||||
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_list_guardrails_v1_carries_in_memory_ids(mocker, mock_in_memory_handler):
|
||||
"""
|
||||
#35256: the v1 list (UI fallback) must stamp each config guardrail with the
|
||||
in-memory guardrail_id (matched by name) so the UI opens /guardrails/<id>/info
|
||||
rather than /guardrails/undefined/info.
|
||||
"""
|
||||
from litellm.proxy.guardrails.guardrail_endpoints import list_guardrails
|
||||
|
||||
mock_proxy_config = mocker.Mock()
|
||||
mock_proxy_config.config = {
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": "Test Config Guardrail",
|
||||
"litellm_params": {"guardrail": "bedrock", "mode": "pre_call"},
|
||||
"guardrail_info": {"description": "x"},
|
||||
}
|
||||
]
|
||||
}
|
||||
mocker.patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config)
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
|
||||
mock_in_memory_handler,
|
||||
)
|
||||
|
||||
response = await list_guardrails()
|
||||
|
||||
assert len(response.guardrails) == 1
|
||||
assert response.guardrails[0].guardrail_id == "test-config-guardrail"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_from_db(mocker, mock_prisma_client):
|
||||
"""Test getting guardrail info from DB"""
|
||||
|
|
|
|||
|
|
@ -367,3 +367,65 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances():
|
|||
finally:
|
||||
for cb_list, snapshot in zip(lists, snapshots):
|
||||
cb_list[:] = snapshot
|
||||
|
||||
|
||||
def test_config_guardrail_id_is_stable_across_boots():
|
||||
"""
|
||||
#35256: a config.yaml guardrail with no explicit guardrail_id must get a
|
||||
deterministic id derived from its name, not a fresh uuid4 per process. Two
|
||||
independent handlers (simulating two boots / replicas) must agree, otherwise
|
||||
UI info lookups 404 after a restart or on a different pod.
|
||||
"""
|
||||
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=False,
|
||||
)
|
||||
|
||||
registry_module.guardrail_initializer_registry["stable_id_test"] = _initializer
|
||||
try:
|
||||
params = {"guardrail": "stable_id_test", "mode": "pre_call"}
|
||||
|
||||
first = InMemoryGuardrailHandler().initialize_guardrail(
|
||||
guardrail={"guardrail_name": "tooling", "litellm_params": dict(params)},
|
||||
source="config",
|
||||
)
|
||||
second = InMemoryGuardrailHandler().initialize_guardrail(
|
||||
guardrail={"guardrail_name": "tooling", "litellm_params": dict(params)},
|
||||
source="config",
|
||||
)
|
||||
|
||||
assert first["guardrail_id"] == second["guardrail_id"]
|
||||
assert first["guardrail_id"] == registry_module.get_stable_config_guardrail_id("tooling")
|
||||
assert first["guardrail_id"] != registry_module.get_stable_config_guardrail_id("other")
|
||||
finally:
|
||||
registry_module.guardrail_initializer_registry.pop("stable_id_test", None)
|
||||
|
||||
|
||||
def test_explicit_config_guardrail_id_is_preserved():
|
||||
"""An operator-provided guardrail_id must win over the derived-from-name id."""
|
||||
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=False,
|
||||
)
|
||||
|
||||
registry_module.guardrail_initializer_registry["explicit_id_test"] = _initializer
|
||||
try:
|
||||
result = InMemoryGuardrailHandler().initialize_guardrail(
|
||||
guardrail={
|
||||
"guardrail_id": "operator-chosen",
|
||||
"guardrail_name": "tooling",
|
||||
"litellm_params": {"guardrail": "explicit_id_test", "mode": "pre_call"},
|
||||
},
|
||||
source="config",
|
||||
)
|
||||
assert result["guardrail_id"] == "operator-chosen"
|
||||
finally:
|
||||
registry_module.guardrail_initializer_registry.pop("explicit_id_test", None)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue