mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-23 22:27:10 +00:00
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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
Claude Haiku 4.5
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
parent
278c9babc6
commit
d15c2d546e
@@ -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,
|
||||
}
|
||||
|
||||
@@ -44,6 +44,7 @@ guardrails:
|
||||
class SupportedGuardrailIntegrations(Enum):
|
||||
APORIA = "aporia"
|
||||
BEDROCK = "bedrock"
|
||||
DYNAMOAI = "dynamoai"
|
||||
GUARDRAILS_AI = "guardrails_ai"
|
||||
LAKERA = "lakera"
|
||||
LAKERA_V2 = "lakera_v2"
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user