mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Every repository handed its `.table` back untyped, so a dozen modules had each grown a private `_PrismaTableActions` Protocol to paper over it. They had drifted: some declared `update` as returning the row, others the row or None, and none agreed on whether `find_many` was covariant Replace all of them with a single `TableActions[RowT_co]` in `litellm/repositories/prisma_protocols.py`, keyed to the prisma row each repository is bound to. Query inputs stay `Mapping[str, object]` so callers keep passing plain dicts, and `find_many` returns `Sequence` so the row type stays covariant Typing the nullable returns honestly surfaced paths that were already crashing. A team admin could never edit or delete a memory entry owned by their team: the write-auth check fed a raw prisma row to a helper that expects the domain model, so `members_with_roles` arrived as plain dicts and the request died as a 500 instead of applying the edit. Non-admin members hit the same 500 in place of the 403 they were owed, so refusal and breakage were indistinguishable. `/v2/model/info?user_models_only=true` dereferenced a missing user row rather than returning the 400 the route already had, three team routes dereferenced a team deleted between the read and the write, and the agent registry dereferenced a missing agent instead of naming it basedpyright drops 2,132 errors, 1,454 of them reportAny and 73 reportExplicitAny. The dashboard's generated types pick up `string[]` where they had `unknown[]` for a team's members, admins and models
681 lines
26 KiB
Python
681 lines
26 KiB
Python
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy.guardrails.guardrail_registry import (
|
|
get_guardrail_initializer_from_hooks,
|
|
GuardrailRegistry,
|
|
InMemoryGuardrailHandler,
|
|
)
|
|
from litellm.types.guardrails import GuardrailEventHooks, Guardrail, LitellmParams
|
|
|
|
|
|
def test_get_guardrail_initializer_from_hooks():
|
|
initializers = get_guardrail_initializer_from_hooks()
|
|
assert "aim" in initializers
|
|
|
|
|
|
def test_guardrail_class_registry():
|
|
from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
|
|
|
|
assert "aim" in guardrail_class_registry
|
|
assert "aporia" in guardrail_class_registry
|
|
|
|
|
|
def test_noma_registry_resolution():
|
|
from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaGuardrail
|
|
from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail
|
|
from litellm.proxy.guardrails.guardrail_registry import (
|
|
guardrail_class_registry,
|
|
guardrail_initializer_registry,
|
|
)
|
|
|
|
assert guardrail_class_registry["noma"] is NomaGuardrail
|
|
assert guardrail_class_registry["noma_v2"] is NomaV2Guardrail
|
|
assert "noma" in guardrail_initializer_registry
|
|
assert "noma_v2" in guardrail_initializer_registry
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"configured, expected",
|
|
[(None, True), (False, False), (True, True)],
|
|
)
|
|
def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected):
|
|
"""
|
|
A guardrail whose constructor sets run_in_parallel=True must keep that default when
|
|
the config omits the key; only an explicit config value may override it. The
|
|
previous code wrote bool(None)==False on every instance, silently disabling the
|
|
opt-in for such guardrails.
|
|
"""
|
|
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=True,
|
|
run_in_parallel=True,
|
|
)
|
|
|
|
registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer
|
|
try:
|
|
params = {"guardrail": "parallel_default_test", "mode": "pre_call"}
|
|
if configured is not None:
|
|
params["run_in_parallel"] = configured
|
|
|
|
handler = InMemoryGuardrailHandler()
|
|
result = handler.initialize_guardrail(
|
|
guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params},
|
|
)
|
|
|
|
stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
|
|
assert stored.run_in_parallel is expected
|
|
finally:
|
|
registry_module.guardrail_initializer_registry.pop("parallel_default_test", None)
|
|
|
|
|
|
def _register_noop_initializer(guardrail_type: str):
|
|
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[guardrail_type] = _initializer
|
|
return registry_module
|
|
|
|
|
|
def _config_guardrail(name: str, guardrail_type: str, guardrail_id=None) -> dict:
|
|
guardrail = {
|
|
"guardrail_name": name,
|
|
"litellm_params": {"guardrail": guardrail_type, "mode": "pre_call"},
|
|
}
|
|
if guardrail_id is not None:
|
|
guardrail["guardrail_id"] = guardrail_id
|
|
return guardrail
|
|
|
|
|
|
def test_config_guardrail_id_is_stable_across_boots():
|
|
"""
|
|
Config guardrails used to get a fresh uuid4 per process, so ids from a
|
|
previous boot (or another replica) 404'd on /guardrails/{id}/info even
|
|
though the guardrail was alive.
|
|
"""
|
|
registry_module = _register_noop_initializer("stable_id_test")
|
|
try:
|
|
first_boot = InMemoryGuardrailHandler().initialize_guardrail(
|
|
guardrail=_config_guardrail("tooling", "stable_id_test")
|
|
)
|
|
second_boot = InMemoryGuardrailHandler().initialize_guardrail(
|
|
guardrail=_config_guardrail("tooling", "stable_id_test")
|
|
)
|
|
|
|
assert first_boot["guardrail_id"] == second_boot["guardrail_id"]
|
|
finally:
|
|
registry_module.guardrail_initializer_registry.pop("stable_id_test", None)
|
|
|
|
|
|
def test_explicit_config_guardrail_id_wins_over_derived_id():
|
|
registry_module = _register_noop_initializer("explicit_id_test")
|
|
try:
|
|
result = InMemoryGuardrailHandler().initialize_guardrail(
|
|
guardrail=_config_guardrail(
|
|
"tooling", "explicit_id_test", guardrail_id="my-explicit-id"
|
|
)
|
|
)
|
|
|
|
assert result["guardrail_id"] == "my-explicit-id"
|
|
finally:
|
|
registry_module.guardrail_initializer_registry.pop("explicit_id_test", None)
|
|
|
|
|
|
def test_duplicate_config_guardrail_names_get_distinct_stable_ids():
|
|
"""
|
|
Duplicate guardrail_name entries are legitimate (load balancing across
|
|
deployments); each occurrence must keep its own id, stable across boots.
|
|
"""
|
|
registry_module = _register_noop_initializer("dup_name_test")
|
|
try:
|
|
handler = InMemoryGuardrailHandler()
|
|
first = handler.initialize_guardrail(
|
|
guardrail=_config_guardrail("dup", "dup_name_test")
|
|
)
|
|
second = handler.initialize_guardrail(
|
|
guardrail=_config_guardrail("dup", "dup_name_test")
|
|
)
|
|
|
|
rebooted_handler = InMemoryGuardrailHandler()
|
|
rebooted_first = rebooted_handler.initialize_guardrail(
|
|
guardrail=_config_guardrail("dup", "dup_name_test")
|
|
)
|
|
rebooted_second = rebooted_handler.initialize_guardrail(
|
|
guardrail=_config_guardrail("dup", "dup_name_test")
|
|
)
|
|
|
|
assert first["guardrail_id"] != second["guardrail_id"]
|
|
assert first["guardrail_id"] == rebooted_first["guardrail_id"]
|
|
assert second["guardrail_id"] == rebooted_second["guardrail_id"]
|
|
assert len(handler.IN_MEMORY_GUARDRAILS) == 2
|
|
finally:
|
|
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,
|
|
)
|
|
|
|
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
|
|
)
|
|
is True
|
|
)
|
|
assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call
|
|
|
|
|
|
def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail:
|
|
return Guardrail(
|
|
guardrail_id=guardrail_id,
|
|
guardrail_name=name,
|
|
litellm_params=LitellmParams(guardrail=name, mode="pre_call", default_on=False),
|
|
)
|
|
|
|
|
|
def test_reconcile_db_guardrails_drops_stale_db_entries_only():
|
|
"""
|
|
The reconcile pass must drop in-memory entries marked source='db' that are
|
|
missing from the DB result, and never touch source='config' entries.
|
|
Models the multi-pod case where another pod deleted a DB-backed guardrail.
|
|
"""
|
|
handler = InMemoryGuardrailHandler()
|
|
|
|
# Two DB-backed entries on this pod (synced from earlier polling cycles)
|
|
handler.IN_MEMORY_GUARDRAILS["db-keep"] = _make_guardrail("db-keep")
|
|
handler.IN_MEMORY_GUARDRAILS["db-stale"] = _make_guardrail("db-stale")
|
|
handler._sources["db-keep"] = "db"
|
|
handler._sources["db-stale"] = "db"
|
|
|
|
# One config-loaded entry that must survive reconciliation
|
|
handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg")
|
|
handler._sources["cfg"] = "config"
|
|
|
|
# The DB now only contains db-keep — db-stale was deleted on another pod.
|
|
removed = handler.reconcile_db_guardrails(db_guardrail_ids={"db-keep"})
|
|
|
|
assert removed == ["db-stale"]
|
|
assert "db-stale" not in handler.IN_MEMORY_GUARDRAILS
|
|
assert "db-stale" not in handler._sources
|
|
assert "db-keep" in handler.IN_MEMORY_GUARDRAILS
|
|
assert "cfg" in handler.IN_MEMORY_GUARDRAILS
|
|
assert handler._sources["cfg"] == "config"
|
|
|
|
|
|
def test_reconcile_does_not_drop_config_entries_missing_from_db():
|
|
"""A config-only guardrail (no DB row) must never be reconciled away."""
|
|
handler = InMemoryGuardrailHandler()
|
|
handler.IN_MEMORY_GUARDRAILS["cfg-only"] = _make_guardrail("cfg-only")
|
|
handler._sources["cfg-only"] = "config"
|
|
|
|
removed = handler.reconcile_db_guardrails(db_guardrail_ids=set())
|
|
|
|
assert removed == []
|
|
assert "cfg-only" in handler.IN_MEMORY_GUARDRAILS
|
|
|
|
|
|
def test_get_source_returns_marker_set_at_insert():
|
|
handler = InMemoryGuardrailHandler()
|
|
handler.IN_MEMORY_GUARDRAILS["a"] = _make_guardrail("a")
|
|
handler._sources["a"] = "db"
|
|
handler.IN_MEMORY_GUARDRAILS["b"] = _make_guardrail("b")
|
|
handler._sources["b"] = "config"
|
|
|
|
assert handler.get_source("a") == "db"
|
|
assert handler.get_source("b") == "config"
|
|
assert handler.get_source("missing") is None
|
|
|
|
|
|
def test_delete_in_memory_guardrail_clears_source_marker():
|
|
handler = InMemoryGuardrailHandler()
|
|
handler.IN_MEMORY_GUARDRAILS["a"] = _make_guardrail("a")
|
|
handler._sources["a"] = "db"
|
|
|
|
handler.delete_in_memory_guardrail("a")
|
|
|
|
assert "a" not in handler.IN_MEMORY_GUARDRAILS
|
|
assert "a" not in handler._sources
|
|
assert handler.get_source("a") is None
|
|
|
|
|
|
def test_list_config_guardrails_excludes_db_sourced():
|
|
"""LIT-2529: read surfaces union DB rows with config guardrails; db-sourced
|
|
in-memory entries would double-count (or resurrect stale ones), so exclude them."""
|
|
handler = InMemoryGuardrailHandler()
|
|
handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one")
|
|
handler._sources["cfg"] = "config"
|
|
handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one")
|
|
handler._sources["db"] = "db"
|
|
|
|
config_guardrails = handler.list_config_guardrails()
|
|
|
|
assert [g["guardrail_id"] for g in config_guardrails] == ["cfg"]
|
|
|
|
|
|
def test_get_config_guardrail_by_id_returns_config_only():
|
|
"""LIT-2529: the detail/logs fallback must return config-owned guardrails and
|
|
treat a db-sourced (stale) or missing id as a miss."""
|
|
handler = InMemoryGuardrailHandler()
|
|
handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one")
|
|
handler._sources["cfg"] = "config"
|
|
handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one")
|
|
handler._sources["db"] = "db"
|
|
|
|
assert handler.get_config_guardrail_by_id("cfg")["guardrail_name"] == "config-one"
|
|
assert handler.get_config_guardrail_by_id("db") is None
|
|
assert handler.get_config_guardrail_by_id("missing") is None
|
|
|
|
|
|
def test_initialize_guardrail_early_return_updates_source_marker():
|
|
"""
|
|
When initialize_guardrail is called for a guardrail that already exists
|
|
in memory, the early-return path must still honor the caller's source.
|
|
Otherwise a racing polling tick that placed a DB entry in memory first
|
|
would leave a later config-init call wrongly marked as 'db' (or vice
|
|
versa), and the entry would be reconciled with the wrong classification.
|
|
"""
|
|
handler = InMemoryGuardrailHandler()
|
|
# Simulate a polling tick already placing the entry as DB-backed.
|
|
handler.IN_MEMORY_GUARDRAILS["collide"] = _make_guardrail("collide", name="bedrock")
|
|
handler._sources["collide"] = "db"
|
|
|
|
# Config init re-visits the same id (e.g., hot-reload, or UUID collision).
|
|
g = Guardrail(
|
|
guardrail_id="collide",
|
|
guardrail_name="bedrock",
|
|
litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call", default_on=False),
|
|
)
|
|
handler.initialize_guardrail(guardrail=g, source="config")
|
|
|
|
assert handler.get_source("collide") == "config"
|
|
|
|
# And the symmetric direction: db sync should override an entry left
|
|
# marked as 'config' from a stale init path.
|
|
handler.initialize_guardrail(guardrail=g, source="db")
|
|
assert handler.get_source("collide") == "db"
|
|
|
|
|
|
def test_sync_guardrail_from_db_marks_source_db_when_unchanged():
|
|
"""
|
|
sync_guardrail_from_db must enforce source='db' even when params are
|
|
unchanged, so a config entry whose UUID happens to collide with a later
|
|
DB row gets re-tagged correctly.
|
|
"""
|
|
handler = InMemoryGuardrailHandler()
|
|
g = _make_guardrail("collide")
|
|
handler.IN_MEMORY_GUARDRAILS["collide"] = g
|
|
handler._sources["collide"] = "config"
|
|
|
|
handler.sync_guardrail_from_db(g)
|
|
|
|
assert handler.get_source("collide") == "db"
|
|
|
|
|
|
def _db_litellm_params() -> dict:
|
|
"""
|
|
Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params
|
|
is a raw dict (not a LitellmParams), holding only the keys originally stored,
|
|
a non-schema extra key, and plain-string enum values.
|
|
"""
|
|
return {
|
|
"guardrail": "litellm_content_filter",
|
|
"mode": "pre_call",
|
|
"default_on": True,
|
|
"version": 2,
|
|
"blocked_words": [{"keyword": "secret", "action": "BLOCK"}],
|
|
}
|
|
|
|
|
|
def test_unchanged_db_params_do_not_register_as_changed():
|
|
"""
|
|
A DB poll returns litellm_params as a raw dict while the in-memory copy is a
|
|
LitellmParams whose model_dump() fills every field default and coerces enums.
|
|
The two shapes must compare equal when the config is identical; otherwise
|
|
every poll cycle re-initializes the guardrail indefinitely.
|
|
"""
|
|
handler = InMemoryGuardrailHandler()
|
|
raw = _db_litellm_params()
|
|
gid = "11111111-1111-1111-1111-111111111111"
|
|
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
|
|
guardrail_id=gid,
|
|
guardrail_name="cf",
|
|
litellm_params=LitellmParams(**raw),
|
|
)
|
|
|
|
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=dict(raw))
|
|
assert handler._has_guardrail_params_changed(gid, new) is False
|
|
|
|
|
|
def test_changed_db_params_register_as_changed():
|
|
"""Normalizing both sides must still surface a genuine config change."""
|
|
handler = InMemoryGuardrailHandler()
|
|
raw = _db_litellm_params()
|
|
gid = "22222222-2222-2222-2222-222222222222"
|
|
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
|
|
guardrail_id=gid,
|
|
guardrail_name="cf",
|
|
litellm_params=LitellmParams(**raw),
|
|
)
|
|
|
|
changed = {**raw, "blocked_words": [{"keyword": "different", "action": "BLOCK"}]}
|
|
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=changed)
|
|
assert handler._has_guardrail_params_changed(gid, new) is True
|
|
|
|
|
|
def test_unnormalizable_db_params_register_as_changed_without_raising():
|
|
"""
|
|
A DB row whose litellm_params fail LitellmParams validation must not crash the
|
|
poll loop. The comparison falls back to treating the guardrail as changed so it
|
|
re-initializes (and surfaces the bad row in logs) rather than propagating the
|
|
validation error up through the polling cycle.
|
|
"""
|
|
handler = InMemoryGuardrailHandler()
|
|
raw = _db_litellm_params()
|
|
gid = "55555555-5555-5555-5555-555555555555"
|
|
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
|
|
guardrail_id=gid,
|
|
guardrail_name="cf",
|
|
litellm_params=LitellmParams(**raw),
|
|
)
|
|
|
|
malformed = {**raw, "default_on": "not-a-bool-xyz"}
|
|
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=malformed)
|
|
assert handler._has_guardrail_params_changed(gid, new) is True
|
|
|
|
|
|
def _all_callback_lists():
|
|
import litellm
|
|
|
|
return [
|
|
litellm.callbacks,
|
|
litellm.success_callback,
|
|
litellm.failure_callback,
|
|
litellm._async_success_callback,
|
|
litellm._async_failure_callback,
|
|
]
|
|
|
|
|
|
def test_delete_in_memory_guardrail_removes_callback_from_all_lists():
|
|
"""
|
|
Request handling promotes guardrail callbacks from litellm.callbacks into the
|
|
success/failure/async lists. delete_in_memory_guardrail must purge the callback
|
|
from every list, otherwise a re-initialized guardrail leaves its old instance
|
|
stranded in those lists and instances accumulate.
|
|
"""
|
|
handler = InMemoryGuardrailHandler()
|
|
callback = CustomGuardrail(
|
|
guardrail_name="cf-delete",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
gid = "33333333-3333-3333-3333-333333333333"
|
|
handler.IN_MEMORY_GUARDRAILS[gid] = _make_guardrail(gid, "cf-delete")
|
|
handler._sources[gid] = "db"
|
|
handler.guardrail_id_to_custom_guardrail[gid] = callback
|
|
|
|
lists = _all_callback_lists()
|
|
snapshots = [list(cb_list) for cb_list in lists]
|
|
try:
|
|
for cb_list in lists:
|
|
cb_list.append(callback)
|
|
|
|
handler.delete_in_memory_guardrail(gid)
|
|
|
|
for cb_list in lists:
|
|
assert callback not in cb_list
|
|
finally:
|
|
for cb_list, snapshot in zip(lists, snapshots):
|
|
cb_list[:] = snapshot
|
|
|
|
|
|
def test_repeated_db_sync_does_not_accumulate_runner_instances():
|
|
"""
|
|
End-to-end regression for the OOM: across repeated DB polls (with the config
|
|
genuinely changing each cycle to force re-initialization), exactly one live
|
|
guardrail instance must exist across all callback lists. On the unfixed code
|
|
the stale instance lingers in the success/failure lists and the distinct count
|
|
climbs above one.
|
|
"""
|
|
import litellm
|
|
|
|
handler = InMemoryGuardrailHandler()
|
|
gid = "44444444-4444-4444-4444-444444444444"
|
|
name = "cf-accum"
|
|
|
|
def db_guardrail(word: str) -> Guardrail:
|
|
params = {
|
|
**_db_litellm_params(),
|
|
"blocked_words": [{"keyword": word, "action": "BLOCK"}],
|
|
}
|
|
return Guardrail(guardrail_id=gid, guardrail_name=name, litellm_params=params)
|
|
|
|
def promote_into_request_lists() -> None:
|
|
manager = litellm.logging_callback_manager
|
|
for callback in list(litellm.callbacks):
|
|
manager.add_litellm_success_callback(callback)
|
|
manager.add_litellm_failure_callback(callback)
|
|
manager.add_litellm_async_success_callback(callback)
|
|
manager.add_litellm_async_failure_callback(callback)
|
|
|
|
def distinct_runner_instances() -> int:
|
|
seen = set()
|
|
for callback in litellm.logging_callback_manager._get_all_callbacks():
|
|
if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == name:
|
|
seen.add(id(callback))
|
|
return len(seen)
|
|
|
|
lists = _all_callback_lists()
|
|
snapshots = [list(cb_list) for cb_list in lists]
|
|
try:
|
|
for cycle in range(5):
|
|
handler.sync_guardrail_from_db(db_guardrail(f"word-{cycle}"))
|
|
promote_into_request_lists()
|
|
|
|
assert distinct_runner_instances() == 1
|
|
finally:
|
|
for cb_list, snapshot in zip(lists, snapshots):
|
|
cb_list[:] = snapshot
|
|
|
|
|
|
def _judge_guardrail(guardrail_id: str) -> Guardrail:
|
|
return Guardrail(
|
|
guardrail_id=guardrail_id,
|
|
guardrail_name="quality-judge",
|
|
litellm_params={
|
|
"guardrail": "llm_as_a_judge",
|
|
"mode": "post_call",
|
|
"judge_model": "my-judge-alias",
|
|
"overall_threshold": 80,
|
|
"on_failure": "log",
|
|
"criteria": [{"name": "helpfulness", "weight": 100, "description": "helpful?"}],
|
|
},
|
|
)
|
|
|
|
|
|
def test_db_synced_judge_guardrail_uses_lazy_router_provider():
|
|
"""A judge guardrail created/synced through a DB path must resolve the active
|
|
Router lazily at call time (issue: UI-created guardrails failed open because the
|
|
Router was captured at construction; a guardrail created before the Router
|
|
existed captured None and never recovered). Asserting the default provider is
|
|
wired guarantees the instance reads the live global rather than a stale value."""
|
|
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
|
|
LLMAsAJudgeGuardrail,
|
|
_default_router_provider,
|
|
)
|
|
|
|
handler = InMemoryGuardrailHandler()
|
|
|
|
lists = _all_callback_lists()
|
|
snapshots = [list(cb_list) for cb_list in lists]
|
|
try:
|
|
handler.sync_guardrail_from_db(_judge_guardrail("judge-db"))
|
|
|
|
instance = handler.guardrail_id_to_custom_guardrail["judge-db"]
|
|
assert isinstance(instance, LLMAsAJudgeGuardrail)
|
|
assert instance._router_provider is _default_router_provider
|
|
finally:
|
|
for cb_list, snapshot in zip(lists, snapshots):
|
|
cb_list[:] = snapshot
|
|
|
|
|
|
def test_reinitialized_judge_guardrail_uses_lazy_router_provider():
|
|
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
|
|
LLMAsAJudgeGuardrail,
|
|
_default_router_provider,
|
|
)
|
|
|
|
handler = InMemoryGuardrailHandler()
|
|
|
|
lists = _all_callback_lists()
|
|
snapshots = [list(cb_list) for cb_list in lists]
|
|
try:
|
|
handler.reinitialize_guardrail(_judge_guardrail("judge-reinit"), source="db")
|
|
|
|
instance = handler.guardrail_id_to_custom_guardrail["judge-reinit"]
|
|
assert isinstance(instance, LLMAsAJudgeGuardrail)
|
|
assert instance._router_provider is _default_router_provider
|
|
finally:
|
|
for cb_list, snapshot in zip(lists, snapshots):
|
|
cb_list[:] = snapshot
|
|
|
|
|
|
class TestScanOnlyToolResultsInitRefusal:
|
|
"""A guardrail whose role filtering never scans tool results must be rejected at
|
|
initialization when configured with scan_only_tool_results, instead of booting a
|
|
proxy that silently scans nothing on every request."""
|
|
|
|
def _initialize(self, name: str, params: dict):
|
|
lists = _all_callback_lists()
|
|
snapshots = [list(cb_list) for cb_list in lists]
|
|
try:
|
|
return InMemoryGuardrailHandler().initialize_guardrail(
|
|
guardrail={"guardrail_name": name, "litellm_params": params},
|
|
)
|
|
finally:
|
|
for cb_list, snapshot in zip(lists, snapshots):
|
|
cb_list[:] = snapshot
|
|
|
|
def test_panw_prisma_airs_with_scan_only_tool_results_is_rejected(self):
|
|
with pytest.raises(ValueError, match="never scans tool results"):
|
|
self._initialize(
|
|
"panw-scan-only-combo",
|
|
{
|
|
"guardrail": "panw_prisma_airs",
|
|
"mode": "pre_call",
|
|
"api_key": "test-key",
|
|
"profile_name": "test-profile",
|
|
"scan_only_tool_results": True,
|
|
},
|
|
)
|
|
|
|
def test_bedrock_latest_role_with_scan_only_tool_results_is_rejected(self):
|
|
with pytest.raises(ValueError, match="never scans tool results"):
|
|
self._initialize(
|
|
"bedrock-latest-role-scan-only-combo",
|
|
{
|
|
"guardrail": "bedrock",
|
|
"mode": "pre_call",
|
|
"guardrailIdentifier": "gr-1",
|
|
"guardrailVersion": "1",
|
|
"experimental_use_latest_role_message_only": True,
|
|
"scan_only_tool_results": True,
|
|
},
|
|
)
|
|
|
|
def test_bedrock_without_latest_role_accepts_scan_only_tool_results(self):
|
|
result = self._initialize(
|
|
"bedrock-scan-only-ok",
|
|
{
|
|
"guardrail": "bedrock",
|
|
"mode": "pre_call",
|
|
"guardrailIdentifier": "gr-1",
|
|
"guardrailVersion": "1",
|
|
"scan_only_tool_results": True,
|
|
},
|
|
)
|
|
assert result is not None
|
|
|
|
def test_prompt_security_default_tool_filtering_rejects_scan_only_tool_results(self, monkeypatch):
|
|
monkeypatch.delenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", raising=False)
|
|
with pytest.raises(ValueError, match="never scans tool results"):
|
|
self._initialize(
|
|
"prompt-security-scan-only-combo",
|
|
{
|
|
"guardrail": "prompt_security",
|
|
"mode": "pre_call",
|
|
"api_key": "test-key",
|
|
"api_base": "https://ps.example.com",
|
|
"scan_only_tool_results": True,
|
|
},
|
|
)
|
|
|
|
def test_prompt_security_check_tool_results_accepts_scan_only_tool_results(self, monkeypatch):
|
|
monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true")
|
|
result = self._initialize(
|
|
"prompt-security-scan-only-ok",
|
|
{
|
|
"guardrail": "prompt_security",
|
|
"mode": "pre_call",
|
|
"api_key": "test-key",
|
|
"api_base": "https://ps.example.com",
|
|
"scan_only_tool_results": True,
|
|
},
|
|
)
|
|
assert result is not None
|
|
|
|
def test_skip_tool_message_with_scan_only_tool_results_is_rejected(self):
|
|
with pytest.raises(ValueError, match="skip_tool_message_in_guardrail are enabled together"):
|
|
self._initialize(
|
|
"bedrock-skip-tool-scan-only-combo",
|
|
{
|
|
"guardrail": "bedrock",
|
|
"mode": "pre_call",
|
|
"guardrailIdentifier": "gr-1",
|
|
"guardrailVersion": "1",
|
|
"skip_tool_message_in_guardrail": True,
|
|
"scan_only_tool_results": True,
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_guardrail_in_db_raises_when_row_missing():
|
|
prisma_client = MagicMock()
|
|
prisma_client.db.litellm_guardrailstable.update = AsyncMock(return_value=None)
|
|
|
|
with pytest.raises(
|
|
Exception,
|
|
match=r"^Error updating guardrail in DB: Guardrail not found, passed guardrail_id=missing-guardrail$",
|
|
):
|
|
await GuardrailRegistry().update_guardrail_in_db(
|
|
guardrail_id="missing-guardrail",
|
|
guardrail=Guardrail(
|
|
guardrail_name="missing-guardrail",
|
|
litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call"),
|
|
),
|
|
prisma_client=prisma_client,
|
|
)
|