fix: guardrails working template policy

This commit is contained in:
Harshit28j 2026-03-16 13:29:06 +05:30
parent cd37ee1459
commit 17831f45a6
5 changed files with 697 additions and 372 deletions

View file

@ -805,7 +805,13 @@ class CustomGuardrail(CustomLogger):
"""
Update the guardrails litellm params in memory
"""
for key, value in vars(litellm_params).items():
# Handle both dict and Pydantic model/object
items = (
litellm_params.items()
if isinstance(litellm_params, dict)
else vars(litellm_params).items()
)
for key, value in items:
setattr(self, key, value)
def get_guardrails_messages_for_call_type(

View file

@ -335,11 +335,25 @@ async def create_guardrail(
f"Immediate sync: Successfully initialized guardrail '{guardrail_name}' (ID: {guardrail_id})"
)
except Exception as init_error:
verbose_proxy_logger.warning(
verbose_proxy_logger.error(
f"Immediate sync: Failed to initialize guardrail '{guardrail_name}' (ID: {guardrail_id}) in memory: {init_error}"
)
# Rollback: remove the ghost row from DB
try:
await GUARDRAIL_REGISTRY.delete_guardrail_from_db(
guardrail_id=guardrail_id, prisma_client=prisma_client
)
except Exception:
verbose_proxy_logger.error("Failed to rollback guardrail DB entry")
raise HTTPException(
status_code=422,
detail=f"Guardrail saved but failed to initialize: {init_error}",
)
return result
except HTTPException as e:
raise e
except Exception as e:
verbose_proxy_logger.exception(f"Error adding guardrail to db: {e}")
raise HTTPException(status_code=500, detail=str(e))
@ -444,9 +458,30 @@ async def update_guardrail(
f"Immediate sync: Successfully updated guardrail '{guardrail_name}' (ID: {guardrail_id})"
)
except Exception as update_error:
verbose_proxy_logger.warning(
verbose_proxy_logger.error(
f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}"
)
# Rollback: restore previous guardrail data in DB
try:
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
prisma_client=prisma_client,
)
# Re-initialize the old guardrail in memory
IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
)
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail DB entry after update failure"
)
raise HTTPException(
status_code=422,
detail=f"Guardrail update failed, rolled back: {update_error}",
)
return result
except HTTPException as e:
@ -518,9 +553,13 @@ async def delete_guardrail(
f"Immediate sync: Successfully removed guardrail '{guardrail_name}' (ID: {guardrail_id}) from memory"
)
except Exception as delete_error:
verbose_proxy_logger.warning(
verbose_proxy_logger.error(
f"Immediate sync: Failed to remove guardrail '{guardrail_name}' (ID: {guardrail_id}) from memory: {delete_error}"
)
raise HTTPException(
status_code=422,
detail=f"Guardrail deleted from DB but failed to remove from memory: {delete_error}",
)
return result
except HTTPException as e:
@ -1101,9 +1140,13 @@ async def patch_guardrail(
f"Immediate sync: Successfully updated guardrail '{guardrail_name}' (ID: {guardrail_id})"
)
except Exception as update_error:
verbose_proxy_logger.warning(
verbose_proxy_logger.error(
f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}"
)
raise HTTPException(
status_code=422,
detail=f"Guardrail patched in DB but failed to update in memory: {update_error}",
)
return result
except HTTPException as e:

View file

@ -544,23 +544,21 @@ class InMemoryGuardrailHandler:
self, guardrail_id: str, guardrail: Guardrail
) -> None:
"""
Update a guardrail in memory
Update a guardrail in memory by deleting and re-initializing.
- updates the guardrail in memory
- updates the guardrail params in litellm.callback_manager
Re-initialization is necessary because guardrails like
``ContentFilterGuardrail`` compile patterns at init time.
Simply patching attributes via ``setattr`` would leave stale
compiled state and skip validation of new patterns.
"""
self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail
# Delete old callback and in-memory references
self.delete_in_memory_guardrail(guardrail_id)
custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get(
guardrail_id
)
if custom_guardrail_callback:
updated_litellm_params = cast(
LitellmParams, guardrail.get("litellm_params", {})
)
custom_guardrail_callback.update_in_memory_litellm_params(
litellm_params=updated_litellm_params
)
# Ensure the guardrail_id is set so initialize_guardrail uses it
guardrail["guardrail_id"] = guardrail_id
# Re-initialize (validates patterns, compiles regexes, registers callback)
self.initialize_guardrail(guardrail=guardrail)
def delete_in_memory_guardrail(self, guardrail_id: str) -> None:
"""

View file

@ -659,8 +659,10 @@ async def get_policy_templates(
"yes",
)
if use_local:
return _load_policy_templates_from_local_backup()
templates = _load_policy_templates_from_local_backup()
return _filter_templates_by_available_patterns(templates)
templates = []
try:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -671,13 +673,51 @@ async def get_policy_templates(
)
response = await async_client.get(POLICY_TEMPLATES_GITHUB_URL)
if response.status_code == 200:
return response.json()
templates = response.json()
except Exception as e:
verbose_proxy_logger.debug(
"Failed to fetch policy templates from GitHub, using local backup: %s", e
)
return _load_policy_templates_from_local_backup()
if not templates:
templates = _load_policy_templates_from_local_backup()
return _filter_templates_by_available_patterns(templates)
def _filter_templates_by_available_patterns(templates: list) -> list:
"""Filter out templates that reference unavailable prebuilt patterns."""
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import (
PREBUILT_PATTERNS,
)
available_patterns = set(PREBUILT_PATTERNS.keys())
filtered_templates = []
for template in templates:
is_valid = True
guardrail_definitions = template.get("guardrailDefinitions", [])
for gd in guardrail_definitions:
litellm_params = gd.get("litellm_params", {})
if litellm_params.get("guardrail") == "litellm_content_filter":
patterns = litellm_params.get("patterns", [])
for p in patterns:
if p.get("pattern_type") == "prebuilt":
pattern_name = p.get("pattern_name")
if pattern_name and pattern_name not in available_patterns:
is_valid = False
break
if not is_valid:
break
if is_valid:
filtered_templates.append(template)
else:
verbose_proxy_logger.debug(
f"Filtering out policy template '{template.get('id')}' because it references unavailable patterns"
)
return filtered_templates
class EnrichTemplateRequest(BaseModel):