diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 47c6cb7e01a..ab524aec89c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -56,6 +56,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): api_endpoint: Optional[str] = None, **kwargs, ): + # Set supported event hooks if not already provided + if "event_hook" not in kwargs: + kwargs["event_hook"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ] + # Initialize parent classes first super().__init__(**kwargs) VertexBase.__init__(self)