mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Implement Bedrock Guardrail apply_guardrail endpoint support (#15892)
* Add bedrock support for apply gaurdrails * Add bedrock support doc * remove unused variable * remove unused variable
This commit is contained in:
Sameer Kankute
GitHub
parent
b9585b1db5
commit
c638f45213
3 changed files with 333 additions and 3 deletions
|
|
@ -3,13 +3,49 @@ import TabItem from '@theme/TabItem';
|
|||
|
||||
# /guardrails/apply_guardrail
|
||||
|
||||
Use this endpoint to directly call a guardrail configured on your LiteLLM instance. This is useful when you have services that need to directly call a guardrail.
|
||||
Use this endpoint to directly call a guardrail configured on your LiteLLM instance. This is useful when you have services that need to directly call a guardrail.
|
||||
|
||||
## Supported Guardrail Types
|
||||
|
||||
This endpoint supports various guardrail types including:
|
||||
- **Presidio** - PII detection and masking
|
||||
- **Bedrock** - AWS Bedrock guardrails for content moderation
|
||||
- **Lakera** - AI safety guardrails
|
||||
- **Custom guardrails** - User-defined guardrails
|
||||
|
||||
## Configuration
|
||||
|
||||
### Bedrock Guardrail Configuration
|
||||
|
||||
To use Bedrock guardrails with the apply_guardrail endpoint, configure your guardrail in your LiteLLM config.yaml:
|
||||
|
||||
```yaml
|
||||
guardrails:
|
||||
- guardrail_name: "bedrock-content-guard"
|
||||
litellm_params:
|
||||
guardrail: bedrock
|
||||
mode: "pre_call"
|
||||
guardrailIdentifier: "your-guardrail-id" # Your actual Bedrock guardrail ID
|
||||
guardrailVersion: "DRAFT" # or your version number
|
||||
aws_region_name: "us-east-1" # Your AWS region
|
||||
aws_role_name: "your-role-arn" # Your AWS role with Bedrock permissions
|
||||
default_on: true
|
||||
```
|
||||
|
||||
**Required AWS Setup:**
|
||||
1. Create a Bedrock guardrail in AWS Console
|
||||
2. Get the guardrail ID and version
|
||||
3. Ensure your AWS credentials have Bedrock permissions
|
||||
4. Configure the guardrail in your LiteLLM config
|
||||
|
||||
|
||||
## Usage
|
||||
---
|
||||
|
||||
In this example `mask_pii` is the guardrail name configured on LiteLLM.
|
||||
<Tabs>
|
||||
<TabItem value="presidio" label="Presidio PII Guardrail" default>
|
||||
|
||||
In this example `mask_pii` is a Presidio guardrail configured on LiteLLM.
|
||||
|
||||
```bash showLineNumbers title="Example calling the endpoint"
|
||||
curl -X POST 'http://localhost:4000/guardrails/apply_guardrail' \
|
||||
|
|
@ -23,6 +59,27 @@ curl -X POST 'http://localhost:4000/guardrails/apply_guardrail' \
|
|||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="bedrock" label="Bedrock Guardrail">
|
||||
|
||||
In this example `bedrock-content-guard` is a Bedrock guardrail configured on LiteLLM.
|
||||
|
||||
```bash showLineNumbers title="Example calling the endpoint"
|
||||
curl -X POST 'http://localhost:4000/guardrails/apply_guardrail' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-H 'Authorization: Bearer your-api-key' \
|
||||
-d '{
|
||||
"guardrail_name": "bedrock-content-guard",
|
||||
"text": "This is potentially harmful content that should be blocked",
|
||||
"language": "en"
|
||||
}'
|
||||
```
|
||||
|
||||
**Note**: For Bedrock guardrails, the `entities` parameter is not used as Bedrock handles content moderation based on its own policies.
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
|
||||
## Request Format
|
||||
---
|
||||
|
|
@ -59,12 +116,39 @@ The response will contain the processed text after applying the guardrail.
|
|||
|
||||
#### Example Response
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="presidio" label="Presidio Response" default>
|
||||
|
||||
```json
|
||||
{
|
||||
"response_text": "My name is [REDACTED] and my email is [REDACTED]"
|
||||
}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="bedrock" label="Bedrock Response">
|
||||
|
||||
```json
|
||||
{
|
||||
"response_text": "This is potentially harmful content that should be blocked"
|
||||
}
|
||||
```
|
||||
|
||||
**Note**: If Bedrock guardrail blocks the content, the endpoint will return an error with the blocking reason.
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
#### Response Fields
|
||||
- **response_text** (string):
|
||||
The text after applying the guardrail.
|
||||
|
||||
#### Error Responses
|
||||
|
||||
If a guardrail blocks content (e.g., Bedrock guardrail), the endpoint will return an error:
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": "Content blocked by Bedrock guardrail: Content violates policy"
|
||||
}
|
||||
```
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.types.llms.openai import ChatCompletionUserMessage
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
|
|
@ -30,7 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.guardrails import GuardrailEventHooks, PiiEntityType
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
||||
BedrockContentItem,
|
||||
|
|
@ -1085,3 +1086,56 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
verbose_proxy_logger.debug(
|
||||
"Applied masking to choice text content"
|
||||
)
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
text: str,
|
||||
language: Optional[str] = None,
|
||||
entities: Optional[List[PiiEntityType]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Apply Bedrock guardrail to the given text for testing purposes.
|
||||
|
||||
This method allows users to test Bedrock guardrails without making actual LLM calls.
|
||||
It creates a mock request and response to test the guardrail functionality.
|
||||
"""
|
||||
try:
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Guardrail: Applying guardrail"
|
||||
)
|
||||
mock_messages = [ChatCompletionUserMessage(role="user", content=text)]
|
||||
bedrock_response = await self.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=mock_messages,
|
||||
request_data={"messages": mock_messages}
|
||||
)
|
||||
|
||||
if bedrock_response.get("action") == "BLOCKED":
|
||||
raise Exception(f"Content blocked by Bedrock guardrail: {bedrock_response.get('reason', 'Unknown reason')}")
|
||||
|
||||
# Apply any masking that was applied by the guardrail
|
||||
masked_text = text
|
||||
if bedrock_response.get("output") and bedrock_response["output"]:
|
||||
# If the guardrail returned modified content, use that
|
||||
for output_item in bedrock_response["output"]:
|
||||
if output_item.get("text"):
|
||||
masked_text = str(output_item["text"])
|
||||
break
|
||||
elif bedrock_response.get("content") and bedrock_response["content"]:
|
||||
# Fallback to content field if output is not available
|
||||
for content_item in bedrock_response["content"]:
|
||||
if content_item.get("text") and content_item["text"].get("text"):
|
||||
masked_text = str(content_item["text"]["text"])
|
||||
break
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Guardrail: Successfully applied guardrail"
|
||||
)
|
||||
|
||||
return masked_text
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Bedrock Guardrail: Failed to apply guardrail: %s", str(e)
|
||||
)
|
||||
raise Exception(f"Bedrock guardrail failed: {str(e)}")
|
||||
|
|
@ -0,0 +1,192 @@
|
|||
"""
|
||||
Test the Bedrock guardrail apply_guardrail functionality
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from fastapi import HTTPException
|
||||
from litellm.types.guardrails import ApplyGuardrailRequest, ApplyGuardrailResponse
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_apply_guardrail_success():
|
||||
"""Test that Bedrock guardrail apply_guardrail method works correctly"""
|
||||
# Create a BedrockGuardrail instance
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guard-id",
|
||||
guardrailVersion="DRAFT"
|
||||
)
|
||||
|
||||
# Mock the make_bedrock_api_request method
|
||||
with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request:
|
||||
# Mock a successful response from Bedrock
|
||||
mock_response = {
|
||||
"action": "ALLOWED",
|
||||
"content": [
|
||||
{
|
||||
"text": {
|
||||
"text": "This is a test message with some content"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_api_request.return_value = mock_response
|
||||
|
||||
# Test the apply_guardrail method
|
||||
result = await guardrail.apply_guardrail(
|
||||
text="This is a test message with some content",
|
||||
language="en"
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert result == "This is a test message with some content"
|
||||
mock_api_request.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_apply_guardrail_blocked():
|
||||
"""Test that Bedrock guardrail apply_guardrail method handles blocked content"""
|
||||
# Create a BedrockGuardrail instance
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guard-id",
|
||||
guardrailVersion="DRAFT"
|
||||
)
|
||||
|
||||
# Mock the make_bedrock_api_request method
|
||||
with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request:
|
||||
# Mock a blocked response from Bedrock
|
||||
mock_response = {
|
||||
"action": "BLOCKED",
|
||||
"reason": "Content violates policy"
|
||||
}
|
||||
mock_api_request.return_value = mock_response
|
||||
|
||||
# Test the apply_guardrail method should raise an exception
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
text="This is blocked content",
|
||||
language="en"
|
||||
)
|
||||
|
||||
assert "Content blocked by Bedrock guardrail" in str(exc_info.value)
|
||||
assert "Content violates policy" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_apply_guardrail_with_masking():
|
||||
"""Test that Bedrock guardrail apply_guardrail method handles content masking"""
|
||||
# Create a BedrockGuardrail instance
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guard-id",
|
||||
guardrailVersion="DRAFT"
|
||||
)
|
||||
|
||||
# Mock the make_bedrock_api_request method
|
||||
with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request:
|
||||
# Mock a response with masked content
|
||||
mock_response = {
|
||||
"action": "ALLOWED",
|
||||
"content": [
|
||||
{
|
||||
"text": {
|
||||
"text": "This is a test message with [REDACTED] content"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_api_request.return_value = mock_response
|
||||
|
||||
# Test the apply_guardrail method
|
||||
result = await guardrail.apply_guardrail(
|
||||
text="This is a test message with sensitive content",
|
||||
language="en"
|
||||
)
|
||||
|
||||
# Verify the result contains the masked content
|
||||
assert result == "This is a test message with [REDACTED] content"
|
||||
mock_api_request.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_apply_guardrail_api_failure():
|
||||
"""Test that Bedrock guardrail apply_guardrail method handles API failures"""
|
||||
# Create a BedrockGuardrail instance
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guard-id",
|
||||
guardrailVersion="DRAFT"
|
||||
)
|
||||
|
||||
# Mock the make_bedrock_api_request method to raise an exception
|
||||
with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request:
|
||||
mock_api_request.side_effect = Exception("API connection failed")
|
||||
|
||||
# Test the apply_guardrail method should raise an exception
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
text="This is a test message",
|
||||
language="en"
|
||||
)
|
||||
|
||||
assert "Bedrock guardrail failed" in str(exc_info.value)
|
||||
assert "API connection failed" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_apply_guardrail_endpoint_integration():
|
||||
"""Test the full endpoint integration with Bedrock guardrail"""
|
||||
from enterprise.litellm_enterprise.proxy.guardrails.endpoints import apply_guardrail
|
||||
|
||||
# Create a real BedrockGuardrail instance
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guard-id",
|
||||
guardrailVersion="DRAFT"
|
||||
)
|
||||
|
||||
# Mock the guardrail registry
|
||||
with patch("enterprise.litellm_enterprise.proxy.guardrails.endpoints.GUARDRAIL_REGISTRY") as mock_registry:
|
||||
# Mock the make_bedrock_api_request method
|
||||
with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request:
|
||||
# Mock a successful response from Bedrock
|
||||
mock_response = {
|
||||
"action": "ALLOWED",
|
||||
"content": [
|
||||
{
|
||||
"text": {
|
||||
"text": "This is a test message with processed content"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_api_request.return_value = mock_response
|
||||
|
||||
# Configure the registry to return our guardrail
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = guardrail
|
||||
|
||||
# Create the request
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
text="This is a test message with some content",
|
||||
language="en"
|
||||
)
|
||||
|
||||
# Create a mock user API key
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
# Call the endpoint
|
||||
response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
# Verify the response
|
||||
assert isinstance(response, ApplyGuardrailResponse)
|
||||
assert response.response_text == "This is a test message with processed content"
|
||||
mock_api_request.assert_called_once()
|
||||
Loading…
Add table
Reference in a new issue