scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix(guardrails): rebuild the serving worker guardrail on PUT instead of patching it in place

Browse source 
update_in_memory_guardrail now goes through reinitialize_guardrail, the same
delete-and-construct path the DB poller and PATCH already use, whenever the
row name or litellm_params changed. Patching raw DB values over constructor
derived state clobbered normalized URLs, derived api_base values, and resolved
secrets, which 500d the serving worker in the earlier revision. An unchanged
config only refreshes the cached row, and a row the constructor rejects keeps
the previous instance enforcing and raises
This commit is contained in:
mateo-berri 2026-09-02 14:58:01 -07:00
parent cffa202bbf
commit 52f34ff553
18 changed files with 217 additions and 425 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 14075
"limit": 14076
},
"reportArgumentType": {
"limit": 2216
@ -9,7 +9,7 @@
"limit": 319
},
"reportAttributeAccessIssue": {
"limit": 479
"limit": 480
},
"reportCallIssue": {
"limit": 112
@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15303
"limit": 15306
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,22 +99,22 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44360
"limit": 44364
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38346
"limit": 38350
},
"reportUnknownParameterType": {
"limit": 19623
"limit": 19626
},
"reportUnknownVariableType": {
"limit": 29881
"limit": 29890
},
"reportUnnecessaryCast": {
"limit": 110
"limit": 111
},
"reportUnnecessaryComparison": {
"limit": 692
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 825
"limit": 826
},
"reportUntypedBaseClass": {
"limit": 0

100
litellm/integrations/custom_guardrail.py
View file

@ -3,11 +3,9 @@ import copy
import hashlib
import os
import secrets
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
from pydantic import TypeAdapter
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
from litellm._logging import verbose_logger
from litellm.caching import DualCache
@ -125,35 +123,6 @@ def _strict_guardrail_modes_enabled() -> bool:
return True if parsed is None else parsed
def updated_litellm_param(litellm_params: "LitellmParams | Mapping[str, object]", key: str) -> object:
if isinstance(litellm_params, Mapping):
return litellm_params.get(key)
value: Final[object] = getattr(litellm_params, key, None)
return value
GUARDRAIL_MODE_ADAPTER: Final[TypeAdapter[GuardrailEventHooks | list[GuardrailEventHooks] | Mode]] = TypeAdapter(
GuardrailEventHooks | list[GuardrailEventHooks] | Mode
)
def event_hook_as_constructed(
validated_mode: GuardrailEventHooks | list[GuardrailEventHooks] | Mode,
) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:
"""
Return the shape ``__init__`` stores for the same mode: ``LitellmParams``
coerces enum members to plain strings, so a resynced ``event_hook`` must
hold plain strings too or workers end up disagreeing on ``str(event_hook)``.
"""
if isinstance(validated_mode, Mode):
return validated_mode
if isinstance(validated_mode, list):
return cast( # cast-ok: __init__ stores the plain strings LitellmParams.mode carries
list[GuardrailEventHooks], [hook.value for hook in validated_mode]
)
return cast(GuardrailEventHooks, validated_mode.value) # cast-ok: same parity as the list branch
def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None:
"""Extract session_id from request data (litellm_session_id or metadata)."""
session_id = request_data.get("litellm_session_id")
@ -247,7 +216,18 @@ class CustomGuardrail(CustomLogger):
self.only_scan_new_messages: bool = only_scan_new_messages
if supported_event_hooks:
self._validate_or_warn_event_hook(event_hook, supported_event_hooks)
## validate event_hook is in supported_event_hooks
try:
self._validate_event_hook(event_hook, supported_event_hooks)
except ValueError as validation_error:
if _strict_guardrail_modes_enabled():
raise
verbose_logger.warning(
"%s. LITELLM_STRICT_GUARDRAIL_MODES=false; continuing "
"with unsupported event_hook. Set the env var to true "
"(default) to enforce validation and fail at startup.",
validation_error,
)
super().__init__(**kwargs)
def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str:
@ -610,12 +590,12 @@ class CustomGuardrail(CustomLogger):
def _validate_event_hook(
self,
event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | None,
supported_event_hooks: Sequence[GuardrailEventHooks],
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None,
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
def _validate_event_hook_list_is_in_supported_event_hooks(
event_hook: Sequence[GuardrailEventHooks] | Sequence[str],
supported_event_hooks: Sequence[GuardrailEventHooks],
event_hook: list[GuardrailEventHooks] | list[str],
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
for hook in event_hook:
if isinstance(hook, str):
@ -644,23 +624,6 @@ class CustomGuardrail(CustomLogger):
if event_hook not in supported_event_hooks:
raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}")
def _validate_or_warn_event_hook(
self,
event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | None,
supported_event_hooks: Sequence[GuardrailEventHooks],
) -> None:
try:
self._validate_event_hook(event_hook, supported_event_hooks)
except ValueError as validation_error:
if _strict_guardrail_modes_enabled():
raise
verbose_logger.warning(
"%s. LITELLM_STRICT_GUARDRAIL_MODES=false; continuing "
"with unsupported event_hook. Set the env var to true "
"(default) to enforce validation and fail at startup.",
validation_error,
)
@staticmethod
def _get_admin_metadata(data: dict) -> dict:
"""Return merged admin-configured key and team metadata from the request data.
@ -1373,29 +1336,12 @@ class CustomGuardrail(CustomLogger):
# Mask the content
return content_string[:start_index] + mask_string + content_string[end_index:]
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
"""
Update the guardrails litellm params in memory, accepting either a
LitellmParams object or the raw params mapping stored in the DB, and
resync ``event_hook`` when the update carries a new ``mode``. The new
mode is validated against ``supported_event_hooks`` before any state
is mutated, so a rejected update leaves the guardrail untouched.
``None`` values are skipped because both sources serialize every unset
LitellmParams field as ``None``; applying them would clobber
constructor-derived state (e.g. dict defaults) with ``None``.
Update the guardrails litellm params in memory
"""
updated_params: Final[Mapping[str, object]] = (
litellm_params if isinstance(litellm_params, Mapping) else vars(litellm_params)
)
raw_mode: Final = updated_params.get("mode")
new_event_hook: Final = None if raw_mode is None else GUARDRAIL_MODE_ADAPTER.validate_python(raw_mode)
if new_event_hook is not None and self.supported_event_hooks:
self._validate_or_warn_event_hook(new_event_hook, self.supported_event_hooks)
for key, value in updated_params.items():
if value is not None:
setattr(self, key, value)
if new_event_hook is not None:
self.event_hook = event_hook_as_constructed(new_event_hook)
for key, value in vars(litellm_params).items():
setattr(self, key, value)
def get_guardrails_messages_for_call_type(
self, call_type: CallTypes, data: dict | None = None
@ -1424,6 +1370,8 @@ class CustomGuardrail(CustomLogger):
# User/System messages are stored in the "input" key, use litellm transformation to get the messages
#########################################################
if call_type == CallTypes.responses.value or call_type == CallTypes.aresponses.value:
from typing import cast
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)

40
litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py
View file

@ -14,7 +14,6 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
updated_litellm_param,
)
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT,
@ -62,6 +61,15 @@ def _resolved_secret_value(value: object) -> object:
return value
def _updated_param(litellm_params: "LitellmParams | dict", key: str) -> object: # mutable-ok: DB dict
"""Read one param from a Mapping or a pydantic object, including pydantic
extras (cost_tier / price_per_1000_text_records live there), which the base
class ``vars()`` loop never sees."""
if isinstance(litellm_params, Mapping):
return litellm_params.get(key)
return getattr(litellm_params, key, None)
def _resolved_cost_tier(raw: object) -> str | None:
"""Normalize the configured cost_tier to 'free' / 'paid' / None."""
value: Final = _resolved_secret_value(raw)
@ -262,27 +270,29 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
verbose_proxy_logger.warning("Azure Prompt Shield: No user prompt found")
return None
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | dict") -> None: # mutable-ok: DB dict
"""Apply updated params in place, re-resolving billing and credentials.
Pricing is read via ``updated_litellm_param`` (the values are pydantic
extras, and the immediate PUT sync hands this method the raw DB dict).
Pricing and any ``os.environ/`` credential references are validated and
resolved BEFORE any state is mutated, so an invalid update leaves the
running guardrail untouched and a raw reference never overwrites a
resolved credential. Both input shapes flow through the base update so
the event_hook resync applies to each.
Pricing is read via ``_updated_param`` (the values are pydantic extras, and
the immediate PUT sync hands this method the raw DB dict). Pricing and any
``os.environ/`` credential references are validated and resolved BEFORE any
state is mutated, so an invalid update leaves the running guardrail
untouched and a raw reference never overwrites a resolved credential.
"""
cost_tier: Final = _resolved_cost_tier(updated_litellm_param(litellm_params, "cost_tier"))
price: Final = _resolved_price(updated_litellm_param(litellm_params, "price_per_1000_text_records"), cost_tier)
cost_tier: Final = _resolved_cost_tier(_updated_param(litellm_params, "cost_tier"))
price: Final = _resolved_price(_updated_param(litellm_params, "price_per_1000_text_records"), cost_tier)
resolved_credentials: dict[str, object] = {} # mutable-ok: staged before mutation
for cred_key in ("api_key", "api_base"):
cred_value = updated_litellm_param(litellm_params, cred_key)
cred_value = _updated_param(litellm_params, cred_key)
if isinstance(cred_value, str) and cred_value.startswith("os.environ/"):
resolved_credentials[cred_key] = _resolved_secret_value(cred_value)
super().update_in_memory_litellm_params(litellm_params)
for cred_key, cred_value in resolved_credentials.items():
setattr(self, cred_key, cred_value)
if isinstance(litellm_params, Mapping):
for key, value in litellm_params.items():
setattr(self, key, resolved_credentials.get(key, value))
else:
super().update_in_memory_litellm_params(litellm_params)
for cred_key, cred_value in resolved_credentials.items():
setattr(self, cred_key, cred_value)
self.cost_tier = cost_tier
self.price_per_1000_text_records = price

5
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -317,10 +317,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
super().update_in_memory_litellm_params(litellm_params)
extras: Final = litellm_params if isinstance(litellm_params, Mapping) else litellm_params.model_extra
self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(extras))
self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
def _streams_incrementally(self) -> bool:
return not self.streaming_buffer_until_moderated and not self.mask_response_content

25
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -13,7 +13,6 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
DEFAULT_ADVISORY_MESSAGE,
CustomGuardrail,
updated_litellm_param,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_skip_system_message_for_guardrail,
@ -305,7 +304,7 @@ class LakeraAIGuardrail(CustomGuardrail):
breakdown=self.breakdown,
)
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
"""
The base implementation blindly ``setattr``s every field on ``litellm_params``
(including ``on_flagged``/``advisory_system_message``/``payload``/``breakdown``)
@ -314,18 +313,24 @@ class LakeraAIGuardrail(CustomGuardrail):
on_flagged combinations __init__ rejects. Validate the prospective post-update
state *before* mutating, so a rejected update leaves the live instance untouched
instead of raising after it's already been corrupted.
The base setattr also writes ``litellm_params.mode`` onto a new ``self.mode``
attribute rather than the ``self.event_hook`` dispatch actually reads
(LitellmParams has no field literally named ``event_hook``), so without the
explicit sync below a hot reload that changes mode would pass validation but
keep dispatching on the stale event_hook.
"""
raw_on_flagged: Final = updated_litellm_param(litellm_params, "on_flagged")
raw_advisory: Final = updated_litellm_param(litellm_params, "advisory_system_message")
raw_payload: Final = updated_litellm_param(litellm_params, "payload")
raw_breakdown: Final = updated_litellm_param(litellm_params, "breakdown")
new_event_hook: Final = litellm_params.mode or self.event_hook
prospective_payload: Final = litellm_params.payload
prospective_breakdown: Final = litellm_params.breakdown
self._validate_advisory_config(
on_flagged=raw_on_flagged if isinstance(raw_on_flagged, str) and raw_on_flagged else self.on_flagged,
advisory_system_message=raw_advisory if isinstance(raw_advisory, str) else None,
payload=raw_payload if isinstance(raw_payload, bool) else self.payload,
breakdown=raw_breakdown if isinstance(raw_breakdown, bool) else self.breakdown,
on_flagged=litellm_params.on_flagged or self.on_flagged,
advisory_system_message=litellm_params.advisory_system_message,
payload=self.payload if prospective_payload is None else prospective_payload,
breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown,
)
super().update_in_memory_litellm_params(litellm_params=litellm_params)
self.event_hook = new_event_hook
def _validate_advisory_config(
self,

2
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -185,7 +185,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if self.optional_params.get("fail_on_error", True):
raise e from None
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
super().update_in_memory_litellm_params(litellm_params)
self.sanitize_error_detail = self.sanitize_error_detail is not False

50
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -11,7 +11,7 @@
import asyncio
import json
import threading
from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Mapping, Sequence
from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Sequence
from contextlib import asynccontextmanager
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, cast
@ -35,9 +35,7 @@ if TYPE_CHECKING:
from litellm.caching.caching import DualCache
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.integrations.custom_guardrail import (
GUARDRAIL_MODE_ADAPTER,
CustomGuardrail,
event_hook_as_constructed,
log_guardrail_information,
)
from litellm.proxy._types import UserAPIKeyAuth
@ -532,17 +530,17 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return created
@staticmethod
def _coerce_analyze_chunk_size(value: object) -> int:
def _coerce_analyze_chunk_size(value: int | None) -> int:
"""
Validate a configured chunk size, falling back to the default.
Non-positive or non-integer values would either bypass chunking entirely
or degenerate it into per-character splits (silently disabling
detection), so they are replaced by the default; values below 4 bytes
are floored to 4 and the splitter always emits at least one character
per chunk, so the chunked path can never re-enter itself.
Non-positive values would either bypass chunking entirely or degenerate
it into per-character splits (silently disabling detection), so they are
replaced by the default; values below 4 bytes are floored to 4 and the
splitter always emits at least one character per chunk, so the chunked
path can never re-enter itself.
"""
if not isinstance(value, int) or value <= 0:
if not value or value <= 0:
return DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES
return max(value, 4)
@ -1630,28 +1628,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
inputs["texts"] = new_texts
return inputs
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
"""
Update the guardrails litellm params in memory
"""
super().update_in_memory_litellm_params(litellm_params)
self.presidio_analyze_chunk_size_bytes = self._coerce_analyze_chunk_size(self.presidio_analyze_chunk_size_bytes)
self._resync_output_stage_event_hook()
def _resync_output_stage_event_hook(self) -> None:
if self.event_hook == GuardrailEventHooks.logging_only:
return
if self.apply_to_output:
self.event_hook = event_hook_as_constructed(GuardrailEventHooks.post_call)
return
if not self.output_parse_pii:
return
current_hook: Final = self.event_hook
if isinstance(current_hook, str) and current_hook != "post_call":
self.event_hook = event_hook_as_constructed(
GUARDRAIL_MODE_ADAPTER.validate_python((current_hook, GuardrailEventHooks.post_call))
)
elif isinstance(current_hook, list) and "post_call" not in current_hook:
self.event_hook = event_hook_as_constructed(
GUARDRAIL_MODE_ADAPTER.validate_python((*current_hook, GuardrailEventHooks.post_call))
if litellm_params.pii_entities_config:
self.pii_entities_config = litellm_params.pii_entities_config
if litellm_params.presidio_score_thresholds:
self.presidio_score_thresholds = litellm_params.presidio_score_thresholds
if litellm_params.presidio_entities_deny_list:
self.presidio_entities_deny_list = litellm_params.presidio_entities_deny_list
if litellm_params.presidio_analyze_chunk_size_bytes is not None:
# Same validation as __init__: a non-positive value from a guardrail
# update must not silently disable detection via degenerate chunking.
self.presidio_analyze_chunk_size_bytes = self._coerce_analyze_chunk_size(
litellm_params.presidio_analyze_chunk_size_bytes
)

9
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -7,7 +7,6 @@
import json
import os
from collections.abc import Mapping
from typing import Any, Final, Literal
from fastapi import HTTPException
@ -16,7 +15,6 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
updated_litellm_param,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
@ -113,7 +111,7 @@ class QualifireGuardrail(CustomGuardrail):
"only 'block' and 'monitor' are supported."
)
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
"""
The base implementation blindly ``setattr``s every field on ``litellm_params``
(including ``on_flagged``) onto this live instance with no revalidation, so an
@ -123,10 +121,7 @@ class QualifireGuardrail(CustomGuardrail):
the live instance untouched instead of raising after it's already been
corrupted. Mirrors LakeraAIGuardrail's own override of this same method.
"""
raw_on_flagged: Final = updated_litellm_param(litellm_params, "on_flagged")
prospective_on_flagged: Final = (
raw_on_flagged if isinstance(raw_on_flagged, str) and raw_on_flagged else self.on_flagged
)
prospective_on_flagged: Final = litellm_params.on_flagged or self.on_flagged
self._validate_on_flagged(prospective_on_flagged)
super().update_in_memory_litellm_params(litellm_params=litellm_params)

6
litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py
View file

@ -3,7 +3,6 @@ from __future__ import annotations
import asyncio
import json
import random
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn
from urllib.parse import urlsplit
@ -48,7 +47,6 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.guardrails import LitellmParams
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
GUARDRAIL_NAME: Final = "straiker"
@ -331,10 +329,6 @@ class StraikerGuardrail(CustomGuardrail):
self.configured_modes = _configured_modes(self.event_hook)
def update_in_memory_litellm_params(self, litellm_params: LitellmParams | Mapping[str, object]) -> None:
super().update_in_memory_litellm_params(litellm_params)
self.configured_modes = _configured_modes(self.event_hook)
def _webhook_url(self) -> str:
return f"{self.api_base}{WEBHOOK_PATH}"

26
litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
View file

@ -159,7 +159,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
self._compiled_rule_targets = compiled_targets
self._compiled_rule_patterns = compiled_patterns
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: LitellmParams | dict) -> None:
"""Apply updated params in place, rebuilding the compiled rule state.
The base implementation only ``setattr``s raw fields, which would leave
@ -169,19 +169,25 @@ class ToolPermissionGuardrail(CustomGuardrail):
immediate in-memory sync take effect, mirroring the PresidioGuardrail
override of this method.
"""
# ``litellm_params`` may arrive as the raw DB dict (the proxy ``cast()``s
# it to ``LitellmParams`` without converting), so handle both shapes. The
# base ``setattr`` loop is model-only, so apply the dict case here.
previous_rules: Final = self.rules
params: Final[Mapping[str, object]] = (
litellm_params if isinstance(litellm_params, Mapping) else vars(litellm_params)
)
super().update_in_memory_litellm_params(litellm_params)
if isinstance(litellm_params, dict):
params = litellm_params
for key, value in params.items():
setattr(self, key, value)
else:
super().update_in_memory_litellm_params(litellm_params)
params = vars(litellm_params)
# The generic update above sets ``self.rules`` from the incoming value
# (skipping None) but never rebuilds the compiled maps. Rebuild them
# when a rules list is provided; otherwise restore the previous ruleset
# so a non-list value can't silently wipe it. An explicit empty list
# still clears the rules.
# (None on a partial update that omits rules), but never rebuilds the
# compiled maps. Rebuild them when rules are provided; otherwise restore
# the previous ruleset so a partial update doesn't silently wipe it. An
# explicit empty list still clears the rules.
rules: Final = params.get("rules")
if isinstance(rules, list):
if rules is not None:
try:
self._load_rules(rules)
except Exception:

7
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py
View file

@ -4,7 +4,6 @@
#
# +-------------------------------------------------------------+
import os
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Literal, Optional
from fastapi import HTTPException
@ -13,7 +12,6 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
updated_litellm_param,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -104,10 +102,9 @@ class ZscalerAIGuard(CustomGuardrail):
return timeout
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams | Mapping[str, object]") -> None:
def update_in_memory_litellm_params(self, litellm_params: "LitellmParams") -> None:
super().update_in_memory_litellm_params(litellm_params)
raw_timeout: Final = updated_litellm_param(litellm_params, "timeout")
self.timeout = self._resolve_timeout(raw_timeout if isinstance(raw_timeout, (int, float)) else None)
self.timeout = self._resolve_timeout(litellm_params.timeout)
@staticmethod
def _resolve_metadata_value(request_data: dict | None, key: str) -> str | None:

22
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -6,7 +6,7 @@ import os
from collections.abc import Callable, Iterator, Mapping
from datetime import datetime, timezone
from itertools import chain, count
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast
from pydantic import ValidationError
@ -622,19 +622,15 @@ class InMemoryGuardrailHandler:
source: Literal["db", "config"] = "db",
) -> None:
"""
Update a guardrail in memory
- updates the guardrail params in litellm.callback_manager
- stores the guardrail in memory only after the callback update
succeeds, so a failed update stays visible as a diff to the
per-worker DB poller and gets retried instead of going stale
Update a guardrail in memory: a changed name or litellm_params rebuilds the
live callback from the new row (fail-closed: an invalid row keeps the
previous instance and raises), anything else only refreshes the stored row
"""
custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id)
updated_litellm_params: Final = guardrail.get("litellm_params")
if custom_guardrail_callback and updated_litellm_params:
custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params)
self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail
updated_guardrail: Final = cast(Guardrail, {**guardrail, "guardrail_id": guardrail_id})
if self._has_guardrail_params_changed(guardrail_id, updated_guardrail):
self.reinitialize_guardrail(guardrail=updated_guardrail, source=source)
return
self.IN_MEMORY_GUARDRAILS[guardrail_id] = updated_guardrail
self._sources[guardrail_id] = source
def delete_in_memory_guardrail(self, guardrail_id: str) -> None:

2
ruff-strict-budget.json
View file

@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1072
"limit": 1073
},
"TRY002": {
"limit": 524

112
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -9,7 +9,6 @@ from litellm.integrations.custom_guardrail import (
log_guardrail_information,
)
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail
@ -2240,117 +2239,6 @@ class TestRecordsOwnGuardrailInformation:
assert _guardrail_entries(request_data) == []
class TestUpdateInMemoryLitellmParams:
"""A PUT /guardrails update reaches the live callback through
update_in_memory_litellm_params: it must accept both a LitellmParams object
and the raw DB dict, and resync self.event_hook (which dispatch reads) from
the incoming mode instead of only writing a dead self.mode attribute (LIT-6591)."""
def _guardrail(self) -> CustomGuardrail:
return CustomGuardrail(
guardrail_name="update-test",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
)
def test_mode_change_resyncs_event_hook_dispatch(self):
guardrail = self._guardrail()
guardrail.update_in_memory_litellm_params(
LitellmParams(guardrail="update-test", mode="post_call", default_on=True)
)
assert guardrail.event_hook == "post_call"
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
def test_raw_db_dict_copies_params_and_resyncs_event_hook(self):
guardrail = self._guardrail()
guardrail.update_in_memory_litellm_params(
{
"guardrail": "update-test",
"mode": "post_call",
"api_base": "https://guardrail.example.com",
"default_on": True,
}
)
assert guardrail.event_hook == "post_call"
assert getattr(guardrail, "api_base", None) == "https://guardrail.example.com"
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True
@pytest.mark.parametrize(
"mode",
[
"post_call",
["pre_call", "post_call"],
{"default": "post_call", "tags": {"team-a": ["pre_call", "post_call"]}},
],
ids=["str", "list", "mode"],
)
def test_resynced_event_hook_has_the_shape_a_fresh_worker_constructs(self, mode):
"""Other workers rebuild the guardrail from the same DB row through
LitellmParams, which coerces enum members to plain strings; the serving
worker's in-place resync must land on that exact shape, or type-sensitive
readers such as str(self.event_hook) disagree across workers."""
updated = self._guardrail()
updated.update_in_memory_litellm_params({"guardrail": "update-test", "mode": mode, "default_on": True})
constructed = CustomGuardrail(
guardrail_name="update-test",
event_hook=LitellmParams(guardrail="update-test", mode=mode).mode,
default_on=True,
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
)
assert updated.event_hook == constructed.event_hook
assert type(updated.event_hook) is type(constructed.event_hook)
assert str(updated.event_hook) == str(constructed.event_hook)
if isinstance(updated.event_hook, list):
assert [type(hook) for hook in updated.event_hook] == [type(hook) for hook in constructed.event_hook]
def test_none_values_do_not_clobber_constructor_state(self):
guardrail = self._guardrail()
guardrail.additional_provider_specific_params = {"team": "security"}
guardrail.api_base = "https://guardrail.example.com"
guardrail.update_in_memory_litellm_params(
{
"mode": "post_call",
"api_base": None,
"additional_provider_specific_params": None,
"extra_headers": None,
}
)
assert guardrail.additional_provider_specific_params == {"team": "security"}
assert guardrail.api_base == "https://guardrail.example.com"
assert guardrail.event_hook == "post_call"
def test_strict_mode_rejects_unsupported_mode_without_mutating(self, monkeypatch):
monkeypatch.delenv("LITELLM_STRICT_GUARDRAIL_MODES", raising=False)
guardrail = self._guardrail()
with pytest.raises(ValueError, match="not in the supported event hooks"):
guardrail.update_in_memory_litellm_params(
{"mode": "during_call", "api_base": "https://guardrail.example.com"}
)
assert guardrail.event_hook is GuardrailEventHooks.pre_call
assert getattr(guardrail, "api_base", None) is None
def test_non_strict_mode_warns_and_applies_unsupported_mode(self, monkeypatch):
monkeypatch.setenv("LITELLM_STRICT_GUARDRAIL_MODES", "false")
guardrail = self._guardrail()
guardrail.update_in_memory_litellm_params({"mode": "during_call", "api_base": "https://guardrail.example.com"})
assert guardrail.event_hook == "during_call"
assert getattr(guardrail, "api_base", None) == "https://guardrail.example.com"
class _ApplyOnlyObserver(CustomGuardrail):
"""Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook."""

38
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -17,7 +17,7 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import (
_OPTIONAL_PresidioPIIMasking,
)
from litellm.exceptions import GuardrailRaisedException
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, PiiAction, PiiEntityType
from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType
from litellm.types.utils import Choices, Message, ModelResponse
from litellm.exceptions import BlockedPiiEntityError
@ -3167,42 +3167,6 @@ def test_update_in_memory_coerces_invalid_chunk_size():
assert guardrail.presidio_analyze_chunk_size_bytes == DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES
def test_update_in_memory_output_callback_keeps_forced_post_call():
"""The registry-tracked callback for filter_scope='output' is initialized with a
forced post_call hook regardless of the configured mode; a mode-changing update
must not move it off the response stage (LIT-6591)."""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
apply_to_output=True,
event_hook=GuardrailEventHooks.post_call.value,
)
guardrail.update_in_memory_litellm_params({"guardrail": "presidio", "mode": "pre_call", "default_on": True})
assert guardrail.event_hook == "post_call"
assert type(guardrail.event_hook) is str
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
def test_update_in_memory_output_parse_pii_keeps_post_call_expansion():
"""A guardrail with output_parse_pii must keep running on post_call to unmask the
response after a mode-changing update, mirroring the constructor's expansion."""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
event_hook="pre_call",
)
guardrail.update_in_memory_litellm_params(
LitellmParams(guardrail="presidio", mode="during_call", output_parse_pii=True, default_on=True)
)
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True
assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
def test_split_text_handles_chunk_size_below_char_width():
chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis(
text="\U0001f642\U0001f642", chunk_size_bytes=3, overlap_chars=8

15
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py
View file

@ -431,21 +431,6 @@ async def test_context_mode_omitted_when_event_hook_absent():
assert "mode" not in _posted_payload(g)["context"]
@pytest.mark.asyncio
async def test_context_mode_follows_an_in_memory_mode_update():
g = _make_guardrail(event_hook="pre_call")
g.update_in_memory_litellm_params({"guardrail": "straiker", "mode": "post_call", "default_on": True})
g.async_handler.post.return_value = _mock_response("NONE")
await g.apply_guardrail(
inputs={"texts": ["x"]},
request_data={"model": "m"},
input_type="response",
logging_obj=_logging_obj(),
)
assert g.configured_modes == ["post_call"]
assert _posted_payload(g)["context"]["mode"] == ["post_call"]
@pytest.mark.asyncio
async def test_identity_key_and_team_coalesce_alias_over_id():
g = _make_guardrail()

157
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -154,88 +154,103 @@ def test_duplicate_config_guardrail_names_get_distinct_stable_ids():
registry_module.guardrail_initializer_registry.pop("dup_name_test", None)
def test_update_in_memory_guardrail():
handler = InMemoryGuardrailHandler()
handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(
guardrail_name="test-guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
def _register_mode_following_initializer(guardrail_type: str):
"""Registers like the shipped initializers do: construct, then add the instance to litellm's callbacks."""
import litellm
from litellm.proxy.guardrails import guardrail_registry as registry_module
handler.update_in_memory_guardrail(
"123",
Guardrail(
guardrail_name="test-guardrail",
litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True),
),
)
assert (
handler.guardrail_id_to_custom_guardrail["123"].should_run_guardrail(
data={}, event_type=GuardrailEventHooks.pre_call
def _initializer(litellm_params, guardrail):
callback = CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
event_hook=GuardrailEventHooks(litellm_params.mode),
default_on=True,
)
is True
)
assert handler.guardrail_id_to_custom_guardrail["123"].event_hook == "pre_call"
litellm.logging_callback_manager.add_litellm_callback(callback)
return callback
registry_module.guardrail_initializer_registry[guardrail_type] = _initializer
return registry_module
def test_update_in_memory_guardrail_raw_db_dict_resyncs_event_hook():
"""PUT /guardrails hands this method the raw DB row, whose litellm_params is a
plain dict; the update must still apply and move dispatch to the new mode
instead of raising inside vars() and leaving the worker stale (LIT-6591)."""
handler = InMemoryGuardrailHandler()
handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(
guardrail_name="test-guardrail",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
def _mode_following_db_row(guardrail_id: str, mode: str, description: str = "") -> Guardrail:
"""The raw row GuardrailRegistry.update_guardrail_in_db hands back: litellm_params is a plain dict."""
return Guardrail(
guardrail_id=guardrail_id,
guardrail_name="mode-following",
litellm_params={"guardrail": "mode_following_test", "mode": mode, "default_on": True},
guardrail_info={"description": description},
)
updated_row = {
"guardrail_id": "123",
"guardrail_name": "test-guardrail",
"litellm_params": {"guardrail": "test-guardrail", "mode": "post_call", "default_on": True},
}
handler.update_in_memory_guardrail("123", updated_row)
callback = handler.guardrail_id_to_custom_guardrail["123"]
assert callback.event_hook == "post_call"
assert callback.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True
assert callback.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
assert handler.IN_MEMORY_GUARDRAILS["123"] == updated_row
def _live_instances_named(name: str) -> int:
return sum(1 for cb_list in _all_callback_lists() for cb in cb_list if getattr(cb, "guardrail_name", None) == name)
def test_update_in_memory_guardrail_failed_callback_update_stays_visible_to_poller(monkeypatch):
"""When the callback update raises, IN_MEMORY_GUARDRAILS must keep the old row:
storing the new row first would make the per-worker DB poller see no diff and
never re-initialize, leaving the PUT-serving worker stale until restart."""
monkeypatch.delenv("LITELLM_STRICT_GUARDRAIL_MODES", raising=False)
handler = InMemoryGuardrailHandler()
stale_row = Guardrail(
guardrail_id="123",
guardrail_name="test-guardrail",
litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True),
)
handler.IN_MEMORY_GUARDRAILS["123"] = stale_row
handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(
guardrail_name="test-guardrail",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
supported_event_hooks=[GuardrailEventHooks.pre_call],
)
def test_update_in_memory_guardrail_raw_db_row_mode_change_gates_at_the_new_stage():
registry_module = _register_mode_following_initializer("mode_following_test")
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler = InMemoryGuardrailHandler()
handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call"), source="db")
original = handler.guardrail_id_to_custom_guardrail["123"]
with pytest.raises(ValueError, match="not in the supported event hooks"):
handler.update_in_memory_guardrail(
"123",
{
"guardrail_id": "123",
"guardrail_name": "test-guardrail",
"litellm_params": {"guardrail": "test-guardrail", "mode": "post_call", "default_on": True},
},
)
handler.update_in_memory_guardrail("123", _mode_following_db_row("123", "post_call"))
assert handler.IN_MEMORY_GUARDRAILS["123"] == stale_row
assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call
replacement = handler.guardrail_id_to_custom_guardrail["123"]
assert replacement is not original
assert replacement.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True
assert replacement.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False
assert all(original not in cb_list for cb_list in lists)
assert _live_instances_named("mode-following") == 1
assert handler.IN_MEMORY_GUARDRAILS["123"]["litellm_params"].mode == "post_call"
assert handler.get_source("123") == "db"
finally:
registry_module.guardrail_initializer_registry.pop("mode_following_test", None)
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_update_in_memory_guardrail_unchanged_params_keep_the_live_instance():
registry_module = _register_mode_following_initializer("mode_following_test")
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler = InMemoryGuardrailHandler()
handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call", "old"), source="db")
original = handler.guardrail_id_to_custom_guardrail["123"]
handler.update_in_memory_guardrail("123", _mode_following_db_row("123", "pre_call", "new"))
assert handler.guardrail_id_to_custom_guardrail["123"] is original
assert handler.IN_MEMORY_GUARDRAILS["123"]["guardrail_info"] == {"description": "new"}
assert _live_instances_named("mode-following") == 1
finally:
registry_module.guardrail_initializer_registry.pop("mode_following_test", None)
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_update_in_memory_guardrail_invalid_row_keeps_the_previous_instance_enforcing():
registry_module = _register_mode_following_initializer("mode_following_test")
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler = InMemoryGuardrailHandler()
handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call"), source="db")
with pytest.raises(ValueError, match="not in the supported event hooks"):
handler.update_in_memory_guardrail("123", _mode_following_db_row("123", "during_call"))
restored = handler.guardrail_id_to_custom_guardrail["123"]
assert restored.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True
assert handler.IN_MEMORY_GUARDRAILS["123"]["litellm_params"].mode == "pre_call"
assert _live_instances_named("mode-following") == 1
finally:
registry_module.guardrail_initializer_registry.pop("mode_following_test", None)
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail:

8
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22359
"limit": 22364
},
"LIT002": {
"limit": 26776
"limit": 26777
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1038
"limit": 1039
},
"LIT007": {
"limit": 0
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16505
"limit": 16507
},
"LIT011": {
"limit": 5535