mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix: guardrails working template policy
This commit is contained in:
parent
cd37ee1459
commit
17831f45a6
5 changed files with 697 additions and 372 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue