diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index dfb01a7cd0..28819024be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -383,6 +383,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) ######################################################### if response.status_code == 200: + # check if the response contains an error + if self._check_bedrock_response_for_exception(response=response): + raise self._get_http_exception_for_failed_guardrail(response) # check if the response was flagged _json_response = response.json() redacted_response = _redact_pii_matches(_json_response) @@ -403,6 +406,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return bedrock_guardrail_response + def _check_bedrock_response_for_exception(self, response: httpx.Response) -> bool: + return "Exception" in json.loads(response.content.decode("utf-8")).get( + "Output", {} + ).get("__type", "") + def _get_bedrock_guardrail_response_status( self, response: httpx.Response ) -> Literal["success", "failure"]: @@ -410,9 +418,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Get the status of the bedrock guardrail response. """ if response.status_code == 200: + if self._check_bedrock_response_for_exception(response): + return "failure" return "success" return "failure" + def _get_http_exception_for_failed_guardrail( + self, response: httpx.Response + ) -> HTTPException: + return HTTPException( + status_code=400, + detail={ + "error": "Guardrail application failed.", + "bedrock_guardrail_response": json.loads( + response.content.decode("utf-8") + ).get("Output", {}), + }, + ) + def _get_http_exception_for_blocked_guardrail( self, response: BedrockGuardrailResponse ) -> HTTPException: @@ -562,11 +585,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 2. Update the messages with the guardrail response ########## ######################################################### - data["messages"] = ( - self._update_messages_with_updated_bedrock_guardrail_response( - messages=new_messages, - bedrock_guardrail_response=bedrock_guardrail_response, - ) + data[ + "messages" + ] = self._update_messages_with_updated_bedrock_guardrail_response( + messages=new_messages, + bedrock_guardrail_response=bedrock_guardrail_response, ) ######################################################### @@ -617,11 +640,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 2. Update the messages with the guardrail response ########## ######################################################### - data["messages"] = ( - self._update_messages_with_updated_bedrock_guardrail_response( - messages=new_messages, - bedrock_guardrail_response=bedrock_guardrail_response, - ) + data[ + "messages" + ] = self._update_messages_with_updated_bedrock_guardrail_response( + messages=new_messages, + bedrock_guardrail_response=bedrock_guardrail_response, ) ######################################################### diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index a3f17a1424..19cf1a0277 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1,12 +1,13 @@ """ Unit tests for Bedrock Guardrails """ - +import json import os import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException sys.path.insert(0, os.path.abspath("../../../../../..")) @@ -238,7 +239,6 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): ) as mock_load_creds, patch.object( guardrail, "_prepare_request", return_value=MagicMock() ) as mock_prepare_request: - mock_post.return_value = mock_bedrock_response # Call the method that should log the redacted response @@ -345,7 +345,6 @@ async def test_bedrock_guardrail_original_response_not_modified(): ) as mock_load_creds, patch.object( guardrail, "_prepare_request", return_value=MagicMock() ) as mock_prepare_request: - mock_post.return_value = mock_bedrock_response # Call the method @@ -860,7 +859,6 @@ async def test__redact_pii_matches_comprehensive_coverage(): print("Comprehensive coverage redaction test passed") - @pytest.mark.asyncio async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): """Test that BedrockGuardrail respects aws_bedrock_runtime_endpoint when set""" @@ -1049,3 +1047,77 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" print(f"Parameter precedence test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_200_with_exception_in_output_raises_and_logs_failure(): + """ + When Bedrock returns HTTP 200 but the body contains Output.__type with 'Exception', + the guardrail should: + - raise an HTTPException(400) with the Output payload in detail + - log the request trace with guardrail_status='failure' + """ + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock a Bedrock "success" HTTP status but an Exception embedded in the body + payload = { + "Output": { + "__type": "com.amazonaws#InternalServerException", + "message": "Something went wrong upstream", + }, + "action": "NONE", + } + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.content = json.dumps(payload).encode("utf-8") + mock_resp.text = json.dumps(payload) + mock_resp.json.return_value = payload + + # Minimal request data + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + } + + # Mock creds and request prep + mock_credentials = MagicMock() + mock_credentials.access_key = "ak" + mock_credentials.secret_key = "sk" + mock_credentials.token = None + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, + "_prepare_request", + return_value=MagicMock(url="http://example", headers={}, body=b""), + ), patch.object( + guardrail, "add_standard_logging_guardrail_information_to_request_data" + ) as mock_add_trace: + mock_post.return_value = mock_resp + + with pytest.raises(HTTPException) as excinfo: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + ) + + # 1) Raised HTTPException with 400 status + err = excinfo.value + assert err.status_code == 400 + assert err.detail["error"] == "Guardrail application failed." + + # 2) Detail includes the Output object from the Bedrock body + assert err.detail["bedrock_guardrail_response"] == payload["Output"] + + # 3) Trace logging received a 'failure' status + assert mock_add_trace.called + _, kwargs = mock_add_trace.call_args + assert kwargs["guardrail_status"] == "failure" + # And the JSON passed to tracing is the same response we received + assert kwargs["guardrail_json_response"] == payload