diff --git a/docs/my-website/docs/apply_guardrail.md b/docs/my-website/docs/apply_guardrail.md index 740eb232e134..18fe951c52ac 100644 --- a/docs/my-website/docs/apply_guardrail.md +++ b/docs/my-website/docs/apply_guardrail.md @@ -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. + + + +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' \ }' ``` + + + +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. + + + + ## Request Format --- @@ -59,12 +116,39 @@ The response will contain the processed text after applying the guardrail. #### Example Response + + + ```json { "response_text": "My name is [REDACTED] and my email is [REDACTED]" } ``` + + + +```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. + + + + #### 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" +} +``` diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 99f3e4cf2437..a6dc1d57886c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -23,6 +23,7 @@ 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.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 @@ def _apply_masking_to_model_response( 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)}") \ No newline at end of file diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py new file mode 100644 index 000000000000..fbdf390faf31 --- /dev/null +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -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()