Update bedrock_guardrails.py

This commit is contained in:
= 2025-09-18 16:05:30 -07:00
parent a3f0a3c05f
commit 2bbcf5a851

View file

@ -406,10 +406,38 @@ 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 _check_bedrock_response_for_exception(self, response) -> bool:
"""
Return True if the Bedrock ApplyGuardrail response indicates an exception.
Works with real httpx.Response objects and MagicMock responses used in tests.
"""
payload = None
try:
json_method = getattr(response, "json", None)
if callable(json_method):
payload = json_method()
except Exception:
payload = None
if payload is None:
try:
raw = getattr(response, "content", None)
if isinstance(raw, (bytes, bytearray)):
payload = json.loads(raw.decode("utf-8"))
else:
text = getattr(response, "text", None)
if isinstance(text, str):
payload = json.loads(text)
except Exception:
# Can't parse -> assume no explicit Exception marker
return False
if not isinstance(payload, dict):
return False
return "Exception" in payload.get("Output", {}).get("__type", "")
def _get_bedrock_guardrail_response_status(
self, response: httpx.Response
@ -423,6 +451,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
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: