mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(guardrails): keep no-op PUT and logging_only scan-failure semantics stable
Three regressions from the logging_only_scope feature, fixed while keeping the input/output/both selection working: 1. PUT with byte-identical litellm_params no longer forces a teardown + re-init of the live callback. reject_invalid_logging_only_scope now only validates; re-init still happens exactly when params/name change. Previously a description-only PUT re-appended the callback at the END of litellm.callbacks, reordering guardrails: with a BLOCK guardrail created before a MASK one, the mask started winning and blocked requests started succeeding (400 -> PUT 200 -> 200). Invalid unchanged scopes are still rejected with 422 without touching the live instance. 2. CustomGuardrail._scan_logged_call no longer swallows input-scan exceptions per branch: a raising or BLOCKED input scan aborts the logging_only hook again (one policy call, one verdict) for guardrails that never selected a logging_only_scope. Explicit input/output/both selection keeps selecting which scans run. 3. A failed PATCH now rolls the DB row back to exactly the stored raw litellm_params (a legacy 4-key row stays 4 keys) instead of expanding it to a full LitellmParams dump; pinned with a test. This matches the PUT rollback shape and is the faithful rollback.
This commit is contained in:
parent
9a1b1307cd
commit
73d034b12e
4 changed files with 173 additions and 17 deletions
|
|
@ -1033,18 +1033,12 @@ class CustomGuardrail(CustomLogger):
|
|||
"metadata": scratch_metadata,
|
||||
}
|
||||
if self.logging_only_scope != "output":
|
||||
try:
|
||||
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
|
||||
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
|
||||
if response is None or self.logging_only_scope == "input":
|
||||
return
|
||||
try:
|
||||
await output_translation.process_output_response(
|
||||
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
|
||||
await output_translation.process_output_response(
|
||||
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
|
||||
)
|
||||
|
||||
def supports_scan_only_tool_results(self) -> bool:
|
||||
"""Whether this guardrail can scan tool-result content.
|
||||
|
|
|
|||
|
|
@ -669,6 +669,28 @@ class InMemoryGuardrailHandler:
|
|||
siblings: Final = self.guardrail_id_to_sibling_callbacks.get(guardrail_id, ())
|
||||
return (() if primary is None else (primary,)) + siblings
|
||||
|
||||
def _reject_invalid_logging_only_scope(self, guardrail_id: str, guardrail: Guardrail) -> None:
|
||||
"""
|
||||
Strictly validate logging_only_scope on a row whose params are otherwise
|
||||
unchanged, without rebuilding the live callback.
|
||||
|
||||
API write paths send the whole object, so an invalid scope must still be
|
||||
rejected even when the write changed nothing else. But an unchanged row
|
||||
must not force a teardown + re-append: initialize_guardrail appends the
|
||||
rebuilt callback at the END of litellm.callbacks, so a no-op PUT would
|
||||
reorder guardrails and change which one wins between a BLOCK and a MASK
|
||||
guardrail over the same content.
|
||||
"""
|
||||
params: Final = guardrail.get("litellm_params")
|
||||
if not isinstance(params, (dict, LitellmParams)):
|
||||
return
|
||||
litellm_params: Final = LitellmParams(**params) if isinstance(params, dict) else params
|
||||
guardrail_name: Final = guardrail.get("guardrail_name", "Unknown")
|
||||
for custom_guardrail_callback in self._tracked_callbacks(guardrail_id):
|
||||
error: Final = _logging_only_scope_error(custom_guardrail_callback, guardrail_name, litellm_params)
|
||||
if error is not None:
|
||||
raise ValueError(error)
|
||||
|
||||
def initialize_custom_guardrail(
|
||||
self,
|
||||
guardrail: Guardrail,
|
||||
|
|
@ -737,13 +759,15 @@ class InMemoryGuardrailHandler:
|
|||
previous instance and raises), anything else only refreshes the stored row
|
||||
"""
|
||||
updated_guardrail: Final = cast(Guardrail, {**guardrail, "guardrail_id": guardrail_id})
|
||||
if reject_invalid_logging_only_scope or self._has_guardrail_params_changed(guardrail_id, updated_guardrail):
|
||||
if self._has_guardrail_params_changed(guardrail_id, updated_guardrail):
|
||||
self.reinitialize_guardrail(
|
||||
guardrail=updated_guardrail,
|
||||
source=source,
|
||||
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
|
||||
)
|
||||
return
|
||||
if reject_invalid_logging_only_scope:
|
||||
self._reject_invalid_logging_only_scope(guardrail_id, updated_guardrail)
|
||||
self.IN_MEMORY_GUARDRAILS[guardrail_id] = updated_guardrail
|
||||
self._sources[guardrail_id] = source
|
||||
|
||||
|
|
@ -967,7 +991,7 @@ class InMemoryGuardrailHandler:
|
|||
verbose_proxy_logger.error("Cannot sync guardrail without guardrail_id")
|
||||
return None
|
||||
|
||||
if reject_invalid_logging_only_scope or self._has_guardrail_params_changed(guardrail_id, guardrail):
|
||||
if self._has_guardrail_params_changed(guardrail_id, guardrail):
|
||||
guardrail_name: Final = guardrail.get("guardrail_name", "Unknown")
|
||||
verbose_proxy_logger.info(
|
||||
"Guardrail '%s' (ID: %s) params changed, re-initializing...", guardrail_name, guardrail_id
|
||||
|
|
@ -979,6 +1003,9 @@ class InMemoryGuardrailHandler:
|
|||
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
|
||||
)
|
||||
|
||||
if reject_invalid_logging_only_scope:
|
||||
self._reject_invalid_logging_only_scope(guardrail_id, guardrail)
|
||||
|
||||
# Params unchanged but the entry is still DB-backed; make sure the
|
||||
# source marker reflects that even if it was previously set differently
|
||||
# (e.g. a config entry whose UUID later collided with a DB row).
|
||||
|
|
|
|||
|
|
@ -2654,12 +2654,31 @@ class TestLoggingOnlyApplyGuardrail:
|
|||
|
||||
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
|
||||
|
||||
assert guardrail.calls == [
|
||||
("request", ["flagged content"]),
|
||||
("response", ["general kenobi"]),
|
||||
]
|
||||
# A raising input scan aborts the hook: the response scan must not run,
|
||||
# and exactly one intervened verdict is recorded (base semantics for a
|
||||
# guardrail that never selected a logging_only_scope).
|
||||
assert guardrail.calls == [("request", ["flagged content"])]
|
||||
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert [entry["guardrail_status"] for entry in entries] == ["guardrail_intervened", "guardrail_intervened"]
|
||||
assert [e["guardrail_status"] for e in entries] == ["guardrail_intervened"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_scan_error_aborts_the_response_scan(self):
|
||||
class _FailingObserver(_ApplyOnlyObserver):
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
self.calls.append((input_type, list(inputs.get("texts") or [])))
|
||||
if input_type == "request":
|
||||
raise RuntimeError("guardrail service unavailable")
|
||||
return GenericGuardrailAPIInputs(texts=[])
|
||||
|
||||
guardrail = _FailingObserver()
|
||||
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
|
||||
|
||||
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
|
||||
|
||||
assert guardrail.calls == [("request", ["hello there"])]
|
||||
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert [e["guardrail_status"] for e in entries] == ["guardrail_failed_to_respond"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_type_without_translation_is_skipped(self):
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from collections.abc import Iterable
|
||||
import json
|
||||
from typing import ClassVar, Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -404,6 +405,86 @@ def test_sync_guardrail_from_db_marks_source_db_when_unchanged():
|
|||
assert handler.get_source("collide") == "db"
|
||||
|
||||
|
||||
def test_sync_guardrail_from_db_reject_flag_keeps_callback_order_on_noop_update():
|
||||
"""
|
||||
The PUT endpoint syncs the whole object with reject_invalid_logging_only_scope=True.
|
||||
That strictness must not force a teardown + re-append of an unchanged guardrail:
|
||||
initialize_guardrail appends the rebuilt callback at the END of litellm.callbacks,
|
||||
so a description-only PUT would reorder guardrails and change which one wins
|
||||
between a BLOCK and a MASK guardrail over the same content.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
registry_module = _register_mode_following_initializer("mode_following_test")
|
||||
lists = _all_callback_lists()
|
||||
snapshots = [list(cb_list) for cb_list in lists]
|
||||
sentinel: Final = CustomGuardrail(
|
||||
guardrail_name="order-sentinel",
|
||||
supported_event_hooks=[GuardrailEventHooks.pre_call],
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
try:
|
||||
handler = InMemoryGuardrailHandler()
|
||||
handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call"), source="db")
|
||||
original = handler.guardrail_id_to_custom_guardrail["123"]
|
||||
assert original is not None
|
||||
# Park another callback after the guardrail so a re-append would be visible.
|
||||
litellm.callbacks.append(sentinel)
|
||||
index_before = litellm.callbacks.index(original)
|
||||
|
||||
handler.sync_guardrail_from_db(
|
||||
guardrail=_mode_following_db_row("123", "pre_call", "description-only edit"),
|
||||
reject_invalid_logging_only_scope=True,
|
||||
)
|
||||
|
||||
assert handler.guardrail_id_to_custom_guardrail["123"] is original
|
||||
assert litellm.callbacks.index(original) == index_before
|
||||
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_sync_guardrail_from_db_reject_flag_still_rejects_invalid_unchanged_scope():
|
||||
"""
|
||||
A PUT sends the whole object, so an unchanged row that already carries an
|
||||
invalid logging_only_scope (tolerated at load) must still be rejected on the
|
||||
strict sync path, without rebuilding the live callback.
|
||||
"""
|
||||
registry_module = _register_mode_following_initializer("mode_following_test")
|
||||
lists = _all_callback_lists()
|
||||
snapshots = [list(cb_list) for cb_list in lists]
|
||||
row: Final = Guardrail(
|
||||
guardrail_id="123",
|
||||
guardrail_name="mode-following",
|
||||
litellm_params={
|
||||
"guardrail": "mode_following_test",
|
||||
"mode": "pre_call",
|
||||
"default_on": True,
|
||||
"logging_only_scope": "input",
|
||||
},
|
||||
guardrail_info={},
|
||||
)
|
||||
try:
|
||||
handler = InMemoryGuardrailHandler()
|
||||
handler.initialize_guardrail(guardrail=row, source="db")
|
||||
original = handler.guardrail_id_to_custom_guardrail["123"]
|
||||
assert original is not None
|
||||
assert original.logging_only_scope is None # tolerated at load
|
||||
|
||||
with pytest.raises(ValueError, match="logging_only_scope is set"):
|
||||
handler.sync_guardrail_from_db(guardrail=row, reject_invalid_logging_only_scope=True)
|
||||
|
||||
# Rejected without touching the live instance.
|
||||
assert handler.guardrail_id_to_custom_guardrail["123"] is original
|
||||
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 _db_litellm_params() -> dict:
|
||||
"""
|
||||
Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params
|
||||
|
|
@ -1260,6 +1341,41 @@ async def test_update_guardrail_in_db_raises_when_row_missing():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_guardrail_in_db_persists_raw_sparse_params_verbatim():
|
||||
"""
|
||||
After a rejected PATCH, the endpoint rolls back by writing the stored row's
|
||||
raw litellm_params through update_guardrail_in_db. A raw dict must be
|
||||
persisted exactly as stored — a legacy 4-key row stays a 4-key row — instead
|
||||
of being round-tripped through LitellmParams.model_dump(), which materializes
|
||||
every field default and rewrites a row the admin never wrote.
|
||||
"""
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_guardrailstable.update = AsyncMock(
|
||||
return_value={"guardrail_id": "legacy-row", "guardrail_name": "legacy-one"}
|
||||
)
|
||||
legacy_params: Final = {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "pre_call",
|
||||
"guardrail_name": "legacy-one",
|
||||
"blocked_words": [{"keyword": "x", "action": "BLOCK"}],
|
||||
}
|
||||
|
||||
await GuardrailRegistry().update_guardrail_in_db(
|
||||
guardrail_id="legacy-row",
|
||||
guardrail=Guardrail(
|
||||
guardrail_id="legacy-row",
|
||||
guardrail_name="legacy-one",
|
||||
litellm_params=legacy_params,
|
||||
guardrail_info={},
|
||||
),
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
persisted: Final = prisma_client.db.litellm_guardrailstable.update.call_args.kwargs["data"]
|
||||
assert json.loads(persisted["litellm_params"]) == legacy_params
|
||||
|
||||
|
||||
def test_reinitialize_guardrail_restores_previous_on_failure():
|
||||
"""A reinitialization whose new params make the guardrail constructor raise must
|
||||
restore the previous instance instead of leaving the guardrail silently removed:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue