Merge pull request #14653 from ARajan1084/in-memory-guardrail-fix

fix: In Memory Guardrail fails to update

Author: krrishdholakia

litellm/integrations/custom_guardrail.py

@@ -487,7 +487,8 @@ def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None
         """
         Update the guardrails litellm params in memory
         """
-        pass
+        for key, value in vars(litellm_params).items():
+            setattr(self, key, value)
 
 
 def log_guardrail_information(func):

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -675,5 +675,6 @@ 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)
         if litellm_params.pii_entities_config:
             self.pii_entities_config = litellm_params.pii_entities_config

tests/test_litellm/proxy/guardrails/test_guardrail_registry.py

@@ -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
+    )