mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-20 06:23:46 +00:00
check for AWS exceptions despite a 200 response
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user