check for AWS exceptions despite a 200 response

This commit is contained in:
=
2025-09-18 15:42:36 -07:00
parent c6e37bb0bb
commit a86b9a1808
2 changed files with 109 additions and 14 deletions
@@ -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,
)
#########################################################
@@ -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