From d15c2d546ecf3448b2447b24aea0d7643b6efe3f Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 17 Mar 2026 10:54:29 +0530 Subject: [PATCH] fix: Register DynamoAI guardrail initializer and enum entry (#23752) * fix: Register DynamoAI guardrail initializer and enum entry Fix the "Unsupported guardrail: dynamoai" error by: 1. Adding DYNAMOAI to SupportedGuardrailIntegrations enum 2. Implementing initialize_guardrail() and registries in dynamoai/__init__.py The DynamoAI guardrail was added in PR #15920 but never properly registered in the initialization system. The __init__.py was missing the guardrail_initializer_registry and guardrail_class_registry dictionaries that the dynamic discovery mechanism looks for at module load time. Fixes #22773 Co-Authored-By: Claude Haiku 4.5 * Update litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * test: Add tests for DynamoAI guardrail registration Verifies enum entry, initializer registry, class registry, instance creation, and global registry discovery. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Haiku 4.5 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../guardrail_hooks/dynamoai/__init__.py | 32 +++++++- litellm/types/guardrails.py | 1 + .../guardrail_hooks/test_dynamoai.py | 81 +++++++++++++++++++ 3 files changed, 113 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py index 79f1992da4..f9ebf46a27 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py @@ -1,3 +1,33 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + from .dynamoai import DynamoAIGuardrails -__all__ = ["DynamoAIGuardrails"] +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _dynamoai_callback = DynamoAIGuardrails( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_dynamoai_callback) + + return _dynamoai_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.DYNAMOAI.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.DYNAMOAI.value: DynamoAIGuardrails, +} diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index d5abf5c8fb..27fa27e6da 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -44,6 +44,7 @@ guardrails: class SupportedGuardrailIntegrations(Enum): APORIA = "aporia" BEDROCK = "bedrock" + DYNAMOAI = "dynamoai" GUARDRAILS_AI = "guardrails_ai" LAKERA = "lakera" LAKERA_V2 = "lakera_v2" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py new file mode 100644 index 0000000000..7bc4e951a5 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py @@ -0,0 +1,81 @@ +""" +Tests for DynamoAI guardrail registration and initialization. +""" + +import os +from unittest.mock import patch + +import pytest + + +class TestDynamoAIGuardrailRegistration: + """Tests for DynamoAI guardrail registration in the guardrail system.""" + + def test_supported_guardrail_enum_entry(self): + """Test that DYNAMOAI is in SupportedGuardrailIntegrations enum.""" + from litellm.types.guardrails import SupportedGuardrailIntegrations + + assert hasattr(SupportedGuardrailIntegrations, "DYNAMOAI") + assert SupportedGuardrailIntegrations.DYNAMOAI.value == "dynamoai" + + def test_initialize_guardrail_function_exists(self): + """Test that initialize_guardrail function is properly exported.""" + from litellm.proxy.guardrails.guardrail_hooks.dynamoai import ( + guardrail_initializer_registry, + initialize_guardrail, + ) + + assert initialize_guardrail is not None + assert "dynamoai" in guardrail_initializer_registry + + def test_guardrail_class_registry_exists(self): + """Test that guardrail_class_registry is properly exported.""" + from litellm.proxy.guardrails.guardrail_hooks.dynamoai import ( + guardrail_class_registry, + ) + from litellm.proxy.guardrails.guardrail_hooks.dynamoai.dynamoai import ( + DynamoAIGuardrails, + ) + + assert "dynamoai" in guardrail_class_registry + assert guardrail_class_registry["dynamoai"] == DynamoAIGuardrails + + def test_initialize_guardrail_creates_instance(self): + """Test that initialize_guardrail creates a DynamoAIGuardrails instance.""" + from litellm.proxy.guardrails.guardrail_hooks.dynamoai import ( + initialize_guardrail, + ) + from litellm.proxy.guardrails.guardrail_hooks.dynamoai.dynamoai import ( + DynamoAIGuardrails, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="dynamoai", + mode="pre_call", + api_key="test-key", + api_base="https://test.dynamo.ai", + ) + + guardrail = { + "guardrail_name": "test-dynamoai-guard", + } + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ) as mock_add: + result = initialize_guardrail(litellm_params, guardrail) + + assert isinstance(result, DynamoAIGuardrails) + assert result.api_key == "test-key" + assert result.api_base == "https://test.dynamo.ai" + assert result.guardrail_name == "test-dynamoai-guard" + mock_add.assert_called_once_with(result) + + def test_dynamoai_in_global_registry(self): + """Test that dynamoai is discoverable in the global guardrail registry.""" + from litellm.proxy.guardrails.guardrail_registry import ( + guardrail_initializer_registry, + ) + + assert "dynamoai" in guardrail_initializer_registry