Implement Bedrock Guardrail apply_guardrail endpoint support

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.
+<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"
+}
+```

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -23,14 +23,15 @@
 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,
     httpxSpecialProvider,
 )
 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)}")

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()