-
-
Notifications
You must be signed in to change notification settings - Fork 9.9k
fix: Register DynamoAI guardrail initializer and enum entry #23752
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
a02962c
62ab632
4b29152
ed213ba
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
|
Comment on lines
+11
to
+23
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No test covering the new registration The PR fixes the "Unsupported guardrail: dynamoai" error by wiring the Adding a focused assertion (similar to def test_dynamoai_registry_resolution():
from litellm.proxy.guardrails.guardrail_registry import (
guardrail_class_registry,
guardrail_initializer_registry,
)
from litellm.proxy.guardrails.guardrail_hooks.dynamoai.dynamoai import DynamoAIGuardrails
assert guardrail_class_registry["dynamoai"] is DynamoAIGuardrails
assert "dynamoai" in guardrail_initializer_registryThis mirrors what already exists for Rule Used: What: Ensure that any PR claiming to fix an issue ... (source) |
||
|
|
||
|
|
||
| guardrail_initializer_registry = { | ||
| SupportedGuardrailIntegrations.DYNAMOAI.value: initialize_guardrail, | ||
| } | ||
|
|
||
|
|
||
| guardrail_class_registry = { | ||
| SupportedGuardrailIntegrations.DYNAMOAI.value: DynamoAIGuardrails, | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| """ | ||
| Tests for DynamoAI guardrail registration and initialization. | ||
| """ | ||
|
|
||
| import os | ||
| from unittest.mock import patch | ||
|
|
||
| import pytest | ||
|
Harshit28j marked this conversation as resolved.
|
||
|
|
||
|
|
||
| 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 | ||
Uh oh!
There was an error while loading. Please reload this page.