mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
added test for updating in memory guardrails
This commit is contained in:
parent
36e01b7881
commit
40fa415c89
1 changed files with 33 additions and 0 deletions
|
|
@ -1,6 +1,9 @@
|
|||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.guardrails.guardrail_registry import (
|
||||
get_guardrail_initializer_from_hooks,
|
||||
InMemoryGuardrailHandler,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Guardrail, LitellmParams
|
||||
|
||||
|
||||
def test_get_guardrail_initializer_from_hooks():
|
||||
|
|
@ -15,3 +18,33 @@ def test_guardrail_class_registry():
|
|||
print(f"guardrail_class_registry: {guardrail_class_registry}")
|
||||
assert "aim" in guardrail_class_registry
|
||||
assert "aporia" in guardrail_class_registry
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue